GPUModelRunner, the input batch, the attention backend

Here is what the tree looks like:

$ wc -l vllm/v1/worker/gpu_model_runner.py
8008 vllm/v1/worker/gpu_model_runner.py
$ wc -l vllm/v1/worker/gpu/model_runner.py
2024 vllm/v1/worker/gpu/model_runner.py
$ find vllm/v1/worker/gpu -name '*.py' | xargs wc -l | tail -1
   19467 total

Two files named model_runner.py, both defining a class named GPUModelRunner, both accepting a SchedulerOutput and returning a ModelRunnerOutput. One is an 8,008-line monolith. The other is the entry point of a 19,467-line package. The package's README.md is four lines long and reads, in full:

# [Experimental] Model Runner V2

This directory contains the new model runner which is under active development.
Ping [Woosuk Kwon](https://github.com/WoosukKwon) for any changes.

That README is misleading, and the mismatch is the single most useful thing this chapter can tell you. The selection policy in vllm/config/vllm.py defaults to the "experimental" runner for eligible dense configurations after all guards, including the reference Llama-3-8B setup. A family name alone does not certify eligibility. The 8,008-line file is the fallback, reached by MoE architectures not on an opt-in list, by hybrid and attention-free models, and by any configuration that trips one of the V2 feature gates.

And the log tells you nothing. There is no line that says which runner was constructed. There are only two negative warnings, emitted when V2 was wanted and refused. Silence means V2 — or means the model was never eligible in the first place. You cannot distinguish those two silences from the log.

§11.1 owns the rewrite narrative and the repo map. §11.3 ends where this chapter starts: with a SchedulerOutput in hand. This chapter owns the code that turns that object into GPU tensors, launches a forward, and hands back sampled token ids — in both runners.

Whichever runner you are on, one step is the same seven-stage pipeline. The scheduler has already decided who runs and how many tokens each. The runner's entire job is to turn that per-request dictionary into a small number of flat, contiguous GPU tensors whose layout the attention kernels expect, run the model over them once, and scatter the result back out per request.

The reason this is hard — the reason it takes 8,008 lines in one design and a 19,467-line package in the other — is that the batch is ragged. Request 0 may be prefilling 512 tokens while requests 1 through 3 each decode a single token. The forward pass sees one flat sequence of 515 tokens with no request boundaries in it. Every boundary the kernels need has to be reconstructed from side tensors.

The two runners differ in where each stage runs. V1 computes stages 2 and 3 largely in NumPy on the CPU, into pinned buffers it then copies to the GPU. V2 computes them in Triton kernels on the GPU, reading persistent GPU-resident state. That single choice cascades into almost every other difference between them.

Flatten a ragged batch and you lose two independent pieces of information, and you need both back.

The query axis — how many tokens of the flat buffer belong to each request. This is query_start_loc, an exclusive-prefix-sum of shape [num_reqs + 1]. Request $i$ owns flat positions the half-open interval [query_start_loc[i], query_start_loc[i+1]).

The key axis — how many keys each request's queries must attend over. This is seq_lens, shape [num_reqs], equal to num_computed_tokens + num_scheduled_tokens. It is not derivable from query_start_loc: a decoding request contributes one query but attends over two thousand keys.

vLLM builds them side by side. In V1 the arithmetic is NumPy:

        # Get request indices.
        # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
        req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)

        # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]
        # self.query_pos.np[:10]: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
        cu_num_tokens = self._get_cumsum_and_arange(
            num_scheduled_tokens, self.query_pos.np
        )

        # Get positions.
        positions_np = (
            self.input_batch.num_computed_tokens_cpu[req_indices]
            + self.query_pos.np[: cu_num_tokens[-1]]
        )

Three arrays fall out of that comment block. req_indices maps every flat token back to its request row. query_pos is the token's offset within its request's scheduled chunk. And positions — the absolute index in the sequence, which is what RoPE and the KV write both need — is just num_computed_tokens[req] + query_pos.

Then the same three quantities are recomputed on the GPU, because the GPU copies are the ones the kernels read:

        self.num_scheduled_tokens.np[:num_reqs] = num_scheduled_tokens
        self.num_scheduled_tokens.copy_to_gpu(num_reqs)
        num_scheduled_tokens_gpu = self.num_scheduled_tokens.gpu[:num_reqs]
        self.positions[:total_num_scheduled_tokens] = (
            self.num_computed_tokens[req_indices_gpu].to(torch.int64)
            + self.query_pos.gpu[:total_num_scheduled_tokens]
        )
        self.seq_lens[:num_reqs] = (
            self.num_computed_tokens[:num_reqs] + num_scheduled_tokens_gpu
        )
        self.seq_lens[num_reqs:].fill_(0)

        self.input_batch.block_table.compute_slot_mapping(
            num_reqs,
            self.query_start_loc.gpu[: num_reqs + 1],
            self.positions[:total_num_scheduled_tokens],
        )

