How vLLM's `torch.compile` Backend Handles Dynamic Batch Sizes — A Code Walkthrough

LLM serving has a constraint that training rarely cares about: the batch and sequence dimensions vary per request, but you can’t afford to JIT-compile a fresh kernel each time the size changes. A naïve torch.compile(model, dynamic=False) would specialise per shape — fast at steady state once you’ve warmed every shape, but cold-starts forever. A dynamic=True compile collapses everything into a single graph parameterised by SymInts — works always, but the kernels carry shape arithmetic that hand-tuned static-shape kernels don’t, and Triton autotune can’t pick shape-specific tile configs.

vLLM’s compile backend (used since v1’s CompilationMode.VLLM_COMPILE) gets the best of both: one Dynamo trace, N Inductor specialisations, O(1) shape dispatch at runtime. The pattern is small and elegant, and once you’ve seen it, it’s portable to other dynamic-shape workloads (training with packed sequences, beam search, etc.).

This post walks through the code in order. All links are GitHub permalinks at vLLM commit 6d09769 so they don’t rot.

The two-stage idea

                         ┌──────────────────────────────┐
                         │  user call: model(input)     │
                         └──────────────┬───────────────┘
                                        ▼
                   ┌─────────────────────────────────────────┐
                   │  Dynamo @ dynamic=True                  │
                   │   ─ traces forward                      │
                   │   ─ promotes seq_len to SymInt s29      │
                   │   ─ emits ONE FX graph                  │
                   └──────────────┬──────────────────────────┘
                                  ▼
        ┌─────────────────── VllmBackend.__call__ ───────────────────┐
        │                                                              │
        │  for size in compile_sizes:    # e.g. [256, 512, 1024]       │
        │     args = create_concrete_args(graph, size)                 │
        │       # SymInt s29 → int(size) in every placeholder          │
        │     compiled[size] = inductor_compile(graph, args)            │
        │       # static-shape kernels — Triton autotune sees concrete  │
        │       # tile dims, no SymInt arithmetic in the loop body      │
        │                                                              │
        │  return PiecewiseBackend(graph, compiled, sym_shape_indices) │
        └──────────────────────────────────────────────────────────────┘
                                  ▼
                     ┌────────────────────────────┐
                     │ runtime: dispatch by shape │
                     │   shape = args[sym_idx]    │
                     │   compiled[bucket(shape)]  │
                     └────────────────────────────┘

Key insight: Dynamo’s symbolic capture is decoupled from Inductor’s compilation. Dynamo runs once, doesn’t recompile when shape changes. Inductor runs N times, once per anticipated shape, each with concrete dims. The result of each Inductor run is a fully shape-static kernel with no SymInt baggage.

Step 1 — capture once with dynamic=True

Nothing custom here. vLLM just wires its backend in via the standard torch.compile API:

torch.compile(model, backend=VllmBackend(...), dynamic=True)

Dynamo traces the model forward, sees the seq_len dim varies, lifts it to a SymInt (s29), and emits a single FX graph whose placeholder nodes carry SymInt-valued shapes in their meta["example_value"].

VllmBackend.__call__ receives (gm, example_inputs). The example inputs are FakeTensors with SymInt shapes. The class lives at vllm/compilation/backends.py:806:

class VllmBackend:
    """The compilation backend for `torch.compile` with vLLM.
    ...
    The major work of this backend is to split the graph into
    piecewise graphs, and pass them to the piecewise backend.
    """
    vllm_config: VllmConfig
    compilation_config: CompilationConfig
    _called: bool = False
    graph: fx.GraphModule
    split_gm: fx.GraphModule  # the stitching graph for all the piecewise graphs
    piecewise_graphs: list[SplitItem]

(The “piecewise” part is a separate graph-partitioning step — VllmBackend splits the captured FX graph at known boundaries like attention so each subgraph can be compiled independently. For this post we’ll focus on what each piece does, not how the splitting itself works.)

Step 2 — find the symbolic-shape inputs

Each piecewise subgraph has one or more SymInt placeholders. To dispatch at runtime we need to know which positional arg holds the symbolic shape. vLLM finds them with a one-liner at vllm/compilation/backends.py:747:

sym_shape_indices = [
    i for i, x in enumerate(args) if isinstance(x, torch.SymInt)
]

That’s it. Walk the example inputs, collect indices where the value is a torch.SymInt. Those positions are guaranteed to hold shape-valued ints at runtime (Dynamo passes them in alongside the actual tensors).

Step 3 — concretise the FX graph for each target size

This is the trick that makes Inductor produce static kernels. vLLM’s helper at vllm/compilation/piecewise_backend.py:37-76:

