GPUModelRunner.execute_model : point d'entrée du forward d'un step

源码版本v0.25.1

Responsabilités

GPUModelRunner(gpu_model_runner.py:440-442) est l'objet, dans le worker, qui exécute réellement le forward. Worker.execute_model lui passe le SchedulerOutput, et en un step il fait tout : _update_states synchronise l'état des requêtes, _prepare_inputs calcule les logits_indices, _determine_batch_execution_and_padding décide quel cudagraph utiliser, maybe_create_ubatch_slices découpe les micro-batchs, _build_attention_metadata construit les métadonnées d'attention, _preprocess assemble les entrées du modèle, _model_forward lance le modèle, et au final on renvoie soit ModelRunnerOutput (rang PP final), soit IntermediateTensors (rang PP intermédiaire).

La relation avec le worker est une composition : Worker.__init__ ne crée pas directement le runner ; il passe par model_runner.load_model dans load_model(gpu_worker.py:406-413). Le constructeur se contente de stocker la configuration, calculer max_num_tokens/max_num_reqs, créer le sampler, et décider selon la config s'il faut envelopper le modèle d'un UBatchWrapper(gpu_model_runner.py:5350-5379).

execute_model et sample_tokens forment un binôme : quand l'async scheduling ou le split du sampler Ray est activé, execute_model termine le forward sans échantillonner, mais stocke les états intermédiaires (logits, etc.) dans ExecuteModelState(gpu_model_runner.py:424-437), renvoie None, puis attend l'appel de sample_tokens pour _sample(gpu_model_runner.py:4456-4492). C'est la clé qui permet à v1 de faire chevaucher « le prochain step de l'ordonnanceur » et « l'échantillonnage ».

Motivation de conception

• Machine à états en deux temps : le NamedTuple ExecuteModelState emballe explicitement tous les états intermédiaires entre forward et sample (logits, hidden_states, spec_decode_metadata, cudagraph_stats, slot_mappings, etc.)(gpu_model_runner.py:424-437). À la fin de execute_model, on fait self.execute_model_state = ExecuteModelState(...), et sample_tokens le dépaquette(gpu_model_runner.py:4470-4483) — pas de champs épars.
• Dispatch cudagraph selon la forme du batch : _determine_batch_execution_and_padding(gpu_model_runner.py:3836-3948) calcule uniform_decode, has_lora, has_encoder_output, puis cudagraph_dispatcher.dispatch(num_tokens, has_lora, uniform_decode, ...) choisit NONE/PIECEWISE/FULL(gpu_model_runner.py:3881-3893) : une même forme peut emprunter des graphes différents selon la sémantique.
• Coordination DP avant découpage : coordinate_batch_across_dp(gpu_model_runner.py:3907-3918) décide should_ubatch et num_tokens_across_dp entre DP ranks, pour garantir que tous découpent les micro-batchs de la même forme — sinon le replay cudagraph de l'ubatch se désalignerait.
• Modèle enveloppable : à la fin de load_model, selon cudagraph_mode et use_ubatching, self.model est remplacé par BreakableCUDAGraphWrapper / CUDAGraphWrapper / UBatchWrapper(gpu_model_runner.py:5353-5379) ; _model_forward ne fait qu'appeler self.model(...), indépendamment de l'enveloppe.
• Préparation des entrées isolée : _prepare_inputs(gpu_model_runner.py:1914) calcule logits_indices et spec_decode_metadata ; _preprocess(gpu_model_runner.py:3449) assemble input_ids/positions/inputs_embeds/intermediate_tensors/model_kwargs, séparé du forward pour que profile et dummy_run puissent le réutiliser.
• KV scales en eager au premier tour : quand calculate_kv_scales=True, on force cudagraph_mode à NONE(gpu_model_runner.py:4311-4314) car l'opération dynamique de calcul des scales KV ne peut pas entrer dans le graphe ; ce n'est qu'après le premier forward qu'on autorise cudagraph.

Fichiers clés

