The forward context: how layers find per-step data

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

The idea#

A model file calls self.attn(q, k, v) and nothing else (How a vLLM model file is written (llama.py)). But attention needs more: the KV cache tensor of this layer, where to write the new keys and values, and the shape of the batch. vLLM passes these through two dictionaries keyed by the layer name, the prefix (of a layer) such as model.layers.0.self_attn.attn:

• The static_forward_context maps each name to its layer object. A layer adds itself when it is built. At start-up, the runner binds each layer's KV cache tensor to the object, as layer.kv_cache.
• The forward context exists only during one forward pass. The model runner sets it with the attention metadata per layer name, the slot mapping per layer name, the CUDA graph mode and a BatchDescriptor.

The attention custom op gets the tensors q, k, v and output, plus the layer name. Inside, it looks up the layer, the KV cache, the metadata and the slot mapping. Why so indirect? For torch.compile, the op is one opaque node with plain inputs (PyTorch for serving: custom ops and torch.compile). The graph does not depend on Python objects that change every step, so vLLM can compile it once and split it at these ops.

An everyday comparison. Think of a theater with a prop table. Each actor has a spot on the table, labeled with the actor's name. Before each show, the stage manager puts that night's props on each labeled spot. During the show, an actor walks to the table and takes the props under the label. The script never lists the props. The comparison stops here: in vLLM, the label also leads to the actor's own permanent locker, the KV cache, which nobody refills between shows.

The picture#

Follow the name model.layers.0.self_attn.attn through the three panels. It is a key in every dictionary. The violet parts happen once at start-up; the blue parts happen in every step.

In the code#

Step 1: the layer registers itself by name#

        # For cuda-alike (CUDA and ROCM) and cpu platforms, we control how
        # torch.compile works by registering the attention as one giant
        # opaque custom op. For other platforms, we directly call them
        # and let torch.compile handle them.
        self.use_direct_call = not current_platform.opaque_attention_op()

        compilation_config = vllm_config.compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

1. Line 440. use_direct_call is false on platforms that want "one giant opaque custom op" for attention. CUDA, ROCm, CPU and XPU return true from opaque_attention_op() (vllm/platforms/cuda.py:595-596, vllm/platforms/rocm.py:1231-1232, vllm/platforms/cpu.py:640-641, vllm/platforms/xpu.py:609-610); the base Platform returns false (vllm/platforms/interface.py:1286-1290).
2. Lines 443-445. The prefix is the key. A second layer with the same prefix raises ValueError: Duplicate layer name.

The dict is a field of the compilation config. Its docstring says: "Map from layer name to layer objects that need to be accessed outside model code" (vllm/config/compilation.py:757-760). Attention is not the only user. MLA attention (vllm/model_executor/layers/attention/mla_attention.py:620), the MoE runner (vllm/model_executor/layers/fused_moe/runner/moe_runner.py:68), Mamba layers (vllm/model_executor/layers/mamba/mamba_mixer2.py:528) and LoRA linear layers (vllm/lora/layers/base_linear.py:102) also register.

The runner reads this dict to find all attention layers. get_layers_from_vllm_config filters it by type (vllm/config/vllm.py:3596-3617), and the KV cache spec of the model comes from these layers (vllm/v1/worker/gpu/attn_utils.py:138-155).

Step 2: the runner binds the KV cache tensor#

When the KV cache is allocated, init_kv_cache gets compilation_config.static_forward_context as its first argument (vllm/v1/worker/gpu/model_runner.py:794-797). It allocates one backing buffer and makes a view per layer, then binds each view to its layer:

def bind_kv_cache_to_layers(
    kv_caches: dict[str, torch.Tensor],
    forward_context: dict[str, Attention],
    num_attn_module: int = 1,
    kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
) -> None:
    """Bind layer caches and share ReplaySSM trackers in model-layer order."""
    # Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks
    # its raw allocation into the per-layer view(s) it needs (e.g. Mamba
    # splits conv/ssm), so the kv_caches dict can hold a single tensor per
    # layer for the KV connector to register.
    for layer_name, kv_cache in kv_caches.items():
        forward_context[layer_name].bind_kv_cache(kv_cache)

The default bind_kv_cache just stores the view: self.kv_cache = kv_cache (vllm/model_executor/layers/attention_layer_base.py:27-33). Until then, the layer holds an empty placeholder tensor (vllm/model_executor/layers/attention/attention.py:461-464). The binding happens once; the tensor stays the same for the whole life of the server. How the views are laid out is on Attention backends: backend, metadata builder and implementation.

