Tour: One forward pass

1execute_model starts and applies the plan#

We are at stop 5 of Tour: One engine step. The uni executor called Worker.execute_model, which hands the SchedulerOutput to the model runner (vllm/v1/worker/gpu_worker.py:1350-1353). Model Runner V2 keeps every running request in a fixed row, its request slot. So it first applies the changes of the plan to those rows:

        if not dummy_run:
            # Update the request states.
            self.update_pp_decode_requests()
            self.finish_requests(scheduler_output)
            self.free_states(scheduler_output)
            self.add_requests(scheduler_output)
            self.update_requests(scheduler_output)
            self.block_tables.apply_staged_writes()
            if self.aux_output_connector is not None:
                # Register this step before the GPU forward.
                self.aux_output_connector.begin_step(
                    scheduler_output.aux_output_connector_metadata
                )
            if scheduler_output.total_num_scheduled_tokens == 0:
                # No need to run the model.
                empty_output = self.kv_connector.no_forward(scheduler_output)
                return self._merge_ec_connector_no_forward(
                    scheduler_output, empty_output
                )

• R3 is new. add_requests gives it a slot, stores its 298 prompt ids, overwrites its block table row, and registers its sampling parameters with the sampler (vllm/v1/worker/gpu/model_runner.py:1214-1220, vllm/v1/worker/gpu/model_runner.py:1238-1245).
• R1 and R2 run already. update_requests writes their computed-token counts into the CPU mirror. It appends new block ids only when the scheduler gave new blocks (vllm/v1/worker/gpu/model_runner.py:1257-1269).
• Block tables. apply_staged_writes() copies all staged rows to the GPU tables in one go (line 1760).

The step has 300 tokens, so it does not return early at line 1766. The persistent rows and the staged writes are explained on The model runner (Model Runner V2).

What the data looks like now:

req_states.req_id_to_index = {"R1": 1023, "R2": 1022, "R3": 1021}
num_computed_tokens_np[[1023, 1022, 1021]] = [540, 129, 0]
prefill_len[1021] = 298                       # R3 has a prompt to fill
block table row 1021 = 19 block ids           # ceil(298 / 16) = 19 (derived)

2Sort the batch, then pick the CUDA graph mode#

Next, gather_batch_req_state puts the requests in batch order: requests with draft tokens first, then decodes, then the others by size (vllm/v1/worker/gpu/model_runner.py:2471-2483). R1 and R2 have 1 token each, which equals the decode query length, so the order is R1, R2, R3. The function also maps each batch row to its slot (idx_mapping_np) and finds that R3 is still in its prefill (vllm/v1/worker/gpu/model_runner.py:1310-1326).

A batch with a real prefill is not eligible for a decode graph, so the uniform decode token count is None (vllm/v1/worker/gpu/model_runner.py:1343-1373, vllm/v1/worker/utils.py:967-975). Then the runner asks the CUDA graph manager for a descriptor, before it builds any input:

        batch_desc, dp_sync_state = dispatch_cg_and_sync_dp(
            self.cudagraph_manager,
            num_reqs,
            num_toks,
            uniform_tok_count,
            self.dp_size,
            self.dp_rank,
            max_query_len=max_query_len,
            need_eager=is_profile or skip_compiled,
            num_active_loras=num_active_loras,
            parallel_config=self.parallel_config,
            allow_ubatching=(
                self.ubatch_runner is not None and not skip_attn_for_dummy_run
            ),
            uniform_decode=uniform_tok_count == self.decode_query_len,
        )

With DP = 1 this is one call of CudaGraphManager.dispatch (vllm/v1/worker/gpu/dp_utils.py:265-279). The manager keeps a list of candidates for every token count up to 512. For 300 tokens, the candidates are the graphs of the next capture size, 304: the FULL decode graph first, then the PIECEWISE graph (vllm/v1/worker/gpu/cudagraph_utils.py:380-394).

    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, 304 tokens. Its uniform_token_count is 1, but our batch has None. Line 128 fails.
2. PIECEWISE, 304 tokens. It has no uniform count and no request limit, and 304 ≥ 300. It matches.