Note self.seq_lens[num_reqs:].fill_(0). Padded rows must be zeroed, because a full CUDA graph replays over a fixed-size buffer and a stale seq_len from three steps ago in row 7 would make an attention kernel walk a block table row that has since been freed. This is the recurring hazard of the whole design: persistent buffers are correct only if every unused slot is explicitly neutralised.

The slot mapping is the bridge

§2.2 established the block table and the slot-mapping concept. Here is the function that computes it, once per step, on the GPU:

    def compute_slot_mapping(
        self,
        num_reqs: int,
        query_start_loc: torch.Tensor,
        positions: torch.Tensor,
    ) -> None:
        num_tokens = positions.shape[0]
        if self.slot_mapping_mode == SlotMappingMode.NONE:
            # Mamba/GDN groups consume the block table as recurrent state
            # indices and do not use per-token slot mappings.
            return
        assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT

        _COMPUTE_SLOT_MAPPING_KERNEL(
            num_reqs,
            num_tokens,
            self.max_num_batched_tokens,
            query_start_loc,
            positions,
            self.block_table.gpu,
            self.block_table.gpu.stride(0),
            self.block_size,
            self.slot_mapping.gpu,
            self.kv_cache_block_size,
            self.blocks_per_kv_block,
            self.dcp_world_size,
            self.dcp_rank,
            self.cp_kv_cache_interleave_size,
        )

The Triton kernel body at vllm/v1/worker/block_table.py:L412-L476 is the whole of paged attention's addressing in twenty lines. Strip context parallelism (set TOTAL_CP_WORLD_SIZE = 1) and it reduces to: block_indices = pos // block_size; block_numbers = block_table[req, block_indices]; slot_ids = block_numbers * block_size + pos % block_size. Every token gets one int64 naming the physical row of the flat KV cache tensor that its K and V will be written into.

Two details matter. First, the last thread block pads the tail of the buffer with PAD_ID up to max_num_tokens, so a CUDA graph replaying at a larger padded shape marks padded entries with a negative no-write sentinel; the consuming cache kernel must mask them rather than index a live slot. Second, BLOCKS_PER_KV_BLOCK is where §2.2's kernel_block_size subdivision lands: if the allocator hands out 32-token blocks and the kernel wants 16, BlockTable.__init__ at vllm/v1/worker/block_table.py:L89-L110 sets blocks_per_kv_block = 2 and every allocator block becomes two kernel blocks.

Worked arithmetic: batch of 4, Llama-3-8B, block_size 16

Four requests, all decoding one token, with num_computed_tokens of 1000, 512, 37, and 2047. Every quantity below is arithmetic from the formulas above.

query_start_loc = [0, 1, 2, 3, 4] as int32[5]. positions = [1000, 512, 37, 2047] as int64[4] — note the dtype asymmetry, visible in the buffer allocations at vllm/v1/worker/gpu_model_runner.py:L818-L827: positions are int64 because they index into the token-id table, while query_start_loc and seq_lens are int32 because that is what FlashAttention's varlen API wants. logits_indices = query_start_loc[1:] - 1 = [0, 1, 2, 3].

The block-table copy is the other cost worth sizing. At max_model_len = 8192 and block_size = 16 a row is 512 int32 entries, 2,048 bytes. With max_num_reqs = 1024 the full CPU-side table is 2 MiB. commit_block_table(num_reqs) at vllm/v1/worker/block_table.py:L231-L232 copies only the live prefix — 8 KiB for our four requests. That prefix-only copy is exactly what forces the persistent batch to stay dense, which is the subject of the next section, and it is the cost V2 removes entirely by keeping the table on the GPU and shipping only diffs.

Answer this before you read either file, or you will read the wrong one. The property is VllmConfig.use_v2_model_runner at vllm/config/vllm.py:L648-L700. It consults the env var first — VLLM_USE_V2_MODEL_RUNNER, declared at vllm/envs.py:L2041-L2044 as a tri-state bool | None, so unset means "use policy", not "off". Then it forces V2 on for prefill context parallelism, for the dspark speculative method, for two DFlash draft shapes, and for diffusion models. Then it reaches the model check:

    def _is_default_v2_model_runner_model(self) -> bool:
        model_config = self.model_config
        if model_config is None:
            return False

        architectures = getattr(model_config, "architectures", [])
        default_architectures = default_v2_model_runner_architectures()
        is_default_v2_architecture = any(
            arch in default_architectures for arch in architectures
        )

        if getattr(model_config, "is_hybrid", False) and (
            not is_default_v2_architecture
        ):
            return False

        if getattr(model_config, "is_attention_free", False):
            return False
        return is_default_v2_architecture or not model_config.is_moe

