GPUModelRunner.execute_model:一つのステップの forward はここから入る

源码版本v0.25.1

役割

GPUModelRunner(gpu_model_runner.py:440-442)は worker プロセス内で本当に forward を走らせるオブジェクトです。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(...) を呼ぶだけで外側を気にしません。
• 入力準備は単独メソッド:_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 を計算する動的操作はグラフに入れられないため、初回 forward 後にだけ 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 を算出し、トークンを 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 バッファのクリア、オプションの ngram_gpu で scheduler_output を copy、KV 転送プリエンプション処理、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 転送グループを登録。

データフロー

一ステップ入ると、execute_model はまず永続状態を update し、その後バッチ形状と cudagraph モードを決定し、ubatch を切り、最後に forward を走らせます。この部分はステップ全体の中盤(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 が追加 VRAM を占有するのを防ぎます。

まとめ

GPUModelRunner は v1 が GPU 上で一つのステップを走らせる中核です。「バッチ形状決定 → 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