Principle:Mlc ai Mlc llm JIT Model Preparation

Knowledge Sources	MLC-LLM TVM Documentation
Domains	Deep_Learning, Model_Serving, Compiler_Optimization
Last Updated	2026-02-09 00:00 GMT

Overview

Just-in-time (JIT) model preparation is the technique of compiling machine learning model libraries on demand at runtime rather than ahead of time, with deterministic caching for reuse across subsequent invocations.

Description

In large language model serving systems, models must be compiled into hardware-specific libraries before inference can begin. Traditionally, this compilation step is performed ahead of time (AOT) as a separate build phase, producing platform-specific shared objects (e.g., .so files on Linux, .dll on Windows) or archive bundles (e.g., .tar for mobile platforms). While AOT compilation ensures that the serving path incurs no compilation overhead, it introduces friction in the deployment workflow: every combination of model architecture, quantization scheme, optimization level, and target device requires a distinct pre-built artifact.

JIT model preparation addresses this by deferring compilation to the first time a particular configuration is requested. The core workflow proceeds as follows:

1. Configuration Fingerprinting: A deterministic hash is computed from the full set of compilation parameters, including model type, quantization scheme, model configuration overrides, optimization flags, and target device. This hash uniquely identifies the compiled artifact.
2. Cache Lookup: The system checks whether a compiled library matching the hash already exists in a persistent cache directory. If a cache hit occurs, the pre-compiled library is loaded directly, incurring negligible overhead.
3. On-Demand Compilation: On a cache miss, the system invokes the full model compilation pipeline as a subprocess, producing the compiled library in a temporary directory. Upon successful compilation, the artifact is atomically moved to the cache directory.
4. Policy Control: The behavior is governed by a policy setting (e.g., MLC_JIT_POLICY) that supports modes such as ON (compile on miss, use cache on hit), OFF (never compile, raise an error), REDO (always recompile), and READONLY (use cache only, error on miss).

This approach combines the deployment convenience of dynamic compilation with the performance benefits of caching, ensuring that repeated invocations of the same model configuration do not incur redundant compilation.

Usage

JIT model preparation is appropriate in the following scenarios:

• Development and Prototyping: When experimenting with different model configurations, quantization schemes, or optimization levels, JIT compilation eliminates the need for a separate build step.
• First-Time Deployment: When deploying a model to a new hardware target where pre-compiled libraries are not yet available.
• Multi-Model Serving: In systems that serve multiple models with varying configurations, JIT allows on-demand preparation of each model without pre-building every variant.
• CI/CD Pipelines: Automated testing and deployment workflows benefit from JIT by reducing the number of pre-built artifacts that must be maintained.

JIT compilation is less suitable for latency-critical cold starts in production environments where the initial compilation delay is unacceptable. In such cases, the READONLY policy can enforce cache-only behavior after an initial warm-up phase.

Theoretical Basis

Content-Addressable Caching

The caching mechanism is based on content-addressable storage. A cryptographic hash function (MD5 in this case) maps the full compilation configuration to a fixed-length digest. This digest serves as the file name for the cached artifact. The properties of this approach are:

• Determinism: Identical configurations always produce the same hash, guaranteeing cache hits for repeated requests.
• Collision Resistance: Different configurations are extremely unlikely to produce the same hash, preventing incorrect library reuse.
• Composability: The hash input is a JSON-serialized dictionary of all relevant parameters, making it straightforward to extend with new configuration axes.

Compilation as Subprocess

The compilation itself is delegated to a subprocess invocation of the same Python runtime. This design provides process-level isolation, ensuring that compilation failures or resource exhaustion do not corrupt the serving process. The subprocess approach also enables parallelism, as multiple JIT compilations can proceed concurrently for different model configurations.

Pseudocode

function jit(model_path, overrides, device):
    config = load_model_config(model_path)
    hash_key = {
        model_config: config,
        overrides: serialize(overrides),
        opt: optimization_flags,
        device: device,
        model_type: config.model_type,
        quantization: config.quantization
    }
    hash_value = md5(json_serialize(hash_key))
    cached_path = CACHE_DIR / f"{hash_value}.{lib_suffix}"

    if exists(cached_path) and policy in [ON, READONLY]:
        return load(cached_path)

    if policy == READONLY:
        raise Error("No cached library found")

    compile_model(model_path, overrides, device, temp_output)
    atomic_move(temp_output, cached_path)
    return load(cached_path)

Related Pages

Implemented By

• Implementation:Mlc_ai_Mlc_llm_Jit