Read the last line slowly. not model_config.is_moe — a dense model is eligible by default, with no allow-list entry required. The allow-list, DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES at vllm/config/vllm.py:L69-L82, contains ten entries and every one of them is a MoE architecture (DeepSeek V2/V3.2/V4, GLM-MoE-DSA, GraniteMoe, Inkling, Kimi-K3, LongCat, Qwen2-MoE). Its purpose is to opt specific MoE models into V2, not to gate dense models. Llama-3-8B is not on that list and does not need to be.

Eligibility is not the end. Two more gates can send you back to V1, and they are the only things that log:

        if not HAS_TRITON:
            logger.warning_once(
                "Model Runner V2 requires Triton; using the V1 model runner instead."
            )
            return False

        unsupported = self._get_v2_model_runner_unsupported_features()
        if unsupported:
            logger.warning_once(
                "Model Runner V2 does not yet support %s; using the V1 model "
                "runner instead.",
                ", ".join(unsupported),
            )
            return False

        return True

_get_v2_model_runner_unsupported_features() at vllm/config/vllm.py:L2438-L2543 is the honest inventory of what V2 cannot yet do, and it is worth reading in full before you choose a configuration: stock torch.compile mode, sequence parallelism with TP>1, pipeline parallelism under external_launcher, ngram/ngram_gpu speculation, dual batch overlap, elastic expert parallelism, custom logits processors, and KV-sharing fast prefill. Any one of those flips you to V1 with a single warning_once.

Downstream, the worker picks the class:

        # Construct the model runner
        if self.use_v2_model_runner:
            from vllm.v1.worker.gpu.model_runner import (
                GPUModelRunner as GPUModelRunnerV2,
            )

            # HACK(woosuk): This is a temporary fix to avoid type errors.
            self.model_runner: GPUModelRunner = GPUModelRunnerV2(  # type: ignore
                self.vllm_config, self.device
            )
        else:
            from vllm.v1.worker.gpu_model_runner import (
                GPUModelRunner as GPUModelRunnerV1,
            )

            self.model_runner = GPUModelRunnerV1(self.vllm_config, self.device)

And the choice leaks past the worker: use_v2_model_runner is read at four sites in vllm/v1/core/sched/scheduler.py (L307, L1202, L1234, L1521), once in vllm/v1/core/sched/async_scheduler.py:L46, and once in vllm/v1/attention/backends/flashinfer.py:L940, where it decides whether the FlashInfer builder pins its host buffers. The runner choice is not confined to the worker process's inner loop; it changes scheduler behaviour and backend memory policy too.

Both runners face the same question: where does per-request state live between steps? §1.3 owns the concept — the persistent batch exists because rebuilding a 2 MiB block table in Python every 5 ms is unaffordable. The two runners answer it differently, and the difference is the cleanest illustration of what the rewrite is actually about.

V1: slot-indexed, dense, and it moves rows

vllm/v1/worker/gpu_input_batch.py is 1,155 lines defining two things: CachedRequestState (L34-L89), a plain Python mirror of everything about a request, and InputBatch (L92-L1155), a bank of parallel arrays indexed by a batch slot.

The [R, M] token table is the memory hazard. At max_num_reqs = 1024 and max_model_len = 128k it is 512 MiB of host RAM, which is why the source carries a TODO(woosuk) at L130 saying exactly that.

The invariant that makes everything else work is density: live requests occupy rows 0..num_reqs-1 with no holes, because every downstream consumer slices [:num_reqs]. When a request finishes mid-batch it leaves a hole, and condense() closes it:

    def condense(self) -> None:
        """Slide non-empty requests down into lower, empty indices.

        Any consecutive empty indices at the very end of the list are not
        filled.

        Returns:
          swaps: list of (from,to) swap tuples for moved requests
          empty_req_indices: indices not filled by condensation
        """
        num_reqs = self.num_reqs

        if not (empty_req_indices := self.batch_update_builder.removed):
            # All removed requests were replaced by added requests, or else no
            # requests were removed at all. No condense() needed
            return