So the step runs as a piecewise CUDA graph of 304 tokens. The capture sizes on this GPU go up to min(1024 × 1 × 2, 512) = 512, in steps of 16 above 256 (vllm/config/vllm.py:2440-2449, vllm/config/vllm.py:2566-2570). A batch above 512 tokens has no candidate and runs eagerly (lines 547-553 there). The rule and its reasons are on torch.compile and CUDA graphs in vLLM.

What the data looks like now:

req_ids        = ["R1", "R2", "R3"]
idx_mapping_np = [1023, 1022, 1021]
has_prefill    = True;  uniform_tok_count = None
batch_desc     = BatchExecutionDescriptor(cg_mode=PIECEWISE, num_tokens=304, num_reqs=None, ...)

3prepare_inputs builds the input tensors on the GPU#

prepare_inputs fills the persistent input buffers. Only two small arrays come from the CPU: idx_mapping, copied as int64 (vllm/v1/worker/gpu/model_runner.py:1393-1395), and query_start_loc:

        # Get query_start_loc.
        # num_reqs_padded is None for PIECEWISE graphs (no request padding needed)
        num_reqs_padded = batch_desc.num_reqs or num_reqs
        query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32)
        query_start_loc_np[0] = 0
        np.cumsum(num_scheduled_tokens_np, out=query_start_loc_np[1 : num_reqs + 1])
        # Pad for full CUDA graph mode.
        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
        query_start_loc_np[num_reqs + 1 :] = num_tokens
        query_start_loc = self.input_buffers.query_start_loc
        async_tensor_h2d(query_start_loc_np, out=query_start_loc)

The PIECEWISE descriptor has no request count, so num_reqs_padded stays 3 (line 1448). The token count is padded to 304 (vllm/v1/worker/gpu/model_runner.py:1383). Without draft tokens, each request gets exactly one logits row (vllm/v1/worker/gpu/model_runner.py:1401-1412). Then three GPU kernels read the request rows through idx_mapping:

        # Get prefill tokens if any.
        if batch_req_state.has_prefill:
            prepare_prefill_inputs(
                self.input_buffers.input_ids,
                self.req_states.next_prefill_tokens,
                idx_mapping,
                query_start_loc,
                self.req_states.all_token_ids.gpu,
                self.req_states.prefill_len.gpu,
                self.req_states.num_computed_tokens.gpu,
            )

        # Prepare positions and seq_lens.
        prepare_pos_seq_lens(
            idx_mapping,
            query_start_loc,
            self.req_states.num_computed_tokens.gpu,
            self.input_buffers.positions,
            self.input_buffers.seq_lens,
        )
        seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]

        # Some input token ids are directly read from the last sampled tokens
        # and draft tokens. Also, get the logits indices to sample tokens from.
        logits_indices = combine_sampled_and_draft_tokens(
            self.input_buffers.input_ids,
            idx_mapping,
            self.req_states.last_sampled_tokens,
            query_start_loc,
            seq_lens,
            self.req_states.prefill_len.gpu,
            self.req_states.draft_tokens,
            cu_num_logits,
            total_num_logits,
            self.model_state.num_new_sampled_tokens_per_step,
        )

1. Lines 1470-1479. The batch has a prefill, so a kernel copies R3's 298 prompt ids from all_token_ids into input_ids.
2. Lines 1482-1488. A kernel writes the positions and seq_lens = num_computed_tokens + query_len of each request.
3. Lines 1493-1504. A kernel writes the last sampled token of R1 and R2 into input_ids, from last_sampled_tokens on the GPU. It also returns logits_indices, the last row of each request.

The CPU does not know the last tokens of R1 and R2 yet; the step before writes them on the GPU (engine-step tour, stop 6). The varlen layout is the one of Input metadata↗.

What the data looks like now (derived from our numbers):

query_start_loc int32 [4]   = [0, 1, 2, 300]           # entries 4..1024 of the buffer are also 300
input_ids       int32 [304] = [last token of R1, last token of R2, R3 prompt ids 0..297, 4 padding rows]
positions       int64 [304] = [540, 129, 0, 1, ..., 297, 4 padding rows]
seq_lens        int32 [3]   = [541, 130, 298]
logits_indices  int64 [3]   = [0, 1, 299]
is_padding      bool  [304] = 300 x False, 4 x True

