Compilation, CUDA graphs, and custom ops

Active contributors: youkaichao, Lucas Wilkinson, Wentao Ye.

Purpose

vLLM's hot path is the model forward. To remove Python overhead, vLLM uses torch.compile (with custom passes), captures replayable CUDA graphs, and dispatches into hand-written kernels through CustomOp. Together these pieces decide what the GPU actually executes per step.

Directory layout

vllm/compilation/
├── backends.py             # The compile backend (50 KB)
├── caching.py              # Cache compiled artifacts to disk
├── codegen.py              # Inductor codegen overrides
├── compiler_interface.py   # `compile_fx` shim
├── cuda_graph.py           # CUDAGraphStat, capture/replay
├── decorators.py           # `@support_torch_compile`, `@cache_compiled`, ...
├── monitor.py              # Compile-time monitoring
├── partition_rules.py      # Splits the FX graph for piecewise capture
├── piecewise_backend.py    # Piecewise CUDA-graph backend
├── wrapper.py              # Wraps a model in compile + graph capture
├── base_static_graph.py
├── counter.py              # Compile-step counter
└── passes/                 # Custom Inductor passes (fusion, memory, etc.)

vllm/v1/cudagraph_dispatcher.py  # Per-step CUDA-graph selection

vllm/_custom_ops.py             # ~3,000-line Python surface for custom CUDA ops
vllm/_aiter_ops.py              # ROCm AITER op surface
vllm/_xpu_ops.py                # XPU op surface
vllm/model_executor/custom_op.py # CustomOp base class

csrc/                           # Native kernels (~82 K lines)
csrc/cumem_allocator.cpp        # Per-block CUDA memory allocator

Compilation modes

CompilationConfig.compilation_mode (vllm/config/compilation.py):

Mode	Behavior
NO_COMPILATION	Pure eager forward. Slowest, most flexible.
STOCK_TORCH_COMPILE	Wrap the model in torch.compile with default settings.
VLLM_COMPILE (default)	vLLM's tuned torch.compile invocation: custom passes, fused attention, donated buffers, FX partitioning.

The compile backend (vllm/compilation/backends.py) extends Inductor with passes from vllm/compilation/passes/ (e.g., RMSNorm + linear fusion, attention rewriting, custom-op rewriting, moe-fusion).

CUDA graph modes

CompilationConfig.cuda_graph_mode (vllm/config/compilation.py::CUDAGraphMode):

Mode	Behavior
NONE	No graphs. Always eager.
PIECEWISE	Partition the FX graph at attention boundaries; capture deterministic chunks; fall back to eager for the rest.
FULL	Capture the entire forward as one graph. Lowest overhead, strictest shape constraints.
FULL_AND_PIECEWISE	Use full graphs for hot shapes; piecewise for the long tail.

The dispatcher (vllm/v1/cudagraph_dispatcher.py) decides which captured graph (if any) to replay each step based on the current batch shape, attention metadata, and active LoRAs.

Custom ops

A CustomOp (vllm/model_executor/custom_op.py) is a torch.nn.Module that selects between several backend implementations (eager Python, native CUDA, Triton, ROCm AITER, XPU). For each backend you implement a forward_<backend> method; the dispatch happens in forward() based on current_platform and config flags.

Custom ops are registered with torch.library so that torch.compile can pattern-match and replace them. The Python surface for the underlying CUDA kernels is in vllm/_custom_ops.py (~3,000 lines) — every csrc/ symbol surfaces through there.

# vllm/model_executor/layers/layernorm.py — typical pattern
class RMSNorm(CustomOp):
    def forward_native(self, x): ...
    def forward_cuda(self, x):
        return ops.rms_norm(x, self.weight, self.epsilon)
    def forward_rocm(self, x):
        return aiter_ops.rms_norm(x, self.weight, self.epsilon)

Pass framework

vllm/compilation/passes/ is a small directory of FX-graph passes invoked by the compile backend. Examples:

• Attention rewriting — fuses pre-attention RMSNorm + QKV projection.
• MoE fusion — combines router + dispatch + experts where possible.
• Memory passes — donated buffers, in-place ops, KV-write fusion.

Each pass is opt-in via CompilationConfig.pass_config (PassConfig).

Caching

Compiled artifacts are cached to disk so that the second vllm serve of the same model boots fast. The cache key includes:

• Model hash, dtype, parallelism layout
• Compilation mode + pass config
• CUDA / NCCL / Triton versions
• Compiled-out config flags

vllm/compilation/caching.py and vllm/utils/cache.py implement the cache. The default location is ~/.cache/vllm/torch_compile_cache/.

CUDA memory

csrc/cumem_allocator.cpp is vLLM's custom CUDA caching allocator. It exposes virtual-memory–backed blocks so that:

• KV cache eviction can madvise(DONTNEED) instead of freeing+re-mallocating
• Sleep/wake can release physical pages while keeping VA reservations
• Page-level fragmentation stays bounded

The allocator is enabled via VLLM_USE_CUDA_MEMORY_POOL (default for V1).

Key source files

File	Purpose
vllm/compilation/backends.py	compile_fx backend
vllm/compilation/cuda_graph.py	Capture / replay primitives
vllm/compilation/piecewise_backend.py	Piecewise capture
vllm/compilation/decorators.py	@support_torch_compile, @cache_compiled, etc.
vllm/compilation/passes/	FX passes
vllm/v1/cudagraph_dispatcher.py	Per-step graph selection
vllm/model_executor/custom_op.py	CustomOp base
vllm/_custom_ops.py	Python surface for csrc kernels
csrc/cumem_allocator.cpp	Custom CUDA allocator

Entry points for modification

• Add a kernel: write it in csrc/, surface a Python wrapper in vllm/_custom_ops.py, register with torch.library, wrap in a CustomOp.
• Add a compile pass: drop a module in vllm/compilation/passes/, register in PassConfig.
• Tune CUDA graph capture: adjust CUDAGraphMode defaults in CompilationConfig (or pass --compilation-config '{"cuda_graph_mode": "...", "pass_config": {...}}').

For where the compiled forward is actually executed, see Executors and workers.