The loop body, L734-L830, is the cost: for each hole it copies a slice of token_ids_cpu, a slice of is_token_ids, five scalar arrays, a block_table.move_row(), the LoRA mapping, six sampling scalars, the RNG generator, the allowed-token mask row, and the bad-words dict entry — and appends a (from, to, UNIDIRECTIONAL) tuple so logits processors can rewrite their own indices. Add a field to InputBatch and forget to add it to condense() and you get a silent correctness bug that only appears when a request finishes out of order.

V2: fixed slots, a free list, and a gather

V2 deletes the problem. vllm/v1/worker/gpu/states.py is 133 lines. A request gets a permanent row for its lifetime:

        self.req_id_to_index: dict[str, int] = {}
        self.index_to_req_id: dict[int, str] = {}
        self.free_indices = list(range(max_num_reqs))
# ...
    def add_request(
        self,
        req_id: str,
        prompt_len: int,
        all_token_ids: list[int],
        num_computed_tokens: int,
        max_tokens: int,
    ) -> None:
        assert len(self.free_indices) > 0, "No free indices"
        req_idx = self.free_indices.pop()
# ...
    def remove_request(self, req_id: str) -> int | None:
        """Return the freed slot index, or None if the request was not found."""
        req_idx = self.req_id_to_index.pop(req_id, None)
        if req_idx is None:
            return None
        self.index_to_req_id.pop(req_idx, None)
        self.free_indices.append(req_idx)
        return req_idx

No condense(). No move. The state table is allowed to be sparse, and the per-step batch is produced by a gather: gather_batch_req_state() at vllm/v1/worker/gpu/model_runner.py:L1057-L1109 builds idx_mapping_np, an intp array of length num_reqs mapping batch position to state row, and every kernel afterwards takes that array as an argument. The design document states the rationale directly: "Assign each request a permanent row for its active lifetime... This removes the need for CachedRequestState and simplifies bookkeeping" (docs/design/model_runner_v2.md:L36-L42).

Batch order is not arrival order either. sort_batch_req_ids() at vllm/v1/worker/gpu/model_runner.py:L2012-L2024 sorts verification/decode first, then short extends, then prefills, with a comment that split_decodes_and_prefills relies on decode-like requests leading. In V1 the equivalent is _may_reorder_batch() at L1169-L1193, which physically swaps rows of the persistent batch to satisfy the same backend requirement. Same constraint; one sorts a list of ids, the other permutes a bank of arrays.

The cost V2 pays is that the state table must be sized for max_num_reqs rows even when few are live, and that all_token_ids — potentially several GB — is backed by UVA rather than device memory (vllm/v1/worker/gpu/states.py:L31-L38), so kernels reach across PCIe to read it. The cost it avoids is per-step block-table copies: StagedWriteTensor at vllm/v1/worker/gpu/buffer_utils.py:L114-L205 keeps the base tensor on the GPU, accumulates ragged CPU diffs in three Python lists, packs them, and applies them with a single Triton kernel launch in apply_write().

§3.4 owns the backend abstraction, CommonAttentionMetadata, AttentionCGSupport, and the metadata-built-once-per-step design. This section shows only where the runners invoke it.

Both build one CommonAttentionMetadata per KV-cache group and call builder.build() once per attention group within it, then fan the single resulting object out to every layer in that group:

        for attn_group in attn_groups[i]:
            attn_metadata_builder = attn_group.get_metadata_builder(0)
            if for_cudagraph_capture:
                metadata = attn_metadata_builder.build_for_cudagraph_capture(
                    common_attn_metadata
                )
            else:
                attn_metadata_extra_kwargs = (
                    model_specific_attn_metadata.get_extra_attn_kwargs(
                        attn_metadata_builder,
                        num_reqs,
                    )
                    if model_specific_attn_metadata is not None
                    else {}
                )
                metadata = attn_metadata_builder.build(
                    common_prefix_len=0,
                    common_attn_metadata=common_attn_metadata,
                    **attn_metadata_extra_kwargs,
                )
            for layer_name in attn_group.layer_names:
                attn_metadata[layer_name] = metadata

For Llama-3-8B that is one build() call whose result is stored under 32 layer names. V1's equivalent, _build_attention_metadata() at vllm/v1/worker/gpu_model_runner.py:L2355-L2691, is 336 lines and does more: it also handles cascade-attention prefix lengths, micro-batch splitting for dual-batch overlap, a memoisation cache keyed on (KVCacheSpec, builder_type) that lets hybrid groups reuse a build via builder.update_block_table(), and per-drafter block-table capture. V2's is 90 lines because none of those features exist in it yet.

Dispatch: replay or eager

§8.1 owns the bucket ladder and capture. At the runner level the decision has three inputs: the padded token count, whether the batch is a uniform decode, and how many LoRAs are active.