The kernels do not write the 4 padding rows. is_padding marks them when VLLM_MOE_SKIP_PADDING is on, which is the default (vllm/v1/worker/gpu/model_runner.py:1385-1388, vllm/envs.py:1640), and the input buffers have fixed addresses, which the CUDA graphs need (vllm/v1/worker/gpu/input_batch.py:19-31).

4prepare_attn: block tables and slot mappings#

    def prepare_attn(
        self, input_batch: InputBatch
    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
        if self.pcp_manager is not None:
            return self.pcp_manager.prepare_attn(input_batch)

        # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
        block_tables = self.block_tables.gather_block_tables(
            input_batch.idx_mapping,
            num_reqs_padded=input_batch.num_reqs_after_padding,
            num_computed_tokens=self.req_states.num_computed_tokens.gpu,
        )
        # Slot mappings: [num_kv_cache_groups, num_tokens_padded].
        # Kernel pads beyond num_tokens with PAD_SLOT_ID.
        slot_mappings = self.block_tables.compute_slot_mappings(
            input_batch.idx_mapping,
            input_batch.query_start_loc,
            input_batch.positions,
            num_tokens_padded=input_batch.num_tokens_after_padding,
        )
        return block_tables, slot_mappings

Llama-3-8B has one KV cache group, so both results have one entry. gather_block_tables copies the block table rows of slots 1023, 1022 and 1021 into a second tensor, in batch order. compute_slot_mappings gives each token its slot in the KV cache: block_number × block_size + offset (vllm/v1/worker/gpu/block_table.py:409). The 4 padding rows get PAD_SLOT_ID = -1 (vllm/v1/worker/gpu/block_table.py:257-260, vllm/v1/attention/backends/utils.py:53).

What the data looks like now (block size 16, derived):

block_tables[0]  int32 [3, max_num_blocks]    # rows of R1, R2, R3
slot_mappings    int64 [1, 304]
  R1, position 540 -> block index 33, offset 12  (540 = 33 x 16 + 12)
  R2, position 129 -> block index 8,  offset 1   (129 = 8 x 16 + 1)
  R3, positions 0..297 -> block indexes 0..18, the last one holds 10 tokens
  rows 300..303 -> -1

The block ids themselves come from the scheduler's block pool (The KV cache manager: blocks, the block pool and memory sizing). The runner only mirrors them.

5The attention metadata: one object for 32 layers#

Back in execute_model, the runner makes a dict from layer name to slot mapping (vllm/v1/worker/gpu/model_runner.py:1924-1928, vllm/v1/worker/gpu/attn_utils.py:447-455). Then the default model state builds the attention metadata (vllm/v1/worker/gpu/model_runner.py:1940-1951). In PIECEWISE and eager mode it uses the unpadded sizes, 3 requests and 300 tokens (vllm/v1/worker/gpu/model_states/default.py:179-186). For each KV cache group, build_attn_metadata makes one CommonAttentionMetadata:

        common_attn_metadata = CommonAttentionMetadata(
            query_start_loc=query_start_loc_gpu,
            query_start_loc_cpu=query_start_loc_cpu,
            seq_lens=seq_lens,
            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
            max_seq_len=max_seq_len,
            num_reqs=num_reqs,
            num_actual_tokens=num_tokens,
            max_query_len=max_query_len,
            block_table_tensor=block_table,
            slot_mapping=slot_mapping,
            causal=group_causal,
            dcp_local_seq_lens=dcp_local_seq_lens,
            dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound,
            positions=positions,
            is_prefilling=group_is_prefilling,
            mm_req_doc_ranges=mm_req_doc_ranges,
            rswa_prefix_lens=rswa_prefix_lens,
            req_idx=req_idx,
            _token_to_req_indices_cache=token_to_req_indices,
            **common_attn_metadata_extra_kwargs,
        )

The metadata builder of the FlashAttention backend turns it into FlashAttentionMetadata. It copies most fields straight from the common metadata (vllm/v1/attention/backends/flash_attn.py:812-829). All 32 layers share one backend and one KV cache spec, so they form one attention group. The same object goes under all 32 layer names (vllm/v1/worker/gpu/attn_utils.py:574-580).

What the data looks like now:

CommonAttentionMetadata(num_reqs=3, num_actual_tokens=300,
    query_start_loc=[0, 1, 2, 300], seq_lens=[541, 130, 298],
    max_query_len=298, max_seq_len=541,
    block_table_tensor=int32 [3, max_num_blocks], slot_mapping=int64 [304], causal=True)
attn_metadata = {"model.layers.0.self_attn.attn": m, ..., "model.layers.31.self_attn.attn": m}

max_query_len and max_seq_len come from CPU arrays, so this stop needs no GPU-to-CPU copy (vllm/v1/worker/gpu/model_states/default.py:191-199). The three classes of a backend are explained on Attention backends: backend, metadata builder and implementation.

6set_forward_context, and the call into the model#

The model inputs are the ids and the positions; inputs_embeds and intermediate_tensors are None, and the model state adds nothing for Llama (vllm/v1/worker/gpu/model_runner.py:1984-1992). The metadata does not go into the call. The runner puts it into the forward context around the call:

        else:
            # 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,
            ):
                self.kv_connector.pre_forward(scheduler_output, input_batch)
                if ubatch_state is not None:
                    assert self.ubatch_runner is not None
                    model_output = self.ubatch_runner.run(
                        self.model, model_inputs, ubatch_state
                    )
                elif batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
                    # Run the PIECEWISE graph (compiled PW cudagraph or breakable
                    # cudagraph, chosen inside run_pw_graph). cg_mode is only
                    # PIECEWISE after the cudagraph manager exists.
                    assert self.cudagraph_manager is not None
                    model_output = self.cudagraph_manager.run_pw_graph(
                        self.model, model_inputs
                    )
                else:
                    # Eager (NONE): call the raw model directly.
                    model_output = self.model(**model_inputs)

