vLLM's op IR: where the inference engine meets the compiler

vLLM is worth understanding not as "the thing that serves Llama fast" but as a case study in a specific tension: an inference engine has to be a compiler target and a hand-tuned-kernel dispatcher at the same time. Recently vLLM grew a small op-level IR to resolve that tension explicitly. Model code calls named IR ops; torch.compile (Dynamo) traces them into an FX graph where they survive as opaque nodes; and at runtime, each op dispatches to the best available kernel for the current platform.

The two-faced inference engine

At its core, vLLM is a model-serving engine. It does not train. It takes trained weights and runs forward passes for inference at the highest throughput and lowest cost it can manage. Its founding idea was PagedAttention: managing the KV cache the way an OS manages virtual memory, with non-contiguous pages and an indirection table, which enables continuous batching without significant memory fragmentation.

In a data center it occupies the layer between the runtime and the orchestration tier:

[Applications] [Agents] [RAG]                     ← your product
[API Gateway / Router (LiteLLM, Envoy)]           ← auth, multi-model, billing
[Cluster Orchestration (K8s, Ray Serve, llm-d)]   ← N replicas, KV-aware routing, P/D split
──────────────────────────────────────────────────────────
>>> vLLM engine <<<                               ← scheduler, KV cache, model exec, sampling
[PyTorch] [CUDA/ROCm] [NCCL] [Triton] [FlashAttn] ← runtime
[GPUs/TPUs] [NVLink] [RDMA]                       ← hardware

The part that matters: vLLM's model-execution path is migrating to be torch.compile-centric. During the autoregressive decode phase, where the engine generates one token at a time and each operation is a fast, memory-bound matrix-vector multiplication, eager-mode PyTorch's Python overhead and kernel launch latencies become significant bottlenecks. Capturing the model graph via torch.compile and CUDA graphs is mandatory to eliminate that overhead. But once you commit to ahead-of-time compilation, you inherit every problem a compiler frontend has, and that is where the IR comes from.

When the compiler meets reality

This duality, acting as a compiler target while relying on opaque kernels, creates immediate friction. Consider a standard operation like RMSNorm. In vLLM, it has to exist in at least these forms simultaneously:

a pure-PyTorch reference (correct, slow, runs anywhere),
one or more hand-written CUDA kernels,
ROCm and other backend variants,
a fused variant (fused_add_rms_norm) that folds the residual add in and writes in place,
quantized variants.

Now layer on three requirements that pull in different directions:

You must retain graph-level fusion. Inductor's whole value is fusing pointwise/reduction chains. If you inject a completely opaque custom op into the middle of the graph, Inductor's generic code generator cannot see inside it, forcing it to route around the op and break fusion at the boundary.
You must be able to swap implementations by hardware and by argument shape/dtype, at runtime, on the hot path, cheaply.
Every fast path must be checkable against the reference, within numerical tolerance, or you will ship a kernel that is subtly wrong at one shape out of millions.

Without a unifying abstraction, all of this lives as if current_platform.is_cuda() and dtype == ... branches sprinkled through model code, invisible to the compiler and impossible to test uniformly. The IR is the seam that pulls those three concerns out of the model and into one object.

A dialect in disguise

vllm/ir is not an AST or an LLVM/MLIR-style IR with its own blocks, regions, or SSA values. Rather, it is an op registry built on PyTorch's torch.library custom-op machinery. It acts as an "IR" only in the sense that it strictly defines the node semantics that populate the FX graph. If you think in compiler terms, it is a dialect: a set of named, stable, non-decomposed ops with reference semantics, plus the metadata required to lower and verify them.

Here is the registration:

lib = vllm_ir_torch_lib  # Library("vllm_ir", "FRAGMENT")
lib.define(self.name + self._schema_str)
# CompositeExplicitAutograd is not decomposed
# by ATen IR normalization in AOTAutograd
lib.impl(self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd")
lib._register_fake(self.name, self._fake_call)

That comment is the whole point. The CompositeExplicitAutograd dispatch key is chosen specifically so AOTAutograd does not decompose the op during ATen IR normalization. The reference implementation is written in plain PyTorch, and if it were composite-decomposed, Inductor would see the constituent pointwise/reduction ops and the identity "this is an RMSNorm" would be gone. By registering it explicitly, the op survives into the FX graph as a single opaque node.

Because Inductor cannot fuse this opaque node automatically, vLLM runs custom compilation passes prior to lowering. These passes pattern-match the high-level semantic nodes (e.g., rms_norm followed by add) and rewrite them into fused_add_rms_norm. This shifts fusion to a graph rewrite over the dialect, rather than relying on the backend or hand-coding it in the model.

So each op has a dual personality: to the compiler, one opaque, schema-typed node with a fake/meta function for shape and dtype propagation; to the runtime, a dispatcher over N implementations.

Traffic control on the hot path

Once the graph is captured and compiled, execution drops into the runtime where every microsecond matters. Here, an op holds impls: dict[str, IrOpImpl]. "native" and "unfused" are reserved; everything else is a named provider, a CUDA kernel, a Triton kernel, a fused variant. Selection is two-tiered:

def dispatch(self, *args, **kwargs) -> "IrOpImpl":
    """This function is on the hot path (op dispatch), must be fast."""
    for impl in self._priority_impls:
        if impl.supports_args(*args, **kwargs):
            return impl
    ...

Two distinct predicates, deliberately separated:

supported: a static property. Does this platform/library even have the kernel? Evaluated once. The docs explicitly forbid putting env-var or global-state logic here.
supports_args: a dynamic check against the actual tensors (dtype, shape, alignment). This runs per call.

The invariant the dispatcher enforces, that the last impl in any priority list must accept all args, is exactly the "total function at the bottom of the lattice" discipline you would want in a lowering pipeline.

There is also an escape hatch: enable_torch_wrap. When wrapping is off, __call__ skips the torch.ops dispatch entirely and calls _inner_call directly. The motivation is simple: to avoid torch dispatch overhead in eager mode, and to avoid forcing a lowering step on platforms that do not use Inductor. The abstraction is free when you opt out.

Battle scars: strict schemas and hidden mutations

Beyond the hot path, building a compiler integration at this scale reveals subtle edge cases. Two details in the IR design address common production issues:

Every provider must have byte-identical schema to the reference. IrOpImpl.__init__ re-infers the schema from the impl function and rejects any mismatch, names, types, and defaults. It even validates that a supports_args predicate has the same parameter names and defaults as the native signature, so the dispatch hot path can forward *args positionally without rebinding. This is the registry refusing to let two implementations of "the same op" silently diverge in interface.
Inplace is modeled as a separate overload, with forced functionalization on the default path. Fused kernels want to reuse activation memory; the functional graph the compiler reasons about must not have hidden mutation. The IR splits this: IrOpInplace creates a .maybe_inplace overload whose schema is inferred with mutates_args=op.activations. The default overload stays functional by cloning activation inputs before calling an inplace impl.

Mental model

An inference engine is a bilingual diplomat. It speaks compiler (opaque nodes in an FX graph) and it speaks kernel (a dispatcher over N implementations). The IR is the interpreter that keeps both languages honest, and it refuses to let either side silently drift.

Why this matters

The design decisions in vllm/ir are the kind of thing you will recognize if you have ever fought AOTAutograd's decomposition table or tried to keep a custom kernel alive through torch.compile. It is the pattern of treating the engine as a compiler target, not just a kernel launcher, and it is where inference engineering is heading: the engine is becoming a compiler, and the compiler is becoming the engine.

vLLM's op IR is the seam that pulls compiler concerns out of the model and into one object: fusion, dispatch, and verification, unified.

Back to the blog