2314 lines
107 KiB
Python
2314 lines
107 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""AR GPU Model Runner for vLLM-Omni.
|
|
|
|
Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
|
|
and also outputs sampled tokens.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import threading
|
|
from collections.abc import Callable, Mapping, Sequence
|
|
from contextlib import nullcontext
|
|
from copy import copy
|
|
from dataclasses import replace
|
|
from typing import Any, NamedTuple, cast
|
|
|
|
import numpy as np
|
|
import torch
|
|
from vllm.config import CUDAGraphMode
|
|
from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
|
|
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
|
|
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
|
|
from vllm.forward_context import set_forward_context
|
|
from vllm.logger import init_logger
|
|
from vllm.sequence import IntermediateTensors
|
|
from vllm.utils.platform_utils import is_pin_memory_available
|
|
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
|
|
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
|
|
from vllm.v1.spec_decode.dflash import DFlashProposer
|
|
from vllm.v1.spec_decode.draft_model import DraftModelProposer
|
|
from vllm.v1.spec_decode.eagle import EagleProposer
|
|
from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
|
|
from vllm.v1.spec_decode.gemma4 import Gemma4Proposer
|
|
from vllm.v1.structured_output.utils import apply_grammar_bitmask
|
|
from vllm.v1.utils import record_function_or_nullcontext
|
|
from vllm.v1.worker.gpu_model_runner import (
|
|
AsyncGPUModelRunnerOutput,
|
|
)
|
|
from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
|
|
from vllm.v1.worker.utils import is_residual_scattered_for_sp
|
|
|
|
from vllm_omni.data_entry_keys import flatten_payload
|
|
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
|
|
from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output
|
|
from vllm_omni.model_executor.duplex_sampling import DuplexSamplingRunnerMixin
|
|
from vllm_omni.outputs import OmniModelRunnerOutput
|
|
from vllm_omni.utils.mm_outputs import (
|
|
build_mm_cpu,
|
|
partition_payload_list,
|
|
snapshot_mm_payload,
|
|
)
|
|
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
|
|
from vllm_omni.worker.omni_connector_model_runner_mixin import (
|
|
OmniConnectorModelRunnerMixin,
|
|
needs_omni_connector,
|
|
)
|
|
from vllm_omni.worker.output.payload_build import build_omni_mm_payload
|
|
from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
|
|
from vllm_omni.worker.sampling_utils import (
|
|
call_model_sampler,
|
|
clamp_prompt_ids_to_penalty_padding,
|
|
sanitize_min_tokens_stop_ids,
|
|
)
|
|
from vllm_omni.worker.sparse_audio import resolve_sparse_mm_routing
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _to_cpu_contiguous(tensor: torch.Tensor) -> torch.Tensor:
|
|
tensor = tensor.detach()
|
|
if tensor.device.type == "cpu":
|
|
return tensor.contiguous()
|
|
return tensor.to("cpu").contiguous()
|
|
|
|
|
|
def _clone_cuda_tensor_payload(value: Any, sources: list[torch.Tensor]) -> Any:
|
|
"""Clone CUDA tensors on the current stream before async CPU copies.
|
|
|
|
The clone protects async Omni output snapshots from CUDA graph output
|
|
buffers that may be reused by subsequent decode steps. CPU tensors are
|
|
cloned synchronously because they are already host-owned snapshots.
|
|
"""
|
|
if isinstance(value, torch.Tensor):
|
|
if value.device.type == "cuda":
|
|
cloned = value.detach().clone()
|
|
sources.append(cloned)
|
|
return cloned
|
|
return value.detach().clone()
|
|
if isinstance(value, dict):
|
|
return {k: _clone_cuda_tensor_payload(v, sources) for k, v in value.items()}
|
|
if isinstance(value, (list, tuple)):
|
|
if (
|
|
value
|
|
and all(
|
|
isinstance(v, torch.Tensor)
|
|
and v.is_cuda
|
|
and v.dtype == value[0].dtype
|
|
and v.device == value[0].device
|
|
and v.layout == torch.strided
|
|
and v.is_contiguous()
|
|
for v in value
|
|
)
|
|
# Keep frame/state packing bounded; large payloads use individual snapshots.
|
|
and sum(v.numel() * v.element_size() for v in value) <= 1024 * 1024
|
|
):
|
|
packed = torch.cat([v.detach().reshape(-1) for v in value])
|
|
sources.append(packed)
|
|
return _PackedTensorPayload(packed, [v.shape for v in value], isinstance(value, tuple))
|
|
items = [_clone_cuda_tensor_payload(v, sources) for v in value]
|
|
return tuple(items) if isinstance(value, tuple) else items
|
|
return value
|
|
|
|
|
|
def _copy_tensor_payload_to_cpu(value: Any, pin_memory: bool) -> Any:
|
|
if isinstance(value, _PackedTensorPayload):
|
|
flat = _copy_tensor_payload_to_cpu(value.tensor, pin_memory)
|
|
items = []
|
|
offset = 0
|
|
for shape in value.shapes:
|
|
count = shape.numel()
|
|
items.append(flat[offset : offset + count].view(shape))
|
|
offset += count
|
|
return tuple(items) if value.is_tuple else items
|
|
if isinstance(value, torch.Tensor):
|
|
if value.device.type != "cuda":
|
|
return value
|
|
cpu = torch.empty_like(value, device="cpu", pin_memory=pin_memory)
|
|
cpu.copy_(value, non_blocking=True)
|
|
return cpu
|
|
if isinstance(value, dict):
|
|
return {k: _copy_tensor_payload_to_cpu(v, pin_memory) for k, v in value.items()}
|
|
if isinstance(value, list):
|
|
return [_copy_tensor_payload_to_cpu(v, pin_memory) for v in value]
|
|
if isinstance(value, tuple):
|
|
return tuple(_copy_tensor_payload_to_cpu(v, pin_memory) for v in value)
|
|
return value
|
|
|
|
|
|
class _PackedTensorPayload(NamedTuple):
|
|
tensor: torch.Tensor
|
|
shapes: list[torch.Size]
|
|
is_tuple: bool
|
|
|
|
|
|
class _AsyncCPUPayloadSnapshot:
|
|
def __init__(
|
|
self,
|
|
payload: Any,
|
|
ready_event: torch.cuda.Event | None,
|
|
cuda_sources: list[torch.Tensor],
|
|
) -> None:
|
|
self.payload = payload
|
|
self._ready_event = ready_event
|
|
self._cuda_sources = cuda_sources
|
|
self._waited = False
|
|
|
|
def wait(self) -> None:
|
|
if self._waited:
|
|
return
|
|
if self._ready_event is not None:
|
|
self._ready_event.synchronize()
|
|
self._cuda_sources.clear()
|
|
self._waited = True
|
|
|
|
|
|
def _snapshot_tensor_payload_to_cpu_async(
|
|
value: Any,
|
|
*,
|
|
copy_stream: torch.cuda.Stream,
|
|
pin_memory: bool,
|
|
) -> _AsyncCPUPayloadSnapshot:
|
|
cuda_sources: list[torch.Tensor] = []
|
|
cloned = _clone_cuda_tensor_payload(value, cuda_sources)
|
|
if not cuda_sources:
|
|
return _AsyncCPUPayloadSnapshot(cloned, None, cuda_sources)
|
|
|
|
source_stream = torch.cuda.current_stream()
|
|
ready_event = torch.cuda.Event()
|
|
with torch.cuda.stream(copy_stream):
|
|
copy_stream.wait_stream(source_stream)
|
|
cpu_payload = _copy_tensor_payload_to_cpu(cloned, pin_memory)
|
|
# An executor may discard an unconsumed output (e.g. cancellation).
|
|
# Keep the allocator from reusing its source storage before D2H ends.
|
|
for source in cuda_sources:
|
|
source.record_stream(copy_stream)
|
|
ready_event.record(copy_stream)
|
|
return _AsyncCPUPayloadSnapshot(cpu_payload, ready_event, cuda_sources)
|
|
|
|
|
|
class _OmniOutputTensorSnapshot(NamedTuple):
|
|
hidden_states: torch.Tensor
|
|
staged_hidden_states_cpu: torch.Tensor | None
|
|
multimodal_outputs: Any
|
|
async_payload: _AsyncCPUPayloadSnapshot | None = None
|
|
|
|
|
|
class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model_runner_output_builder: Callable[[], OmniModelRunnerOutput],
|
|
cuda_device: torch.device | int | str | None = None,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
sampled_token_ids = kwargs.pop("sampled_token_ids")
|
|
logprobs_tensors = kwargs.pop("logprobs_tensors")
|
|
invalid_req_indices = kwargs.pop("invalid_req_indices")
|
|
async_output_copy_stream = kwargs.pop("async_output_copy_stream")
|
|
vocab_size = kwargs.pop("vocab_size")
|
|
num_nans = kwargs.pop("num_nans", None)
|
|
# Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault
|
|
# for EP all2all fault tolerance (PR #43637). Omni doesn't use this
|
|
# feature but must consume the kwarg to prevent TypeError from stray
|
|
# kwargs and initialize the attribute so super().get_output() works.
|
|
kwargs.pop("check_ep_fault", False)
|
|
if kwargs:
|
|
raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}")
|
|
|
|
self._model_runner_output: OmniModelRunnerOutput | None = None
|
|
self._invalid_req_indices = invalid_req_indices
|
|
|
|
self.async_copy_ready_event = torch.Event()
|
|
self._sampled_token_ids = sampled_token_ids
|
|
self.vocab_size = vocab_size
|
|
self._logprobs_tensors = logprobs_tensors
|
|
self._has_fault: torch.Tensor | None = None
|
|
# Upstream b1e12d142d (PR #51304) added device-side NaN-in-logits
|
|
# counts (num_nans) to AsyncGPUModelRunnerOutput. Omni keeps the
|
|
# counts on the async copy stream and lets super().get_output()
|
|
# populate num_nans_in_logits from the CPU copy.
|
|
self._num_nans = num_nans
|
|
|
|
default_stream = torch.cuda.current_stream()
|
|
with torch.cuda.stream(async_output_copy_stream):
|
|
async_output_copy_stream.wait_stream(default_stream)
|
|
# Keep sampled-token feedback identical to upstream async
|
|
# scheduling. This tensor drives the next decode step, so avoid
|
|
# changing its host-copy allocation semantics while building Omni
|
|
# output asynchronously.
|
|
self.sampled_token_ids_cpu = self._sampled_token_ids.to("cpu", non_blocking=True)
|
|
self._logprobs_tensors_cpu = self._logprobs_tensors.to_cpu_nonblocking() if self._logprobs_tensors else None
|
|
self._num_nans_cpu = self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None else None
|
|
self.async_copy_ready_event.record()
|
|
|
|
self._model_runner_output_builder: Callable[[], OmniModelRunnerOutput] | None = model_runner_output_builder
|
|
self._background_exception: BaseException | None = None
|
|
self._background_thread: threading.Thread | None = None
|
|
self._cuda_device = cuda_device
|
|
self._background_thread = threading.Thread(
|
|
target=self._build_output_in_background,
|
|
daemon=True,
|
|
name="omni-async-output-builder",
|
|
)
|
|
self._background_thread.start()
|
|
|
|
def _build_model_runner_output_once(self) -> None:
|
|
if self._model_runner_output is not None:
|
|
return
|
|
with record_function_or_nullcontext("omni_async_output:get_output/build_model_runner_output"):
|
|
assert self._model_runner_output_builder is not None
|
|
self._model_runner_output = self._model_runner_output_builder()
|
|
self._model_runner_output_builder = None
|
|
|
|
def _build_output_in_background(self) -> None:
|
|
try:
|
|
if self._cuda_device is not None:
|
|
torch.cuda.set_device(self._cuda_device)
|
|
self._build_model_runner_output_once()
|
|
except BaseException as exc: # noqa: BLE001 - re-raised by get_output().
|
|
self._background_exception = exc
|
|
|
|
def get_output(self) -> OmniModelRunnerOutput:
|
|
background_thread = getattr(self, "_background_thread", None)
|
|
if background_thread is not None:
|
|
background_thread.join()
|
|
self._background_thread = None
|
|
background_exception = getattr(self, "_background_exception", None)
|
|
if background_exception is not None:
|
|
raise background_exception
|
|
self._build_model_runner_output_once()
|
|
# Upstream AsyncGPUModelRunnerOutput.get_output() accesses
|
|
# self._has_fault for EP all2all fault tolerance (PR #43637).
|
|
# Ensure the attribute exists even when __init__ was bypassed
|
|
# (e.g. unit tests using object.__new__).
|
|
if not hasattr(self, "_has_fault"):
|
|
self._has_fault = None
|
|
# Upstream b1e12d142d (PR #51304) also touches _num_nans/_num_nans_cpu
|
|
# in get_output(). Guard them the same way for object.__new__ tests.
|
|
if not hasattr(self, "_num_nans"):
|
|
self._num_nans = None
|
|
if not hasattr(self, "_num_nans_cpu"):
|
|
self._num_nans_cpu = None
|
|
with record_function_or_nullcontext("omni_async_output:get_output/finalize_async_sampled_tokens"):
|
|
return super().get_output()
|
|
|
|
|
|
def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]:
|
|
"""Convert a flattened payload to strictly ``dict[str, torch.Tensor]``.
|
|
|
|
Non-tensor scalars (int, float) are wrapped with ``torch.tensor()``.
|
|
Values that cannot be safely converted are dropped with a warning.
|
|
This enforces the tensor-only invariant required by the
|
|
``OmniEngineCoreOutput.multimodal_output`` wire field and msgspec
|
|
serialization.
|
|
"""
|
|
result: dict[str, torch.Tensor] = {}
|
|
for key, val in payload.items():
|
|
# Sparse per-request conditioning uses None for requests with no
|
|
# update. Absence is expected, not an unsupported wire value.
|
|
if val is None:
|
|
continue
|
|
if isinstance(val, torch.Tensor):
|
|
result[key] = val
|
|
elif isinstance(val, (int, float, bool)):
|
|
result[key] = torch.tensor(val)
|
|
elif isinstance(val, (list, tuple)):
|
|
try:
|
|
result[key] = torch.tensor(val)
|
|
except (ValueError, TypeError, RuntimeError):
|
|
logger.warning(
|
|
"Dropping non-tensorizable multimodal output key '%s' (type=%s) from wire payload.",
|
|
key,
|
|
type(val).__name__,
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"Dropping non-tensor multimodal output key '%s' (type=%s) from wire payload.",
|
|
key,
|
|
type(val).__name__,
|
|
)
|
|
return result
|
|
|
|
|
|
class ExecuteModelState(NamedTuple):
|
|
scheduler_output: SchedulerOutput
|
|
logits: torch.Tensor | None
|
|
spec_decode_metadata: Any
|
|
spec_decode_common_attn_metadata: Any
|
|
hidden_states: torch.Tensor
|
|
hidden_states_cpu: torch.Tensor | None
|
|
sample_hidden_states: torch.Tensor
|
|
aux_hidden_states: list[torch.Tensor] | None
|
|
ec_connector_output: Any
|
|
cudagraph_stats: Any
|
|
# OMNI: multimodal_outputs field for omni-specific multimodal handling
|
|
multimodal_outputs: Any
|
|
# slot_mappings for attention/drafter (aligned with upstream v1 API)
|
|
slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None
|
|
# OMNI: prefix-cache step id; the step context saved under it must be
|
|
# consumed exactly once (materialize or discard_step) in sample_tokens.
|
|
prefix_cache_step_id: int | None = None
|
|
|
|
|
|
class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin):
|
|
"""Autoregressive GPU model runner that returns hidden states per request.
|
|
|
|
Follows the v0.12 two-phase execute/sample flow from GPUModelRunner, and
|
|
reuses Omni hooks for additional_information / multimodal outputs. This
|
|
class only overrides sample_tokens to expose hidden states + multimodal
|
|
outputs per request while keeping Async output semantics.
|
|
"""
|
|
|
|
execute_model_state: ExecuteModelState | None
|
|
kv_extracted_req_ids: list[str] | None
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
|
|
# each model stage has their own hidden size
|
|
self.hidden_size = self.model_config.hf_text_config.hidden_size
|
|
# The width of embeddings received by this stage. It can differ from
|
|
# the stage's internal hidden size (e.g., Qwen2.5-Omni Talker receives
|
|
# 3584-wide Thinker states and projects them to a 896-wide hidden state).
|
|
self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
|
|
self.inputs_embeds = self._make_buffer(
|
|
self.max_num_tokens,
|
|
self.inputs_embeds_size,
|
|
dtype=self.dtype,
|
|
numpy=False,
|
|
)
|
|
# Initialize KV cache manager (preserve vllm_config fallback behavior)
|
|
self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
|
|
self._async_chunk = getattr(self.model_config, "async_chunk", False)
|
|
# Connector init is derived from pipeline capabilities and active
|
|
# producer/KV/connector configuration. A stage needs it when it is:
|
|
# - a full-payload consumer, or
|
|
# - a producer (custom_process_next_stage_input_func is set, meaning
|
|
# it sends inter-stage data to a downstream consumer via connector), or
|
|
# - explicitly configured with a connector role via stage_connector_config.
|
|
if needs_omni_connector(self.model_config):
|
|
self.init_omni_connectors(
|
|
model_config=self.model_config,
|
|
kv_transfer_manager=self.kv_transfer_manager,
|
|
)
|
|
self._downstream_payload_cache: dict[str, bool] = {}
|
|
self._init_duplex_sampling_state()
|
|
|
|
def load_model(self, *args, **kwargs) -> None:
|
|
super().load_model(*args, **kwargs)
|
|
self._resolve_duplex_sampling_hook(force=True)
|
|
|
|
def _make_buffer(self, *size, dtype, numpy=True):
|
|
# Prevent ray from pinning the buffer due to large size
|
|
from vllm_omni.distributed.ray_utils.utils import (
|
|
calculate_total_bytes,
|
|
maybe_disable_pin_memory_for_ray,
|
|
)
|
|
|
|
total_bytes = calculate_total_bytes(size, dtype)
|
|
|
|
# Use the context manager to temporarily disable pinning if needed
|
|
with maybe_disable_pin_memory_for_ray(self, total_bytes):
|
|
return super()._make_buffer(*size, dtype=dtype, numpy=numpy)
|
|
|
|
def _build_model_sampler_output_token_ids(self) -> list[list[int]]:
|
|
"""Build decoded-token history for custom model samplers.
|
|
|
|
vLLM only populates sampling_metadata.output_token_ids when penalties or
|
|
logits processors require it. CosyVoice3's custom RAS sampler also
|
|
depends on this history, so we reconstruct it directly from the input
|
|
batch for prefer_model_sampler models.
|
|
"""
|
|
req_output_token_ids = getattr(self.input_batch, "req_output_token_ids", [])
|
|
req_ids = list(getattr(self.input_batch, "req_ids", []))
|
|
output_token_ids = [list(req_output_token_ids[idx] or []) for idx in range(len(req_ids))]
|
|
|
|
sampled_token_ids_cpu = getattr(self.input_batch, "sampled_token_ids_cpu", None)
|
|
async_copy_ready_event = getattr(self.input_batch, "async_copy_ready_event", None)
|
|
prev_req_id_to_index = getattr(self.input_batch, "prev_req_id_to_index", None)
|
|
if sampled_token_ids_cpu is None or not output_token_ids or prev_req_id_to_index is None:
|
|
return output_token_ids
|
|
|
|
sampled_token_ids: list[list[int]] | None = None
|
|
for index, req_id in enumerate(req_ids):
|
|
prev_index = prev_req_id_to_index.get(req_id)
|
|
if prev_index is None:
|
|
continue
|
|
req_history = output_token_ids[index]
|
|
if not req_history or req_history[-1] != -1:
|
|
continue
|
|
if sampled_token_ids is None:
|
|
assert async_copy_ready_event is not None
|
|
async_copy_ready_event.synchronize()
|
|
sampled_token_ids = sampled_token_ids_cpu.tolist()
|
|
new_ids = list(sampled_token_ids[prev_index])
|
|
if not new_ids:
|
|
continue
|
|
num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1)
|
|
first_placeholder = req_history.index(-1)
|
|
num_placeholders = len(req_history) - first_placeholder
|
|
num_to_replace = min(num_sampled_ids, num_placeholders)
|
|
req_history[first_placeholder : first_placeholder + num_to_replace] = new_ids[:num_to_replace]
|
|
|
|
for index, req_history in enumerate(output_token_ids):
|
|
if -1 in req_history:
|
|
output_token_ids[index] = req_history[: req_history.index(-1)]
|
|
|
|
return output_token_ids
|
|
|
|
def _sampling_metadata_for_model_sampler(self, sampling_metadata):
|
|
if getattr(self.model, "skips_model_sampler_output_token_history", False):
|
|
return sampling_metadata
|
|
output_token_ids = self._build_model_sampler_output_token_ids()
|
|
if output_token_ids != sampling_metadata.output_token_ids:
|
|
return replace(sampling_metadata, output_token_ids=output_token_ids)
|
|
return sampling_metadata
|
|
|
|
def _update_states(self, scheduler_output: SchedulerOutput) -> Callable | None:
|
|
deferred_state_corrections_fn = super()._update_states(scheduler_output)
|
|
self._update_duplex_sampling_states(scheduler_output)
|
|
self._maybe_apply_duplex_window_reanchor(scheduler_output)
|
|
return deferred_state_corrections_fn
|
|
|
|
def _maybe_apply_duplex_window_reanchor(self, scheduler_output: SchedulerOutput | None = None) -> None:
|
|
"""Apply in-place KV reanchor and rotation on worker before model forward."""
|
|
reanchor_hook = getattr(getattr(self, "model", None), "apply_duplex_kv_reanchor", None)
|
|
if callable(reanchor_hook):
|
|
try:
|
|
reanchor_hook(self, scheduler_output=scheduler_output)
|
|
except TypeError:
|
|
reanchor_hook(self)
|
|
return
|
|
|
|
helper = getattr(self, "_duplex_window_helper", None)
|
|
if helper is not None and hasattr(helper, "maybe_apply_reanchor"):
|
|
try:
|
|
helper.maybe_apply_reanchor(self, scheduler_output=scheduler_output)
|
|
except TypeError:
|
|
helper.maybe_apply_reanchor(self)
|
|
return
|
|
|
|
# Fallback for dynamic runner inspection without hardcoding model classes
|
|
if not hasattr(self, "input_batch") or self.input_batch is None:
|
|
return
|
|
num_reqs = getattr(self.input_batch, "num_reqs", 0)
|
|
req_ids = self.input_batch.req_ids[:num_reqs]
|
|
has_reanchor = False
|
|
for req_id in req_ids:
|
|
info = self.model_intermediate_buffer.get(req_id)
|
|
if isinstance(info, dict) and isinstance(info.get("duplex"), dict):
|
|
if "stage0_reanchor" in info["duplex"]:
|
|
has_reanchor = True
|
|
break
|
|
if not has_reanchor:
|
|
return
|
|
|
|
model_module = getattr(getattr(self, "model", None), "__module__", "")
|
|
if "minicpmo_4_5" in model_module:
|
|
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.window_kv import (
|
|
MiniCPMO45DuplexWorkerHelper,
|
|
)
|
|
|
|
MiniCPMO45DuplexWorkerHelper.maybe_apply_reanchor(self, scheduler_output=scheduler_output)
|
|
|
|
def _maybe_apply_stage0_reanchor(self) -> None:
|
|
"""Backward-compatible alias for _maybe_apply_duplex_window_reanchor."""
|
|
self._maybe_apply_duplex_window_reanchor()
|
|
|
|
def _request_final_stage_id(self, req_id: str) -> int | None:
|
|
info = self.model_intermediate_buffer.get(req_id)
|
|
if not isinstance(info, dict):
|
|
req_state = self.requests.get(req_id)
|
|
info = getattr(req_state, "additional_information_cpu", None)
|
|
if not isinstance(info, dict):
|
|
return None
|
|
val = info.get("omni_final_stage_id")
|
|
if val is None:
|
|
return None
|
|
try:
|
|
return int(val)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
def _request_needs_downstream_stage_payload(self, req_id: str) -> bool:
|
|
cached = self._downstream_payload_cache.get(req_id)
|
|
if cached is not None:
|
|
return cached
|
|
final_stage_id = self._request_final_stage_id(req_id)
|
|
if final_stage_id is None:
|
|
# Conservative default while the marker is missing: keep the
|
|
# payload, but do NOT memoize - the marker arrives via
|
|
# `model_intermediate_buffer`, which may be unpopulated on the
|
|
# first call (memoizing here pinned the request to True forever,
|
|
# never refreshing once the marker landed).
|
|
return True
|
|
needs_payload = final_stage_id > 0
|
|
self._downstream_payload_cache[req_id] = needs_payload
|
|
return needs_payload
|
|
|
|
def _resolve_pooler_payload_req_ids(self, req_ids_output_copy: list[str]) -> tuple[str, list[str]]:
|
|
downstream_req_ids = [rid for rid in req_ids_output_copy if self._request_needs_downstream_stage_payload(rid)]
|
|
engine_output_type = (self.vllm_config.model_config.engine_output_type or "").lower()
|
|
if self._client_multimodal_output_keys():
|
|
downstream_req_ids = req_ids_output_copy
|
|
# Single-stage AR TTS models (e.g. VoxCPM2) finish on this stage but still
|
|
# need multimodal payloads for final audio postprocess/output.
|
|
elif engine_output_type == "audio" and not downstream_req_ids:
|
|
downstream_req_ids = req_ids_output_copy
|
|
return engine_output_type, downstream_req_ids
|
|
|
|
def capture_model(self) -> int:
|
|
result = super().capture_model()
|
|
self._capture_talker_mtp_graphs()
|
|
return result
|
|
|
|
def shutdown(self) -> None:
|
|
"""Release omni-specific GPU resources before upstream shutdown.
|
|
|
|
Order of operations (must match upstream's expectation):
|
|
1. Unfreeze Python GC so model weights are collected immediately
|
|
when self.model is set to None (upstream Worker.init_device
|
|
calls gc.freeze() / freeze_gc_heap()).
|
|
2. Destroy omni-specific CUDA graphs (talker MTP) so references to
|
|
model parameters are released before self.model = None.
|
|
3. Clear GPU-side buffers (input_ids, inputs_embeds) and per-request
|
|
caches that may hold GPU tensor references.
|
|
4. Call CUDAGraphWrapper.clear_all_graphs() unconditionally (not just
|
|
on ROCm) to ensure all CUDA graphs including talker MTP are
|
|
released before model weight teardown.
|
|
5. Call BreakableCUDAGraphWrapper.clear_all_graphs() as well, to
|
|
match the upstream ROCm-only pattern but also protect CUDA.
|
|
6. Delegate to upstream GPUModelRunner.shutdown() which sets
|
|
self.model = None, clears KV caches, resets workspace, etc.
|
|
|
|
This prevents abrupt GPU memory release during EngineCore subprocess
|
|
exit that can trigger GPU OOM signals when the parent process
|
|
concurrently cleans up its own GPU state.
|
|
"""
|
|
# 1. Unfreeze GC so model weights and GPU tensors are collected
|
|
# immediately when references are dropped (upstream Worker.shutdown
|
|
# also does this before any teardown).
|
|
gc.unfreeze()
|
|
|
|
# 2. Destroy talker MTP CUDA graph wrapper to release captured graphs.
|
|
if hasattr(self, "talker_mtp") and self.talker_mtp is not None:
|
|
self.talker_mtp = None
|
|
self.has_talker_mtp = False
|
|
|
|
# 3. Clear GPU-side buffers (small tensors, but every MiB helps).
|
|
if hasattr(self, "input_ids") and self.input_ids is not None:
|
|
self.input_ids = None
|
|
if hasattr(self, "inputs_embeds") and self.inputs_embeds is not None:
|
|
self.inputs_embeds = None
|
|
|
|
# 4. Clear per-request caches that may hold GPU tensor references.
|
|
if hasattr(self, "_downstream_payload_cache"):
|
|
self._downstream_payload_cache.clear()
|
|
if hasattr(self, "model_intermediate_buffer"):
|
|
self.model_intermediate_buffer.clear()
|
|
self._clear_duplex_sampling()
|
|
|
|
# 5. Release all CUDA graphs unconditionally (upstream only does this
|
|
# on ROCm; on CUDA the graphs are only freed by Python GC during
|
|
# interpreter shutdown, which is too late to prevent memory spikes).
|
|
from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
|
|
from vllm.compilation.cuda_graph import CUDAGraphWrapper
|
|
|
|
CUDAGraphWrapper.clear_all_graphs()
|
|
BreakableCUDAGraphWrapper.clear_all_graphs()
|
|
|
|
# 6. Stop the prefix-cache committer before the CUDA context and the
|
|
# tensors it copies from are torn down.
|
|
cache = self.omni_prefix_cache
|
|
if cache is not None:
|
|
cache.shutdown()
|
|
self.omni_prefix_cache = None
|
|
|
|
# 7. Delegate to upstream shutdown (model = None, KV caches, workspace).
|
|
super().shutdown()
|
|
|
|
def _capture_talker_mtp_graphs(self) -> None:
|
|
from vllm.compilation.cuda_graph import CUDAGraphWrapper
|
|
|
|
if not self.has_talker_mtp or not isinstance(self.talker_mtp, CUDAGraphWrapper):
|
|
return
|
|
|
|
from vllm.compilation.monitor import set_cudagraph_capturing_enabled
|
|
from vllm.distributed.parallel_state import graph_capture
|
|
|
|
capture_sizes = self.compilation_config.cudagraph_capture_sizes
|
|
num_warmups = self.compilation_config.cudagraph_num_of_warmups
|
|
capture_sizes = sorted(capture_sizes, reverse=True)
|
|
logger.info("Capturing talker_mtp graphs for sizes %s", capture_sizes)
|
|
|
|
set_cudagraph_capturing_enabled(True)
|
|
try:
|
|
with torch.inference_mode(), graph_capture(device=self.device):
|
|
for bsz in capture_sizes:
|
|
_, batch_desc, _, _, _ = self._determine_batch_execution_and_padding(
|
|
num_tokens=bsz,
|
|
num_reqs=bsz,
|
|
num_scheduled_tokens_np=np.ones(bsz, dtype=np.int32),
|
|
max_num_scheduled_tokens=1,
|
|
use_cascade_attn=False,
|
|
)
|
|
n = batch_desc.num_tokens
|
|
ids = self.talker_mtp_input_ids.gpu[:n]
|
|
emb = self.talker_mtp_inputs_embeds.gpu[:n]
|
|
hid = self.last_talker_hidden.gpu[:n]
|
|
ts = self.text_step.gpu[:n]
|
|
|
|
for _ in range(num_warmups):
|
|
with set_forward_context(
|
|
None,
|
|
self.vllm_config,
|
|
cudagraph_runtime_mode=CUDAGraphMode.NONE,
|
|
batch_descriptor=batch_desc,
|
|
):
|
|
self.talker_mtp(ids, emb, hid, ts)
|
|
|
|
with set_forward_context(
|
|
None,
|
|
self.vllm_config,
|
|
cudagraph_runtime_mode=CUDAGraphMode.FULL,
|
|
batch_descriptor=batch_desc,
|
|
):
|
|
self.talker_mtp(ids, emb, hid, ts)
|
|
torch.accelerator.synchronize()
|
|
|
|
logger.info("Captured talker_mtp graphs for %d sizes", len(capture_sizes))
|
|
except RuntimeError as e:
|
|
raise RuntimeError(
|
|
f"talker_mtp graph capture failed for a model that declared talker_mtp_graph_safe=True: {e}"
|
|
) from e
|
|
finally:
|
|
set_cudagraph_capturing_enabled(False)
|
|
|
|
def _get_runner_assisted_full_attention_metadata_request(
|
|
self,
|
|
*,
|
|
req_ids: Sequence[str],
|
|
num_reqs: int,
|
|
num_reqs_padded: int,
|
|
num_scheduled_tokens_np: np.ndarray,
|
|
num_computed_tokens_cpu: torch.Tensor | np.ndarray,
|
|
max_num_scheduled_tokens: int,
|
|
) -> RunnerAssistedFullAttentionMetadataRequest | None:
|
|
# Models without this hook keep the normal runner path. VoxCPM2 uses
|
|
# it to ask the runner for padded FULL attention metadata while keeping
|
|
# graph policy and lifecycle in the model layer.
|
|
hook = getattr(self.model, "get_runner_assisted_full_attention_metadata_request", None)
|
|
if not callable(hook):
|
|
return None
|
|
num_computed_tokens = num_computed_tokens_cpu
|
|
if hasattr(num_computed_tokens, "numpy"):
|
|
num_computed_tokens = num_computed_tokens.numpy()
|
|
request = hook(
|
|
req_ids=req_ids,
|
|
num_reqs=num_reqs,
|
|
num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
|
|
num_computed_tokens=num_computed_tokens,
|
|
max_num_scheduled_tokens=max_num_scheduled_tokens,
|
|
)
|
|
if request is None:
|
|
return None
|
|
if not isinstance(request, RunnerAssistedFullAttentionMetadataRequest):
|
|
raise TypeError(
|
|
"runner-assisted full attention metadata hook must return "
|
|
"RunnerAssistedFullAttentionMetadataRequest or None, got "
|
|
f"{type(request).__name__}"
|
|
)
|
|
padded_num_reqs = max(
|
|
num_reqs_padded,
|
|
min(int(request.num_reqs_padded), self.scheduler_config.max_num_seqs),
|
|
)
|
|
return RunnerAssistedFullAttentionMetadataRequest(
|
|
num_reqs_padded=padded_num_reqs,
|
|
for_cudagraph_capture=bool(request.for_cudagraph_capture),
|
|
)
|
|
|
|
def _refresh_runner_assisted_full_attention_metadata_buffers(
|
|
self,
|
|
*,
|
|
num_reqs: int,
|
|
num_reqs_padded: int,
|
|
num_scheduled_tokens_np: np.ndarray,
|
|
) -> None:
|
|
num_computed_tokens = self.input_batch.num_computed_tokens_cpu[:num_reqs]
|
|
if hasattr(num_computed_tokens, "numpy"):
|
|
num_computed_tokens = num_computed_tokens.numpy()
|
|
np.add(
|
|
num_computed_tokens,
|
|
num_scheduled_tokens_np,
|
|
out=self.optimistic_seq_lens_cpu[:num_reqs].numpy(),
|
|
)
|
|
self.optimistic_seq_lens_cpu[num_reqs:num_reqs_padded].fill_(0)
|
|
self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True)
|
|
|
|
cum_num_tokens = self._get_cumsum_and_arange(
|
|
num_scheduled_tokens_np,
|
|
self.query_pos.np,
|
|
)
|
|
self.query_start_loc.np[0] = 0
|
|
self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
|
|
self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1])
|
|
self.query_start_loc.copy_to_gpu()
|
|
|
|
self.input_batch.block_table.commit_block_table(num_reqs_padded)
|
|
|
|
def _set_runner_assisted_full_attention_metadata_context(
|
|
self,
|
|
*,
|
|
enabled: bool,
|
|
num_reqs: int = 0,
|
|
) -> bool:
|
|
hook = getattr(self.model, "set_runner_assisted_full_attention_metadata_context", None)
|
|
if not callable(hook):
|
|
return False
|
|
hook(
|
|
enabled=enabled,
|
|
num_reqs=num_reqs,
|
|
)
|
|
return True
|
|
|
|
def _prepare_prefix_cache_pooler_payload_sources(
|
|
self,
|
|
*,
|
|
staged_hidden_states_cpu: torch.Tensor | None,
|
|
needs_scheduled_hidden_payload: bool,
|
|
req_ids: list[str],
|
|
step_id: int | None,
|
|
) -> tuple[torch.Tensor | None, dict[str, torch.Tensor] | None, dict | None]:
|
|
"""Prefix-cache payload sources for the pooler payload build.
|
|
|
|
req_ids must be the step's snapshot, not ``input_batch.req_ids``: under
|
|
async omni output this runs a step late, when the live batch has moved on.
|
|
"""
|
|
if not get_pp_group().is_last_rank:
|
|
raise RuntimeError("Omni prefix-cache merge is only valid on the last pipeline parallel rank.")
|
|
hidden_states_cpu = None
|
|
if needs_scheduled_hidden_payload:
|
|
if staged_hidden_states_cpu is None:
|
|
raise RuntimeError("Prefix-cache hidden-state payload requires staged CPU hidden states.")
|
|
hidden_states_cpu = staged_hidden_states_cpu
|
|
if step_id is None:
|
|
# save_outputs did not run this step (no last-rank AR output), so
|
|
# there is no step context to consume; serve the fresh slice.
|
|
if hidden_states_cpu is None and staged_hidden_states_cpu is not None:
|
|
hidden_states_cpu = staged_hidden_states_cpu
|
|
return hidden_states_cpu, None, None
|
|
combined_hidden, combined_mm = self._prefix_cache_materialize(step_id, list(req_ids))
|
|
return hidden_states_cpu, combined_hidden, combined_mm
|
|
|
|
def _build_omni_pooler_payload(
|
|
self,
|
|
*,
|
|
rid: str,
|
|
idx: int,
|
|
start: int,
|
|
end: int,
|
|
hidden_states_cpu: torch.Tensor | None,
|
|
req_hidden_states_cpu: dict[str, torch.Tensor] | None,
|
|
combined_hidden_states: dict[str, torch.Tensor] | None,
|
|
combined_multimodal_outputs: dict | None,
|
|
mm_cpu: dict[str, object] | None,
|
|
audio_sparse_output: bool,
|
|
sparse_mm_index: dict[str, int],
|
|
hidden_seq_len: int,
|
|
scheduled_seq_len: int,
|
|
) -> dict[str, object]:
|
|
payload: dict[str, object] = {}
|
|
if not audio_sparse_output:
|
|
if req_hidden_states_cpu is not None and combined_hidden_states is None:
|
|
req_hidden_states = req_hidden_states_cpu[rid]
|
|
else:
|
|
req_hidden_states = self._resolve_req_hidden_states(
|
|
hidden_states_cpu,
|
|
combined_hidden_states,
|
|
rid,
|
|
start,
|
|
end,
|
|
)
|
|
if req_hidden_states is not None:
|
|
payload["hidden"] = req_hidden_states
|
|
|
|
mm_payload = build_omni_mm_payload(
|
|
combined_multimodal_outputs=combined_multimodal_outputs,
|
|
mm_cpu=mm_cpu,
|
|
rid=rid,
|
|
idx=idx,
|
|
start=start,
|
|
end=end,
|
|
audio_sparse_output=audio_sparse_output,
|
|
sparse_mm_index=sparse_mm_index,
|
|
hidden_seq_len=hidden_seq_len,
|
|
scheduled_seq_len=scheduled_seq_len,
|
|
)
|
|
payload.update(mm_payload)
|
|
return payload
|
|
|
|
def _merge_model_kv_transfer_metadata(
|
|
self,
|
|
finished_reqs: dict[str, dict[str, Any]],
|
|
) -> dict[str, dict[str, Any]]:
|
|
"""Merge model-provided KV-transfer metadata copy-on-write.
|
|
|
|
The outer mapping is always new; a per-request dict is shallow-copied
|
|
(with a fresh merged ``custom_metadata``) only when the model supplies
|
|
metadata for it, and passes through unchanged otherwise.
|
|
`scheduler_output` (and its per-request dicts) can be shared with the
|
|
engine-core process, so nothing here mutates them — the ngram block
|
|
in `execute_model` `replace()`-copies for exactly the same reason. The
|
|
merged mapping is only consumed by
|
|
`kv_transfer_manager.handle_finished_requests_kv_transfer`.
|
|
"""
|
|
merged: dict[str, dict[str, Any]] = {}
|
|
for req_id, data in finished_reqs.items():
|
|
try:
|
|
# NOTE: seq_len is the same as num_computed_tokens_cpu in current
|
|
# async scheduling, since both exclude async placeholders. We use
|
|
# seq_len since we control it, just in case upstream async scheduler
|
|
# semantics change in the future.
|
|
model_meta = self.model.get_kv_transfer_metadata(
|
|
req_id,
|
|
num_computed_tokens=data.get("seq_len"),
|
|
)
|
|
except Exception:
|
|
# A swallowed failure here used to ship the KV transfer WITHOUT
|
|
# the model's custom metadata and surface much later as a
|
|
# corrupt/incomplete transfer on the consumer stage. Fail at
|
|
# the point of the defect instead.
|
|
logger.exception("Failed to get custom KV-transfer metadata from model for %s", req_id)
|
|
raise
|
|
if model_meta:
|
|
data = {
|
|
**data,
|
|
"custom_metadata": {**(data.get("custom_metadata") or {}), **model_meta},
|
|
}
|
|
merged[req_id] = data
|
|
return merged
|
|
|
|
def _call_prepare_runner_inputs(
|
|
self,
|
|
prepare_runner_inputs: Callable[..., tuple[torch.Tensor, torch.Tensor]],
|
|
*,
|
|
req_ids: list[str],
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
inputs_embeds: torch.Tensor | None,
|
|
num_computed_tokens: np.ndarray,
|
|
num_scheduled_tokens: np.ndarray,
|
|
input_ids_buffer: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Keep opt-in request metadata outside captured model forwards."""
|
|
if getattr(self.model, "accepts_runner_sampling_extra_args", False):
|
|
return prepare_runner_inputs(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
inputs_embeds=inputs_embeds,
|
|
req_ids=req_ids,
|
|
num_computed_tokens=num_computed_tokens,
|
|
num_scheduled_tokens=num_scheduled_tokens,
|
|
input_ids_buffer=input_ids_buffer,
|
|
sampling_extra_args=[
|
|
params.extra_args if params and params.extra_args else {}
|
|
for params in (self.requests[rid].sampling_params for rid in req_ids)
|
|
],
|
|
discard_mask=self.discard_request_mask.np[: len(req_ids)],
|
|
)
|
|
return prepare_runner_inputs(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
inputs_embeds=inputs_embeds,
|
|
req_ids=req_ids,
|
|
num_computed_tokens=num_computed_tokens,
|
|
num_scheduled_tokens=num_scheduled_tokens,
|
|
input_ids_buffer=input_ids_buffer,
|
|
)
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None:
|
|
if self.execute_model_state is not None:
|
|
raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
|
|
|
|
if not getattr(self, "_warmup_state_cleared", False):
|
|
self._warmup_state_cleared = True
|
|
if hasattr(self.model, "_clear_warmup_state"):
|
|
self.model._clear_warmup_state()
|
|
|
|
# Exactly once per real scheduler_output, before _update_states
|
|
# (dummy/warmup runs never reach execute_model).
|
|
self._prefix_cache_step_begin(scheduler_output)
|
|
|
|
# [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
|
|
finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
|
|
if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"):
|
|
finished_reqs = self._merge_model_kv_transfer_metadata(finished_reqs)
|
|
self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
|
|
finished_reqs=finished_reqs,
|
|
kv_caches=self.kv_caches,
|
|
block_size=self.cache_config.block_size,
|
|
cache_dtype=str(self.cache_config.cache_dtype),
|
|
request_id_resolver=self._resolve_global_request_id,
|
|
)
|
|
|
|
if hasattr(self, "_omni_connector"):
|
|
for request in getattr(scheduler_output, "pending_input_registrations", []):
|
|
self.register_chunk_recv(request)
|
|
self.recv_full_payload_inputs(scheduler_output)
|
|
if self._pending_full_payload_send:
|
|
flush_ids = set(getattr(scheduler_output, "finished_req_ids", set()))
|
|
flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests})
|
|
if flush_ids:
|
|
self.flush_full_payload_outputs(flush_ids)
|
|
|
|
# If ngram_gpu is used, we need to copy the scheduler_output to avoid
|
|
# the modification has influence on the scheduler_output in engine core process.
|
|
# The replace is much faster than deepcopy.
|
|
if self.speculative_config is not None and self.speculative_config.use_ngram_gpu():
|
|
num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy()
|
|
spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy()
|
|
scheduler_output = replace(
|
|
scheduler_output,
|
|
num_scheduled_tokens=num_scheduled_tokens_copy,
|
|
scheduled_spec_decode_tokens=spec_decode_tokens_copy,
|
|
)
|
|
|
|
if has_kv_transfer_group():
|
|
kv_connector_metadata = scheduler_output.kv_connector_metadata
|
|
if kv_connector_metadata is not None:
|
|
get_kv_transfer_group().handle_preemptions(kv_connector_metadata)
|
|
|
|
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
|
|
with (
|
|
record_function_or_nullcontext("gpu_model_runner: preprocess"),
|
|
self.synchronize_input_prep(),
|
|
):
|
|
# Update persistent batch states.
|
|
deferred_state_corrections_fn = self._update_states(scheduler_output)
|
|
|
|
# Notify model of finished requests for state cleanup
|
|
if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
|
|
self.model.on_requests_finished(scheduler_output.finished_req_ids)
|
|
|
|
if has_ec_transfer() and not get_ec_transfer().is_consumer:
|
|
with self.maybe_get_ec_connector_output(
|
|
scheduler_output,
|
|
encoder_cache=self.encoder_cache,
|
|
) as ec_connector_output:
|
|
self._execute_mm_encoder(scheduler_output)
|
|
|
|
kv_ids = self.kv_extracted_req_ids
|
|
self.kv_extracted_req_ids = None
|
|
|
|
output = make_empty_encoder_model_runner_output(scheduler_output)
|
|
if kv_ids:
|
|
output = copy(output)
|
|
output.kv_extracted_req_ids = kv_ids
|
|
return self.attach_omni_connector_output(output)
|
|
|
|
# `<= 0`: upstream can schedule a negative span, which is truthy (#5196).
|
|
if num_scheduled_tokens <= 0:
|
|
if (
|
|
self.parallel_config.distributed_executor_backend == "external_launcher"
|
|
and self.parallel_config.data_parallel_size > 1
|
|
):
|
|
self._dummy_run(1)
|
|
|
|
# Capture KV extraction results before early return;
|
|
# sample_tokens() is skipped on this path so the IDs
|
|
# would otherwise be silently overwritten next step.
|
|
kv_ids = self.kv_extracted_req_ids
|
|
self.kv_extracted_req_ids = None
|
|
|
|
if not has_kv_transfer_group():
|
|
output = EMPTY_MODEL_RUNNER_OUTPUT
|
|
else:
|
|
output = self.kv_connector_no_forward(scheduler_output, self.vllm_config)
|
|
|
|
if kv_ids:
|
|
output = copy(output)
|
|
output.kv_extracted_req_ids = kv_ids
|
|
|
|
return self.attach_omni_connector_output(output)
|
|
|
|
if self.cache_config.kv_sharing_fast_prefill:
|
|
assert not self.num_prompt_logprobs, (
|
|
"--kv-sharing-fast-prefill produces incorrect "
|
|
"logprobs for prompt tokens, tokens, please disable "
|
|
"it when the requests need prompt logprobs"
|
|
)
|
|
|
|
num_reqs = self.input_batch.num_reqs
|
|
req_ids = self.input_batch.req_ids
|
|
tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
|
|
num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
|
|
max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
|
|
num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
|
|
|
|
logits_indices, spec_decode_metadata, max_num_sampled_tokens = self._prepare_inputs(
|
|
scheduler_output, num_scheduled_tokens_np
|
|
)
|
|
|
|
cascade_attn_prefix_lens = None
|
|
# Disable cascade attention when using microbatching (DBO)
|
|
if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
|
|
# Pre-compute cascade attention prefix lengths
|
|
cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
|
|
num_scheduled_tokens_np,
|
|
self.input_batch.num_computed_tokens_cpu[:num_reqs],
|
|
scheduler_output.num_common_prefix_blocks,
|
|
)
|
|
|
|
(
|
|
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),
|
|
)
|
|
num_tokens_padded = batch_desc.num_tokens
|
|
num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
|
|
num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu[:num_reqs]
|
|
|
|
runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request(
|
|
req_ids=req_ids[:num_reqs],
|
|
num_reqs=num_reqs,
|
|
num_reqs_padded=num_reqs_padded,
|
|
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
|
num_computed_tokens_cpu=num_computed_tokens_cpu,
|
|
max_num_scheduled_tokens=max_num_scheduled_tokens,
|
|
)
|
|
runner_assisted_full_attn_capture = False
|
|
if runner_assisted_full_attn_request is not None:
|
|
num_reqs_padded = runner_assisted_full_attn_request.num_reqs_padded
|
|
runner_assisted_full_attn_capture = runner_assisted_full_attn_request.for_cudagraph_capture
|
|
num_tokens_padded = max(
|
|
num_tokens_padded,
|
|
num_reqs_padded * max_num_scheduled_tokens,
|
|
)
|
|
runner_assisted_full_attn = runner_assisted_full_attn_request is not None
|
|
|
|
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
|
|
should_ubatch,
|
|
num_scheduled_tokens_np,
|
|
num_tokens_padded,
|
|
num_reqs_padded,
|
|
self.parallel_config.num_ubatches,
|
|
)
|
|
pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn
|
|
|
|
use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
|
|
ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
|
|
|
|
# True if any attention backend handles KV cache update separately
|
|
# from forward() (i.e., forward_includes_kv_cache_update=False). When true,
|
|
# slot_mappings must use padded dimensions to match the key/value tensors.
|
|
from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec
|
|
|
|
has_separate_kv_update = not all(
|
|
all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id])
|
|
for id, spec in enumerate(self.kv_cache_config.kv_cache_groups)
|
|
if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec)
|
|
)
|
|
|
|
slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
|
|
num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded,
|
|
num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs),
|
|
num_tokens_unpadded=num_tokens_unpadded,
|
|
ubatch_slices=ubatch_slices_padded,
|
|
)
|
|
if runner_assisted_full_attn:
|
|
self._refresh_runner_assisted_full_attention_metadata_buffers(
|
|
num_reqs=num_reqs,
|
|
num_reqs_padded=num_reqs_padded,
|
|
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
|
)
|
|
|
|
attn_num_reqs = num_reqs_padded if runner_assisted_full_attn else num_reqs
|
|
attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
|
|
num_tokens=num_tokens_unpadded,
|
|
num_tokens_padded=num_tokens_padded if pad_attn else None,
|
|
num_reqs=attn_num_reqs,
|
|
num_reqs_padded=num_reqs_padded if pad_attn else None,
|
|
max_query_len=max_num_scheduled_tokens,
|
|
ubatch_slices=ubatch_slices_attn,
|
|
logits_indices=logits_indices,
|
|
max_num_sampled_tokens=max_num_sampled_tokens,
|
|
use_spec_decode=use_spec_decode,
|
|
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
|
|
cascade_attn_prefix_lens=cascade_attn_prefix_lens,
|
|
slot_mappings=slot_mappings_by_group,
|
|
for_cudagraph_capture=runner_assisted_full_attn_capture,
|
|
)
|
|
self._maybe_attach_attention_metadata_extensions(
|
|
attn_metadata=attn_metadata,
|
|
num_reqs=num_reqs,
|
|
num_reqs_padded=num_reqs_padded,
|
|
max_query_len=max_num_scheduled_tokens,
|
|
pad_attn=pad_attn,
|
|
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
|
for_cudagraph_capture=runner_assisted_full_attn_capture,
|
|
)
|
|
|
|
(
|
|
input_ids,
|
|
inputs_embeds,
|
|
positions,
|
|
intermediate_tensors,
|
|
model_kwargs,
|
|
ec_connector_output,
|
|
) = self._preprocess(scheduler_output, num_tokens_padded, intermediate_tensors)
|
|
|
|
# Let the model adjust inputs before forward (e.g. restore input_ids
|
|
# for multimodal position detection, fix decode position offsets).
|
|
prepare_runner_inputs = getattr(self.model, "prepare_runner_inputs", None)
|
|
if callable(prepare_runner_inputs):
|
|
input_ids, positions = self._call_prepare_runner_inputs(
|
|
prepare_runner_inputs,
|
|
req_ids=req_ids[:num_reqs],
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
inputs_embeds=inputs_embeds,
|
|
num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs],
|
|
num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
|
|
input_ids_buffer=self.input_ids.gpu[:num_tokens_padded],
|
|
)
|
|
|
|
runner_assisted_context_enabled = False
|
|
if runner_assisted_full_attn:
|
|
runner_assisted_context_enabled = self._set_runner_assisted_full_attention_metadata_context(
|
|
enabled=True,
|
|
num_reqs=num_reqs,
|
|
)
|
|
runner_assisted_full_attn = runner_assisted_context_enabled
|
|
|
|
# Run the model.
|
|
# Use persistent buffers for CUDA graphs.
|
|
# When spec decode is enabled, defer connector finalization
|
|
# (wait_for_save + clear metadata) until after draft model runs.
|
|
defer_kv_connector_finalize = self.speculative_config is not None
|
|
try:
|
|
with (
|
|
nullcontext(),
|
|
set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_tokens_padded,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
cudagraph_runtime_mode=(CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode),
|
|
batch_descriptor=batch_desc,
|
|
ubatch_slices=ubatch_slices_padded,
|
|
slot_mapping=slot_mappings, # OMNI: required for KV cache operations
|
|
),
|
|
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,
|
|
sampling_metadata=self.input_batch.sampling_metadata,
|
|
logits_index=logits_indices,
|
|
sampler=self.sampler,
|
|
)
|
|
|
|
# [Omni] Map pending ropes metadata to req_ids.
|
|
flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None)
|
|
if callable(flush_pending_metadata):
|
|
flush_pending_metadata(req_ids[:num_reqs])
|
|
|
|
# [Omni] Hand the model the batch's req_ids in logits order, for
|
|
# models that gate logits per request. Only valid without spec
|
|
# decode: there logits_indices carries several rows per request,
|
|
# so row i no longer corresponds to req_ids[i].
|
|
if spec_decode_metadata is None:
|
|
set_batch_req_ids = getattr(self.model, "set_batch_req_ids", None)
|
|
if callable(set_batch_req_ids):
|
|
set_batch_req_ids(req_ids[:num_reqs])
|
|
finally:
|
|
if runner_assisted_context_enabled:
|
|
self._set_runner_assisted_full_attention_metadata_context(enabled=False)
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: postprocess"):
|
|
if self.use_aux_hidden_state_outputs:
|
|
# True when EAGLE 3 is used.
|
|
hidden_states, aux_hidden_states = model_output
|
|
else:
|
|
# Common case.
|
|
hidden_states = model_output
|
|
aux_hidden_states = None
|
|
|
|
# A pooling stage skips mm extraction + the omni prefix cache (both assume an AR
|
|
# [n_tok, hidden] output) and exits through the shared _pool() branch below.
|
|
if self.is_pooling_model:
|
|
multimodal_outputs = None
|
|
else:
|
|
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output)
|
|
hidden_states_cpu = None
|
|
# Prefix-cache write: freeze + submit; policy gating (full
|
|
# hidden, skip/deferred keys) happens inside the manager.
|
|
prefix_cache_step_id = self._prefix_cache_save_step(
|
|
hidden_states,
|
|
multimodal_outputs,
|
|
num_tokens_unpadded=num_tokens_unpadded,
|
|
num_tokens_padded=num_tokens_padded,
|
|
)
|
|
if (
|
|
prefix_cache_step_id is not None
|
|
and not self._model_needs_full_prefix_hidden_states()
|
|
and self._model_omni_pooler_payload_include_hidden()
|
|
):
|
|
# Opt-out models keep the scheduled-slice CPU view for the
|
|
# pooler payload (runner-side bypass, legacy parity).
|
|
hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous()
|
|
|
|
if not self.broadcast_pp_output:
|
|
# Common case.
|
|
if not get_pp_group().is_last_rank:
|
|
# Return the intermediate tensors.
|
|
assert isinstance(hidden_states, IntermediateTensors)
|
|
self.kv_connector_output = kv_connector_output
|
|
return hidden_states
|
|
|
|
if self.is_pooling_model:
|
|
# Return the pooling output.
|
|
return self._pool(
|
|
hidden_states,
|
|
num_scheduled_tokens,
|
|
num_scheduled_tokens_np,
|
|
kv_connector_output,
|
|
)
|
|
|
|
sample_hidden_states = hidden_states[logits_indices.to(hidden_states.device)]
|
|
# Try with sampling_metadata first; fall back to without for models that don't support it
|
|
try:
|
|
logits = self.model.compute_logits(
|
|
sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
|
|
)
|
|
except TypeError:
|
|
logits = self.model.compute_logits(sample_hidden_states)
|
|
else:
|
|
# Rare case.
|
|
assert not self.is_pooling_model
|
|
|
|
sample_hidden_states = hidden_states[logits_indices.to(hidden_states.device)]
|
|
if not get_pp_group().is_last_rank:
|
|
all_gather_tensors = {
|
|
"residual": not is_residual_scattered_for_sp(self.vllm_config, num_tokens_padded)
|
|
}
|
|
get_pp_group().send_tensor_dict(
|
|
hidden_states.tensors,
|
|
all_gather_group=get_tp_group(),
|
|
all_gather_tensors=all_gather_tensors,
|
|
)
|
|
logits = None
|
|
else:
|
|
# Try with sampling_metadata first; fall back to without for models that don't support it
|
|
try:
|
|
logits = self.model.compute_logits(
|
|
sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
|
|
)
|
|
except TypeError:
|
|
logits = self.model.compute_logits(sample_hidden_states)
|
|
|
|
model_output_broadcast_data: dict[str, Any] = {}
|
|
if logits is not None:
|
|
model_output_broadcast_data["logits"] = logits.contiguous()
|
|
|
|
broadcasted = get_pp_group().broadcast_tensor_dict(
|
|
model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
|
|
)
|
|
assert broadcasted is not None
|
|
logits = broadcasted["logits"]
|
|
|
|
self.execute_model_state = ExecuteModelState(
|
|
scheduler_output,
|
|
logits,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
hidden_states,
|
|
hidden_states_cpu,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
ec_connector_output,
|
|
cudagraph_stats,
|
|
multimodal_outputs,
|
|
slot_mappings, # OMNI: pass slot_mappings for drafter
|
|
prefix_cache_step_id,
|
|
)
|
|
self.kv_connector_output = kv_connector_output
|
|
|
|
if deferred_state_corrections_fn:
|
|
deferred_state_corrections_fn()
|
|
|
|
return None
|
|
|
|
def _sample(
|
|
self,
|
|
logits: torch.Tensor | None,
|
|
spec_decode_metadata: Any,
|
|
):
|
|
"""Sample tokens, preferring the model's own sampler when declared.
|
|
|
|
Model-interface contract for `prefer_model_sampler` models: a custom
|
|
`model.sample(logits, sampling_metadata)` may return ``None`` to
|
|
DECLINE the step, and the runner then falls back to the default
|
|
sampler. Declarers rely on this deliberately for empty-logits or
|
|
inapplicable-stage steps (e.g. cosyvoice3, glm_tts), so the
|
|
fallthrough is load-bearing, not an error path; it is logged once per
|
|
model for visibility.
|
|
"""
|
|
sampling_metadata = self.input_batch.sampling_metadata
|
|
if spec_decode_metadata is None:
|
|
model_sample = getattr(self.model, "sample", None)
|
|
self.input_batch.update_async_output_token_ids()
|
|
if logits is not None and callable(model_sample) and getattr(self.model, "prefer_model_sampler", False):
|
|
# Apply logit bias (min_tokens, allowed_token_ids) before
|
|
# the custom model sampler — the standard GPU sampler does
|
|
# this internally, but prefer_model_sampler bypasses it.
|
|
if hasattr(self.sampler, "logit_bias_state"):
|
|
self.sampler.logit_bias_state.apply_logit_bias(
|
|
logits,
|
|
self.input_batch.expanded_idx_mapping,
|
|
self.input_batch.idx_mapping_np,
|
|
self.input_batch.positions[self.input_batch.logits_indices],
|
|
)
|
|
prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata)
|
|
self._apply_duplex_sampling(logits, prepared_sampling_metadata)
|
|
sampler_output = call_model_sampler(
|
|
self.model,
|
|
model_sample,
|
|
logits,
|
|
prepared_sampling_metadata,
|
|
input_batch=self.input_batch,
|
|
requests=getattr(self, "requests", None),
|
|
)
|
|
if sampler_output is not None:
|
|
return sampler_output
|
|
# Contract: None => fall back to the default sampler (see
|
|
# docstring). If a custom sampler returns None by accident,
|
|
# its tokens silently come from the default sampler - the
|
|
# one-time log is the visibility hook for that case.
|
|
logger.warning_once(
|
|
"prefer_model_sampler model %s returned None from sample(); falling back to the default sampler.",
|
|
type(self.model).__name__,
|
|
)
|
|
return self.sampler(
|
|
logits=logits,
|
|
sampling_metadata=sampling_metadata,
|
|
)
|
|
|
|
return super()._sample(logits, spec_decode_metadata)
|
|
|
|
@staticmethod
|
|
def _resolve_req_hidden_states(
|
|
hidden_states_cpu: torch.Tensor | None,
|
|
combined_hidden_states: dict[str, torch.Tensor] | None,
|
|
rid: str,
|
|
start: int,
|
|
end: int,
|
|
) -> torch.Tensor | None:
|
|
if combined_hidden_states is not None:
|
|
# We always have all request IDs for prefix cache, even for
|
|
# partial cache misses, so this should never happen.
|
|
if rid not in combined_hidden_states:
|
|
raise RuntimeError("Request IDs in the batch are missing from the merged states!")
|
|
return combined_hidden_states[rid]
|
|
# Prefix caching is disabled. hidden_states_cpu may legitimately be
|
|
# None (e.g. sparse audio output or no scheduled hidden payload);
|
|
# callers must omit the "hidden" key in that case.
|
|
if hidden_states_cpu is None:
|
|
return None
|
|
return hidden_states_cpu[start:end]
|
|
|
|
def _run_post_sample_talker_mtp(
|
|
self,
|
|
*,
|
|
req_ids: list[str],
|
|
valid_sampled_token_ids: list[list[int]],
|
|
sampled_token_ids: torch.Tensor,
|
|
invalid_req_indices: list[int],
|
|
sample_hidden_states: torch.Tensor,
|
|
multimodal_outputs: Any,
|
|
) -> Any:
|
|
"""Run an opt-in talker MTP with the token sampled in this step.
|
|
|
|
Some frame-synchronous talkers need the current temporal hidden state
|
|
and its newly sampled text token to produce the same frame's codec
|
|
codes. A resumable ``max_tokens=1`` segment has no later decode
|
|
preprocess in which to do that work, so those models expose
|
|
``post_sample_talker_mtp`` instead.
|
|
"""
|
|
model = getattr(self, "model", None)
|
|
hook = getattr(model, "post_sample_talker_mtp", None)
|
|
if not callable(hook):
|
|
return multimodal_outputs
|
|
|
|
selected_indices: list[int] = []
|
|
selected_req_ids: list[str] = []
|
|
selected_req_infos: list[dict[str, Any]] = []
|
|
duplex_indices: list[int] = []
|
|
invalid_indices = set(invalid_req_indices)
|
|
use_async_scheduling = bool(getattr(self, "use_async_scheduling", False))
|
|
for idx, req_id in enumerate(req_ids):
|
|
req_info = self.model_intermediate_buffer.get(req_id)
|
|
duplex = req_info.get("duplex") if isinstance(req_info, dict) else None
|
|
if not isinstance(duplex, dict) or duplex.get("data_plane") is not True:
|
|
continue
|
|
duplex_indices.append(idx)
|
|
if use_async_scheduling:
|
|
if idx in invalid_indices:
|
|
continue
|
|
else:
|
|
token_ids = valid_sampled_token_ids[idx] if idx < len(valid_sampled_token_ids) else []
|
|
if not token_ids:
|
|
continue
|
|
if len(token_ids) != 1:
|
|
raise ValueError(
|
|
"post_sample_talker_mtp requires exactly one sampled token "
|
|
f"per request, got {len(token_ids)} for {req_id}"
|
|
)
|
|
selected_indices.append(idx)
|
|
selected_req_ids.append(req_id)
|
|
selected_req_infos.append(cast(dict[str, Any], req_info))
|
|
|
|
if not selected_indices:
|
|
return multimodal_outputs
|
|
if sampled_token_ids.ndim != 2 or sampled_token_ids.shape[0] < len(req_ids):
|
|
raise ValueError(
|
|
"post_sample_talker_mtp sampled-token rows do not match requests: "
|
|
f"shape={tuple(sampled_token_ids.shape)}, requests={len(req_ids)}"
|
|
)
|
|
if sampled_token_ids.shape[1] != 1:
|
|
raise ValueError(
|
|
"post_sample_talker_mtp does not support speculative token rows: "
|
|
f"shape={tuple(sampled_token_ids.shape)}"
|
|
)
|
|
|
|
empty_audio = torch.empty(0, dtype=torch.long, device=sample_hidden_states.device)
|
|
merged = dict(multimodal_outputs) if isinstance(multimodal_outputs, dict) else {}
|
|
existing_codes = merged.get("codes")
|
|
codes_payload = dict(existing_codes) if isinstance(existing_codes, dict) else {}
|
|
existing_audio = codes_payload.get("audio")
|
|
if isinstance(existing_audio, list) and len(existing_audio) == len(req_ids):
|
|
per_request_audio = list(existing_audio)
|
|
elif (
|
|
isinstance(existing_audio, torch.Tensor)
|
|
and existing_audio.ndim > 0
|
|
and existing_audio.shape[0] == len(req_ids)
|
|
):
|
|
per_request_audio = [existing_audio[idx : idx + 1] for idx in range(len(req_ids))]
|
|
else:
|
|
per_request_audio = [empty_audio for _ in req_ids]
|
|
for idx in duplex_indices:
|
|
per_request_audio[idx] = empty_audio
|
|
|
|
indices = torch.tensor(
|
|
selected_indices,
|
|
dtype=torch.long,
|
|
device=sample_hidden_states.device,
|
|
)
|
|
input_ids = sampled_token_ids.index_select(0, indices).reshape(-1).to(torch.long)
|
|
codes = hook(
|
|
input_ids=input_ids,
|
|
hidden_states=sample_hidden_states.index_select(0, indices),
|
|
req_ids=selected_req_ids,
|
|
req_infos=selected_req_infos,
|
|
)
|
|
if not isinstance(codes, torch.Tensor) or codes.ndim != 2:
|
|
raise TypeError(f"post_sample_talker_mtp must return a rank-2 Tensor, got {type(codes).__name__}")
|
|
if codes.shape[0] != len(selected_req_ids):
|
|
raise ValueError(
|
|
f"post_sample_talker_mtp row count does not match requests: {codes.shape[0]} != {len(selected_req_ids)}"
|
|
)
|
|
for row, (idx, req_id) in enumerate(zip(selected_indices, selected_req_ids, strict=True)):
|
|
request_codes = codes[row : row + 1]
|
|
per_request_audio[idx] = request_codes
|
|
self._update_intermediate_buffer(
|
|
req_id,
|
|
{"codes": {"audio": request_codes}},
|
|
)
|
|
|
|
# Empty rows are intentional: they suppress replay of the previous
|
|
# frame while a chunked prefill has not sampled a token yet.
|
|
codes_payload["audio"] = per_request_audio
|
|
merged["codes"] = codes_payload
|
|
return merged
|
|
|
|
def _build_multimodal_outputs(
|
|
self,
|
|
per_req_payloads: Sequence[dict[str, object] | None] | None,
|
|
) -> list[dict[str, torch.Tensor] | None] | None:
|
|
"""Build per-request multimodal output payloads (dedicated channel).
|
|
|
|
Reuses the per-request payloads assembled by the pooler-payload loop
|
|
in sample_tokens() (which already handles prefix-cache merging,
|
|
sparse audio output, and partial downstream batches) so the wire
|
|
channel stays consistent with the full-payload accumulation path.
|
|
Enforces the tensor-only invariant required by msgspec: scalars and
|
|
lists are wrapped into tensors, and anything that cannot be safely
|
|
converted is dropped.
|
|
"""
|
|
if self.vllm_config.model_config.engine_output_type == "text" and not self._client_multimodal_output_keys():
|
|
return None
|
|
if per_req_payloads is None:
|
|
return None
|
|
wire_payloads: list[dict[str, torch.Tensor] | None] = []
|
|
for payload in per_req_payloads:
|
|
if not payload:
|
|
wire_payloads.append(None)
|
|
else:
|
|
wire_payloads.append(_ensure_tensor_values(payload))
|
|
if all(item is None for item in wire_payloads):
|
|
return None
|
|
return wire_payloads
|
|
|
|
def _get_query_start_loc_cpu(self, *, snapshot: bool) -> Any:
|
|
query_start_loc_cpu = self.query_start_loc.cpu
|
|
if callable(query_start_loc_cpu):
|
|
query_start_loc_cpu = query_start_loc_cpu()
|
|
if not snapshot:
|
|
return query_start_loc_cpu
|
|
if isinstance(query_start_loc_cpu, torch.Tensor):
|
|
return query_start_loc_cpu.detach().cpu().clone()
|
|
if isinstance(query_start_loc_cpu, np.ndarray):
|
|
return query_start_loc_cpu.copy()
|
|
if isinstance(query_start_loc_cpu, list):
|
|
return list(query_start_loc_cpu)
|
|
return query_start_loc_cpu
|
|
|
|
@staticmethod
|
|
def _snapshot_scheduler_output_for_async_omni_output(
|
|
scheduler_output: SchedulerOutput,
|
|
) -> SchedulerOutput:
|
|
updates: dict[str, Any] = {}
|
|
for attr in ("num_scheduled_tokens", "scheduled_spec_decode_tokens"):
|
|
val = getattr(scheduler_output, attr, None)
|
|
if isinstance(val, dict):
|
|
updates[attr] = val.copy()
|
|
elif isinstance(val, list):
|
|
updates[attr] = list(val)
|
|
if not updates:
|
|
return scheduler_output
|
|
try:
|
|
return replace(scheduler_output, **updates)
|
|
except TypeError:
|
|
return scheduler_output
|
|
|
|
@staticmethod
|
|
def _model_omni_flag(model: Any, name: str, default: bool = False) -> bool:
|
|
return bool(getattr(model, name, default)) if model is not None else default
|
|
|
|
def _runner_model_omni_flag(self, name: str, default: bool = False) -> bool:
|
|
return self._model_omni_flag(getattr(self, "model", None), name, default)
|
|
|
|
def _client_multimodal_output_keys(self) -> tuple[str, ...]:
|
|
raw = getattr(
|
|
getattr(self, "model", None),
|
|
"omni_client_multimodal_output_keys",
|
|
(),
|
|
)
|
|
if not isinstance(raw, tuple) or any(not isinstance(key, str) or not key for key in raw):
|
|
raise TypeError("omni_client_multimodal_output_keys must be a tuple of non-empty strings")
|
|
if len(raw) != len(set(raw)):
|
|
raise ValueError("omni_client_multimodal_output_keys must not contain duplicates")
|
|
return raw
|
|
|
|
def _model_omni_pooler_payload_include_hidden(self) -> bool:
|
|
return self._pooler_payload_include_hidden_flag
|
|
|
|
def _should_defer_full_payload_d2h(self) -> bool:
|
|
"""Keep opted-in full payloads on device until request completion."""
|
|
return (
|
|
self.omni_prefix_cache is None
|
|
and self._runner_model_omni_flag("omni_payload_at_request_end")
|
|
and self._should_accumulate_full_payload_output()
|
|
)
|
|
|
|
def _build_omni_step_outputs(
|
|
self,
|
|
pooler_inter: Sequence[dict[str, object] | None] | None,
|
|
pooler_client: Sequence[dict[str, object] | None] | None,
|
|
*,
|
|
defer_full_payload_d2h: bool,
|
|
) -> tuple[
|
|
list[dict[str, torch.Tensor] | None] | None,
|
|
list[dict[str, torch.Tensor] | None] | None,
|
|
]:
|
|
if defer_full_payload_d2h:
|
|
# The connector-side full-payload accumulator owns the snapshots.
|
|
# Nothing from an intermediate stage needs to cross the worker/core
|
|
# boundary on each decode step.
|
|
return None, None
|
|
inter_stage_outputs = self._build_multimodal_outputs(pooler_inter)
|
|
multimodal_outputs = (
|
|
inter_stage_outputs if pooler_client is pooler_inter else self._build_multimodal_outputs(pooler_client)
|
|
)
|
|
return inter_stage_outputs, multimodal_outputs
|
|
|
|
def _should_use_async_omni_output(self) -> bool:
|
|
if not self.use_async_scheduling:
|
|
return False
|
|
if self.speculative_config is not None:
|
|
return False
|
|
|
|
model_config = getattr(self, "model_config", None)
|
|
if model_config is None:
|
|
model_config = getattr(getattr(self, "vllm_config", None), "model_config", None)
|
|
if not bool(getattr(model_config, "async_chunk", False)) and not self._model_omni_flag(
|
|
getattr(self, "model", None), "supports_async_whole_payload"
|
|
):
|
|
return False
|
|
|
|
model = getattr(self, "model", None)
|
|
if not self._model_omni_flag(model, "use_async_omni_output"):
|
|
return False
|
|
if self._model_omni_flag(model, "has_postprocess") and not self._model_omni_flag(
|
|
model, "eager_omni_postprocess_before_async_output"
|
|
):
|
|
return False
|
|
|
|
return True
|
|
|
|
def _build_omni_async_snapshot_payload(
|
|
self,
|
|
*,
|
|
hidden_states: torch.Tensor,
|
|
staged_hidden_states_cpu: torch.Tensor | None,
|
|
multimodal_outputs: Any,
|
|
) -> dict[str, Any]:
|
|
payload: dict[str, Any] = {"multimodal_outputs": multimodal_outputs}
|
|
if self._model_omni_pooler_payload_include_hidden():
|
|
payload["hidden_states"] = hidden_states
|
|
payload["staged_hidden_states_cpu"] = staged_hidden_states_cpu
|
|
return payload
|
|
|
|
def _snapshot_omni_output_tensors_for_async_output(
|
|
self,
|
|
*,
|
|
use_async_omni_output: bool,
|
|
hidden_states: torch.Tensor,
|
|
staged_hidden_states_cpu: torch.Tensor | None,
|
|
multimodal_outputs: Any,
|
|
) -> _OmniOutputTensorSnapshot:
|
|
if not use_async_omni_output:
|
|
return _OmniOutputTensorSnapshot(
|
|
hidden_states=hidden_states,
|
|
staged_hidden_states_cpu=staged_hidden_states_cpu,
|
|
multimodal_outputs=multimodal_outputs,
|
|
)
|
|
|
|
with record_function_or_nullcontext("omni_async_output:snapshot_cpu_payload"):
|
|
async_payload_snapshot = _snapshot_tensor_payload_to_cpu_async(
|
|
self._build_omni_async_snapshot_payload(
|
|
hidden_states=hidden_states,
|
|
staged_hidden_states_cpu=staged_hidden_states_cpu,
|
|
multimodal_outputs=multimodal_outputs,
|
|
),
|
|
copy_stream=self._get_or_create_omni_payload_copy_stream(),
|
|
# NOTE: vLLM v0.24.0's GPUModelRunner no longer exposes a
|
|
# ``self.pin_memory`` attribute (it uses a module-level
|
|
# ``PIN_MEMORY`` constant instead), so the old
|
|
# ``getattr(self, "pin_memory", False)`` silently fell back to
|
|
# False. That allocated the async D2H snapshot destination in
|
|
# *pageable* host memory, which turns ``copy_(non_blocking=True)``
|
|
# into a fully synchronous, stream-stalling copy (~240 ms/step
|
|
# on the 17.5k-token Thinker prefill). Resolve pinning from the
|
|
# platform helper so the copy is a true async cudaMemcpyAsync.
|
|
pin_memory=is_pin_memory_available(),
|
|
)
|
|
|
|
payload = async_payload_snapshot.payload
|
|
hidden_states_snapshot = payload.get("hidden_states")
|
|
if hidden_states_snapshot is None:
|
|
# Models that omit hidden from the async snapshot only need
|
|
# multimodal payloads (for example, talker codes.audio).
|
|
hidden_states_snapshot = hidden_states[:0]
|
|
|
|
return _OmniOutputTensorSnapshot(
|
|
hidden_states=hidden_states_snapshot,
|
|
staged_hidden_states_cpu=payload.get("staged_hidden_states_cpu"),
|
|
multimodal_outputs=payload["multimodal_outputs"],
|
|
async_payload=async_payload_snapshot,
|
|
)
|
|
|
|
def _maybe_run_eager_omni_postprocess_before_async_output(
|
|
self,
|
|
*,
|
|
hidden_states: torch.Tensor,
|
|
multimodal_outputs: Any,
|
|
num_scheduled_tokens_np: np.ndarray,
|
|
scheduler_output: SchedulerOutput,
|
|
req_ids_output_copy: list[str],
|
|
query_start_loc_cpu: Any,
|
|
) -> bool:
|
|
"""Apply model postprocess on live GPU tensors before async payload D2H."""
|
|
model = getattr(self, "model", None)
|
|
if not self._model_omni_flag(model, "has_postprocess"):
|
|
return False
|
|
if not self._model_omni_flag(model, "eager_omni_postprocess_before_async_output"):
|
|
return False
|
|
|
|
_, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
|
|
if not downstream_req_ids:
|
|
return False
|
|
|
|
with record_function_or_nullcontext("omni_output_builder:eager_postprocess"):
|
|
self._process_additional_information_updates(
|
|
hidden_states,
|
|
multimodal_outputs,
|
|
num_scheduled_tokens_np,
|
|
scheduler_output,
|
|
None,
|
|
None,
|
|
req_ids_filter=set(downstream_req_ids),
|
|
req_ids=req_ids_output_copy,
|
|
query_start_loc_cpu=query_start_loc_cpu,
|
|
)
|
|
return True
|
|
|
|
def _get_or_create_omni_payload_copy_stream(self) -> torch.cuda.Stream:
|
|
stream = getattr(self, "_omni_payload_copy_stream", None)
|
|
if stream is None:
|
|
stream = torch.cuda.Stream()
|
|
self._omni_payload_copy_stream = stream
|
|
return stream
|
|
|
|
def _build_omni_model_runner_output_from_snapshot(
|
|
self,
|
|
*,
|
|
scheduler_output: SchedulerOutput,
|
|
hidden_states: torch.Tensor,
|
|
staged_hidden_states_cpu: torch.Tensor | None,
|
|
multimodal_outputs: Any,
|
|
req_ids_output_copy: list[str],
|
|
req_id_to_index_output_copy: dict[str, int],
|
|
valid_sampled_token_ids: list[list[int]],
|
|
logprobs_lists: Any,
|
|
prompt_logprobs_dict: dict[str, Any],
|
|
num_nans_in_logits: Any,
|
|
kv_connector_output: Any,
|
|
ec_connector_output: Any,
|
|
cudagraph_stats: Any,
|
|
kv_extracted_req_ids: list[str] | None,
|
|
num_scheduled_tokens_np: np.ndarray,
|
|
query_start_loc_cpu: Any,
|
|
postprocess_already_applied: bool = False,
|
|
prefix_cache_step_id: int | None = None,
|
|
) -> OmniModelRunnerOutput:
|
|
combined_hidden_states = None
|
|
combined_multimodal_outputs = None
|
|
|
|
engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
|
|
downstream_req_ids, sparse_mm_index, audio_sparse_output = resolve_sparse_mm_routing(
|
|
engine_output_type=engine_output_type,
|
|
req_ids_output_copy=req_ids_output_copy,
|
|
downstream_req_ids=downstream_req_ids,
|
|
multimodal_outputs=multimodal_outputs,
|
|
)
|
|
|
|
needs_pooler_payload = len(downstream_req_ids) > 0
|
|
downstream_req_id_set = set(downstream_req_ids)
|
|
defer_full_payload_d2h = needs_pooler_payload and self._should_defer_full_payload_d2h()
|
|
hidden_states_cpu = None
|
|
req_hidden_states_cpu: dict[str, torch.Tensor] | None = None
|
|
include_hidden_payload = self._model_omni_pooler_payload_include_hidden()
|
|
needs_scheduled_hidden_payload = (
|
|
include_hidden_payload
|
|
and needs_pooler_payload
|
|
and (self.omni_prefix_cache is None or not self._model_needs_full_prefix_hidden_states())
|
|
)
|
|
if not needs_pooler_payload and prefix_cache_step_id is not None:
|
|
# No consumer for this step's merge: consume the step context by
|
|
# id (exactly-once contract). The cache write still lands.
|
|
assert self.omni_prefix_cache is not None
|
|
self.omni_prefix_cache.discard_step(prefix_cache_step_id)
|
|
prefix_cache_step_id = None
|
|
if self.omni_prefix_cache is None and needs_scheduled_hidden_payload and not audio_sparse_output:
|
|
num_valid_tokens = min(
|
|
int(scheduler_output.total_num_scheduled_tokens),
|
|
int(hidden_states.shape[0]),
|
|
)
|
|
if len(downstream_req_ids) == len(req_ids_output_copy):
|
|
with record_function_or_nullcontext("omni_output_builder:hidden_d2h/scheduled"):
|
|
hidden_states_cpu = _to_cpu_contiguous(hidden_states[:num_valid_tokens])
|
|
else:
|
|
req_hidden_states_cpu = {}
|
|
with record_function_or_nullcontext("omni_output_builder:hidden_d2h/per_request"):
|
|
for rid in downstream_req_ids:
|
|
idx = req_id_to_index_output_copy[rid]
|
|
start = int(query_start_loc_cpu[idx])
|
|
sched = int(num_scheduled_tokens_np[idx])
|
|
end = start + sched
|
|
req_hidden_states_cpu[rid] = _to_cpu_contiguous(hidden_states[start:end])
|
|
|
|
# NOTE: pooler_output here is used only for the full-payload accumulation
|
|
# path (accumulate_full_payload_output) and is NOT passed on the wire via
|
|
# OmniModelRunnerOutput.pooler_output (which is set to None below).
|
|
# The actual multimodal wire transport uses multimodal_outputs instead.
|
|
pooler_output: list[dict[str, object]] | None = None
|
|
# Token-only steps have no pooler payload to construct. Keep the full
|
|
# path for cached payloads and model postprocess, even without MM input.
|
|
has_mm_payload = multimodal_outputs is not None and (
|
|
not isinstance(multimodal_outputs, Mapping) or len(multimodal_outputs) > 0
|
|
)
|
|
needs_payload_work = (
|
|
include_hidden_payload
|
|
or has_mm_payload
|
|
or self.omni_prefix_cache is not None
|
|
or (not postprocess_already_applied and self._runner_model_omni_flag("has_postprocess"))
|
|
)
|
|
if needs_pooler_payload and needs_payload_work:
|
|
hidden_seq_len = int(hidden_states.shape[0])
|
|
scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens)
|
|
mm_cpu = None
|
|
if self.omni_prefix_cache is not None:
|
|
(
|
|
hidden_states_cpu,
|
|
combined_hidden_states,
|
|
combined_multimodal_outputs,
|
|
) = self._prepare_prefix_cache_pooler_payload_sources(
|
|
staged_hidden_states_cpu=staged_hidden_states_cpu,
|
|
needs_scheduled_hidden_payload=needs_scheduled_hidden_payload,
|
|
step_id=prefix_cache_step_id,
|
|
req_ids=req_ids_output_copy,
|
|
)
|
|
if combined_multimodal_outputs is None:
|
|
flat_mm = flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs
|
|
if defer_full_payload_d2h:
|
|
with record_function_or_nullcontext("omni_output_builder:snapshot_mm_payload"):
|
|
mm_cpu = snapshot_mm_payload(flat_mm)
|
|
else:
|
|
with record_function_or_nullcontext("omni_output_builder:build_mm_cpu"):
|
|
mm_cpu = build_mm_cpu(flat_mm)
|
|
|
|
with record_function_or_nullcontext("omni_output_builder:process_additional_information"):
|
|
if not postprocess_already_applied:
|
|
self._process_additional_information_updates(
|
|
hidden_states,
|
|
multimodal_outputs,
|
|
num_scheduled_tokens_np,
|
|
scheduler_output,
|
|
combined_hidden_states,
|
|
combined_multimodal_outputs,
|
|
req_ids_filter=downstream_req_id_set,
|
|
req_ids=req_ids_output_copy,
|
|
query_start_loc_cpu=query_start_loc_cpu,
|
|
)
|
|
|
|
pooler_output = []
|
|
with record_function_or_nullcontext("omni_output_builder:build_pooler_payloads"):
|
|
for rid in req_ids_output_copy:
|
|
if rid not in downstream_req_id_set:
|
|
pooler_output.append({})
|
|
continue
|
|
idx = req_id_to_index_output_copy[rid]
|
|
start = int(query_start_loc_cpu[idx])
|
|
sched = int(num_scheduled_tokens_np[idx])
|
|
end = start + sched
|
|
payload = self._build_omni_pooler_payload(
|
|
rid=rid,
|
|
idx=idx,
|
|
start=start,
|
|
end=end,
|
|
hidden_states_cpu=hidden_states_cpu,
|
|
req_hidden_states_cpu=req_hidden_states_cpu,
|
|
combined_hidden_states=combined_hidden_states,
|
|
combined_multimodal_outputs=combined_multimodal_outputs,
|
|
mm_cpu=mm_cpu,
|
|
audio_sparse_output=audio_sparse_output,
|
|
sparse_mm_index=sparse_mm_index,
|
|
hidden_seq_len=hidden_seq_len,
|
|
scheduled_seq_len=scheduled_seq_len,
|
|
)
|
|
pooler_output.append(flatten_payload(payload))
|
|
|
|
pooler_inter: Sequence[dict[str, object] | None] | None
|
|
pooler_client: Sequence[dict[str, object] | None] | None
|
|
pooler_output = pooler_output or []
|
|
if self._async_chunk and stage_sends_async_output(self.model_config):
|
|
pooler_inter, pooler_client = partition_payload_list(pooler_output)
|
|
else:
|
|
# Connector-less stages expose the same payload through the
|
|
# orchestrator bridge; non-async stages preserve legacy behavior.
|
|
pooler_inter, pooler_client = pooler_output, pooler_output
|
|
client_output_keys = self._client_multimodal_output_keys()
|
|
if client_output_keys:
|
|
allowed = frozenset(client_output_keys)
|
|
pooler_client = [
|
|
{key: value for key, value in payload.items() if key in allowed} for payload in pooler_output
|
|
]
|
|
|
|
if pooler_inter and self._should_accumulate_full_payload_output():
|
|
with record_function_or_nullcontext("omni_output_builder:accumulate_full_payload_output"):
|
|
for i, rid in enumerate(req_ids_output_copy):
|
|
req_state = self.requests.get(rid)
|
|
if req_state is not None and pooler_inter[i]:
|
|
self.accumulate_full_payload_output(rid, pooler_inter[i], req_state)
|
|
|
|
with record_function_or_nullcontext("omni_output_builder:build_multimodal_outputs"):
|
|
inter_stage_outputs, multimodal_outputs = self._build_omni_step_outputs(
|
|
pooler_inter,
|
|
pooler_client,
|
|
defer_full_payload_d2h=defer_full_payload_d2h,
|
|
)
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
|
|
output = OmniModelRunnerOutput(
|
|
req_ids=req_ids_output_copy,
|
|
req_id_to_index=req_id_to_index_output_copy,
|
|
sampled_token_ids=valid_sampled_token_ids,
|
|
logprobs=logprobs_lists,
|
|
prompt_logprobs_dict=prompt_logprobs_dict,
|
|
pooler_output=None,
|
|
multimodal_outputs=multimodal_outputs,
|
|
inter_stage_outputs=inter_stage_outputs,
|
|
kv_connector_output=kv_connector_output,
|
|
ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
|
|
num_nans_in_logits=num_nans_in_logits,
|
|
cudagraph_stats=cudagraph_stats,
|
|
)
|
|
output.kv_extracted_req_ids = kv_extracted_req_ids
|
|
return output
|
|
|
|
@torch.inference_mode()
|
|
def sample_tokens(
|
|
self,
|
|
grammar_output: GrammarOutput | None,
|
|
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
|
|
kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None)
|
|
self.kv_extracted_req_ids = None
|
|
|
|
if self.execute_model_state is None:
|
|
kv_connector_output = self.kv_connector_output
|
|
self.kv_connector_output = None
|
|
# receive sampled token ids from the last PP rank.
|
|
if self.use_async_scheduling and not get_pp_group().is_last_rank:
|
|
self._pp_receive_prev_sampled_token_ids_to_input_batch()
|
|
# In case of PP with kv transfer, we need to pass through the
|
|
# kv_connector_output
|
|
return self.attach_omni_connector_output(
|
|
OmniModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
|
|
)
|
|
|
|
# Unpack ephemeral state.
|
|
(
|
|
scheduler_output,
|
|
logits,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
hidden_states,
|
|
staged_hidden_states_cpu,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
ec_connector_output,
|
|
cudagraph_stats,
|
|
multimodal_outputs,
|
|
slot_mappings, # OMNI: unpack slot_mappings for drafter
|
|
prefix_cache_step_id,
|
|
) = self.execute_model_state
|
|
self.execute_model_state = None
|
|
|
|
# Apply structured output bitmasks if present.
|
|
if grammar_output is not None:
|
|
apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits)
|
|
|
|
# Correct padding values of prompt_token_ids to match the logits vocabulary size.
|
|
if logits is not None and not self.input_batch.sampling_metadata.no_penalties:
|
|
smd = self.input_batch.sampling_metadata
|
|
if smd.prompt_token_ids is not None:
|
|
logits_vocab = logits.shape[-1]
|
|
if self.input_batch.vocab_size > logits_vocab:
|
|
smd.prompt_token_ids = clamp_prompt_ids_to_penalty_padding(smd.prompt_token_ids, logits_vocab)
|
|
|
|
# Drop min-tokens stop ids the head cannot emit (e.g. the text
|
|
# tokenizer EOS folded into all_stop_token_ids on a narrow codec
|
|
# talker head); they would index_put_ out of bounds (#4962).
|
|
if logits is not None:
|
|
sanitize_min_tokens_stop_ids(
|
|
self.input_batch.sampling_metadata.logitsprocs,
|
|
logits.shape[-1],
|
|
)
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: sample"):
|
|
sampler_output = self._sample(logits, spec_decode_metadata)
|
|
|
|
self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output)
|
|
if self.use_async_scheduling:
|
|
pp = get_pp_group()
|
|
# Earlier PP ranks posted the matching async receive in the
|
|
# connector-only branch above (inherited from vLLM). Without this
|
|
# send their next `_prepare_input_ids` waits forever.
|
|
# For torchrun external_launcher PP mode with broadcast_pp_output=True,
|
|
# PP outputs have been broadcast to all ranks at logits computation,
|
|
# so the sampled token ids must not be sent again.
|
|
if not self.broadcast_pp_output and pp.world_size > 1 and pp.is_last_rank:
|
|
self._pp_broadcast_prev_sampled_token_ids(sampler_output.sampled_token_ids)
|
|
|
|
self._draft_token_ids = None
|
|
self._draft_token_req_ids = None
|
|
self.valid_sampled_token_count_gpu = None
|
|
self.input_batch.prev_sampled_token_ids = None
|
|
|
|
def propose_draft_token_ids(sampled_token_ids):
|
|
assert spec_decode_common_attn_metadata is not None
|
|
with record_function_or_nullcontext("gpu_model_runner: draft"):
|
|
self._draft_token_ids = self.propose_draft_token_ids(
|
|
scheduler_output,
|
|
sampled_token_ids,
|
|
self.input_batch.sampling_metadata,
|
|
hidden_states,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
slot_mappings, # OMNI: pass slot_mappings to drafter (upstream v1 API)
|
|
)
|
|
self._copy_draft_token_ids_to_cpu(scheduler_output)
|
|
|
|
spec_config = self.speculative_config
|
|
propose_drafts_after_bookkeeping = False
|
|
if spec_config is not None:
|
|
input_fits_in_drafter = self._input_fits_in_drafter(spec_decode_common_attn_metadata)
|
|
use_gpu_toks = (
|
|
spec_config.use_eagle() or spec_config.uses_draft_model() or spec_config.uses_extract_hidden_states()
|
|
) and not spec_config.disable_padded_drafter_batch
|
|
if use_gpu_toks:
|
|
assert isinstance(
|
|
self.drafter,
|
|
EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer,
|
|
)
|
|
sampled_token_ids = sampler_output.sampled_token_ids
|
|
if input_fits_in_drafter:
|
|
propose_draft_token_ids(sampled_token_ids)
|
|
elif self.valid_sampled_token_count_event is not None:
|
|
assert spec_decode_common_attn_metadata is not None
|
|
next_token_ids, valid_sampled_tokens_count = self.drafter.prepare_next_token_ids_padded(
|
|
sampled_token_ids,
|
|
self.requests,
|
|
self.input_batch,
|
|
self.discard_request_mask.gpu,
|
|
)
|
|
self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
|
|
# Since we couldn't run the drafter,
|
|
# just use zeros for the draft tokens.
|
|
self._draft_token_ids = torch.zeros(1, device=self.device, dtype=torch.int32).expand(
|
|
len(self.input_batch.req_ids), self.num_spec_tokens
|
|
)
|
|
self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
|
|
else:
|
|
propose_drafts_after_bookkeeping = input_fits_in_drafter
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
|
|
(
|
|
num_nans_in_logits,
|
|
num_nans,
|
|
logprobs_lists,
|
|
valid_sampled_token_ids,
|
|
prompt_logprobs_dict,
|
|
req_ids_output_copy,
|
|
req_id_to_index_output_copy,
|
|
invalid_req_indices,
|
|
) = self._bookkeeping_sync(
|
|
scheduler_output,
|
|
sampler_output,
|
|
logits,
|
|
hidden_states,
|
|
scheduler_output.total_num_scheduled_tokens,
|
|
)
|
|
|
|
post_sample_mm = getattr(getattr(self, "model", None), "post_sample_multimodal_outputs", None)
|
|
if callable(post_sample_mm):
|
|
multimodal_outputs = post_sample_mm(
|
|
req_ids=req_ids_output_copy,
|
|
invalid_req_indices=invalid_req_indices,
|
|
multimodal_outputs=multimodal_outputs,
|
|
)
|
|
|
|
multimodal_outputs = self._run_post_sample_talker_mtp(
|
|
req_ids=req_ids_output_copy,
|
|
valid_sampled_token_ids=valid_sampled_token_ids,
|
|
sampled_token_ids=sampler_output.sampled_token_ids,
|
|
invalid_req_indices=invalid_req_indices,
|
|
sample_hidden_states=sample_hidden_states,
|
|
multimodal_outputs=multimodal_outputs,
|
|
)
|
|
|
|
if propose_drafts_after_bookkeeping:
|
|
# ngram and other speculative decoding methods use the sampled
|
|
# tokens on the CPU, so they are run after bookkeeping.
|
|
propose_draft_token_ids(valid_sampled_token_ids)
|
|
|
|
# Finalize KV connector (wait_for_save + clear metadata) after
|
|
# draft model runs. Deferred from target model forward to allow
|
|
# draft model to also save its KV cache.
|
|
if self.speculative_config is not None:
|
|
self.finalize_kv_connector()
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: eplb"):
|
|
self.eplb_step()
|
|
|
|
# kv_connector_output may be modified during drafting
|
|
kv_connector_output = self.kv_connector_output
|
|
self.kv_connector_output = None
|
|
|
|
use_async_omni_output = self._should_use_async_omni_output()
|
|
num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
|
|
if num_scheduled_tokens_np is None:
|
|
num_scheduled_tokens_np = np.array(
|
|
[scheduler_output.num_scheduled_tokens[rid] for rid in req_ids_output_copy],
|
|
dtype=np.int32,
|
|
)
|
|
else:
|
|
num_scheduled_tokens_np = np.asarray(num_scheduled_tokens_np, dtype=np.int32)
|
|
if use_async_omni_output:
|
|
num_scheduled_tokens_np = num_scheduled_tokens_np.copy()
|
|
|
|
# Inline materialization finishes before the next execute_model can
|
|
# mutate step metadata. Only the background builder needs snapshots.
|
|
query_start_loc_cpu = self._get_query_start_loc_cpu(snapshot=use_async_omni_output)
|
|
scheduler_output_snapshot = scheduler_output
|
|
valid_sampled_token_ids_snapshot = valid_sampled_token_ids
|
|
logprobs_lists_snapshot = logprobs_lists
|
|
prompt_logprobs_dict_snapshot = prompt_logprobs_dict
|
|
num_nans_in_logits_snapshot = num_nans_in_logits
|
|
if use_async_omni_output:
|
|
scheduler_output_snapshot = self._snapshot_scheduler_output_for_async_omni_output(scheduler_output)
|
|
valid_sampled_token_ids_snapshot = [list(token_ids) for token_ids in valid_sampled_token_ids]
|
|
logprobs_lists_snapshot = copy(logprobs_lists) if logprobs_lists is not None else None
|
|
prompt_logprobs_dict_snapshot = dict(prompt_logprobs_dict) if prompt_logprobs_dict is not None else {}
|
|
num_nans_in_logits_snapshot = (
|
|
dict(num_nans_in_logits) if isinstance(num_nans_in_logits, dict) else num_nans_in_logits
|
|
)
|
|
# _bookkeeping_sync already detached these from the mutable input batch.
|
|
req_ids_output_snapshot = req_ids_output_copy
|
|
req_id_to_index_output_snapshot = req_id_to_index_output_copy
|
|
omni_postprocess_already_applied = False
|
|
if use_async_omni_output:
|
|
omni_postprocess_already_applied = self._maybe_run_eager_omni_postprocess_before_async_output(
|
|
hidden_states=hidden_states,
|
|
multimodal_outputs=multimodal_outputs,
|
|
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
|
scheduler_output=scheduler_output,
|
|
req_ids_output_copy=req_ids_output_copy,
|
|
query_start_loc_cpu=query_start_loc_cpu,
|
|
)
|
|
output_tensor_snapshot = self._snapshot_omni_output_tensors_for_async_output(
|
|
use_async_omni_output=use_async_omni_output,
|
|
hidden_states=hidden_states,
|
|
staged_hidden_states_cpu=staged_hidden_states_cpu,
|
|
multimodal_outputs=multimodal_outputs,
|
|
)
|
|
|
|
# Runs a TP collective, so it must stay on the main thread: the builder
|
|
# below runs on the async output thread, which would race execute_model().
|
|
with record_function_or_nullcontext("omni_async_output:get_omni_connector_output"):
|
|
omni_connector_output = self.get_omni_connector_output()
|
|
|
|
def output_builder() -> OmniModelRunnerOutput:
|
|
if output_tensor_snapshot.async_payload is not None:
|
|
with record_function_or_nullcontext("omni_async_output:wait_cpu_payload"):
|
|
output_tensor_snapshot.async_payload.wait()
|
|
mm_snapshot = output_tensor_snapshot.multimodal_outputs
|
|
finalize_snapshot = getattr(
|
|
getattr(self, "model", None), "finalize_multimodal_outputs_from_cpu_snapshot", None
|
|
)
|
|
if callable(finalize_snapshot):
|
|
mm_snapshot = finalize_snapshot(mm_snapshot)
|
|
with record_function_or_nullcontext("omni_output_builder:total"):
|
|
output = self._build_omni_model_runner_output_from_snapshot(
|
|
scheduler_output=scheduler_output_snapshot,
|
|
hidden_states=output_tensor_snapshot.hidden_states,
|
|
staged_hidden_states_cpu=output_tensor_snapshot.staged_hidden_states_cpu,
|
|
multimodal_outputs=mm_snapshot,
|
|
req_ids_output_copy=req_ids_output_snapshot,
|
|
req_id_to_index_output_copy=req_id_to_index_output_snapshot,
|
|
valid_sampled_token_ids=valid_sampled_token_ids_snapshot,
|
|
logprobs_lists=logprobs_lists_snapshot,
|
|
prompt_logprobs_dict=prompt_logprobs_dict_snapshot,
|
|
num_nans_in_logits=num_nans_in_logits_snapshot,
|
|
kv_connector_output=kv_connector_output,
|
|
ec_connector_output=ec_connector_output,
|
|
cudagraph_stats=cudagraph_stats,
|
|
kv_extracted_req_ids=kv_extracted_req_ids,
|
|
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
|
query_start_loc_cpu=query_start_loc_cpu,
|
|
postprocess_already_applied=omni_postprocess_already_applied,
|
|
prefix_cache_step_id=prefix_cache_step_id,
|
|
)
|
|
output.omni_connector_output = omni_connector_output
|
|
return output
|
|
|
|
if not use_async_omni_output:
|
|
output = output_builder()
|
|
|
|
if not self.use_async_scheduling:
|
|
return output
|
|
with record_function_or_nullcontext("gpu_model_runner: AsyncGPUModelRunnerOutput"):
|
|
async_output_cls = OmniAsyncGPUModelRunnerOutput if use_async_omni_output else AsyncGPUModelRunnerOutput
|
|
async_output_kwargs = dict(
|
|
sampled_token_ids=sampler_output.sampled_token_ids,
|
|
logprobs_tensors=sampler_output.logprobs_tensors,
|
|
invalid_req_indices=invalid_req_indices,
|
|
async_output_copy_stream=self.async_output_copy_stream,
|
|
vocab_size=self.input_batch.vocab_size,
|
|
num_nans=num_nans,
|
|
)
|
|
if use_async_omni_output:
|
|
async_output = async_output_cls(
|
|
model_runner_output_builder=output_builder,
|
|
cuda_device=self.device,
|
|
**async_output_kwargs,
|
|
)
|
|
else:
|
|
async_output = async_output_cls(
|
|
model_runner_output=output,
|
|
**async_output_kwargs,
|
|
)
|
|
with record_function_or_nullcontext("gpu_model_runner: set_async_sampled_token_ids"):
|
|
# Save ref of sampled_token_ids CPU tensor if the batch contains
|
|
# any requests with sampling params that require output ids.
|
|
self.input_batch.set_async_sampled_token_ids(
|
|
async_output.sampled_token_ids_cpu,
|
|
async_output.async_copy_ready_event,
|
|
)
|
|
|
|
return async_output
|
|
|
|
def _resolve_global_request_id(self, req_id: str) -> str:
|
|
"""Resolve global request ID from request state."""
|
|
req_state = self.requests.get(req_id)
|
|
if not req_state:
|
|
return req_id
|
|
|
|
add_info = self.model_intermediate_buffer.get(req_id, {})
|
|
global_id = add_info.get("global_request_id")
|
|
if global_id:
|
|
if isinstance(global_id, list) and global_id:
|
|
global_id = global_id[0]
|
|
if isinstance(global_id, bytes):
|
|
return global_id.decode("utf-8")
|
|
return str(global_id)
|
|
return req_id
|