"Uniform decode" is the predicate that makes a full-graph replay legal, and it is four lines at vllm/v1/worker/gpu_model_runner.py:L4000-L4007: max_num_scheduled_tokens == uniform_decode_query_len and num_tokens == max_num_scheduled_tokens * num_reqs. Both halves matter — the second rules out a batch where one request happens to be scheduled at the decode query length while others are not. One prefill in the batch and the whole step falls to piecewise or eager.

V1 then calls self.cudagraph_dispatcher.dispatch(...) inside _determine_batch_execution_and_padding() (vllm/v1/worker/gpu_model_runner.py:L4054-L4163), receives a (CUDAGraphMode, BatchDescriptor) pair, and passes the mode into the forward context. The actual replay happens implicitly: set_forward_context(..., cudagraph_runtime_mode=cudagraph_mode, ...) at L4546-L4560 wraps a plain self._model_forward(...), and the graph wrapper inside the compiled model reads the mode out of the context.

V2 makes it explicit:

        # Run model.
        if batch_desc.cg_mode == CUDAGraphMode.FULL:
            # Use explicit cudagraph replay for FULL mode.
            # NOTE(woosuk): Here, we don't need to pass the input tensors,
            # because they are already copied to the CUDA graph input buffers.
            assert self.cudagraph_manager is not None
            self.kv_connector.pre_forward(scheduler_output)
            model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
        else:
            # For piecewise and eager mode, just call model().

No arguments. The inputs were written into the graph's static buffers during prepare_inputs(), because self.input_buffers.input_ids and self.input_buffers.positions are the captured buffers. This is what the design doc means by "V1's CUDA graph handling is implicit and hard to reason about" (docs/design/model_runner_v2.md:L188-L190). The matching selection logic is a 30-line dispatch() at vllm/v1/worker/gpu/cudagraph_utils.py:L382-L410 that walks a priority-ordered candidate list keyed on (num_tokens, effective_loras) and returns a BatchExecutionDescriptor with cg_mode = NONE if nothing matches.

"2,024 lines replaces 8,008" is the wrong summary. The V2 tree is 19,467 lines — more code, not less. What changed is coupling.

V1 declares class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin) at vllm/v1/worker/gpu_model_runner.py:L501-L503; V2 declares class GPUModelRunner(LoRAModelRunnerMixin) at vllm/v1/worker/gpu/model_runner.py:L159 and holds self.kv_connector and self.ec_connector as objects instead. Behaviour that was inherited became behaviour that is called, and the import count went up precisely because the dependencies became visible.

Input-batch management

V1: one class, InputBatch, owns persistent state and is the sampler's input. That coupling is what forces condense(); the design doc names it as the root cause — "V1 uses persistent state tensors directly as model and sampler inputs, which imposes strict layout and ordering requirements" (docs/design/model_runner_v2.md:L25-L29). V2 splits it in three: RequestState (persistent, sparse, GPU), InputBuffers (the fixed-size CUDA-graph buffers), and InputBatch (a frozen per-step dataclass of views and slices, vllm/v1/worker/gpu/input_batch.py:L41-L115).

Sampling

§6.1 owns the operators. Structurally: V1 calls into vllm/v1/sample/ (4,624 lines) with a SamplingMetadata object rebuilt by InputBatch.refresh_metadata() whenever the batch changes. V2 has its own gpu/sample/ (2,744 lines, twelve files: gumbel.py, penalties.py, min_p.py, logit_bias.py, bad_words.py, logprob.py, prompt_logprob.py, thinking_budget.py, trace_replay.py, states.py, output.py, sampler.py) and passes the InputBatch itself: sampler_output = self.sampler(logits, input_batch) at vllm/v1/worker/gpu/model_runner.py:L1361. Per-request sampling state is dereferenced inside the kernels through idx_mapping rather than being expanded to per-logit shape first.

Speculative decoding

V1 keeps drafters in vllm/v1/spec_decode/ and dispatches with a chain of isinstance checks against EagleProposer, DFlashProposer, Gemma4Proposer, Step3p5MTPProposer and others inside _build_attention_metadata() — see vllm/v1/worker/gpu_model_runner.py:L2648-L2668. V2 has a Speculator interface with one package per family (eagle/, mtp/, dflash/, dflash2/, dspark/, gemma4/, multi_module_mtp/, autoregressive/) and no isinstance ladder in the runner.

