torch.compile and CUDA graphs in vLLM

Code version: vLLM @ 3e44c0d · 2026-10-22

The idea#

A decoder step launches hundreds of small kernels. Two tools cut this cost (PyTorch for serving: custom ops and torch.compile has the PyTorch basics):

• torch.compile captures the model's Python code as one graph and lets Inductor fuse the small ops into fewer kernels. vLLM does it once per process, with no guards, and keeps the token count as a dynamic dimension.
• CUDA graphs record the kernel launches of one batch shape at start-up and replay them later with one call (CUDA graph capture and replay↗).

Attention is the hard part. Its kernels depend on the batch layout, and some backends cannot run inside a graph for every layout. So vLLM cuts the compiled graph at the attention ops, the splitting ops. Each piece between two attention ops gets its own piecewise CUDA graph, and attention runs eagerly between them. For pure decode batches, it also captures a full CUDA graph of the whole forward pass, attention included.

The optimization level sets the defaults. At the default -O2, the compilation mode is VLLM_COMPILE and the CUDA graph mode is FULL_AND_PIECEWISE. Each step, Model Runner V2 looks at the batch and picks FULL, PIECEWISE or NONE (eager).

An everyday comparison. Think of a factory line with fixed robot programs. The engineers write one program per standard product size. They leave a manual station in the middle, for the step that changes with every order. Small standard orders run one complete program with no stop (FULL). Other orders run the programs with a stop at the manual station (PIECEWISE). A giant order runs fully by hand (eager). The comparison stops here: the programs are recorded once at start-up, for a fixed list of sizes. An order is padded up to the next size.

The picture#

Panel 1 is the one-time compile. Panel 2 shows what runs in a PIECEWISE step: blue pieces replay graphs, amber attention pieces run Python and the backend kernel. Panel 3 shows why FULL and PIECEWISE graphs do not interfere: a wrapper replays only when the forward context says PIECEWISE.

In the code#

Step 1: optimization levels and their defaults#

The OptimizationLevel enum has four values; O2 is the default (vllm/config/vllm.py:131-143, vllm/config/vllm.py:434). The CLI accepts -O2, -O 2, -O=2 and --optimization-level 2 (vllm/utils/argparse_utils.py:395-406).

The compile mode is not in the level tables. VllmConfig.__post_init__ sets it: above O0, an unset mode becomes VLLM_COMPILE (vllm/config/vllm.py:1806-1810). Then the level's table fills only the fields that are still None, so your own settings win (vllm/config/vllm.py:1005-1019, vllm/config/vllm.py:1833-1834).

Later checks can lower the CUDA graph mode again. A pooling model gets PIECEWISE instead of FULL graphs; an encoder-decoder model gets FULL_DECODE_ONLY (vllm/config/vllm.py:1932-1954). A KV connector can ask for PIECEWISE (vllm/config/vllm.py:1957-1980). At KV cache init, the attention backends can lower it once more (Attention backends: backend, metadata builder and implementation).

Step 2: @support_torch_compile marks what to compile#

A model class opts in with the decorator; in Llama it sits on LlamaModel (How a vLLM model file is written (llama.py)). Without explicit dynamic_arg_dims, dimension 0 of every argument annotated as a tensor or IntermediateTensors is dynamic: the token count (vllm/compilation/decorators.py:214-232). The decorator adds TorchCompileWithNoGuardsWrapper as a base class (vllm/compilation/decorators.py:360-367). It skips compilation when the mode is NONE or STOCK_TORCH_COMPILE, when the class is marked with ignore_torch_compile, or when enable_if says no (vllm/compilation/decorators.py:409-418).

The wrapper builds the compiled callable once:

        with aot_context:
            self._compiled_callable = torch.compile(
                compiled_ptr,
                fullgraph=True,
                dynamic=False,
                backend=backend,
                options=options,
            )

