GPUModelRunner.execute_model: el forward de un step entra aquí

源码版本v0.25.1

Responsabilidades

GPUModelRunner (gpu_model_runner.py:440-442) es el objeto del worker que de verdad ejecuta el forward. Worker.execute_model le pasa el SchedulerOutput y en un mismo paso hace todo: _update_states sincroniza el estado de las peticiones, _prepare_inputs calcula logits_indices, _determine_batch_execution_and_padding decide qué cudagraph usar, maybe_create_ubatch_slices corta los micro-batches, _build_attention_metadata construye los metadatos de attention, _preprocess ensambla las entradas del modelo, _model_forward ejecuta el modelo, y al final devuelve ModelRunnerOutput (último rank PP) o IntermediateTensors (rank intermedio PP).

La relación con el worker es de composición: Worker.__init__ no lo crea directamente, sino que lo hace en load_model vía model_runner.load_model (gpu_worker.py:406-413); el constructor se limita a guardar configuración, calcular max_num_tokens/max_num_reqs, construir el sampler y, según la configuración, decidir si envolver el modelo con un UBatchWrapper (gpu_model_runner.py:5350-5379).

execute_model y sample_tokens forman un par: con async scheduling o con el sampler desdoblado en Ray, execute_model corre el forward sin muestrear; mete los logits y demás estados intermedios en un ExecuteModelState (gpu_model_runner.py:424-437) en espera, devuelve None y, cuando se invoca sample_tokens, se ejecuta _sample (gpu_model_runner.py:4456-4492). Esta es la clave de v1 para que el "siguiente paso del scheduler" pueda correr en paralelo con el "muestreo".

Motivación de diseño

• State machine en dos tramos: el NamedTuple ExecuteModelState empaqueta de forma explícita todos los valores intermedios entre forward y sample (logits, hidden_states, spec_decode_metadata, cudagraph_stats, slot_mappings, etc.) (gpu_model_runner.py:424-437); al final de execute_model se asigna self.execute_model_state = ExecuteModelState(...) y sample_tokens lo desempaqueta (gpu_model_runner.py:4470-4483), evitando campos dispersos.
• Modo cudagraph despachado por la forma del batch: _determine_batch_execution_and_padding (gpu_model_runner.py:3836-3948) calcula uniform_decode, has_lora, has_encoder_output y luego cudagraph_dispatcher.dispatch(num_tokens, has_lora, uniform_decode, ...) elige NONE/PIECEWISE/FULL (gpu_model_runner.py:3881-3893), de modo que una misma forma con semántica distinta usa un grafo distinto.
• Coordinación entre DP antes de cortar el batch: coordinate_batch_across_dp (gpu_model_runner.py:3907-3918) decide entre todos los ranks DP el should_ubatch y el num_tokens_across_dp, para que todos los ranks corten micro-batches con la misma forma; de lo contrario el replay del cudagraph de un ubatch quedaría desalineado.
• El modelo se puede envolver: al final de load_model, según cudagraph_mode y use_ubatching, self.model se reemplaza por BreakableCUDAGraphWrapper / CUDAGraphWrapper / UBatchWrapper (gpu_model_runner.py:5353-5379); _model_forward se limita a self.model(...) sin importar el envoltorio exterior.
• Preparación de entradas aislada en su método: _prepare_inputs (gpu_model_runner.py:1914) calcula logits_indices y spec_decode_metadata; _preprocess (gpu_model_runner.py:3449) ensambla input_ids/positions/inputs_embeds/intermediate_tensors/model_kwargs, separado del forward, de modo que profile y dummy_run también pueden reutilizarlo.
• KV scales en la primera ronda en modo eager: con calculate_kv_scales=True se fuerza cudagraph_mode = NONE (gpu_model_runner.py:4311-4314), ya que la operación dinámica para calcular la KV scale no puede entrar en el grafo; solo tras el primer forward se permite usar cudagraph.

Archivos clave