1. Lines 2035-2039. The BatchDescriptor holds the padded token count, 304, and the LoRA flags.
2. Lines 2041-2053. set_forward_context installs the metadata dict, the slot mapping dict, the mode PIECEWISE and the descriptor for the duration of the with block (vllm/forward_context.py:345-360).
3. Lines 2060-2067. In PIECEWISE mode, run_pw_graph calls the model object directly, because breakable CUDA graphs are off for Llama (vllm/v1/worker/gpu/cudagraph_utils.py:578-583). They turn on by default only for the architectures in a fixed list, and LlamaForCausalLM is not in it (vllm/config/vllm.py:75-107, vllm/config/vllm.py:824-833).

self.model is LlamaForCausalLM. Its forward only calls the inner LlamaModel with the four arguments (vllm/model_executor/models/llama.py:534-544). The page The forward context: how layers find per-step data explains the two dicts and why the attention ops read them by name.

7The compiled LlamaModel: embedding and layer loop#

LlamaModel carries @support_torch_compile, with dimension 0 of the ids and positions dynamic (vllm/model_executor/models/llama.py:337-346). The first call, in the profile run at start-up, compiled the model (vllm/v1/worker/gpu_worker.py:680-697, torch.compile and CUDA graphs in vLLM). So this call goes to the compiled code (vllm/compilation/decorators.py:533-551). The compiled graph is cut at the attention ops into 33 pieces. Each piece sits in a CUDAGraphWrapper. A wrapper replays its graph only when the forward context says PIECEWISE (vllm/compilation/cuda_graph.py:242-256):

        # Sync offloader before replay - ensures any external dependencies
        # from pre-capture prefetches are satisfied.
        get_offloader().sync_prev_onload()
        entry.cudagraph.replay()
        return entry.output

So the GPU work of this stop is 33 graph replays for the 304-token descriptor, with 32 eager attention calls between them (stop 10). The source code that was traced is this:

        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.embed_input_ids(input_ids)
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

        remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)

        aux_hidden_states: list[torch.Tensor] = []
        if get_pp_group().is_first_rank:
            self._maybe_add_hidden_state(
                aux_hidden_states, self.start_layer, hidden_states, residual
            )
        for idx, layer in enumerate(
            islice(self.layers, self.start_layer, self.end_layer),
            start=self.start_layer,
        ):
            hidden_states, residual = layer(
                positions, hidden_states, residual, **extra_layer_kwargs
            )