fullgraph=True forbids graph breaks; dynamic=False plus the marked dimensions give one graph with a symbolic token count. Outside STOCK_TORCH_COMPILE all guards are dropped, so Dynamo never traces again after the first call (vllm/compilation/wrapper.py:65-70, vllm/compilation/wrapper.py:126-146). The backend comes from init_backend: a VllmBackend in VLLM_COMPILE mode (vllm/config/compilation.py:1083-1134).

The first call happens in the first forward pass of the worker, the profile run with max_num_batched_tokens dummy tokens (vllm/v1/worker/gpu/model_runner.py:1022-1027). The worker runs it inside its memory profiling, and a comment says that this run "compiles the model" (vllm/v1/worker/gpu_worker.py:654-655, vllm/v1/worker/gpu_worker.py:681-697). So on a start without a cache, the compile time is part of the memory profiling time.

Step 3: VllmBackend hashes and caches#

Dynamo calls VllmBackend with the traced graph. First it decides the compile cache folder:

        env_factors = envs.compile_factors()
        env_hash = hash_factors(env_factors)
        # Compute config/compiler/code hashes once and reuse
        config_hash = vllm_config.compute_hash()
        compiler_hash = self.compiler_manager.compute_hash(vllm_config)
        forward_code_files = list(sorted(self.compilation_config.traced_files))

        logger.debug(
            "Traced files (to be considered for compilation cache):\n%s",
            lazy(lambda: "\n".join(forward_code_files)),
        )
        hash_content = []
        for filepath in forward_code_files:
            if filepath == "<string>":
                # This means the function was dynamically generated, with
                # e.g. exec(). We can't actually check these.
                continue
            hash_content.append(filepath)
            try:
                with open(filepath) as f:
                    hash_content.append(f.read())
            except (OSError, UnicodeDecodeError):
                logger.warning("Failed to read file %s", filepath)
                continue
        code_hash = hashlib.sha256("\n".join(hash_content).encode()).hexdigest()
        # Clear after consumption
        self.compilation_config.traced_files.clear()
        if not self.compilation_config.cache_dir:
            # no provided cache dir, generate one based on the known factors
            # that affects the compilation. if none of the factors change,
            # the cache dir will be the same so that we can reuse the compiled
            # graph.
            factors = [env_hash, config_hash, code_hash, compiler_hash]
            # Use SHA-256 for cache key hashing to be consistent across
            # compute_hash functions. Truncate for a short cache dir name.
            hash_key = hashlib.sha256(str(factors).encode()).hexdigest()[:10]
            cache_dir = os.path.join(
                envs.VLLM_CACHE_ROOT, "torch_compile_cache", hash_key
            )
            self.compilation_config.cache_dir = cache_dir

1. Env factors. Every known vLLM environment variable, minus a list of ignored ones such as VLLM_PORT and VLLM_CACHE_ROOT (vllm/envs.py:2351-2381).
2. Config hash. VllmConfig.compute_hash(): the vLLM version and the hashes of the model, cache, parallel, scheduler, compilation and other configs, cut to 10 hex digits (vllm/config/vllm.py:459-501, vllm/config/vllm.py:571-574). The scheduler part holds max_num_batched_tokens, because LoRA buffers and Inductor's index width depend on it, and max_num_seqs, because some static per-request buffers depend on it (vllm/config/scheduler.py:271-284).
3. Code hash. The source of every file that Dynamo traced (lines 1064-1083).
4. Compiler hash. From the Inductor adaptor (vllm/compilation/backends.py:159-160).

The folder is $VLLM_CACHE_ROOT/torch_compile_cache/<10 hex digits>/rank_<rank>_<dp>/<prefix>, where the prefix is backbone for the main model (vllm/compilation/backends.py:1103-1107, vllm/compilation/backends.py:804). VLLM_CACHE_ROOT defaults to ~/.cache/vllm when XDG_CACHE_HOME is not set (vllm/envs.py:710-718). vLLM computes the hash folder only when compilation_config.cache_dir is empty (vllm/compilation/backends.py:1086-1098). VLLM_DISABLE_COMPILE_CACHE=1 turns the cache off (vllm/compilation/compiler_interface.py:204-219).

