GPUModelRunner.execute_model:一个 step 的前向从这进

源码版本v0.25.1

职责

GPUModelRunner(gpu_model_runner.py:440-442)是 worker 进程里真正跑前向的对象。Worker.execute_model 把 SchedulerOutput 喂给它,它在这一步里干完:_update_states 同步请求状态、_prepare_inputs 算 logits_indices、_determine_batch_execution_and_padding 决定走哪种 cudagraph、maybe_create_ubatch_slices 切微批、_build_attention_metadata 建 attention 元数据、_preprocess 拼模型输入、_model_forward 跑模型,最后要么返回 ModelRunnerOutput(末级 PP rank)、要么返回 IntermediateTensors(中间 PP rank)。

它跟 worker 是组合关系:Worker.__init__ 不直接创建,而是在 load_model 里走 model_runner.load_model(gpu_worker.py:406-413);构造函数本身只是存配置、算 max_num_tokens/max_num_reqs、建 sampler、根据配置决定要不要 UBatchWrapper 包模型(gpu_model_runner.py:5350-5379)。

execute_model 和 sample_tokens 是一对:开了 async scheduling 或 Ray sampler 拆分时,execute_model 跑完 forward 不采样,而是把 logits 等中间状态塞进 ExecuteModelState(gpu_model_runner.py:424-437)暂存,返回 None,等 sample_tokens 被调到再 _sample(gpu_model_runner.py:4456-4492)。这是 v1 让「调度的下一步」可以跟「采样」并行起来的关键。

设计动机

• 状态机分两段:ExecuteModelState 这个 NamedTuple 把 forward 和 sample 之间的所有中间值(logits、hidden_states、spec_decode_metadata、cudagraph_stats、slot_mappings 等)显式打包(gpu_model_runner.py:424-437),execute_model 末尾 self.execute_model_state = ExecuteModelState(...),进 sample_tokens 再解包(gpu_model_runner.py:4470-4483),避免散落字段。
• cudagraph 模式按批形状派发:_determine_batch_execution_and_padding(gpu_model_runner.py:3836-3948)算 uniform_decode、has_lora、has_encoder_output,再 cudagraph_dispatcher.dispatch(num_tokens, has_lora, uniform_decode, ...) 选 NONE/PIECEWISE/FULL(gpu_model_runner.py:3881-3893),同形状不同语义走不同图。
• DP 间先协调再切批:coordinate_batch_across_dp(gpu_model_runner.py:3907-3918)跨 DP rank 决定 should_ubatch 和 num_tokens_across_dp,保证所有 rank 切成相同形状的微批,否则 ubatch 的 cudagraph replay 会错位。
• 模型可以被 wrapper 包:load_model 末尾根据 cudagraph_mode 和 use_ubatching 把 self.model 换成 BreakableCUDAGraphWrapper / CUDAGraphWrapper / UBatchWrapper(gpu_model_runner.py:5353-5379),_model_forward 只调 self.model(...),不关心外面包了什么。
• input prep 单独成方法:_prepare_inputs(gpu_model_runner.py:1914)算 logits_indices 和 spec_decode_metadata,_preprocess(gpu_model_runner.py:3449)拼 input_ids/positions/inputs_embeds/intermediate_tensors/model_kwargs,跟 forward 分开,profile 和 dummy_run 也能复用。
• KV scales 首轮 eager:calculate_kv_scales=True 时强行把 cudagraph_mode 改回 NONE(gpu_model_runner.py:4311-4314),因为算 KV scale 的动态操作没法进图,第一轮前向后才允许走 cudagraph。

关键文件