• GPUModelRunner.__init__:440-563 — stocke la config, calcule max_num_tokens/max_num_reqs, crée le Sampler, MULTIMODAL_REGISTRY, cudagraph_batch_sizes, cudagraph_dispatcher = CudagraphDispatcher(vllm_config).
• ExecuteModelState:424-437 — NamedTuple entre forward et sample, enregistre logits, hidden_states, spec_decode_metadata, cudagraph_stats, etc.
• _update_states:1152 — synchronise les différences entre requêtes nouvelles/reprises du scheduler_output vers l'état persistant de input_batch, renvoie deferred_state_corrections_fn.
• _prepare_inputs:1914 — calcule logits_indices et spec_decode_metadata, décide comment coller les tokens aux logits.
• _determine_batch_execution_and_padding:3836-3948 — calcule uniform_decode/has_lora/has_encoder_output, appelle cudagraph_dispatcher.dispatch pour choisir CUDAGraphMode, multi-DP rank via coordinate_batch_across_dp.
• execute_model entry:4070-4105 — entrée : nettoie les buffers routed_experts, optionnellement copie le scheduler_output pour ngram_gpu, gère la préemption KV transfer, entrée profiler.
• 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, supporte le découpage ubatch (ubatch_slices=ubatch_slices_attn).
• _model_forward:4335-4359 — set_forward_context(...) enveloppe self.model(input_ids=..., positions=..., intermediate_tensors=..., inputs_embeds=..., **model_kwargs) — la vraie ligne qui lance le modèle.
• sample_tokens entry:4456-4492 — dépaquette execute_model_state, optionnellement apply_grammar_bitmask, appelle self._sample(logits, spec_decode_metadata).
• _sample:3596-3624 — sans spec decode, appelle directement self.sampler(...), sinon passe par rejection_sampler.
• capture_model:6647-6712 — capture en une fois tous les cudagraphs au démarrage : grands shapes en premier, set_cudagraph_capturing_enabled(True) + graph_capture(device=self.device), puis lock_workspace().
• initialize_kv_cache:7405-7462 — selon KVCacheConfig, crée attn_groups, les metadata builders, initialize_kv_cache_tensors, enregistre le KV transfer group.

Flux de données

À l'arrivée d'un step, execute_model met d'abord à jour l'état persistant, détermine la forme du batch et le mode cudagraph, découpe l'ubatch, puis lance le forward. Voici le tronçon central (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),
)

Ensuite, on découpe les ubatch slices, set_forward_context enveloppe le forward modèle, et toutes les métadonnées d'attention, le runtime mode cudagraph, les ubatch slices se récupèrent depuis le 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,
    )

Le corps de _model_forward se résume à une ligne self.model(...)(gpu_model_runner.py:3807-3813) ; le replay cudagraph / la concurrence ubatch se passent dans le wrapper de self.model. Le rang PP intermédiaire récupère des IntermediateTensors qu'il renvoie directement ; le rang final stocke hidden_states dans ExecuteModelState, en attendant que sample_tokens échantillonne et lance propose_draft_token_ids.

Limites et échecs

• Retour anticipé sur batch vide : quand num_scheduled_tokens == 0, on return EMPTY_MODEL_RUNNER_OUTPUT(gpu_model_runner.py:4122-4138) sans lancer de forward ; en external_launcher + DP>1, on fait d'abord un _dummy_run(1) pour synchroniser, afin d'éviter une divergence entre DP ranks.
• Détection de réentrance de la machine à états : à l'entrée de execute_model, si self.execute_model_state is not None, on lève « sample_tokens() must be called after execute_model() returns None »(gpu_model_runner.py:4075-4079).
• KV scales forcés en eager au premier tour : voir motivation ci-dessus(gpu_model_runner.py:4311-4314) ; après le premier tour, calculate_kv_scales = False.
• Encoder-decoder en eager au premier step : quand has_encoder_input = is_encoder_decoder and num_encoder_reqs > 0, skip_compiled=True fait que le forward context saute la compilation(gpu_model_runner.py:4318-4321).
• ngram_gpu modifie scheduler_output : quand use_ngram_gpu(), on copy() num_scheduled_tokens et scheduled_spec_decode_tokens puis replace(...)(gpu_model_runner.py:4087-4099), pour éviter qu'une modification in-place ne pollue le scheduler_output côté engine.
• is_graph_capturing pendant la capture cudagraph : _dummy_run(..., is_graph_capturing=True) fait que _prepare_inputs emprunte la branche for_cudagraph_capture=True(gpu_model_runner.py:5936), en remplissant max_seqlen_k avec la vraie longueur d'encoder.
• Capture cudagraph grands shapes en premier : capture_model parcourt en ordre inverse cudagraph_dispatcher.get_capture_descs()(gpu_model_runner.py:6671-6679), les grands graphes d'abord pour que les petits réutilisent le pool mémoire.
• Shutdown nettoie les graphes : CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs()(gpu_model_runner.py:6413-6415) nettoient les graphes historiques lors d'un re-profile, sinon une capture répétée consommerait de la VRAM supplémentaire.

Résumé

GPUModelRunner est le cœur d'un step sur GPU en v1. Il enchaîne « détermination de la forme de batch → dispatch cudagraph → découpe ubatch → métadonnées d'attention → forward modèle → emballage d'état » en un pipeline, puis via execute_model + sample_tokens décale l'ordonnancement du sampling. Le modèle proprement dit est enveloppé dans CUDAGraphWrapper ou UBatchWrapper, et le forward se résume à self.model(...). Pour la manière dont le worker l'appelle, voir /worker/worker ; pour les détails de capture/replay de cudagraph et ubatch, voir /worker/ubatch-cudagraph ; pour la couche exécuteur, voir /executor/executor.

Voir la documentation officielle : Documentation vLLM · README