On PyTorch 2.10 or newer, vLLM also uses ahead-of-time (AOT) compilation by default: VLLM_USE_AOT_COMPILE is 1 unless the compile cache is off (vllm/envs.py:378-387). The decorator then first tries to load a saved compiled function from torch_compile_cache/torch_aot_compile/<hash>/rank_<rank>_<dp>/model. After a fresh compile it saves one there (vllm/compilation/decorators.py:556-605, vllm/compilation/decorators.py:731-756).

So both folders can exist side by side. On a start without a cache, Dynamo traces the model, and VllmBackend compiles the pieces into its hash folder. Then the decorator saves the AOT file. On the next start, a loaded AOT file skips the Dynamo trace (vllm/compilation/decorators.py:595-605). From the saved graph, VllmBackend runs again and finds its compiled pieces in its own folder (vllm/compilation/caching.py:351-380). With VLLM_USE_MEGA_AOT_ARTIFACT (default on with AOT on PyTorch 2.12 or newer), the AOT file holds the Inductor code too (vllm/envs.py:390-397, vllm/compilation/caching.py:326-349).

Step 4: split_graph cuts at the splitting ops#

Without Inductor graph partition (the default at every level, vllm/config/vllm.py:311), the backend splits the FX graph itself:

        if self.compilation_config.use_inductor_graph_partition:
            # Let Inductor decide partitioning; avoid FX-level pre-splitting.
            fx_split_ops: list[str] = []
        else:
            fx_split_ops = self.compilation_config.splitting_ops or []

        self.split_gm, self.piecewise_graphs = split_graph(graph, fx_split_ops)

The default list is the attention ops plus the two KV cache update ops (vllm/config/compilation.py:1175-1213):

• vllm::unified_attention_with_output, vllm::unified_mla_attention_with_output, the Mamba and linear-attention ops and some model-specific ops: 17 names in _attention_ops (vllm/config/compilation.py:768-788);
• vllm::unified_kv_cache_update and vllm::unified_mla_kv_cache_update, because the string argument "prevents Inductor from reusing piecewise graphs" (vllm/config/compilation.py:1186-1187).

split_graph walks the nodes in order. A splitting op starts a new subgraph, and consecutive splitting ops stay in one subgraph (vllm/compilation/backends.py:589-617). Only the pieces that are not splitting graphs are compiled (vllm/compilation/backends.py:1235-1239). The attention subgraphs stay plain FX code that calls the custom ops, so they run eagerly.

If attention-quant fusion is on (Inductor graph partition off), vLLM must keep attention inside the graph. If you did not set splitting_ops, it sets [], and a mode with piecewise graphs becomes FULL (vllm/config/compilation.py:1260-1278).

Step 5: PiecewiseBackend and CUDAGraphWrapper#

PiecewiseCompileInterpreter runs the split graph with fake tensors. For each piece to compile it makes a PiecewiseBackend, wraps it, and puts the wrapper in place of the piece (vllm/compilation/backends.py:765-795). The backend compiles the piece with Inductor once per compile range. The default range list has one range, from 1 to max_num_batched_tokens; all-reduce and RoPE/KV-cache fusions can add an extra endpoint (vllm/config/vllm.py:2637-2645, vllm/compilation/piecewise_backend.py:244-276). At run time it picks the compiled code of the range that holds the token count (vllm/compilation/piecewise_backend.py:357-379).