• Line 418. On TP = 1, the embedding is a plain lookup (vllm/model_executor/layers/vocab_parallel_embedding.py:522-525). It turns input_ids [304] into hidden_states [304, 4096] BF16.
• Line 419. There is no residual yet. Layer 0 starts it (stop 8).
• Lines 432-438. The loop runs all 32 layers, because with PP = 1 this rank owns layers 0 to 31. Each layer returns (hidden_states, residual).

Piece 0 holds the embedding and the start of layer 0. Each later piece holds the end of one layer and the start of the next; the facts are on torch.compile and CUDA graphs in vLLM. The diagram below shows one layer.

8Inside a decoder layer: fused add + RMSNorm#

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)

        # Fully Connected
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual

The layer keeps the residual outside the norm. Layer 0 has no residual, so it calls a plain RMSNorm (lines 320-322). Layers 1 to 31 call input_layernorm(hidden_states, residual), which adds and normalizes in one call (line 324). post_attention_layernorm does the same after attention (line 328).

RMSNorm is a CustomOp. With Inductor compilation, the default of vllm serve, vLLM adds "none" to the custom op list, so the op is disabled (vllm/config/vllm.py:1819-1826). A disabled op uses forward_native (vllm/model_executor/custom_op.py:189-192). For RMSNorm, that method calls an IR op:

    def forward_native(
        self,
        x: torch.Tensor,
        residual: torch.Tensor | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        """PyTorch-native implementation equivalent to forward()."""
        if residual is None:
            return ir.ops.rms_norm(
                x,
                self.weight.data if self.pass_weight else None,
                self.variance_epsilon,
                self.variance_size_override,
            )
        else:
            return ir.ops.fused_add_rms_norm.maybe_inplace(
                x,
                residual,
                self.weight.data if self.pass_weight_add else None,
                self.variance_epsilon,
                self.variance_size_override,
            )

The IR op is one node in the traced graph. During compilation, a lowering pass replaces each IR node with the implementation that the priority list selects (vllm/compilation/passes/ir/lowering_pass.py:41-67, vllm/compilation/passes/pass_manager.py:122-124). On CUDA with Inductor, the list for fused_add_rms_norm is ["native"] (vllm/platforms/cuda.py:716-732). So the native PyTorch body runs, compiled by Inductor inside the piece, and not the vllm_c kernel. The native body adds in FP32, keeps the sum as the new residual, and normalizes (vllm/ir/ops/layernorm.py:43-62).

What the data looks like now: x [304, 4096] BF16 (normalized) and residual [304, 4096] BF16 (the sum). The two switches are explained on CustomOp, out-of-tree ops and IR op providers.

9QKVParallelLinear and the rotary embedding#

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
        attn_output = self.attn(q, k, v)
        output, _ = self.o_proj(attn_output)
        return output

1. Line 229. qkv_proj is a QKVParallelLinear. With TP = 1 its weight is [6144, 4096]: 32 query heads plus 2 × 8 KV heads, times 128 (The tensor-parallel layers (linear.py and vocab_parallel_embedding.py)). The forward is one GEMM through the quantize method, and no collective (vllm/model_executor/layers/linear.py:616-634). For BF16 on CUDA, the default GEMM is torch.nn.functional.linear (vllm/model_executor/layers/utils.py:83-89, vllm/model_executor/layers/utils.py:689-691).
2. Line 230. The split uses the per-rank sizes: q_size = 32 × 128 = 4096, kv_size = 8 × 128 = 1024 (vllm/model_executor/models/llama.py:156-161).
3. Line 231. The rotary embedding rotates q and k by the positions of stop 3. It is a CustomOp too (vllm/model_executor/layers/rotary_embedding/base.py:14-15), so by default it runs forward_native: a lookup in cos_sin_cache by position, then the rotation of each head (vllm/model_executor/layers/rotary_embedding/base.py:160-219). The model's config.json on the Hugging Face Hub has no rope_scaling, so get_rope takes the default branch and builds a plain RotaryEmbedding (vllm/model_executor/layers/rotary_embedding/__init__.py:63-65, vllm/model_executor/layers/rotary_embedding/__init__.py:147-155).

What the data looks like now: qkv [304, 6144] → q [304, 4096], k [304, 1024], v [304, 1024], all BF16. RoPE keeps the shapes. RoPE itself is explained in Inside a transformer↗.

10Attention: two custom ops, a lookup by name, and FlashAttention#

self.attn(q, k, v) calls Attention.forward. It makes the output buffer [304, 4096] and views the inputs per head: q as [304, 32, 128], k and v as [304, 8, 128] (vllm/model_executor/layers/attention/attention.py:525-540). CUDA wants one opaque op for attention, so use_direct_call is false (vllm/model_executor/layers/attention/attention.py:436-440), and the layer calls two registered ops:

        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)