• GPUModelRunner.__init__:440-563 — 存配置、算 max_num_tokens/max_num_reqs、建 Sampler、MULTIMODAL_REGISTRY、cudagraph_batch_sizes、cudagraph_dispatcher = CudagraphDispatcher(vllm_config)。
• ExecuteModelState:424-437 — forward 和 sample 之间的 NamedTuple,记录 logits、hidden_states、spec_decode_metadata、cudagraph_stats 等。
• _update_states:1152 — 把 scheduler_output 里新请求/继续请求的差异同步到 input_batch 持久状态,返回 deferred_state_corrections_fn。
• _prepare_inputs:1914 — 算 logits_indices 和 spec_decode_metadata,决定怎么把 token 拼到 logits。
• _determine_batch_execution_and_padding:3836-3948 — 算 uniform_decode/has_lora/has_encoder_output,调 cudagraph_dispatcher.dispatch 选 CUDAGraphMode,DP 多 rank 走 coordinate_batch_across_dp。
• execute_model entry:4070-4105 — 入口:清 routed_experts buffer、可选 ngram_gpu copy scheduler_output、KV transfer 抢占处理、profile 入口。
• 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,支持 ubatch 切片(ubatch_slices=ubatch_slices_attn)。
• _model_forward:4335-4359 — set_forward_context(...) 包住 self.model(input_ids=..., positions=..., intermediate_tensors=..., inputs_embeds=..., **model_kwargs),真正跑模型的一行。
• sample_tokens entry:4456-4492 — 解包 execute_model_state,可选 apply_grammar_bitmask,调 self._sample(logits, spec_decode_metadata)。
• _sample:3596-3624 — 没 spec decode 直接 self.sampler(...),否则走 rejection_sampler。
• capture_model:6647-6712 — 启动期一次性抓所有 cudagraph:大形状先抓、set_cudagraph_capturing_enabled(True) + graph_capture(device=self.device)、抓完 lock_workspace()。
• initialize_kv_cache:7405-7462 — 按 KVCacheConfig 建 attn_groups、metadata builders、initialize_kv_cache_tensors,注册 KV transfer group。

数据流

一个 step 进来,execute_model 先 update 持久状态,然后决定批形状和 cudagraph 模式,再切 ubatch,最后跑 forward。这段是整个 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),
)

接着切 ubatch slices,set_forward_context 包住模型 forward,所有 attention metadata、cudagraph runtime mode、ubatch slices 都从 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 本体就一行 self.model(...)(gpu_model_runner.py:3807-3813),cudagraph replay / ubatch 并发都在 self.model 的 wrapper 里发生。PP 中间 rank 拿到 IntermediateTensors 直接返回,末级 rank 把 hidden_states 暂存到 ExecuteModelState,等 sample_tokens 来采样并 propose_draft_token_ids。

边界与失败

• 空批早返回:num_scheduled_tokens == 0 时直接 return EMPTY_MODEL_RUNNER_OUTPUT(gpu_model_runner.py:4122-4138),不开 forward;external_launcher + DP>1 时还会先 _dummy_run(1) 同步,避免跨 DP rank 失配。
• 状态机重入检测:execute_model 入口看到 self.execute_model_state is not None 直接抛「sample_tokens() must be called after execute_model() returns None」(gpu_model_runner.py:4075-4079)。
• KV scales 首轮强制 eager:见上面设计动机(gpu_model_runner.py:4311-4314),第一轮后 calculate_kv_scales = False。
• encoder-decoder 首步 eager:has_encoder_input = is_encoder_decoder and num_encoder_reqs > 0 时,skip_compiled=True 让 forward context 跳编译(gpu_model_runner.py:4318-4321)。
• ngram_gpu 改 scheduler_output:use_ngram_gpu() 时把 num_scheduled_tokens 和 scheduled_spec_decode_tokens 都 copy() 再 replace(...)(gpu_model_runner.py:4087-4099),避免 in-place 改动污染 engine 那侧的 scheduler_output。
• cudagraph 捕获期 is_graph_capturing:_dummy_run(..., is_graph_capturing=True) 时 _prepare_inputs 走 for_cudagraph_capture=True 分支(gpu_model_runner.py:5936),用真实 encoder 长度填 max_seqlen_k。
• cudagraph capture 大形状先抓:capture_model 按 cudagraph_dispatcher.get_capture_descs() 倒序遍历(gpu_model_runner.py:6671-6679),大图先抓小图复用内存池。
• shutdown 清图:CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs()(gpu_model_runner.py:6413-6415)在重新 profile 时清掉历史图,否则重复 capture 会占额外显存。

小结

GPUModelRunner 是 v1 在 GPU 上跑一个 step 的核心。它把「批形状决定 → cudagraph 派发 → ubatch 切分 → attention metadata → 模型 forward → 状态打包」这一串串成一条流水,然后通过 execute_model + sample_tokens 两段调用让调度和采样可以错开。模型本体被 CUDAGraphWrapper 或 UBatchWrapper 包住,真正跑就是 self.model(...) 一行。worker 怎么调它见 /worker/worker,cudagraph 和 ubatch 的捕获/重放细节见 /worker/ubatch-cudagraph,执行器层见 /executor/executor。

对照官方资料:vLLM 文档 · README。