The wrapper is added only when the CUDA graph mode has a piecewise part and Inductor graph partition is off (vllm/compilation/backends.py:656-707):

    def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
        if not is_forward_context_available():
            # No forward context means we are outside the normal
            # inference path (e.g. a vision encoder forward pass).
            # Just run the underlying function without cudagraphs.
            return self.runnable(*args, **kwargs)

        forward_context = get_forward_context()
        batch_descriptor = forward_context.batch_descriptor
        cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode

        if (
            cudagraph_runtime_mode == CUDAGraphMode.NONE
            or cudagraph_runtime_mode != self.runtime_mode
        ):
            # CUDAGraphMode.NONE could mean the profile run, a warmup run, or
            # running without cudagraphs.
            # We do not trigger capture/replay if the runtime mode is not
            # matches. This enables properly dispatching to the correct
            # CUDAGraphWrapper when nesting multiple instances with different
            # runtime modes.
            return self.runnable(*args, **kwargs)

The wrapper reads the forward context. It replays only when the run-time mode equals its own mode, PIECEWISE. Then it looks up a graph by the BatchDescriptor: the padded token count and the LoRA flags (vllm/forward_context.py:29-57). A missing entry is captured on the spot, which is allowed only during the capture phase (vllm/compilation/cuda_graph.py:258-279).

Step 6: Inductor post-grad passes#

VllmBackend adds vLLM's own passes to Inductor. PostGradPassManager.configure picks them from PassConfig (vllm/compilation/passes/pass_manager.py:139-223). The flags start as None and get their values from the level table, as functions of the config (vllm/config/vllm.py:296-315):

Take a BF16 Llama with TP = 1 on CUDA, at O2. The CustomOps are off by default (CustomOp, out-of-tree ops and IR op providers), and the CUDA IR priority of rms_norm starts with native when Inductor compiles (vllm/platforms/cuda.py:716-731). So of these flags only eliminate_noops is on (derived from the table). The pass manager always adds its IR lowering and clean-up passes (vllm/compilation/passes/pass_manager.py:119-134). With TP = 2 on an H100 and FlashInfer installed, fuse_allreduce_rms joins it. vLLM logs the list as "Enabled custom fusions: ..." (vllm/config/compilation.py:301-320).

Step 7: CUDA graph modes and capture sizes#

CUDAGraphMode has three run-time modes and two pairs (vllm/config/compilation.py:53-63). A pair has one mode for uniform decode batches and one for mixed batches (vllm/config/compilation.py:65-69):

The capture sizes are padded token counts. When you do not set max_cudagraph_capture_size, the maximum is min(max_num_seqs × decode_query_len × 2, 512), or 1024 on data-center Blackwell (compute capability 10.x) (vllm/config/vllm.py:2440-2449). It is also capped at max_num_batched_tokens (vllm/config/vllm.py:2530-2531). The list is 1, 2, 4, then steps of 8 below 256, then steps of 16 (vllm/config/vllm.py:2557-2570). max_num_batched_tokens joins the list if it is below the maximum (lines 2571-2576).

Step 8: Model Runner V2 captures the graphs#

The runner makes a ModelCudaGraphManager at KV cache init (vllm/v1/worker/gpu/model_runner.py:764-772). Its _init_candidates turns the mode and the sizes into a list of BatchExecutionDescriptors to capture (vllm/v1/worker/gpu/cudagraph_utils.py:256-394). After warm-up, the worker calls capture_model(), unless --enforce-eager is set (vllm/v1/worker/gpu_worker.py:983-985). The manager captures all PIECEWISE descriptors first, then all FULL ones, the largest size first in each mode (vllm/v1/worker/gpu/cudagraph_utils.py:447-449, vllm/v1/worker/gpu/cudagraph_utils.py:376-378):

• PIECEWISE. A warm-up run, then a run in the forward context with mode PIECEWISE. Each piece's wrapper records its graph (vllm/v1/worker/gpu/cudagraph_utils.py:464-479).
• FULL. A warm-up run, fresh inputs, then torch.cuda.graph(...) around the whole forward with the mode NONE inside. The piece wrappers pass through, so the outer graph records everything (vllm/v1/worker/gpu/cudagraph_utils.py:480-512).