These two ops are splitting ops, so this part runs eagerly between piece i and piece i + 1. Both ops resolve the layer name in the forward context of stop 6:

    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. unified_kv_cache_update. FlashAttention writes the cache in a separate op (vllm/v1/attention/backends/flash_attn.py:367). The op calls impl.do_kv_cache_update with the slot mapping of stop 4 (vllm/model_executor/layers/attention/attention.py:734-748). That method scatters the new keys and values into this layer's cache view by slot (vllm/v1/attention/backends/flash_attn.py:1443-1464). A comment there says that the slot mapping is not padded (lines 1450-1454). This comment is stale: here the slot mapping has 304 entries, and rows 300 to 303 hold -1 (stop 4).
2. unified_attention_with_output. It takes the empty tensor of the first op as a dummy argument, so the compiler keeps the order. Then it calls impl.forward(self, query, key, value, kv_cache, attn_metadata, output=output) (vllm/model_executor/layers/attention/attention.py:779-796).

FlashAttentionImpl.forward splits the cache view [num_blocks, 8, 16, 256] into a key half and a value half (vllm/v1/attention/backends/flash_attn.py:1254-1255). It reads the batch layout from the metadata, then calls the kernel with the first num_actual_tokens rows only:

            _FA4_DENSE_ATTENTION_KERNEL(
                q=query[:num_actual_tokens],
                k=key_cache,
                v=value_cache,
                out=output[:num_actual_tokens],
                cu_seqlens_q=cu_seqlens_q,
                max_seqlen_q=max_seqlen_q,
                seqused_k=seqused_k,
                max_seqlen_k=max_seqlen_k,
                softmax_scale=self.scale,
                causal=causal,
                alibi_slopes=self.alibi_slopes,
                window_size=sliding_window_size,
                block_table=block_table,

The arguments come from stop 5: cu_seqlens_q is query_start_loc, seqused_k is seq_lens, and the block table maps each request to its blocks (vllm/v1/attention/backends/flash_attn.py:1279-1284). The wrapper object ends in flash_attn_varlen_func (vllm/v1/attention/backends/flash_attn.py:275-283, vllm/v1/attention/backends/flash_attn.py:113-116). The kernel itself is out of scope. For our batch it computes causal attention for 300 query rows: R1 reads 541 keys, R2 130 keys and R3 298 keys. Attention itself is explained in Attention↗.

What the data looks like now: the KV cache of this layer holds the new keys and values at the 300 slots. output [304, 4096] holds the attention result in its first 300 rows. In this step, get_attention_context runs 64 times: two ops in each of 32 layers.

11o_proj and the MLP#

o_proj starts piece i + 1. It is a row-parallel linear layer layer with weight [4096, 4096]. It calls the all-reduce only with TP > 1, so on our one GPU it is one GEMM (vllm/model_executor/layers/linear.py:1797-1806). Then post_attention_layernorm does the second fused add + RMSNorm (stop 8), and the MLP runs:

    def forward(self, x):
        x, _ = self.gate_up_proj(x)
        x = maybe_fused_act_quant(self.act_fn, x, self.down_proj)
        x, _ = self.down_proj(x)
        return x

1. Line 119. gate_up_proj is a column-parallel linear layer MergedColumnParallelLinear with weight [28672, 4096]: gate and up stacked. Output [304, 28672].
2. Line 120. maybe_fused_act_quant fuses the activation with a quantization of the next input only when down_proj wants a quantized input. For BF16 it just calls act_fn(x) (vllm/model_executor/layers/fusion/fused_act_quant.py:206-213). SiluAndMul is a disabled CustomOp, so it runs forward_native: silu(x[..., :d]) * x[..., d:] (vllm/model_executor/layers/activation.py:140-144). Output [304, 14336].
3. Line 121. down_proj is row-parallel with weight [4096, 14336]. Output [304, 4096], again without an all-reduce.

The layer returns (hidden_states, residual), both [304, 4096], to the loop of stop 7. Stops 8 to 11 repeat for each of the 32 layers. The layer classes and their slices are on The tensor-parallel layers (linear.py and vocab_parallel_embedding.py); the weights in them were loaded at start-up ("Weight loading: from checkpoint files to GPU parameters").

12The final norm; the hidden states are kept#

After layer 31, the last piece applies the final norm. It is the last fused add + RMSNorm, and the new residual is not used (vllm/model_executor/models/llama.py:452-457). The model returns hidden_states [304, 4096], not logits: compute_logits is a method of the outer class and stays outside the compiled graph (How a vLLM model file is written (llama.py)).

Back in the runner, the forward context closes. The runner does not return the hidden states. It stores them with the input batch and the metadata, and returns None:

        finished_req_ids = scheduler_output.finished_req_ids
        self.execute_model_state = ExecuteModelState(
            input_batch=input_batch,
            attn_metadata=attn_metadata,
            slot_mappings_by_layer=slot_mappings_by_layer,
            hidden_states=hidden_states,
            aux_hidden_states=aux_hidden_states,
            dp_sync_state=dp_sync_state,
            finished_req_ids=finished_req_ids,
            ec_connector_output=ec_connector_output,
            cudagraph_stats=cudagraph_stats,
            num_spec_tokens_to_schedule=scheduler_output.num_spec_tokens_to_schedule,
        )

        if not self.is_last_pp_rank:
            # Non-last PP rank: return IntermediateTensors for sending.
            assert output_intermediate_tensors is not None
            assert self.pp_handler is not None
            return self.pp_handler.relay_aux_hidden_states(
                model_inputs["intermediate_tensors"], output_intermediate_tensors
            )
        return None

The Python code returns as soon as all kernels are queued. The GPU may still run the first layers. The uni executor wraps None in a finished future, and the engine core goes on (engine-step tour, stops 7 and 8).

What the data looks like now:

execute_model_state = ExecuteModelState(input_batch=..., attn_metadata={32 names: m},
                                        hidden_states=bf16 [304, 4096], ...)   # 2,490,368 bytes

13sample_tokens: logits for three rows#

The engine core calls sample_tokens(None) (no grammar bitmask). The runner takes the stored state, and on the last PP rank it calls self.sample (vllm/v1/worker/gpu/model_runner.py:2117-2131, vllm/v1/worker/gpu/model_runner.py:2163-2165). Batch-sharded sampling is off, so the plain path runs:

        else:
            sample_hidden_states = hidden_states[input_batch.logits_indices]
            logits = self.model.compute_logits(sample_hidden_states)

hidden_states[logits_indices] keeps rows 0, 1 and 299: the last row of each request. The other 301 rows never become logits. compute_logits calls the LogitsProcessor with the LM head (vllm/model_executor/models/llama.py:546-551):

    def _get_logits(
        self,
        hidden_states: torch.Tensor,
        lm_head: VocabParallelEmbedding,
        embedding_bias: torch.Tensor | None,
        skip_gather: bool = False,
    ) -> torch.Tensor | None:
        # Get the logits for the next tokens.
        logits = self._apply_head(lm_head, hidden_states, embedding_bias)
        if skip_gather:
            return logits

        # Gather logits for TP
        if lm_head.tp_size > 1:
            logits = self._gather_logits(logits)

        # Remove paddings in vocab (if any).
        if logits is not None:
            logits = logits[..., : self.org_vocab_size]
        return logits

The head runs in the model dtype by default for a generation model (vllm/config/model.py:2540-2545, vllm/model_executor/layers/logits_processor.py:143-146). With TP = 1 there is no gather (line 193). The vocabulary 128256 is a multiple of 64, so the slice at line 198 cuts nothing (The tensor-parallel layers (linear.py and vocab_parallel_embedding.py)).

What the data looks like now: sample_hidden_states [3, 4096] and logits [3, 128256] BF16 (769,536 bytes).

14The sampler and the AsyncOutput#

Without draft tokens, the plain Sampler runs (vllm/v1/worker/gpu/model_runner.py:1663-1665). Our requests send no sampling fields, so the server defaults come from the model's generation_config.json (The OpenAI-compatible API). For this model that file sets temperature 0.6 and top-p 0.9 (from the Hugging Face Hub, not from the vLLM repository). Temperature 0.6 needs processing, so the sampler copies the logits to FP32 and applies the temperature (vllm/v1/worker/gpu/sample/states.py:67-72, vllm/v1/worker/gpu/sample/sampler.py:238-265). Then it picks the sampling method:

        # Don't use a fused sampler if no requests use top_k/top_p, if there are
        # any greedy requests or per-request seeds, or if post-processed
        # logprobs need to be returned for any requests.
        fused_sampler_eligible = not (
            (top_k is None and top_p is None)
            or (return_logprobs and self.logprobs_mode in PROCESSED_LOGPROBS_MODES)
            or self.sampling_states.any_greedy(idx_mapping_np)
            or self.sampling_states.any_explicit_seed(idx_mapping_np)
        )
        use_fused_sampler = (
            self.use_flashinfer or self.use_xpu_sampler
        ) and fused_sampler_eligible

        return self._sample_random(
            processed_logits,
            expanded_idx_mapping,
            idx_mapping_np,
            pos,
            top_k,
            top_p,
            has_topp_only_rows,
            use_fused_sampler,
        )

• Top-p is set, and no request is greedy, seeded or asks for processed logprobs. So the batch is eligible for the fused sampler.
• use_flashinfer needs two things. First, the sampling mask is off, which is the default. Second, FlashInfer's sampler supports the GPU: CUDA, a supported compute capability, more than 16 SMs, and a working FlashInfer JIT build. VLLM_USE_FLASHINFER_SAMPLER defaults to on (vllm/v1/worker/gpu/sample/sampler.py:88-90, vllm/envs.py:886-890, vllm/v1/sample/ops/topk_topp_sampler.py:118-150). An H100 has more than 100 SMs. Then FlashInfer's fused top-k/top-p sampler draws the tokens (vllm/v1/worker/gpu/sample/sampler.py:339-343).
• Otherwise, top-k/top-p masking and the Gumbel sampler run (lines 346-359).

R3's sequence length equals its prompt length, so R3 is not in a chunked prefill and samples its first token. All three rows count (vllm/v1/worker/gpu/sample/sampler.py:190-199). The runner then starts the copy to the CPU before it updates its own state, and returns the AsyncOutput:

        # Start async output copy here so that it can overlap with speculator proposal.
        async_output = AsyncOutput(
            model_runner_output=model_runner_output,
            sampler_output=sampler_output,
            num_sampled_tokens=num_sampled,
            main_stream=self.main_stream,
            copy_stream=self.output_copy_stream,
            check_ep_fault=self.check_ep_fault,
            pending_aux_output=pending_aux_output,
        )

The copy waits for the main stream on a second stream, copies into pinned CPU memory, and records copy_event (vllm/v1/worker/gpu/async_utils.py:137-171). postprocess_sampled writes the three new tokens into last_sampled_tokens on the GPU, for the next step (vllm/v1/worker/gpu/model_runner.py:2244-2255). The function returns async_output without a wait (vllm/v1/worker/gpu/model_runner.py:2311). The engine core calls get_output() later, which is the only wait for the GPU (vllm/v1/worker/gpu/async_utils.py:173-184; engine-step tour, stop 10). Sampling methods are explained in Sampling↗.

What the data looks like now:

sampled_token_ids int64 [3, 1] = [[t_R1], [t_R2], [t_R3]]     # on the GPU
num_sampled             [3]    = [1, 1, 1]
last_sampled_tokens[[1023, 1022, 1021]] = [t_R1, t_R2, t_R3]  # inputs of the next step