Step 3: the runner sets the forward context each step#

For PIECEWISE and eager steps, Model Runner V2 wraps the model call:

            # For piecewise and eager mode, just call model().
            batch_descriptor = BatchDescriptor(
                num_tokens=input_batch.num_tokens_after_padding,
                has_lora=self.lora_config is not None,
                num_active_loras=batch_desc.num_active_loras,
            )

            with set_forward_context(
                attn_metadata,
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
                cudagraph_runtime_mode=batch_desc.cg_mode,
                num_tokens_across_dp=dp_sync_state
                and dp_sync_state.num_tokens_across_dp,
                batch_descriptor=batch_descriptor,
                ubatch_slices=ubatch_slices,
                slot_mapping=slot_mappings_by_layer,
                skip_compiled=skip_compiled,
                is_padding=input_batch.is_padding,
            ):

set_forward_context builds a ForwardContext and installs it in a module-level variable for the duration of the with block:

@dataclass
class ForwardContext:
    # copy from vllm_config.compilation_config.static_forward_context
    no_compile_layers: dict[str, Any]
    attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]]
    slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]]
    """
    Type Dict[str, AttentionMetadata] for v1, map from layer_name of each
    attention layer to its attention metadata
    Type List[Dict[str, AttentionMetadata]] for DBO. List of size two, one
    for each microbatch.
    Set dynamically for each forward pass
    """
    # set dynamically for each forward pass
    dp_metadata: DPMetadata | None = None
    # determine the cudagraph style at runtime to be FULL, PIECEWISE, or NONE.
    # by default NONE, no cudagraph is used.
    cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE
    batch_descriptor: BatchDescriptor | None = None

• no_compile_layers is the static dict itself, not a copy: create_forward_context passes the same object (vllm/forward_context.py:247-248).
• attn_metadata and slot_mapping are dicts from layer name to value. They are built fresh each step (The model runner (Model Runner V2), Step 6).
• cudagraph_runtime_mode and batch_descriptor tell the CUDA graph wrappers inside the compiled model what to do (torch.compile and CUDA graphs in vLLM).

The context manager restores the previous context on exit (vllm/forward_context.py:262-274). Any code inside the forward pass reads it with get_forward_context() (vllm/forward_context.py:208-214).

Step 4: the attention layer calls a custom op with its name#

        else:
            # Skip this if sharing KV cache with an earlier attention layer.
            encoded = _encode_layer_name(self.layer_name)
            if (
                not self.attn_backend.forward_includes_kv_cache_update
                and self.kv_sharing_target_layer_name is None
                and key is not None
                and value is not None
            ):
                kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
                    key, value, encoded
                )
            torch.ops.vllm.unified_attention_with_output(
                query,
                key,
                value,
                output,
                encoded,
                kv_cache_dummy_dep=kv_cache_dummy_dep,
            )
        return output.view(-1, hidden_size)

1. Line 563. The name is wrapped by _encode_layer_name. On PyTorch 2.11 or newer with VLLM_USE_LAYERNAME=1 (the default), it becomes a LayerName object; else it stays a str (vllm/utils/torch_utils.py:990-995, vllm/utils/torch_utils.py:1045-1047).
2. Lines 564-572. Only when the backend does not write the KV cache inside its own forward, a first op writes k and v into the cache. It returns an empty tensor, the "dummy dependency".
3. Lines 573-580. The attention op takes the dummy as an argument, so the compiler cannot move it before the cache write (vllm/model_executor/layers/attention/attention.py:779-782). It returns nothing and writes into output.

At this commit, the FlashAttention, FlashInfer and Triton backends set forward_includes_kv_cache_update = False, so they use the separate update op (vllm/v1/attention/backends/flash_attn.py:367, vllm/v1/attention/backends/flashinfer.py:708, vllm/v1/attention/backends/triton_attn.py:347). The base class default is True (vllm/v1/attention/backend.py:69-70).