• GPUModelRunner.__init__:440-563 — guarda configuración, calcula max_num_tokens/max_num_reqs, construye Sampler, MULTIMODAL_REGISTRY, cudagraph_batch_sizes, cudagraph_dispatcher = CudagraphDispatcher(vllm_config).
• ExecuteModelState:424-437 — NamedTuple entre forward y sample; registra logits, hidden_states, spec_decode_metadata, cudagraph_stats, etc.
• _update_states:1152 — vuelca las diferencias de peticiones nuevas/continuadas del scheduler_output al estado persistente del input_batch; devuelve deferred_state_corrections_fn.
• _prepare_inputs:1914 — calcula logits_indices y spec_decode_metadata y decide cómo concatenar los tokens a los logits.
• _determine_batch_execution_and_padding:3836-3948 — calcula uniform_decode/has_lora/has_encoder_output, llama a cudagraph_dispatcher.dispatch para elegir CUDAGraphMode, y con DP multi-rank pasa por coordinate_batch_across_dp.
• execute_model entry:4070-4105 — entrada: limpia el buffer de routed_experts, opcionalmente copia el scheduler_output en ngram_gpu, gestiona la preempt de KV transfer, entrada de 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, soporta cortes ubatch (ubatch_slices=ubatch_slices_attn).
• _model_forward:4335-4359 — envuelve set_forward_context(...) alrededor de self.model(input_ids=..., positions=..., intermediate_tensors=..., inputs_embeds=..., **model_kwargs), la línea que de verdad ejecuta el modelo.
• sample_tokens entry:4456-4492 — desempaqueta execute_model_state, opcionalmente apply_grammar_bitmask, invoca self._sample(logits, spec_decode_metadata).
• _sample:3596-3624 — sin spec decode llama directamente self.sampler(...); en caso contrario pasa por rejection_sampler.
• capture_model:6647-6712 — captura de una sola vez todos los cudagraphs en el arranque: primero captura las formas grandes, set_cudagraph_capturing_enabled(True) + graph_capture(device=self.device), y al terminar lock_workspace().
• initialize_kv_cache:7405-7462 — según el KVCacheConfig construye los attn_groups, los metadata builders, llama a initialize_kv_cache_tensors y registra el KV transfer group.

Flujo de datos

Cuando entra un step, execute_model primero actualiza el estado persistente, luego decide la forma del batch y el modo cudagraph, parte el ubatch y por último ejecuta el forward. Esta porción es la parte central (preprocess) del step completo:

# 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),
)

Luego corta los ubatch slices, set_forward_context envuelve el forward del modelo y todos los attention metadata, el runtime mode del cudagraph y los ubatch slices se obtienen del 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,
    )

El cuerpo de _model_forward es una sola línea self.model(...) (gpu_model_runner.py:3807-3813); el replay del cudagraph y la concurrencia del ubatch suceden dentro del wrapper de self.model. El rank intermedio PP recibe IntermediateTensors y los devuelve; el último rank guarda hidden_states en ExecuteModelState y espera a que sample_tokens los muestree y a propose_draft_token_ids.

Límites y fallos

• Retorno temprano en batch vacío: con num_scheduled_tokens == 0 devuelve directamente EMPTY_MODEL_RUNNER_OUTPUT (gpu_model_runner.py:4122-4138) sin tocar el forward; con external_launcher + DP>1 además hace primero un _dummy_run(1) para sincronizar y evitar desincronización entre ranks DP.
• Detección de reentrada en la state machine: si a la entrada de execute_model se encuentra self.execute_model_state is not None, lanza directamente "sample_tokens() must be called after execute_model() returns None" (gpu_model_runner.py:4075-4079).
• KV scales forzadas a eager en la primera ronda: véase la motivación de diseño anterior (gpu_model_runner.py:4311-4314); tras la primera ronda calculate_kv_scales = False.
• Encoder-decoder en el primer paso en eager: cuando has_encoder_input = is_encoder_decoder and num_encoder_reqs > 0, se activa skip_compiled=True para que el forward context se salte la compilación (gpu_model_runner.py:4318-4321).
• ngram_gpu modifica scheduler_output: con use_ngram_gpu() se hace copy() tanto de num_scheduled_tokens como de scheduled_spec_decode_tokens y luego replace(...) (gpu_model_runner.py:4087-4099), para no mutar in-place el scheduler_output del lado del engine.
• is_graph_capturing durante la captura de cudagraph: en _dummy_run(..., is_graph_capturing=True), _prepare_inputs toma la rama for_cudagraph_capture=True (gpu_model_runner.py:5936) y rellena max_seqlen_k con la longitud real del encoder.
• La captura de cudagraph graba primero las formas grandes: capture_model recorre en orden inverso cudagraph_dispatcher.get_capture_descs() (gpu_model_runner.py:6671-6679), de modo que los grafos grandes se capturan primero y los pequeños reutilizan el pool de memoria.
• Limpiar grafos en shutdown: CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() (gpu_model_runner.py:6413-6415) limpian los grafos históricos al rehacer profile; si no, capturar de nuevo ocuparía memoria extra.

Resumen

GPUModelRunner es el núcleo de v1 para ejecutar un step en GPU. Encadena "decisión de forma del batch → despacho de cudagraph → partición ubatch → attention metadata → forward del modelo → empaquetado de estado" en una sola tubería, y con las dos llamadas execute_model + sample_tokens desacopla scheduler y sampling. El modelo mismo está envuelto por CUDAGraphWrapper o UBatchWrapper; ejecutar de verdad es una sola línea self.model(...). Para cómo lo invoca el worker, véase /worker/worker; los detalles de captura/replay de cudagraph y ubatch, en /worker/ubatch-cudagraph; la capa del executor, en /executor/executor.

Véase la documentación oficial: vLLM 文档 · README