def create_concrete_args(graph: fx.GraphModule, size: int) -> list[Any]:
    """Create Fake example inputs with symbolic dims replaced by a concrete size."""
    from torch._prims_common import compute_required_storage_length
    from torch._subclasses.fake_tensor import FakeTensorMode
    from torch.fx.experimental.symbolic_shapes import ShapeEnv, is_symbolic

    def concretize(sym_val: Any) -> int:
        """Replace all symbolic variables in a SymInt expression with size."""
        if not is_symbolic(sym_val):
            return int(sym_val)
        expr = sym_val.node.expr
        return int(expr.subs({s: size for s in expr.free_symbols}))

    fake_mode = FakeTensorMode(shape_env=ShapeEnv())

    args: list[Any] = []
    with fake_mode:
        for node in graph.graph.nodes:
            if node.op != "placeholder":
                break
            val = node.meta["example_value"]
            if isinstance(val, torch.SymInt):
                args.append(concretize(val))
            elif isinstance(val, torch.Tensor):
                new_shape = tuple(concretize(d) for d in val.shape)
                new_strides = tuple(concretize(s) for s in val.stride())
                new_storage_offset = concretize(val.storage_offset())
                needed_size = compute_required_storage_length(
                    new_shape, new_strides, new_storage_offset
                )
                t = torch.empty(needed_size, dtype=val.dtype, device=val.device)
                t = t.as_strided(new_shape, new_strides, new_storage_offset)
                args.append(t)
            else:
                args.append(val)
    return args

A few things worth pulling out:

concretize works on SymPy expressions, not just bare SymInts. If a placeholder’s shape is 2 * s29 + 1, concretize substitutes s29 → size and evaluates the expression to a plain int. That handles the case where the same symbolic variable appears in multiple derived-shape placeholders (an attention-mask shape derived from seq_len, for example).

Strides and storage offsets get the same treatment. A FakeTensor’s stride pattern can also depend on SymInts (stride = [1024 * s29, 1024, 1] on a contiguous [1, s29, 1024] tensor). Concretising the shape isn’t enough — Inductor needs strides to match the actual memory layout, otherwise it picks wrong access patterns.

compute_required_storage_length is called explicitly. The default torch.empty(shape) would allocate prod(shape) elements, but a non-contiguous tensor with the same shape might need more (or less) storage. vLLM allocates the exact required storage and then as_strideds into it, so the FakeTensor passed to Inductor has the right storage_offset for backward graph analysis.

The output is a list of fake tensors and concrete ints with no SymInts anywhere. That’s what Inductor wants.

Step 4 — Inductor compile per size

Each RangeEntry represents one bucket. The compile loop at vllm/compilation/piecewise_backend.py:245-277:

def compile_all_ranges(self) -> None:
    """Compile all range entries for this piecewise subgraph up front."""
    for range_entry in self.range_entries.values():
        if range_entry.compiled:
            continue

        self._log_compile_start(range_entry.compile_range)

        if range_entry.compile_range.is_single_size():
            args_list = create_concrete_args(
                self.graph, range_entry.compile_range.start
            )
        else:
            args_list = get_fake_args_from_graph(self.graph)

        range_entry.runnable = self.vllm_backend.compiler_manager.compile(
            self.graph,
            args_list,
            self.vllm_backend.inductor_config,
            self.compilation_config,
            compile_range=range_entry.compile_range,
            graph_index=self.piecewise_compile_index,
            num_graphs=self.total_piecewise_compiles,
            is_encoder=self.vllm_backend.is_encoder,
        )

        range_entry.compiled = True

Two flavours:

• Single-size range (e.g. Range(start=512, end=512)): use create_concrete_args to substitute SymInts → 512 → fully static Inductor compile.
• Multi-size range (e.g. Range(start=256, end=511)): get_fake_args_from_graph reuses the original SymInt-valued example inputs as-is, so Inductor produces a dynamic-shape kernel that handles any size in the range. This is the fallback for “I haven’t enumerated this exact size, but it falls in this bucket.”

The serving config typically declares both: pin a few specific sizes that are common (256, 512, 1024) for shape-static optimisation, and let everything else fall into a wider dynamic range.

compiler_manager.compile is just a thin wrapper around torch._inductor.compile_fx (or torch._inductor.standalone_compile if the build supports it). Output is a callable.

Step 5 — runtime dispatch

This is the prettiest part. Once compilation is done, the PiecewiseBackend instance acts as the actual call target. From vllm/compilation/piecewise_backend.py:358-380:

def __call__(self, *args: Any) -> Any:
    if self.sym_shape_indices:
        runtime_shape = args[self.sym_shape_indices[0]]
        range_entry = self._find_range_for_shape(runtime_shape)
        assert range_entry is not None, (
            f"Shape: {runtime_shape} out of considered ranges: "
            f"{self.compile_ranges}"
        )
    else:
        # All inputs have static shapes; use the only compiled range_entry
        compiled_entries = [re for re in self.range_entries.values() if re.compiled]
        assert len(compiled_entries) == 1, (
            f"Expected exactly one compiled range_entry for static shape "
            f"compilation, but found {len(compiled_entries)}"
        )
        range_entry = compiled_entries[0]

    assert range_entry.compiled, ...
    return range_entry.runnable(*args)

At runtime:

1. Read the SymInt-valued positional arg (args[sym_shape_indices[0]]). At call time it’s a plain int — Dynamo’s symbolic capture has resolved.

2. Look up the range entry containing that size via _find_range_for_shape (a linear scan over a small list — the table is short, no hash-map needed):

   def _find_range_for_shape(self, runtime_shape: int) -> RangeEntry | None:
       if self.compile_sizes is None:
           return None
       if runtime_shape in self.compile_sizes:
           return self.range_entries[Range(start=runtime_shape, end=runtime_shape)]
       else:
           for range in self.compile_ranges:
               if runtime_shape in range:
                   return self.range_entries[range]
       return None
   

3. Call the corresponding compiled artifact.

The dispatch overhead is one Python attribute access + one list scan + one indexed call. At GPU step times of a few hundred microseconds it’s invisible.

Step 6 — make_copy_and_call for cudagraphs

There’s one more wrinkle worth knowing about. Cudagraph capture requires inputs to live at fixed memory addresses — but the model’s caller wants to pass freshly-allocated runtime tensors. vLLM bridges this with a tiny copy-into-static-buffer wrapper at vllm/compilation/backends.py:54-88:

def make_copy_and_call(
    sym_tensor_indices: list[int],
    input_buffers: list[torch.Tensor | None],
    callable_fn: Callable[..., Any],
) -> Callable[..., Any]:
    def copy_and_call(*args: Any) -> Any:
        list_args = list(args)
        for i, index in enumerate(sym_tensor_indices):
            runtime_tensor = list_args[index]
            runtime_shape = runtime_tensor.shape[0]

            # lazy initialization of buffer on first call
            if input_buffers[i] is None:
                input_buffers[i] = runtime_tensor.clone()

            static_tensor = input_buffers[i][:runtime_shape]
            static_tensor.copy_(runtime_tensor)
            list_args[index] = static_tensor
        return callable_fn(*list_args)
    return copy_and_call

For each tensor whose shape varies at runtime: lazily allocate a max-sized buffer on first call, then on every call slice the buffer to runtime_shape and copy the user’s tensor in. The downstream cudagraph always sees the same allocation, just the prefix changes. The slice cost is a single copy_ — bandwidth-bound but tiny relative to the actual compute.

make_copy_and_call only wraps when cudagraphs are enabled — see the wrap_with_cudagraph_if_needed helper at vllm/compilation/backends.py:635-680.

Why this matters beyond serving

The trick generalises. In any workload where:

• you can enumerate a small set of distinguished input shapes (or shape buckets)
• the shape variance breaks Dynamo’s recompile budget or makes Inductor’s symbolic kernels slow
• you don’t want to pre-pad everything to the largest shape (because that’s wasted compute)

…you can apply this pattern. The pieces are plain torch.fx and torch._inductor.compile_fx. There’s nothing vLLM-specific about create_concrete_args or the runtime dispatcher.

I’ve been considering it for training a packed-sequence SFT pipeline where total token count varies per step. We pad to a 2-bucket ladder (40,960 / 65,536) so Dynamo only sees two shapes, but currently still compile dynamic — the resulting kernels pay SymInt-arithmetic overhead. Lifting vLLM’s pattern would let us cut that overhead at the cost of one more Inductor compile. The 5-line concretize helper is the same; the bookkeeping is the same; only the integration point with HF Trainer changes.

Summary

vLLM’s compile backend is a clean separation of concerns:

Stage	What it does	Where to read it
Capture	One Dynamo trace, dynamic SymInt shapes	backends.py:806
Identify	Find SymInt-valued arg positions	backends.py:747
Concretise	Substitute SymInts → ints, build fake tensors	piecewise_backend.py:37
Compile	Inductor compile per concrete size	piecewise_backend.py:245
Dispatch	Read shape, look up artifact, call	piecewise_backend.py:358
Cudagraph	Copy into static buffer	backends.py:54

The whole machine is under 500 LOC. Most of the complexity in vLLM’s compilation/ directory is the configuration, caching, and serialisation around it; the core compile-and-dispatch loop is small. Worth reading in full if you ever need to do something similar — and it’s the most pragmatic example I’ve seen of torch.compile co-existing with shape-varying production workloads.