That reorganisation is also where a feature can exist in one tree and not the other. §6.6 found this; here is the code. The config accepts three methods — RejectionSampleMethod = Literal["standard", "synthetic", "block"] at vllm/config/speculative.py:L80 — but only the V2 sampler implements the third:

        rejection_sample_method = spec_config.rejection_sample_method
        self.use_block_verification: bool = False
        self.synthetic_conditional_rates: torch.Tensor | None = None
        if rejection_sample_method == "synthetic":
            assert spec_config.synthetic_acceptance_rates is not None
            self.synthetic_conditional_rates = torch.tensor(
                unconditional_to_conditional_rates(
                    spec_config.synthetic_acceptance_rates
                ),
                dtype=torch.float32,
                device=device,
            )
        elif rejection_sample_method == "block":
            self.use_block_verification = True

V1's rejection sampler at vllm/v1/sample/rejection_sampler.py:L77-L90 has the synthetic branch and no block branch. SpeculativeConfig.__post_init__ at vllm/config/speculative.py:L1380-L1395 validates only the synthetic parameters. So on the V1 path rejection_sample_method: "block" is accepted, stored, never read, and never warned about. (There is no top-level --rejection-sample-method flag at this SHA: the field is reachable only inside the --speculative-config JSON, vllm/engine/arg_utils.py:L1637-L1640.) Since a dense model defaults to V2, the flag usually works — and then silently stops working the day you switch to a MoE checkpoint that is not on the allow-list.

Multimodal and pooling

V1 handles multimodal inline: _execute_mm_encoder() spans L3077-L3294 and _gather_mm_embeddings() L3300-L3415 of the runner file. V2 moves both into gpu/mm/ (encoder_runner.py, encoder_cache.py, rope.py, lora.py) behind self.model_state.prepare_inputs_embeds(...), with the model-family variation pushed into gpu/model_states/ — nine files including encoder_decoder.py, mamba_hybrid.py, and prompt_embeds.py.

Pooling is where the incompleteness is loudest. gpu/pool/ is 212 lines against V1's full pooling support, and it starts with a hard-coded allow-list and two constructor-time raises:

_SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset(
    {"embed", "classify", "token_embed", "token_classify", "embed&token_classify"}
)
# ...
        selected_task = self.model_config.get_pooling_task(model_tasks)
        if selected_task not in _SUPPORTED_TASKS:
            hint = (
                "Set an explicitly supported task or VLLM_USE_V2_MODEL_RUNNER=0."
                if _SUPPORTED_TASKS.intersection(model_tasks)
                else "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
            )
            raise ValueError(
                "Model Runner V2 supports pooling tasks "
                f"{sorted(_SUPPORTED_TASKS)}, but this model selects "
                f"{selected_task!r} from {list(model_tasks)}. {hint}"
            )

Those two error strings are the most honest documentation of V2's maturity in the repo. The escape hatch is always the same: set VLLM_USE_V2_MODEL_RUNNER=0. Note also that this failure mode is a raise, not a fallback — unlike the config-level gates, an unsupported pooling task takes the engine down at startup rather than quietly reverting.

The four requests from §3, one token each, Llama-3-8B, one KV-cache group, one attention group, full CUDA graphs enabled. Function names in call order.

On V2 (the default for this model)

1. execute_model(scheduler_output) — gpu/model_runner.py:L1416. State update first:
   vllm/v1/worker/gpu/model_runner.py:L1425-L1432 vLLM

           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()

   Five named methods instead of one 376-line _update_states. apply_staged_writes() flushes the block-table diffs — for a pure decode step where nobody crossed a block boundary, zero rows.