The comment gives the order's reason: "PIECEWISE has larger activations so FULL activations should fit in already allocated buffers in the graph pool" (vllm/v1/worker/gpu/cudagraph_utils.py:444-446).

Step 9: the dispatch rule for each step#

execute_model calls dispatch_cg_and_sync_dp before it prepares the inputs (vllm/v1/worker/gpu/model_runner.py:1806-1819). It passes the uniform decode token count. This count is set only when every request has the same query length and the batch has no real prefill (vllm/v1/worker/utils.py:958-976, vllm/v1/worker/gpu/model_runner.py:1335-1372). The manager then takes the candidates for this token count and returns the first compatible one, else NONE:

    def dispatch(
        self,
        num_reqs: int,
        num_tokens: int,
        uniform_token_count: int | None,
        num_active_loras: int,
        max_query_len: int | None = None,
        num_ubatches: int = 1,
    ) -> BatchExecutionDescriptor:
        """Find matching cudagraph descriptor from priority-ordered candidates."""
        effective_loras = self._resolve_effective_loras(num_active_loras)
        key = (num_tokens, effective_loras)
        if self._graphs_captured and num_tokens > 0 and key in self._candidates:
            for desc in self._candidates[key]:
                if _is_compatible(
                    desc,
                    num_reqs,
                    num_tokens,
                    uniform_token_count,
                    effective_loras,
                    max_query_len,
                    num_ubatches,
                ):
                    return desc
        return BatchExecutionDescriptor(
            cg_mode=CUDAGraphMode.NONE,
            num_tokens=num_tokens,
            num_reqs=num_reqs,
            num_active_loras=effective_loras,
            num_ubatches=num_ubatches,
        )

The compatibility test:

def _is_compatible(
    desc: BatchExecutionDescriptor,
    num_reqs: int,
    num_tokens: int,
    uniform_token_count: int | None,
    num_active_loras: int,
    max_query_len: int | None,
    num_ubatches: int,
) -> bool:
    # desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count
    # desc.num_reqs=None means no request padding needed (PIECEWISE)
    # desc.max_query_len=None means the graph does not constrain query length; a
    # caller that does not track max_query_len must not match one that does
    # A graph captured for N microbatches can only serve a batch split N ways.
    return (
        (
            desc.uniform_token_count is None
            or desc.uniform_token_count == uniform_token_count
        )
        and (
            desc.max_query_len is None
            or (max_query_len is not None and desc.max_query_len >= max_query_len)
        )
        and (desc.num_reqs is None or desc.num_reqs >= num_reqs)
        and desc.num_tokens >= num_tokens
        and desc.num_active_loras == num_active_loras
        and desc.num_ubatches == num_ubatches
    )

1. FULL decode graph. Its uniform_token_count is decode_query_len: 1 without speculative decoding, 1 + k with k drafts (vllm/v1/worker/gpu/model_runner.py:442-445). The batch must have the same count, and the graph enough request slots.
2. Varlen decode FULL variant. Only with adaptive verification: a FULL graph for decode batches whose requests have 1 to decode_query_len tokens (vllm/v1/worker/gpu/cudagraph_utils.py:302-318). A batch with a prefill passes max_query_len=None and cannot use it (vllm/v1/worker/gpu/model_runner.py:1780-1785).
3. PIECEWISE. No request count and no uniform count, so any batch up to its token count fits (vllm/v1/worker/gpu/cudagraph_utils.py:352-370).
4. Above the largest size. The token count has no candidates, so the result is NONE (line 553).

With DP > 1, the ranks agree on one padded size with a collective (vllm/v1/worker/gpu/dp_utils.py:278-312).

Step 10: breakable CUDA graphs#