Step 5: get_attention_context resolves the name#

    forward_context: ForwardContext = get_forward_context()
    attn_metadata_raw = forward_context.attn_metadata
    attn_metadata: AttentionMetadata
    if isinstance(attn_metadata_raw, dict):
        attn_metadata = attn_metadata_raw[layer_name]
    elif isinstance(attn_metadata_raw, list):
        # list[dict[str, AttentionMetadata]]: used in speculative decoding
        # where [0] is the base-model (non-speculative) metadata dict.
        attn_metadata = attn_metadata_raw[0][layer_name]
    else:
        attn_metadata = attn_metadata_raw
    attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
    kv_cache = attn_layer.kv_cache
    slot_mapping = forward_context.slot_mapping
    assert isinstance(slot_mapping, dict), (
        f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
    )
    layer_slot_mapping = slot_mapping.get(layer_name)
    return attn_metadata, attn_layer, kv_cache, layer_slot_mapping

1. Line 709. The normal case: attn_metadata is a dict, and the layer's entry is taken.
2. Line 713. A list of dicts: the comment says index 0 holds the base model's metadata in speculative decoding.
3. Lines 716-717. The layer object comes from no_compile_layers, and the KV cache from the attribute bound in Step 2.
4. Line 722. The slot mapping may be None; then unified_kv_cache_update writes nothing (vllm/model_executor/layers/attention/attention.py:736).

Then the two ops call the backend: impl.do_kv_cache_update(...) (vllm/model_executor/layers/attention/attention.py:740-746) and impl.forward(self, query, key, value, kv_cache, attn_metadata, output=...) (vllm/model_executor/layers/attention/attention.py:786-796). The backend classes are on Attention backends: backend, metadata builder and implementation.

Step 6: why a name and not the tensors#

The two ops are registered once, at import, in the vllm library (vllm/model_executor/layers/attention/attention.py:759-764, vllm/model_executor/layers/attention/attention.py:799-803). For the compiler, three facts follow:

• Opaque. Dynamo records one node per call and does not trace into the backend code. The node has only tensors and a name as inputs.
• A splitting point. The op names are in the default splitting ops, so vLLM cuts the graph here. Attention then runs eagerly between piecewise CUDA graphs and can use any batch shape (vllm/config/compilation.py:768-770, vllm/config/compilation.py:1184).
• One graph for 32 layers. A str argument is a constant in the graph. The LayerName class is an opaque type. torch.compile "lifts [it] as a graph input instead of baking it as a constant", which avoids "per-layer recompilation" (vllm/utils/torch_utils.py:1003-1010).

The code also says that the string in unified_kv_cache_update "prevents Inductor from reusing piecewise graphs"; this is why that op is a splitting op too (vllm/config/compilation.py:1186-1213).

Step 7: FULL CUDA graph steps do not read the context#

In FULL mode the runner replays one graph and does not call set_forward_context (vllm/v1/worker/gpu/model_runner.py:2024-2032). No Python code of the model runs, so no layer looks anything up. The comment says the inputs "are already copied to the CUDA graph input buffers" (lines 2026-2027). The context was used at capture time: the capture runs the model inside set_forward_context, with the dummy metadata of that batch shape (vllm/v1/worker/gpu/cudagraph_utils.py:754-771).

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#

• A model call outside a context. get_forward_context() asserts "Forward context is not set. Please use `set_forward_context` to set the forward context." (vllm/forward_context.py:210-213). This happens when you call the model directly in a test or a script.
• Two layers with one prefix. ValueError: Duplicate layer name at model build time (vllm/model_executor/layers/attention/attention.py:443-444).
• A wrong KV sharing target. ValueError "must come before the current layer", "is not a valid Attention layer in the model" or "must be the same type" (vllm/model_executor/layers/attention/attention.py:80-91).
• Batch-size logging slows the server. With VLLM_LOG_BATCHSIZE_INTERVAL >= 0, every forward pass that runs inside set_forward_context with metadata ends with a device synchronize (vllm/forward_context.py:295, vllm/forward_context.py:363-370). In Model Runner V2 these are the eager and PIECEWISE steps; FULL replays have no context. That removes the CPU/GPU overlap of Async scheduling: overlap the CPU and the GPU. The comment above it still says "we use synchronous scheduling right now" (line 365); with the default async scheduling this is no longer true.
• A backend without a separate cache write. If a backend says forward_includes_kv_cache_update = False but has no do_kv_cache_update, the update op fails its assertion (vllm/model_executor/layers/attention/attention.py:737-739).

Why it matters. The forward context is the contract between the model runner, the model code and the attention backends. A model file stays free of serving details, the compiler sees a clean graph, and a new backend only needs to read its metadata. When you write a model or a backend plugin, keep the layer name unique, and read per-step data from the context, never from arguments.