2. gather_batch_req_state() — L1057. Produces req_ids sorted decode-first, num_scheduled_tokens = int32[4] = [1,1,1,1], idx_mapping_np = intp[4].
3. dispatch_cg_and_sync_dp(...) → CudaGraphManager.dispatch() at gpu/cudagraph_utils.py:L382. uniform_token_count = 1, so a FULL-mode BatchExecutionDescriptor for the smallest captured bucket ≥ 4 matches.
4. prepare_inputs(...) — L1110. query_start_loc_np = [0,1,2,3,4] by np.cumsum, copied to the persistent int32[max_num_reqs+1] buffer; tail filled with num_tokens so FlashAttention's non-decreasing requirement holds through the padding. Then prepare_pos_seq_lens(...) launches num_reqs + 1 Triton programs — the extra one zeroes padded seq_lens — writing positions int64[4] = [1000,512,37,2047] and seq_lens int32[4] = [1001,513,38,2048]. Then combine_sampled_and_draft_tokens(...) fills input_ids int32[4] directly from req_states.last_sampled_tokens — the previous step's output, never round-tripped to the CPU — and returns logits_indices int32[4].
5. prepare_attn(input_batch) — L1311. gather_block_tables() materialises int32[num_reqs_padded, max_num_blocks] by gathering rows through idx_mapping; compute_slot_mappings() produces int64[1, num_tokens_padded], real slots at 0..3 and PAD_SLOT_ID after.
6. build_attn_metadata(...) — gpu/attn_utils.py:L591. One CommonAttentionMetadata, one builder.build(), stored under 32 layer names.
7. self.cudagraph_manager.run_fullgraph(batch_desc) — L1643. No arguments. hidden_states comes back as bf16[num_tokens_padded, 4096].
8. sample_tokens(grammar_output) — L1715 → sample() at L1341: hidden_states[logits_indices] gives bf16[4, 4096], compute_logits gives [4, 128256], the Triton sampler returns sampled_token_ids int64[4, 1].
9. postprocess_sampled(...) — L1375 → post_update(...) at gpu/input_batch.py:L604. One kernel advances num_computed_tokens, writes last_sampled_tokens, appends to all_token_ids, bumps total_len, and updates penalty bin counts. All on the GPU; the CPU never learns the token ids in this call.
10. AsyncOutput starts the D2H copy on a side stream while the speculator (if any) proposes. The ModelRunnerOutput is constructed with sampled_token_ids=None and filled when the copy lands.

On V1 (a MoE checkpoint, same batch)

Same seven stages, different mechanics. execute_model at L4288 opens with the async barrier:

        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
        with (
            record_function_or_nullcontext("gpu_model_runner: preprocess"),
            self.synchronize_input_prep(),
        ):
            # Update persistent batch states.
            deferred_state_corrections_fn = self._update_states(scheduler_output)

That synchronize_input_prep() context manager (L3942-L3954) is the "async barrier" the design doc lists as V1 design mistake number three. Everything inside it touches pinned CPU buffers that a previous step's H2D copy may still be reading.

            num_reqs = self.input_batch.num_reqs
            req_ids = self.input_batch.req_ids
            tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
            num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
            max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
            num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens

            logits_indices, spec_decode_metadata, max_num_sampled_tokens = (
                self._prepare_inputs(scheduler_output, num_scheduled_tokens_np)
            )

Note the list comprehension: scheduler_output.num_scheduled_tokens is keyed by request id, but the persistent batch is keyed by row, so V1 reindexes through self.input_batch.req_ids on every step. This is the point where a stale req_ids list — one that condense() failed to update — becomes a wrong-token bug rather than a crash.

Then _prepare_inputs (L2019-L2354) does everything from §3 in NumPy, _determine_batch_execution_and_padding (L4054) picks the mode, _get_slot_mappings (L4204) collects per-group and per-layer views and fills the padded tail with -1, _build_attention_metadata (L2355) builds, _preprocess (L3612) assembles the model kwargs, and _model_forward (L3956) calls self.model(...) under a forward context carrying the graph mode. execute_model then stashes an ExecuteModelState NamedTuple and returns None — sampling is a separate call, sample_tokens(grammar_output) at L4667, so that structured-output bitmasks computed on the CPU can be applied between forward and sample without stalling the launch.

Reading the monolith

If you do land on V1, do not read it front to back. The map:

A useful heuristic for a first pass over execute_model: assume speculative_config is None, uses_mrope is False, cascade_attn_enabled is False, data_parallel_size == 1, lora_config is None, and is_pooling_model is False. Roughly 70% of the branches vanish, and what remains is Figure 1.

Silent config no-ops

The block-verification example needs end-to-end validation and selection evidence. A field absent from a selected sampler is a useful diagnostic lead; it may also be rejected earlier. Grep the field across configuration and both runners, then record the resolved class and observed behavior. Do not infer that every V1-path feature gap is a silent no-op.

Padding is not optional

The failures that are hardest to diagnose come from stale padded rows. Three places in the step explicitly neutralise padding, and all three are load-bearing under full CUDA graphs: query_start_loc.np[num_reqs + 1:].fill(cu_num_tokens[-1]) (V1 L2135-L2137) so the varlen kernel sees a non-decreasing array; seq_lens[num_reqs:].fill_(0) (V1 L2255) so no padded row claims a length; and the slot-mapping kernel's PAD_ID tail fill (block_table.py:L431-L440) so consuming cache kernels can skip writes for negative sentinel entries. V1 additionally re-fills with -1 in _get_slot_mappings (L4250-L4252) with the comment "Needed for reshape_and_cache in full cuda graph mode." Symptom of getting this wrong: correct output at eager, garbage output only at certain batch sizes, only with cudagraph_mode=FULL.

Two runners, one InputBatch name