A breakable CUDA graph gives piecewise graphs without torch.compile. Instead of cutting an FX graph, one capture context runs the whole forward. Each attention op ends the current graph segment, runs eagerly, and a new segment starts (vllm/compilation/breakable_cudagraph.py:3-20). The ops that break are marked with @eager_break_during_capture, for example unified_attention_with_output (vllm/model_executor/layers/attention/attention.py:767-769). The decorator does nothing outside a capture, and nothing during a FULL capture (vllm/compilation/breakable_cudagraph.py:99-118). In MRV2, run_pw_graph uses the BreakableCUDAGraphWrapper instead of the compiled model when this mode is on (vllm/v1/worker/gpu/cudagraph_utils.py:578-583). The environment variable is described as experimental (vllm/envs.py:776-779).

Step 11: the start-up cost, and how to reduce it#

Compilation and capture are the slow part of vllm serve start-up. The logs show each part: "Dynamo bytecode transform time" (vllm/compilation/backends.py:1186-1190), "Compiling a graph for compile range ... takes ... s", once per compile range after its last piece (vllm/compilation/backends.py:415-422), and "Graph capturing finished in ... secs, took ... GiB" (vllm/v1/worker/gpu/model_runner.py:1146-1152). The comment there says capture "usually takes 5~20 seconds".

• Keep the compile cache. The second start with the same config, code and env loads the compiled pieces. Mount ~/.cache/vllm (or VLLM_CACHE_ROOT) as a volume in containers. Changing --max-num-batched-tokens or --max-num-seqs changes the hash.
• -O0: no compile, no CUDA graphs, fastest start, slowest steps. --enforce-eager has the same effect on compile and graphs.
• -O1: compile with PIECEWISE graphs only, so no FULL capture.
• Fewer sizes. Set -cc.max_cudagraph_capture_size or -cc.cudagraph_capture_sizes. Fewer sizes mean shorter capture and less memory, but more padding or eager steps.
• -cc.cudagraph_mode=FULL_DECODE_ONLY: no piecewise graphs; the docstring suggests it for decode instances in prefill/decode disaggregation (vllm/config/compilation.py:628-631).

Deep dive#

The sections below give the full detail. Open them one at a time, or use Open all deep dives at the top of the page.

What can go wrong#

• Slow first start, fast second start. This is the compile cache at work. If every start is slow, something in the hash changes: a different --max-num-batched-tokens, another environment variable, or a lost cache volume.
• Large steps run eagerly. Batches above the largest capture size get no CUDA graph (Step 9). This is normal for long prefill chunks, where launch overhead is small.
• A mode quietly lowered. A backend, a pooling model, an encoder-decoder model or a KV connector can change FULL_AND_PIECEWISE with only a log line (Step 1, Attention backends: backend, metadata builder and implementation).
• A model without a compiled part. If the model has no compiled submodule and breakable graphs are off, piecewise capture is not possible. At KV cache init the runner then lowers the mode to its non-piecewise form, with a warning (vllm/config/compilation.py:1467-1480). The capture code checks the same condition and raises "piecewise CUDA graphs ... unavailable" (vllm/v1/worker/gpu/cudagraph_utils.py:633-641).
• FULL with a weak backend. -cc.cudagraph_mode=FULL with a NEVER backend stops start-up with ValueError (vllm/config/compilation.py:1394-1410).
• Docs that do not match. docs/design/torch_compile.md shows a cache folder torch_compile_cache/1517964802/rank_0_0 (docs/design/torch_compile.md:15); the code uses 10 hex digits and adds the prefix folder. docs/design/cuda_graphs.md describes the dispatcher of Model Runner V1 (docs/design/cuda_graphs.md:98-118); Model Runner V2 uses CudaGraphManager.dispatch.

Why it matters. These settings decide both the start-up time and the speed of every decode step. For a latency-sensitive deployment, keep O2, keep the compile cache on disk, and check in the log that FULL decode graphs survived start-up. For fast restarts during development, use -O0 or --enforce-eager.