GPUModelRunner.execute_model: where a step's forward pass enters

源码版本v0.25.1

Responsibilities

GPUModelRunner(gpu_model_runner.py:440-442) is the object inside the worker process that actually runs the forward pass. Worker.execute_model feeds it a SchedulerOutput, and within a single step it does it all: _update_states syncs request state, _prepare_inputs computes logits_indices, _determine_batch_execution_and_padding picks which cudagraph mode to use, maybe_create_ubatch_slices slices micro-batches, _build_attention_metadata builds attention metadata, _preprocess assembles model inputs, and _model_forward runs the model; finally it returns either ModelRunnerOutput(last PP rank) or IntermediateTensors(intermediate PP rank).

It is composed with the worker rather than owned directly: Worker.__init__ does not create it; instead load_model goes through model_runner.load_model(gpu_worker.py:406-413). The constructor itself only stores config, computes max_num_tokens/max_num_reqs, builds the sampler, and conditionally wraps the model with UBatchWrapper(gpu_model_runner.py:5350-5379).

execute_model and sample_tokens are a pair: when async scheduling or the Ray sampler split is enabled, execute_model runs the forward without sampling, instead packing the logits and other intermediate state into ExecuteModelState(gpu_model_runner.py:424-437) for stashing and returning None, then waits for sample_tokens to be called before running _sample(gpu_model_runner.py:4456-4492). This is the key to v1's ability to run "the scheduler's next step" in parallel with "sampling".

Design motivation

• State machine split into two stages: the ExecuteModelState NamedTuple explicitly packages all intermediate values between forward and sample(logits, hidden_states, spec_decode_metadata, cudagraph_stats, slot_mappings, etc.)(gpu_model_runner.py:424-437); at the end of execute_model self.execute_model_state = ExecuteModelState(...), and entering sample_tokens unpacks it(gpu_model_runner.py:4470-4483), avoiding scattered fields.
• cudagraph mode dispatched by batch shape: _determine_batch_execution_and_padding(gpu_model_runner.py:3836-3948) computes uniform_decode, has_lora, has_encoder_output, then cudagraph_dispatcher.dispatch(num_tokens, has_lora, uniform_decode, ...) picks NONE/PIECEWISE/FULL(gpu_model_runner.py:3881-3893); the same shape with different semantics uses different graphs.
• Coordinate across DP before splitting the batch: coordinate_batch_across_dp(gpu_model_runner.py:3907-3918) decides should_ubatch and num_tokens_across_dp across DP ranks, ensuring all ranks slice the same micro-batch shapes; otherwise ubatch's cudagraph replay would misalign.
• The model can be wrapped: at the end of load_model, depending on cudagraph_mode and use_ubatching, self.model is swapped for BreakableCUDAGraphWrapper / CUDAGraphWrapper / UBatchWrapper(gpu_model_runner.py:5353-5379); _model_forward just calls self.model(...) without caring what wraps it.
• Input prep is its own method: _prepare_inputs(gpu_model_runner.py:1914) computes logits_indices and spec_decode_metadata, and _preprocess(gpu_model_runner.py:3449) assembles input_ids/positions/inputs_embeds/intermediate_tensors/model_kwargs; this is separated from forward so profile and dummy_run can reuse it.
• KV scales go eager on the first round: when calculate_kv_scales=True, cudagraph_mode is forced back to NONE(gpu_model_runner.py:4311-4314), because the dynamic operation that computes KV scale cannot enter the graph; cudagraph is only allowed after the first forward.

Key files

• GPUModelRunner.__init__:440-563 — stores config, computes max_num_tokens/max_num_reqs, builds Sampler, MULTIMODAL_REGISTRY, cudagraph_batch_sizes, cudagraph_dispatcher = CudagraphDispatcher(vllm_config).
• ExecuteModelState:424-437 — NamedTuple between forward and sample, recording logits, hidden_states, spec_decode_metadata, cudagraph_stats, etc.
• _update_states:1152 — syncs the diff of new/continued requests from scheduler_output into input_batch persistent state, returns deferred_state_corrections_fn.
• _prepare_inputs:1914 — computes logits_indices and spec_decode_metadata, decides how to map tokens into logits.
• _determine_batch_execution_and_padding:3836-3948 — computes uniform_decode/has_lora/has_encoder_output, calls cudagraph_dispatcher.dispatch to pick a CUDAGraphMode; multi-DP rank goes through coordinate_batch_across_dp.
• execute_model entry:4070-4105 — entry: clears routed_experts buffer, optional ngram_gpu copy of scheduler_output, KV transfer preemption handling, profile entry.
• execute_model batch prep:4147-4203 — num_scheduled_tokens_np, _prepare_inputs, _determine_batch_execution_and_padding, maybe_create_ubatch_slices.
• build attn metadata:4270-4295 — _get_slot_mappings + _build_attention_metadata, supports ubatch slicing(ubatch_slices=ubatch_slices_attn).
• _model_forward:4335-4359 — set_forward_context(...) wrapping self.model(input_ids=..., positions=..., intermediate_tensors=..., inputs_embeds=..., **model_kwargs), the single line that actually runs the model.
• sample_tokens entry:4456-4492 — unpacks execute_model_state, optionally apply_grammar_bitmask, calls self._sample(logits, spec_decode_metadata).
• _sample:3596-3624 — without spec decode, directly self.sampler(...); otherwise goes through rejection_sampler.
• capture_model:6647-6712 — one-shot capture of all cudagraphs at startup: large shapes first, set_cudagraph_capturing_enabled(True) + graph_capture(device=self.device), then lock_workspace() after capture.
• initialize_kv_cache:7405-7462 — builds attn_groups, metadata builders, initialize_kv_cache_tensors per KVCacheConfig, registers the KV transfer group.

