59db6a4285
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Signed-off-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
669 lines
29 KiB
Python
669 lines
29 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Omni v2 GPU model runner hooks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
|
|
import torch
|
|
from vllm.config.compilation import CUDAGraphMode
|
|
from vllm.forward_context import (
|
|
set_forward_context,
|
|
)
|
|
from vllm.logger import init_logger
|
|
from vllm.v1.core.sched.output import SchedulerOutput
|
|
from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
|
|
from vllm.v1.worker.gpu.input_batch import set_dummy_context
|
|
from vllm.v1.worker.gpu.lora_utils import get_num_active_loras_for_dispatch
|
|
from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
|
|
from vllm.v1.worker.gpu.model_runner import (
|
|
BatchDescriptor,
|
|
BatchExecutionDescriptor,
|
|
ExecuteModelState,
|
|
GPUModelRunner,
|
|
IntermediateTensors,
|
|
build_slot_mappings_by_layer,
|
|
)
|
|
from vllm.v1.worker.gpu.sample.sampler import Sampler
|
|
|
|
from vllm_omni.core.sched.omni_scheduling_coordinator import (
|
|
uses_native_mrv2_data_plane,
|
|
)
|
|
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
|
from vllm_omni.worker.sampling_utils import sanitize_sampling_params_min_tokens_stop_ids
|
|
from vllm_omni.worker_v2.model_states import init_omni_model_state
|
|
from vllm_omni.worker_v2.model_states.intermediate_buffer import (
|
|
_resolve_additional_information,
|
|
)
|
|
from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
|
|
from vllm_omni.worker_v2.omni_data_plane import OmniRunnerDataPlane
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
|
|
|
|
from vllm_omni.worker_v2.native_output_worker import NativeOutputWorker
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_model_state_patch_lock = threading.RLock()
|
|
|
|
|
|
def _needs_capture_tensor_unwrap(model: Any) -> bool:
|
|
return bool(getattr(model, "_returns_tuple", False))
|
|
|
|
|
|
class OmniGPUModelRunner(GPUModelRunner):
|
|
"""Thin layer over v2 ``GPUModelRunner`` for Omni lifecycle hooks."""
|
|
|
|
supports_mm_inputs: bool
|
|
encoder_cache: EncoderCache | None
|
|
_native_output_materializer: NativeOutputWorker | None
|
|
model_state: OmniModelState
|
|
_last_aux_output: Any
|
|
_last_multimodal_outputs: dict[str, Any] | None
|
|
_model_returns_tuple: bool
|
|
|
|
def _configure_cudagraph_output_contract(self) -> None:
|
|
"""Select the CUDA graph output contract declared by the model."""
|
|
self._model_returns_tuple = _needs_capture_tensor_unwrap(self.model)
|
|
self._exclude_full_graph = self._model_returns_tuple or hasattr(self.model, "_last_captured_layers")
|
|
|
|
@staticmethod
|
|
def _prepare_cudagraph_capture_output(model_output: Any) -> Any:
|
|
"""Unwrap model outputs for the upstream PIECEWISE capture warmup."""
|
|
if isinstance(model_output, OmniOutput):
|
|
return model_output.text_hidden_states
|
|
if isinstance(model_output, tuple) and len(model_output) == 2:
|
|
return model_output[0]
|
|
return model_output
|
|
|
|
def _add_legacy_forward_inputs(self, model_inputs: dict[str, Any], input_batch: Any) -> None:
|
|
"""Supply forward-only metadata for pipelines not on the native plane."""
|
|
if self._omni_data_plane is not None:
|
|
return
|
|
model_inputs.setdefault("sampling_metadata", getattr(input_batch, "sampling_metadata", None))
|
|
logits_index = getattr(
|
|
input_batch,
|
|
"logits_indices",
|
|
getattr(input_batch, "logits_index", None),
|
|
)
|
|
model_inputs.setdefault("logits_index", logits_index)
|
|
model_inputs.setdefault("sampler", self.sampler)
|
|
|
|
def _get_mm_embeddings(
|
|
self,
|
|
scheduled_encoder_inputs: dict[str, list[int]],
|
|
input_batch: Any,
|
|
) -> torch.Tensor:
|
|
"""Call the native vLLM multimodal embedding contract."""
|
|
return self.model_state.prepare_inputs_embeds(
|
|
scheduled_encoder_inputs,
|
|
input_batch,
|
|
self.req_states,
|
|
)
|
|
|
|
def _prepare_mm_inputs(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
input_batch: Any,
|
|
*,
|
|
dummy_run: bool,
|
|
) -> tuple[Any, torch.Tensor | None, Any]:
|
|
"""Prepare native vLLM multimodal model inputs."""
|
|
input_ids = input_batch.input_ids
|
|
inputs_embeds = None
|
|
ec_connector_output = None
|
|
if not self.supports_mm_inputs or not self.is_first_pp_rank:
|
|
return input_ids, inputs_embeds, ec_connector_output
|
|
|
|
if dummy_run:
|
|
inputs_embeds = self.model_state.dummy_inputs_embeds(input_batch.num_tokens_after_padding)
|
|
else:
|
|
scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
|
|
if self.lora_config is not None:
|
|
set_active_mm_loras(
|
|
model=self.model,
|
|
lora_manager=self.lora_manager,
|
|
encoder_cache=self.encoder_cache,
|
|
req_id_to_index=self.req_states.req_id_to_index,
|
|
lora_state=self.lora_state,
|
|
scheduled_encoder_inputs=scheduled_encoder_inputs,
|
|
)
|
|
with self.ec_connector.maybe_get_output(scheduler_output) as ec_connector_output:
|
|
inputs_embeds = self._get_mm_embeddings(
|
|
scheduled_encoder_inputs,
|
|
input_batch,
|
|
)
|
|
if inputs_embeds is not None and not self.model.requires_raw_input_tokens:
|
|
input_ids = None
|
|
return input_ids, inputs_embeds, ec_connector_output
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self._validate_parallel_support()
|
|
self._omni_data_plane = (
|
|
OmniRunnerDataPlane(self.vllm_config, self.model_config)
|
|
if uses_native_mrv2_data_plane(
|
|
self.model_config,
|
|
use_v2_model_runner=True,
|
|
)
|
|
else None
|
|
)
|
|
|
|
def _validate_parallel_support(self) -> None:
|
|
parallel_config = self.vllm_config.parallel_config
|
|
if parallel_config.pipeline_parallel_size > 1:
|
|
raise NotImplementedError(
|
|
"Omni Model Runner V2 does not support pipeline parallel; "
|
|
"use model_runner: v1 when pipeline_parallel_size > 1"
|
|
)
|
|
if getattr(parallel_config, "prefill_context_parallel_size", 1) > 1:
|
|
raise NotImplementedError(
|
|
"Omni Model Runner V2 does not support prefill context parallelism; "
|
|
"use model_runner: v1 when prefill_context_parallel_size > 1"
|
|
)
|
|
|
|
def add_requests(self, scheduler_output: SchedulerOutput) -> None:
|
|
# The upstream Sampler stages static parameters only in add_request.
|
|
# With no new requests, its UVA snapshots remain valid and immutable.
|
|
# Keep the upstream path for custom/speculative samplers, whose staged
|
|
# writes may have a different lifecycle.
|
|
if not scheduler_output.scheduled_new_reqs and type(getattr(self, "sampler", None)) is Sampler:
|
|
return
|
|
logits_processor = getattr(self.model, "logits_processor", None)
|
|
logits_vocab = getattr(logits_processor, "vocab_size", None)
|
|
if isinstance(logits_vocab, int) and logits_vocab > 0:
|
|
for request_data in scheduler_output.scheduled_new_reqs:
|
|
sampling_params = request_data.sampling_params
|
|
if sampling_params is not None:
|
|
sanitize_sampling_params_min_tokens_stop_ids(
|
|
sampling_params,
|
|
logits_vocab,
|
|
)
|
|
super().add_requests(scheduler_output)
|
|
|
|
def shutdown(self) -> None:
|
|
# Every owner must release its resources even if an earlier drain fails.
|
|
try:
|
|
worker = getattr(self, "_native_output_materializer", None)
|
|
if worker is not None:
|
|
worker.close()
|
|
finally:
|
|
self._native_output_materializer = None
|
|
try:
|
|
if self._omni_data_plane is not None:
|
|
self._omni_data_plane.close()
|
|
finally:
|
|
super().shutdown()
|
|
|
|
def _prepare_native_data_plane(self, scheduler_output: SchedulerOutput) -> None:
|
|
plane = getattr(self, "_omni_data_plane", None)
|
|
if plane is None:
|
|
return
|
|
for request_data in getattr(scheduler_output, "scheduled_new_reqs", []):
|
|
if str(getattr(request_data, "req_id", "")).startswith("_warmup_"):
|
|
continue
|
|
plane.register_request(request_data)
|
|
plane.register_receivers(list(getattr(scheduler_output, "pending_input_registrations", [])))
|
|
natural_terminal_req_ids = set(getattr(scheduler_output, "data_plane_terminal_req_ids", set()))
|
|
aborted_req_ids = set(getattr(scheduler_output, "finished_req_ids", set())).difference(natural_terminal_req_ids)
|
|
if natural_terminal_req_ids:
|
|
plane.request_terminal(natural_terminal_req_ids)
|
|
if aborted_req_ids:
|
|
plane.abort_requests(aborted_req_ids)
|
|
|
|
def _sync_native_data_plane_payloads(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
) -> None:
|
|
plane = getattr(self, "_omni_data_plane", None)
|
|
if plane is None:
|
|
return
|
|
gpu_keys: set[tuple[str, str]] = getattr(self.model, "gpu_resident_buffer_keys", set())
|
|
for req_id in scheduler_output.num_scheduled_tokens:
|
|
req_idx = self.req_states.req_id_to_index.get(req_id)
|
|
if req_idx is None:
|
|
continue
|
|
payload = plane.pop_local_stage_payload(req_id)
|
|
if isinstance(payload, dict) and payload:
|
|
self.model_state.intermediate_buffer.update(
|
|
req_idx,
|
|
payload,
|
|
gpu_keys,
|
|
)
|
|
|
|
def _uses_native_output_materializer(self) -> bool:
|
|
parallel = getattr(getattr(self, "vllm_config", None), "parallel_config", None)
|
|
# Multi-rank executors already consume outputs asynchronously and choose
|
|
# the replying rank themselves. Preserve that ownership contract.
|
|
return getattr(parallel, "tensor_parallel_size", 1) == 1 and getattr(self, "_omni_data_plane", None) is not None
|
|
|
|
def _materialize_native_output(self, output: Any) -> Any:
|
|
from vllm_omni.worker_v2.native_output_worker import NativeOutputWorker
|
|
|
|
worker = getattr(self, "_native_output_materializer", None)
|
|
if worker is None:
|
|
worker = NativeOutputWorker(getattr(self.vllm_config, "max_concurrent_batches", 2), self.device)
|
|
self._native_output_materializer = worker
|
|
return worker.submit(output, self._omni_data_plane)
|
|
|
|
def _finalize_native_data_plane_output(self, output: Any) -> Any:
|
|
plane = getattr(self, "_omni_data_plane", None)
|
|
if plane is None or output is None:
|
|
return output
|
|
# MRv2's data plane consumes the host-materialized token lists from
|
|
# the runner output. Do not let a future connector path fall back to
|
|
# reading sampler tensors and synchronizing the device here.
|
|
if hasattr(output, "sampled_token_ids_materialized") and getattr(output, "sampled_token_ids", None) is not None:
|
|
if not output.sampled_token_ids_materialized:
|
|
raise RuntimeError("MRv2 native data plane received non-materialized sampled_token_ids")
|
|
inter_stage_outputs = getattr(output, "inter_stage_outputs", None)
|
|
plane.enqueue_outputs(
|
|
req_ids=list(getattr(output, "req_ids", [])),
|
|
inter_stage_outputs=inter_stage_outputs,
|
|
sampled_token_ids=getattr(output, "sampled_token_ids", None),
|
|
)
|
|
output.inter_stage_outputs = None
|
|
output.omni_connector_output = plane.get_omni_connector_output()
|
|
return output
|
|
|
|
def _reserve_native_data_plane_outputs(self, req_ids: list[str]) -> None:
|
|
plane = getattr(self, "_omni_data_plane", None)
|
|
if plane is not None:
|
|
plane.reserve_outputs(req_ids)
|
|
|
|
def _attach_native_data_plane_signals(self, output: Any) -> Any:
|
|
plane = getattr(self, "_omni_data_plane", None)
|
|
if plane is None:
|
|
return output
|
|
return plane.attach_omni_connector_output(output)
|
|
|
|
def load_model(self, *args: Any, **kwargs: Any) -> None:
|
|
import vllm.v1.worker.gpu.model_runner as _mr_module
|
|
|
|
with _model_state_patch_lock:
|
|
_orig = _mr_module.init_model_state
|
|
_mr_module.init_model_state = init_omni_model_state
|
|
try:
|
|
super().load_model(*args, **kwargs)
|
|
finally:
|
|
_mr_module.init_model_state = _orig
|
|
self._last_aux_output = None
|
|
self._last_multimodal_outputs = None
|
|
self._configure_cudagraph_output_contract()
|
|
|
|
# Preprocess models own embedding buffers; encoder_runner sizing would mismatch.
|
|
if getattr(self.model, "has_preprocess", False) and self.supports_mm_inputs:
|
|
self.supports_mm_inputs = False
|
|
self.encoder_cache = None
|
|
|
|
# ------------------------------------------------------------------
|
|
# CUDA Graph: conditionally exclude FULL mode
|
|
# ------------------------------------------------------------------
|
|
|
|
def _exclude_unsupported_full_graphs(self) -> None:
|
|
# vLLM 0.29 keys dispatch candidates by (uniform token count, LoRA count).
|
|
manager = self.cudagraph_manager
|
|
assert manager is not None
|
|
manager._capture_descs.pop(CUDAGraphMode.FULL, None)
|
|
manager._candidates = {
|
|
key: [desc for desc in candidates if desc.cg_mode != CUDAGraphMode.FULL]
|
|
for key, candidates in manager._candidates.items()
|
|
}
|
|
logger.info("Excluded FULL CUDA graph capture for Omni model. PIECEWISE graphs will still be captured.")
|
|
|
|
def capture_model(self) -> int:
|
|
"""Handle CUDA graph capture for Omni models.
|
|
|
|
Tuple-returning models use PIECEWISE graphs; FULL replay requires
|
|
a tensor output.
|
|
|
|
For PIECEWISE capture, the warmup pass runs with
|
|
``CUDAGraphMode.NONE`` which hits ``torch.empty_like(hidden_states)``
|
|
in the cudagraph framework. If the model returns a tuple, that call
|
|
crashes. We temporarily wrap the model's forward to extract only the
|
|
tensor part during capture, then restore the original forward.
|
|
"""
|
|
if self._exclude_full_graph:
|
|
self._exclude_unsupported_full_graphs()
|
|
|
|
# Wrap model forward during capture so tuple returns don't crash
|
|
# torch.empty_like() in the PIECEWISE warmup pass.
|
|
if self._model_returns_tuple:
|
|
original_forward = self.model.forward
|
|
|
|
def _capture_forward(*args: Any, **kwargs: Any) -> Any:
|
|
output = original_forward(*args, **kwargs)
|
|
return self._prepare_cudagraph_capture_output(output)
|
|
|
|
self.model.forward = _capture_forward # type: ignore[assignment]
|
|
try:
|
|
result = super().capture_model()
|
|
finally:
|
|
self.model.forward = original_forward # type: ignore[assignment]
|
|
else:
|
|
result = super().capture_model()
|
|
|
|
capture_mtp = getattr(getattr(self, "model_state", None), "capture_mtp_graphs", None)
|
|
if callable(capture_mtp):
|
|
capture_mtp(self._dispatch_mtp_batch_descriptor)
|
|
return result
|
|
|
|
def _dispatch_mtp_batch_descriptor(self, num_mtp_reqs: int) -> Any:
|
|
capture_sizes = self.model_state._get_mtp_capture_sizes()
|
|
captured_bucket = next(
|
|
(size for size in sorted(capture_sizes) if num_mtp_reqs <= size <= self.scheduler_config.max_num_seqs),
|
|
None,
|
|
)
|
|
if captured_bucket is None:
|
|
return BatchExecutionDescriptor(
|
|
cg_mode=CUDAGraphMode.NONE,
|
|
num_tokens=num_mtp_reqs,
|
|
num_reqs=num_mtp_reqs,
|
|
)
|
|
return self.cudagraph_manager.dispatch(
|
|
captured_bucket,
|
|
captured_bucket,
|
|
1,
|
|
0,
|
|
)
|
|
|
|
def _dispatch_batch_descriptor(
|
|
self,
|
|
*,
|
|
num_reqs: int,
|
|
num_toks: int,
|
|
uniform_tok_count: int | None,
|
|
num_active_loras: int,
|
|
use_eager: bool,
|
|
max_query_len: int,
|
|
):
|
|
return dispatch_cg_and_sync_dp(
|
|
self.cudagraph_manager,
|
|
num_reqs,
|
|
num_toks,
|
|
uniform_tok_count,
|
|
self.dp_size,
|
|
self.dp_rank,
|
|
max_query_len=max_query_len,
|
|
need_eager=use_eager,
|
|
num_active_loras=num_active_loras,
|
|
)
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
dummy_run: bool = False,
|
|
skip_attn_for_dummy_run: bool = False,
|
|
is_profile: bool = False,
|
|
context_len: int = 0,
|
|
valid_dummy_state_slots: bool = False,
|
|
) -> Any:
|
|
if not dummy_run:
|
|
self._prepare_native_data_plane(scheduler_output)
|
|
self.finish_requests(scheduler_output)
|
|
self.free_states(scheduler_output)
|
|
self.add_requests(scheduler_output)
|
|
self.update_requests(scheduler_output)
|
|
self._sync_native_data_plane_payloads(scheduler_output)
|
|
self.block_tables.apply_staged_writes()
|
|
if scheduler_output.total_num_scheduled_tokens == 0:
|
|
empty_output = self.kv_connector.no_forward(scheduler_output)
|
|
return self._attach_native_data_plane_signals(
|
|
self._merge_ec_connector_no_forward(scheduler_output, empty_output)
|
|
)
|
|
|
|
num_reqs = len(scheduler_output.num_scheduled_tokens)
|
|
num_toks = scheduler_output.total_num_scheduled_tokens
|
|
max_query_len = max(scheduler_output.num_scheduled_tokens.values())
|
|
batch_req_state, uniform_tok_count = self.gather_batch_req_state(scheduler_output, dummy_run)
|
|
if batch_req_state is not None:
|
|
num_toks = batch_req_state.num_tokens
|
|
num_active_loras = 0
|
|
if self.lora_config:
|
|
req_ids = list(scheduler_output.num_scheduled_tokens.keys())
|
|
num_active_loras = get_num_active_loras_for_dispatch(
|
|
self.lora_config,
|
|
self.lora_state,
|
|
req_ids,
|
|
dummy_run,
|
|
)
|
|
# Encoder-decoder models: disable compilation when encoder inputs
|
|
# are scheduled (dynamic cross-attention cache updates).
|
|
skip_compiled = self.is_encoder_decoder and bool(scheduler_output.scheduled_encoder_inputs)
|
|
batch_desc, dp_sync = self._dispatch_batch_descriptor(
|
|
num_reqs=num_reqs,
|
|
num_toks=num_toks,
|
|
uniform_tok_count=uniform_tok_count,
|
|
num_active_loras=num_active_loras,
|
|
use_eager=is_profile or skip_compiled,
|
|
max_query_len=max_query_len,
|
|
)
|
|
|
|
if batch_desc.num_tokens == 0:
|
|
empty_output = self.kv_connector.no_forward(scheduler_output)
|
|
return self._attach_native_data_plane_signals(
|
|
self._merge_ec_connector_no_forward(scheduler_output, empty_output)
|
|
)
|
|
|
|
if not dummy_run:
|
|
assert batch_req_state is not None
|
|
input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
|
|
block_tables, slot_mappings = self.prepare_attn(input_batch)
|
|
self.model_state.preprocess_state(
|
|
input_batch,
|
|
block_tables,
|
|
self.kv_cache_config,
|
|
self.req_states.num_computed_tokens.gpu,
|
|
)
|
|
|
|
if self.lora_config:
|
|
lora_inputs = self.lora_state.make_lora_inputs(
|
|
input_batch.req_ids,
|
|
input_batch.idx_mapping_np,
|
|
input_batch.num_scheduled_tokens,
|
|
)
|
|
self._set_active_loras(*lora_inputs)
|
|
else:
|
|
from vllm.v1.worker.gpu.input_batch import InputBatch
|
|
|
|
input_batch = InputBatch.make_dummy(
|
|
batch_desc.num_reqs or num_reqs,
|
|
batch_desc.num_tokens,
|
|
self.input_buffers,
|
|
max_query_len=batch_desc.max_query_len,
|
|
)
|
|
if not skip_attn_for_dummy_run:
|
|
block_tables, slot_mappings = self.prepare_dummy_attn(input_batch, valid_dummy_state_slots)
|
|
if context_len:
|
|
set_dummy_context(
|
|
input_batch,
|
|
self.block_tables,
|
|
context_len,
|
|
self.kv_cache_config.num_blocks,
|
|
self.max_model_len,
|
|
)
|
|
else:
|
|
assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
|
|
"Attention metadata is required for FULL CUDA graph dummy runs"
|
|
)
|
|
block_tables = None
|
|
slot_mappings = None
|
|
|
|
attn_metadata = None
|
|
slot_mappings_by_layer = None
|
|
if not (dummy_run and skip_attn_for_dummy_run):
|
|
assert slot_mappings is not None
|
|
slot_mappings_by_layer = build_slot_mappings_by_layer(slot_mappings, self.kv_cache_config)
|
|
assert block_tables is not None
|
|
attn_metadata = self.model_state.prepare_attn(
|
|
input_batch,
|
|
batch_desc.cg_mode,
|
|
block_tables,
|
|
slot_mappings,
|
|
self.attn_groups,
|
|
self.kv_cache_config,
|
|
for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,
|
|
)
|
|
input_ids, inputs_embeds, ec_connector_output = self._prepare_mm_inputs(
|
|
scheduler_output,
|
|
input_batch,
|
|
dummy_run=dummy_run,
|
|
)
|
|
|
|
model_inputs: dict[str, Any] = {
|
|
"input_ids": input_ids,
|
|
"positions": input_batch.positions,
|
|
"inputs_embeds": inputs_embeds,
|
|
"intermediate_tensors": intermediate_tensors,
|
|
**self.model_state.prepare_inputs(input_batch, self.req_states),
|
|
}
|
|
self._add_legacy_forward_inputs(model_inputs, input_batch)
|
|
# ★ PRE-FORWARD: per-request preprocess + batched MTP.
|
|
# Runs for ALL graph modes (FULL, PIECEWISE, NONE).
|
|
# For FULL graph: OmniModelState provides a static inputs_embeds
|
|
# buffer that was captured by the graph. Preprocess writes
|
|
# in-place to this buffer, and FULL graph replay reads the
|
|
# updated values from the same tensor address.
|
|
if not dummy_run:
|
|
self.model_state.run_preprocess(
|
|
input_batch,
|
|
model_inputs,
|
|
self.req_states,
|
|
self._dispatch_mtp_batch_descriptor,
|
|
)
|
|
|
|
self.eplb.prepare_forward(self.model_config, input_batch.num_tokens)
|
|
|
|
# --- Model forward ---
|
|
if batch_desc.cg_mode == CUDAGraphMode.FULL:
|
|
# FULL graph replay. Preprocess already wrote to the static
|
|
# inputs_embeds buffer above.
|
|
assert self.cudagraph_manager is not None
|
|
self.kv_connector.pre_forward(scheduler_output)
|
|
hidden_states = self.cudagraph_manager.run_fullgraph(batch_desc)
|
|
self._last_aux_output = None
|
|
self._last_multimodal_outputs = None
|
|
if hasattr(self.model, "_last_captured_layers"):
|
|
self.model._last_captured_layers = self._last_aux_output
|
|
else:
|
|
batch_descriptor = BatchDescriptor(
|
|
num_tokens=input_batch.num_tokens_after_padding,
|
|
has_lora=self.lora_config is not None,
|
|
num_active_loras=batch_desc.num_active_loras,
|
|
)
|
|
with set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=input_batch.num_tokens_after_padding,
|
|
cudagraph_runtime_mode=batch_desc.cg_mode,
|
|
num_tokens_across_dp=dp_sync.num_tokens_across_dp if dp_sync is not None else None,
|
|
batch_descriptor=batch_descriptor,
|
|
slot_mapping=slot_mappings_by_layer,
|
|
skip_compiled=skip_compiled,
|
|
is_padding=input_batch.is_padding,
|
|
):
|
|
self.kv_connector.pre_forward(scheduler_output)
|
|
if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
|
|
assert self.cudagraph_manager is not None
|
|
model_output = self.cudagraph_manager.run_pw_graph(
|
|
self.model,
|
|
model_inputs,
|
|
)
|
|
else:
|
|
model_output = self.model(**model_inputs)
|
|
|
|
# Extract hidden_states from model output.
|
|
self._last_aux_output = None
|
|
self._last_multimodal_outputs = None
|
|
if isinstance(model_output, OmniOutput):
|
|
hidden_states = model_output.text_hidden_states
|
|
if model_output.multimodal_outputs:
|
|
self._last_multimodal_outputs = cast(dict[str, Any], model_output.multimodal_outputs)
|
|
elif isinstance(model_output, tuple) and len(model_output) == 2:
|
|
hidden_states, self._last_aux_output = model_output
|
|
if hasattr(self.model, "_last_captured_layers"):
|
|
self.model._last_captured_layers = self._last_aux_output
|
|
else:
|
|
hidden_states = model_output
|
|
if not dummy_run and isinstance(hidden_states, torch.Tensor):
|
|
self.model_state.run_postprocess(hidden_states, input_batch)
|
|
|
|
routed_experts = None
|
|
if not dummy_run and (capturer := self.routed_experts_capturer) is not None:
|
|
assert slot_mappings is not None
|
|
routed_experts = capturer.get_routed_experts(slot_mappings, num_toks)
|
|
|
|
self.execute_model_state = ExecuteModelState(
|
|
input_batch=input_batch,
|
|
attn_metadata=attn_metadata,
|
|
slot_mappings_by_layer=slot_mappings_by_layer,
|
|
hidden_states=hidden_states,
|
|
aux_hidden_states=None,
|
|
finished_req_ids=scheduler_output.finished_req_ids,
|
|
dp_sync=dp_sync,
|
|
ec_connector_output=ec_connector_output,
|
|
routed_experts=routed_experts,
|
|
cudagraph_stats=None,
|
|
)
|
|
|
|
assert isinstance(hidden_states, torch.Tensor)
|
|
return None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Request lifecycle: update intermediate buffer from cached requests
|
|
# ------------------------------------------------------------------
|
|
|
|
def update_requests(self, scheduler_output: SchedulerOutput) -> None:
|
|
"""Merge updated additional_information into intermediate_buffer.
|
|
|
|
In async_chunk mode, chunk_transfer_adapter attaches updated
|
|
additional_information (e.g. thinker_decode_embeddings) to
|
|
OmniCachedRequestData for cached requests every schedule step.
|
|
Upstream GPUModelRunner.update_requests does not handle this
|
|
field, so we merge it into the intermediate buffer here.
|
|
"""
|
|
super().update_requests(scheduler_output)
|
|
|
|
cached = scheduler_output.scheduled_cached_reqs
|
|
addl_info = getattr(cached, "additional_information", None)
|
|
if not addl_info:
|
|
return
|
|
for req_id, info in addl_info.items():
|
|
if info is None:
|
|
continue
|
|
req_idx = self.req_states.req_id_to_index.get(req_id)
|
|
if req_idx is None:
|
|
continue
|
|
resolved = _resolve_additional_information(info)
|
|
if resolved:
|
|
gpu_keys: set[tuple[str, str]] = getattr(self.model, "gpu_resident_buffer_keys", set())
|
|
self.model_state.intermediate_buffer.update(req_idx, resolved, gpu_keys)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Request lifecycle: clean up intermediate buffer on finish
|
|
# ------------------------------------------------------------------
|
|
|
|
def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
|
|
# IMPORTANT: Must query req_id_to_index BEFORE super().finish_requests()
|
|
# because super() calls req_states.remove_request(req_id) which pops the
|
|
# mapping and returns the slot index to free_indices.
|
|
finished = scheduler_output.finished_req_ids
|
|
on_finished = getattr(self.model, "on_requests_finished", None)
|
|
if finished and callable(on_finished):
|
|
# Request-owned state outlives generation slots and chunk boundaries.
|
|
# Notify even if the last chunk already released its runner slot.
|
|
on_finished(finished)
|
|
preempted = scheduler_output.preempted_req_ids
|
|
all_done = finished | preempted if preempted else finished
|
|
for req_id in all_done:
|
|
idx = self.req_states.req_id_to_index.get(req_id)
|
|
if idx is not None:
|
|
self.model_state.remove_request(idx)
|
|
super().finish_requests(scheduler_output)
|