GPUModelRunner.execute_model: der Forward eines Steps betritt hier die Bühne

源码版本v0.25.1

Verantwortung

GPUModelRunner(gpu_model_runner.py:440-442) ist das Objekt im Worker-Prozess, das den Forward tatsächlich ausführt. Worker.execute_model übergibt ihm den SchedulerOutput, der in diesem Schritt alles erledigt: _update_states synchronisiert den Anfragezustand, _prepare_inputs berechnet logits_indices, _determine_batch_execution_and_padding wählt die cudagraph-Variante, maybe_create_ubatch_slices teilt Mikrobatches (ubatch) auf, _build_attention_metadata baut die Attention-Metadaten, _preprocess fügt die Modelleingaben zusammen, _model_forward führt das Modell aus — und am Ende steht entweder ein ModelRunnerOutput (letzter PP-Rank) oder IntermediateTensors (mittlerer PP-Rank).

Das Verhältnis zum Worker ist eine Komposition: Worker.__init__ erzeugt den Runner nicht direkt, sondern ruft in load_model die Funktion model_runner.load_model auf(gpu_worker.py:406-413); der Konstruktor selbst speichert nur die Konfiguration, berechnet max_num_tokens/max_num_reqs, baut den Sampler und entscheidet anhand der Konfiguration, ob das Modell mit einem UBatchWrapper umhüllt wird(gpu_model_runner.py:5350-5379).

execute_model und sample_tokens bilden ein Paar: Wenn asynchrones Scheduling oder eine Ray-Sampler-Aufspaltung aktiv sind, führt execute_model zwar den Forward aus, führt aber kein Sampling durch; stattdessen verpackt es Logits und Zwischenzustände in ExecuteModelState(gpu_model_runner.py:424-437), gibt None zurück und wartet auf den Aufruf von sample_tokens, um dann _sample auszuführen(gpu_model_runner.py:4456-4492). Das ist der Hebel, mit dem v1 den nächsten Scheduling-Schritt parallel zum Sampling ausführen kann.

Entwurfsmotivation

• Zweiphasige Zustandsmaschine:ExecuteModelState (ein NamedTuple) bündelt alle Werte zwischen Forward und Sampling (logits, hidden_states, spec_decode_metadata, cudagraph_stats, slot_mappings usw.) explizit(gpu_model_runner.py:424-437); am Ende von execute_model folgt self.execute_model_state = ExecuteModelState(...), und in sample_tokens wird wieder entpackt(gpu_model_runner.py:4470-4483), um verstreute Felder zu vermeiden.
• cudagraph-Modus wird nach Batch-Form dispatcht:_determine_batch_execution_and_padding(gpu_model_runner.py:3836-3948) berechnet uniform_decode, has_lora, has_encoder_output und ruft dann cudagraph_dispatcher.dispatch(num_tokens, has_lora, uniform_decode, ...) auf, um NONE/PIECEWISE/FULL auszuwählen(gpu_model_runner.py:3881-3893) — gleiche Form, unterschiedliche Semantik nimmt unterschiedliche Graphen.
• DP-Abstimmung vor dem Aufteilen:coordinate_batch_across_dp(gpu_model_runner.py:3907-3918) stimmt should_ubatch und num_tokens_across_dp über alle DP-Ranks ab, damit alle Ranks identisch geformte Mikrobatches erhalten — sonst geraten die cudagraph-Replays des ubatch aus dem Tritt.
• Modell kann gewrappt werden:Am Ende von load_model wird self.model abhängig von cudagraph_mode und use_ubatching durch BreakableCUDAGraphWrapper / CUDAGraphWrapper / UBatchWrapper ersetzt(gpu_model_runner.py:5353-5379); _model_forward ruft nur self.model(...) auf und macht sich die äußere Hülle nicht zunutze.
• Input-Vorbereitung als eigene Methode:_prepare_inputs(gpu_model_runner.py:1914) berechnet logits_indices und spec_decode_metadata; _preprocess(gpu_model_runner.py:3449) baut input_ids/positions/inputs_embeds/intermediate_tensors/model_kwargs zusammen — getrennt vom Forward, sodass auch profile und dummy_run es wiederverwenden können.
• KV-Scales in erster Runde eager:Wenn calculate_kv_scales=True, wird cudagraph_mode hart auf NONE zurückgesetzt(gpu_model_runner.py:4311-4314), weil die dynamischen Operationen zur Berechnung der KV-Scale nicht in den Graphen passen; erst nach der ersten Forward-Runde ist cudagraph wieder erlaubt.

Schlüsseldateien