vllm/v1/worker/gpu_input_batch.py and vllm/v1/worker/gpu/input_batch.py both export a class named InputBatch, and they are unrelated types — one is a mutable 1,155-line state manager, the other a frozen per-step dataclass. A stack trace naming InputBatch tells you nothing until you read the module path in the frame above it. The same is true of block_table.py, BlockTable vs. BlockTables, and of course GPUModelRunner itself.

Determine which runner you are on, without a GPU, before you start reading:

# The policy, straight from config. No GPU needed.
python -c "
from vllm.engine.arg_utils import EngineArgs
cfg = EngineArgs(model='meta-llama/Meta-Llama-3-8B-Instruct').create_engine_config()
print('use_v2_model_runner =', cfg.use_v2_model_runner)
print('unsupported        =', cfg._get_v2_model_runner_unsupported_features())
"

# Force the other one and diff the behaviour you care about.
VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Meta-Llama-3-8B-Instruct

# Confirm the allow-list is MoE-only.
python -c "
from vllm.config.vllm import DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES as A
print(len(A)); print(sorted(A))
"

Then read one method in each runner side by side. The most instructive pair is V1's _prepare_inputs (gpu_model_runner.py:L2019-L2354, 336 lines of NumPy) against V2's prepare_inputs (gpu/model_runner.py:L1110-L1310, 201 lines that mostly launch kernels defined in gpu/input_batch.py:L254-L714). Both compute the same four tensors.

If you have a GPU, the measurement that separates them is CPU-side step time at high batch and small model — the regime where Python input prep, not the forward, is the bottleneck. At batch 1 on H100 the Llama-3-8B decode floor is 4.48 ms (see §0.4), so any input-prep difference is buried. Scale --max-num-seqs up and profile with --enable-layerwise-nvtx-tracing: the NVTX ranges gpu_model_runner: preprocess, : forward, : postprocess and : sample are already emitted by V1's execute_model. I have not run this; treat it as the experiment, not a result.

This CPU oracle includes an empty row, mixed query lengths and padded slots. It checks addressing only; it does not certify any CUDA cache kernel. A negative padding sentinel is filtered before indexing, because NumPy/PyTorch would otherwise interpret -1 as the last element.

import numpy as np
query_lengths = np.array([2, 0, 1, 3])
prefix_lengths = np.array([5, 0, 12, 1])
starts = np.r_[0, np.cumsum(query_lengths)]
table = np.array([[2, 3], [4, 5], [6, 7], [8, 9]])
rows = np.repeat(np.arange(4), query_lengths)
positions = np.concatenate([
    np.arange(p, p + n)
    for p, n in zip(prefix_lengths, query_lengths)
])
slots = table[rows, positions // 8] * 8 + positions % 8
np.testing.assert_array_equal(starts, [0, 2, 2, 3, 6])
np.testing.assert_array_equal(slots, [21, 22, 60, 65, 66, 67])
np.testing.assert_array_equal(prefix_lengths + query_lengths, [7, 0, 13, 4])
padded_slots = np.r_[slots, -1, -1]
valid = padded_slots >= 0
cache = np.full(80, -99)
cache[padded_slots[valid]] = np.arange(8)[valid]
assert cache[-1] == -99
assert np.count_nonzero(cache != -99) == 6

After removing or permuting requests, recompute the row map and move every request-associated state consistently. A correct prefix sum cannot compensate for penalties or RNG state attached to the wrong row.

• docs/design/model_runner_v2.md (206 lines, in-tree at a556f3f) — the rationale document. Nine numbered sections: persistent batch, async-first, removing the async barrier, StagedWriteTensor, GPU-native input prep, the Triton sampler, modularity, dummy_run abuse, explicit CUDA graphs. It opens by saying V1 had "several fundamental design mistakes". Read it before either model_runner.py.
• §11.1 — the V0-to-V1-to-V2 rewrite narrative and the repo map.
• §11.3 — SchedulerOutput's contents and the code that produces it.
• §3.4 — CommonAttentionMetadata, AttentionCGSupport, backend selection, and the nineteen MLA backends.
• §8.1 — the bucket ladder, capture, and what CUDAGraphMode.FULL vs. PIECEWISE actually change.
• §2.2 — block tables, slot mapping, and the kernel_block_size subdivision that BLOCKS_PER_KV_BLOCK implements.
• §6.6 — the speculator families that gpu/spec_decode/ packages, and where the block verification gap was first found.
• vllm/v1/outputs.py:L309-L397 — ModelRunnerOutput, the contract on the way back out, including the with_kv_conn_output_only and with_ec_conn_output constructors used on no-forward steps.