Data flow

When a step comes in, execute_model first updates persistent state, then decides the batch shape and cudagraph mode, then slices ubatches, then runs forward. This is the middle of the step (preprocess):

# vllm/v1/worker/gpu_model_runner.py L4169-L4182
(
    cudagraph_mode,
    batch_desc,
    should_ubatch,
    num_tokens_across_dp,
    cudagraph_stats,
) = self._determine_batch_execution_and_padding(
    num_tokens=num_tokens_unpadded,
    num_reqs=num_reqs,
    num_scheduled_tokens_np=num_scheduled_tokens_np,
    max_num_scheduled_tokens=max_num_scheduled_tokens,
    use_cascade_attn=cascade_attn_prefix_lens is not None,
    num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
)

Then ubatch slices are cut, set_forward_context wraps the model forward, and all attention metadata, cudagraph runtime mode, and ubatch slices are obtained from the forward context:

# vllm/v1/worker/gpu_model_runner.py L4335-L4359
with (
    set_forward_context(
        attn_metadata,
        self.vllm_config,
        num_tokens=num_tokens_padded,
        num_tokens_across_dp=num_tokens_across_dp,
        cudagraph_runtime_mode=cudagraph_mode,
        batch_descriptor=batch_desc,
        ubatch_slices=ubatch_slices_padded,
        slot_mapping=slot_mappings,
        skip_compiled=has_encoder_input,
    ),
    record_function_or_nullcontext("gpu_model_runner: forward"),
    self.maybe_get_kv_connector_output(
        scheduler_output,
        defer_finalize=defer_kv_connector_finalize,
    ) as kv_connector_output,
):
    model_output = self._model_forward(
        input_ids=input_ids,
        positions=positions,
        intermediate_tensors=intermediate_tensors,
        inputs_embeds=inputs_embeds,
        **model_kwargs,
    )

_model_forward itself is the single line self.model(...)(gpu_model_runner.py:3807-3813); cudagraph replay / ubatch concurrency all happen inside the wrapper of self.model. An intermediate PP rank takes the IntermediateTensors and returns directly; the last rank stashes hidden_states in ExecuteModelState and waits for sample_tokens to sample and propose_draft_token_ids.

Boundaries and failure modes

• Empty batch early return: when num_scheduled_tokens == 0, directly return EMPTY_MODEL_RUNNER_OUTPUT(gpu_model_runner.py:4122-4138) and skip forward; with external_launcher + DP>1 it also _dummy_run(1) to sync first, avoiding cross-DP-rank divergence.
• State machine reentry detection: execute_model raises "sample_tokens() must be called after execute_model() returns None" if it sees self.execute_model_state is not None at the entry(gpu_model_runner.py:4075-4079).
• KV scales forced eager on the first round: see design motivation above(gpu_model_runner.py:4311-4314); after the first round calculate_kv_scales = False.
• Encoder-decoder first step eager: when has_encoder_input = is_encoder_decoder and num_encoder_reqs > 0, skip_compiled=True makes the forward context skip compilation(gpu_model_runner.py:4318-4321).
• ngram_gpu modifies scheduler_output: when use_ngram_gpu() is set, both num_scheduled_tokens and scheduled_spec_decode_tokens are copy()'d then replace(...)'d(gpu_model_runner.py:4087-4099), to avoid in-place mutations polluting the engine-side scheduler_output.
• cudagraph capture period is_graph_capturing: during _dummy_run(..., is_graph_capturing=True), _prepare_inputs takes the for_cudagraph_capture=True branch(gpu_model_runner.py:5936), filling max_seqlen_k with the real encoder length.
• cudagraph capture large shapes first: capture_model walks cudagraph_dispatcher.get_capture_descs() in reverse(gpu_model_runner.py:6671-6679); large graphs are captured first so small graphs can reuse the memory pool.
• shutdown clears graphs: CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs()(gpu_model_runner.py:6413-6415) clear historical graphs during re-profiling; otherwise duplicate captures consume extra memory.

Summary

GPUModelRunner is v1's core for running a step on the GPU. It strings "decide batch shape -> dispatch cudagraph -> slice ubatch -> attention metadata -> model forward -> state packaging" into one pipeline, and through the two-stage execute_model + sample_tokens call lets scheduling and sampling overlap. The model body is wrapped by CUDAGraphWrapper or UBatchWrapper, and actually running it is the single line self.model(...). For how the worker invokes it, see /worker/worker; for cudagraph and ubatch capture/replay details, see /worker/ubatch-cudagraph; for the executor layer, see /executor/executor.

See official docs: vLLM docs · README