• GPUModelRunner.__init__:440-563 — speichert Konfiguration, berechnet max_num_tokens/max_num_reqs, baut Sampler, MULTIMODAL_REGISTRY, cudagraph_batch_sizes, cudagraph_dispatcher = CudagraphDispatcher(vllm_config) auf.
• ExecuteModelState:424-437 — NamedTuple zwischen Forward und Sampling, erfasst logits, hidden_states, spec_decode_metadata, cudagraph_stats usw.
• _update_states:1152 — synchronisiert Unterschiede aus scheduler_output (neue/fortgesetzte Anfragen) in den persistenten Zustand des input_batch und liefert deferred_state_corrections_fn zurück.
• _prepare_inputs:1914 — berechnet logits_indices und spec_decode_metadata und entscheidet, wie die Token in die logits einfließen.
• _determine_batch_execution_and_padding:3836-3948 — berechnet uniform_decode/has_lora/has_encoder_output, ruft cudagraph_dispatcher.dispatch auf, um den CUDAGraphMode zu wählen, und führt bei mehreren DP-Ranks coordinate_batch_across_dp aus.
• execute_model entry:4070-4105 — Entry-Point: leert routed_experts-Buffer, optionaler ngram_gpu-Copy des scheduler_output, Vorab-Behandlung des KV-Transfers, Profile-Einstieg.
• 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, unterstützt ubatch-Slices (ubatch_slices=ubatch_slices_attn).
• _model_forward:4335-4359 — umschließt self.model(input_ids=..., positions=..., intermediate_tensors=..., inputs_embeds=..., **model_kwargs) mit set_forward_context(...) — die eine Zeile, in der das Modell wirklich läuft.
• sample_tokens entry:4456-4492 — entpackt execute_model_state, wendet optional apply_grammar_bitmask an und ruft self._sample(logits, spec_decode_metadata) auf.
• _sample:3596-3624 — ohne spec decode direkt self.sampler(...), andernfalls über rejection_sampler.
• capture_model:6647-6712 — nimmt beim Start einmal alle cudagraph-Graphen auf: große Formen zuerst, set_cudagraph_capturing_enabled(True) + graph_capture(device=self.device), danach lock_workspace().
• initialize_kv_cache:7405-7462 — baut anhand von KVCacheConfig die attn_groups, Metadata-Builder und initialize_kv_cache_tensors auf und registriert die KV-Transfer-Gruppe.

Datenfluss

Beim Eintreffen eines Steps aktualisiert execute_model zuerst den persistenten Zustand, entscheidet dann über Batch-Form und cudagraph-Modus, teilt dann den ubatch auf und führt schließlich den Forward aus. Dieser Abschnitt ist die Mitte des gesamten Steps (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),
)

Anschließend werden die ubatch-Slices aufgeteilt; set_forward_context umschließt den Modell-Forward, und alle Attention-Metadaten, der cudagraph-Runtime-Modus und die ubatch-Slices stammen aus dem Forward-Kontext:

# 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 besteht aus der einzigen Zeile self.model(...)(gpu_model_runner.py:3807-3813); das cudagraph-Replay bzw. die ubatch-Parallelität spielen sich im Wrapper von self.model ab. Ein mittlerer PP-Rank liefert die empfangenen IntermediateTensors direkt zurück; der letzte Rank parkt die hidden_states in ExecuteModelState und wartet auf sample_tokens, wo das Sampling und propose_draft_token_ids stattfinden.

Grenzen und Fehler

• Leerer Batch: früher Return:Wenn num_scheduled_tokens == 0, wird direkt return EMPTY_MODEL_RUNNER_OUTPUT ausgeführt(gpu_model_runner.py:4122-4138) — ohne Forward; bei external_launcher + DP>1 erfolgt zuvor ein _dummy_run(1) zur Synchronisation, um ein Aufeinanderzulaufen der DP-Ranks zu vermeiden.
• Re-Entry-Prüfung der Zustandsmaschine:Wenn execute_model beim Eintritt feststellt, dass self.execute_model_state is not None, wird sofort „sample_tokens() must be called after execute_model() returns None“ geworfen(gpu_model_runner.py:4075-4079).
• KV-Scales in erster Runde erzwungen eager:Siehe Entwurfsmotivation oben(gpu_model_runner.py:4311-4314); nach der ersten Runde gilt calculate_kv_scales = False.
• Encoder-Decoder: erster Schritt eager:has_encoder_input = is_encoder_decoder and num_encoder_reqs > 0 setzt skip_compiled=True, sodass der Forward-Kontext die Kompilierung überspringt(gpu_model_runner.py:4318-4321).
• ngram_gpu verändert scheduler_output:Wenn use_ngram_gpu() aktiv ist, werden sowohl num_scheduled_tokens als auch scheduled_spec_decode_tokens zunächst copy() und dann replace(...)(gpu_model_runner.py:4087-4099) — um eine in-place-Modifikation des scheduler_output auf Engine-Seite zu vermeiden.
• Während der cudagraph-Aufnahme gilt is_graph_capturing:_dummy_run(..., is_graph_capturing=True) aktiviert in _prepare_inputs den Zweig for_cudagraph_capture=True(gpu_model_runner.py:5936) und befüllt max_seqlen_k mit der echten Encoder-Länge.
• cudagraph: große Formen zuerst:capture_model iteriert über cudagraph_dispatcher.get_capture_descs() in absteigender Reihenfolge(gpu_model_runner.py:6671-6679), sodass große Graphen zuerst aufgenommen werden und kleine Graphen den Speicher-Pool wiederverwenden können.
• Shutdown löscht Graphen:CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs()(gpu_model_runner.py:6413-6415) entfernen beim erneuten Profiling historische Graphen, da sonst wiederholte Captures zusätzlichen VRAM beanspruchen würden.

Zusammenfassung

GPUModelRunner ist der Kern, mit dem v1 einen einzelnen Schritt auf der GPU ausführt. Er verknüpft die Kette „Batch-Form festlegen → cudagraph-Dispatch → ubatch-Aufteilung → Attention-Metadaten → Modell-Forward → Zustand verpacken“ zu einer Pipeline und macht über die Zweiteilung execute_model + sample_tokens Scheduling und Sampling versetzt zueinander. Das Modell selbst ist von CUDAGraphWrapper oder UBatchWrapper umhüllt; der eigentliche Lauf besteht aus der einen Zeile self.model(...). Wie der Worker ihn aufruft, steht in /worker/worker; die Aufnahme- und Replay-Details von cudagraph und ubatch in /worker/ubatch-cudagraph; die Executor-Ebene in /executor/executor.

Siehe offizielle Dokumentation: vLLM 文档 · README.