vLLM 架构拆解

繁體中文

简体中文
English
日本語
Español
Français
Deutsch

繁體中文

简体中文
English
日本語
Español
Français
Deutsch

Appearance

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

python
# 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 拿到:

python
# 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