2203 lines
111 KiB
Python
2203 lines
111 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
||
import contextlib
|
||
import inspect
|
||
from collections.abc import Callable
|
||
from dataclasses import replace
|
||
from typing import TYPE_CHECKING, Any, cast
|
||
|
||
import numpy as np
|
||
import torch
|
||
from vllm.config import CUDAGraphMode
|
||
from vllm.distributed.parallel_state import get_pp_group
|
||
from vllm.forward_context import set_forward_context
|
||
from vllm.logger import init_logger
|
||
from vllm.model_executor.models.interfaces import supports_mrope
|
||
from vllm.model_executor.models.interfaces_base import VllmModelForPooling
|
||
from vllm.sampling_params import SamplingType
|
||
from vllm.sequence import IntermediateTensors
|
||
from vllm.tracing import instrument
|
||
from vllm.utils.import_utils import LazyLoader
|
||
from vllm.utils.math_utils import cdiv
|
||
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.spec_decode.ngram_proposer_gpu import (
|
||
update_ngram_gpu_tensors_incremental,
|
||
update_scheduler_for_invalid_drafts,
|
||
)
|
||
from vllm.v1.worker.gpu_input_batch import CachedRequestState
|
||
from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata
|
||
from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
|
||
|
||
from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
|
||
from vllm_omni.engine.init_timeline import timed_phase, worker_phase
|
||
from vllm_omni.engine.serialization import deserialize_additional_information
|
||
from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding
|
||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||
from vllm_omni.platforms import current_omni_platform
|
||
|
||
if TYPE_CHECKING:
|
||
from vllm.v1.core.sched.output import SchedulerOutput
|
||
from vllm.v1.outputs import RoutedExpertsLists
|
||
else:
|
||
xgr = LazyLoader("xgr", globals(), "xgrammar")
|
||
xgr_torch_compile = LazyLoader(
|
||
"xgr_torch_compile",
|
||
globals(),
|
||
"xgrammar.kernels.apply_token_bitmask_inplace_torch_compile",
|
||
)
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
def _filter_mrope_kwargs_for_model(model: object, kwargs: dict[str, Any]) -> dict[str, Any]:
|
||
"""Return only M-RoPE kwargs accepted by the model implementation."""
|
||
method = getattr(model, "get_mrope_input_positions")
|
||
try:
|
||
signature = inspect.signature(method)
|
||
except (TypeError, ValueError):
|
||
return kwargs
|
||
|
||
if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()):
|
||
return kwargs
|
||
|
||
accepted = {
|
||
name
|
||
for name, param in signature.parameters.items()
|
||
if name not in {"self", "input_tokens"}
|
||
and param.kind
|
||
in {
|
||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||
inspect.Parameter.KEYWORD_ONLY,
|
||
}
|
||
}
|
||
return {key: value for key, value in kwargs.items() if key in accepted}
|
||
|
||
|
||
class OmniGPUModelRunner(GPUModelRunner):
|
||
intermediate_tensors: IntermediateTensors | None
|
||
|
||
def __init__(self, *args, **kwargs):
|
||
super().__init__(*args, **kwargs)
|
||
self.model_intermediate_buffer: dict[str, dict[str, Any]] = {}
|
||
self._omni_num_scheduled_tokens_np: np.ndarray | None = None
|
||
self._omni_last_model_output: object | None = None
|
||
# The Omni tensor prefix cache will be allocated
|
||
# when we initialize the metadata builders if enabled
|
||
self.omni_prefix_cache = None
|
||
self._sampled_token_ids_cpu_override = None
|
||
self._omni_query_start_loc_model_kwarg = False
|
||
|
||
def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
|
||
override_fn = self._sampled_token_ids_cpu_override
|
||
if callable(override_fn):
|
||
sampled = override_fn(sampled_token_ids)
|
||
if sampled is not None:
|
||
return sampled
|
||
return super()._to_list(sampled_token_ids)
|
||
|
||
def _omni_routed_experts_d2h(self, scheduler_output) -> None:
|
||
"""Issue routed-experts D2H copy matching upstream GPUModelRunner pattern.
|
||
|
||
Upstream does this inline in ``execute_model``:
|
||
buf = self.routed_experts_capturer.get_device_buffer()
|
||
total = scheduler_output.total_num_scheduled_tokens
|
||
self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True)
|
||
self.routed_experts_slot_mapping_cpu[:total].copy_(
|
||
self.routed_experts_slot_mapping_device[:total], non_blocking=True)
|
||
"""
|
||
if not self.routed_experts_initialized:
|
||
return
|
||
buf = self.routed_experts_capturer.get_device_buffer()
|
||
total = scheduler_output.total_num_scheduled_tokens
|
||
self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True)
|
||
if hasattr(self, "routed_experts_slot_mapping_device"):
|
||
self.routed_experts_slot_mapping_cpu[:total].copy_(
|
||
self.routed_experts_slot_mapping_device[:total],
|
||
non_blocking=True,
|
||
)
|
||
|
||
def _omni_extract_routed_experts(self, scheduler_output) -> "RoutedExpertsLists | None":
|
||
"""Extract routed experts matching upstream GPUModelRunner pattern.
|
||
|
||
Upstream (sync path, sample_tokens):
|
||
total = scheduler_output.total_num_scheduled_tokens
|
||
output.routed_experts = RoutedExpertsLists(
|
||
routing_data=self.routed_experts_cpu[:total].numpy(),
|
||
slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(),
|
||
)
|
||
|
||
Returns RoutedExpertsLists (batch-level, with slot_mapping) so that
|
||
downstream schedulers can use slot_mapping to map back to requests.
|
||
"""
|
||
from vllm.v1.outputs import RoutedExpertsLists
|
||
|
||
if not self.routed_experts_initialized:
|
||
return None
|
||
total = scheduler_output.total_num_scheduled_tokens
|
||
if total <= 0:
|
||
return None
|
||
return RoutedExpertsLists(
|
||
routing_data=self.routed_experts_cpu[:total].numpy(),
|
||
slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(),
|
||
)
|
||
|
||
def initialize_metadata_builders(self, kv_cache_config, kernel_block_sizes):
|
||
"""Initialize metadata builders and keep FA3 graph metadata buffers sized.
|
||
|
||
FlashAttentionMetadataBuilder can pre-allocate scheduler_metadata for
|
||
only max_num_seqs + 1 entries while FA3 with split scheduling may need
|
||
max_num_seqs * max_num_splits + 1 entries during CUDA graph capture.
|
||
This runner is shared across Omni models, so preserve the existing
|
||
workaround for non-Higgs models that still use FA3.
|
||
"""
|
||
super().initialize_metadata_builders(kv_cache_config, kernel_block_sizes)
|
||
|
||
for kv_cache_group in self.attn_groups:
|
||
for attn_group in kv_cache_group:
|
||
for builder in attn_group.metadata_builders:
|
||
sm = getattr(builder, "scheduler_metadata", None)
|
||
max_num_splits = getattr(builder, "max_num_splits", 0)
|
||
if sm is not None and max_num_splits > 1:
|
||
required = self.scheduler_config.max_num_seqs * max_num_splits + 1
|
||
if sm.shape[0] < required:
|
||
builder.scheduler_metadata = torch.zeros(
|
||
required,
|
||
dtype=sm.dtype,
|
||
device=sm.device,
|
||
)
|
||
|
||
# Initialize the wrapper for both multimodal output tensors
|
||
# and for hidden states to be passed between stages
|
||
if self.cache_config.enable_prefix_caching:
|
||
self.omni_prefix_cache = OmniTensorPrefixCache(
|
||
num_blocks=kv_cache_config.num_blocks,
|
||
block_size=self.cache_config.block_size,
|
||
hidden_size=self.model_config.get_hidden_size(),
|
||
hs_dtype=self.dtype,
|
||
)
|
||
|
||
@instrument(span_name="Loading (GPU)")
|
||
@timed_phase("load_model")
|
||
def load_model(self, *args, **kwargs) -> None:
|
||
with worker_phase("load_weights"):
|
||
super().load_model(*args, **kwargs)
|
||
model = getattr(self, "model", None)
|
||
override_fn = None
|
||
if bool(getattr(model, "supports_sampled_token_ids_cpu_override", False)):
|
||
candidate = getattr(model, "consume_sampled_token_ids_cpu_override", None)
|
||
if callable(candidate):
|
||
override_fn = candidate
|
||
self._sampled_token_ids_cpu_override = override_fn
|
||
self._omni_query_start_loc_model_kwarg = bool(getattr(model, "supports_omni_query_start_loc", False))
|
||
self._maybe_enable_output_token_ids_for_model_sampler()
|
||
self._init_talker_mtp()
|
||
self._prewarm_attention_capture_workspaces()
|
||
|
||
def _maybe_enable_output_token_ids_for_model_sampler(self) -> None:
|
||
if getattr(self.model, "logitsprocs_need_output_token_ids", False):
|
||
self.input_batch.logitsprocs_need_output_token_ids = True
|
||
|
||
@timed_phase("talker_mtp_init")
|
||
def _init_talker_mtp(self) -> None:
|
||
# TODO move this model specific logic to a separate class
|
||
# TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
|
||
self.has_talker_mtp = False
|
||
talker_mtp = getattr(self.model, "talker_mtp", None)
|
||
if talker_mtp is None:
|
||
return
|
||
self.talker_mtp = talker_mtp # type: ignore[assignment]
|
||
self.has_talker_mtp = True
|
||
cudagraph_mode = self.compilation_config.cudagraph_mode
|
||
assert cudagraph_mode is not None
|
||
has_separate_talker = getattr(self.model, "talker", None) is not None
|
||
talker_mtp_graph_safe = getattr(self.model, "talker_mtp_graph_safe", None)
|
||
# Preserve the legacy separate-talker default only when the model has
|
||
# not declared graph safety. An explicit False must disable wrapping.
|
||
if talker_mtp_graph_safe is None:
|
||
talker_mtp_graph_safe = has_separate_talker
|
||
if cudagraph_mode.has_full_cudagraphs() and talker_mtp_graph_safe:
|
||
graph_wrapper_cls = current_omni_platform.get_graph_wrapper_cls()
|
||
self.talker_mtp = graph_wrapper_cls(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
|
||
# TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.
|
||
hidden_size = int(
|
||
getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size")
|
||
)
|
||
# max_cudagraph_capture_size is None when graphs are disabled (e.g. the CPU platform).
|
||
max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size or 0)
|
||
self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
|
||
self.talker_mtp_inputs_embeds = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
|
||
self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
|
||
self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
|
||
|
||
@timed_phase("attn_prewarm")
|
||
def _prewarm_attention_capture_workspaces(self) -> None:
|
||
capture_sizes = getattr(self.compilation_config, "cudagraph_capture_sizes", None)
|
||
if not capture_sizes:
|
||
return
|
||
from vllm_omni.attention.fish_kvcache_backend import prewarm_fish_kvcache_attn_capture_workspaces
|
||
|
||
prewarm_fish_kvcache_attn_capture_workspaces(
|
||
model_config=self.model_config,
|
||
device=self.device,
|
||
dtype=self.dtype,
|
||
capture_sizes=capture_sizes,
|
||
)
|
||
|
||
def _maybe_attach_attention_metadata_extensions(
|
||
self,
|
||
*,
|
||
attn_metadata: Any,
|
||
num_reqs: int,
|
||
num_reqs_padded: int,
|
||
max_query_len: int,
|
||
pad_attn: bool,
|
||
for_cudagraph_capture: bool = False,
|
||
num_scheduled_tokens_np: np.ndarray | None = None,
|
||
) -> None:
|
||
from vllm_omni.attention.fish_kvcache_backend import maybe_attach_fish_kvcache_seq_lens_upper_bound
|
||
|
||
maybe_attach_fish_kvcache_seq_lens_upper_bound(
|
||
model_config=self.model_config,
|
||
attn_metadata=attn_metadata,
|
||
input_batch=self.input_batch,
|
||
optimistic_seq_lens_cpu=self.optimistic_seq_lens_cpu,
|
||
num_reqs=num_reqs,
|
||
num_reqs_padded=num_reqs_padded,
|
||
max_query_len=max_query_len,
|
||
pad_attn=pad_attn,
|
||
for_cudagraph_capture=for_cudagraph_capture,
|
||
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
||
)
|
||
|
||
def _build_model_sampler_output_token_ids(self) -> list[list[int]]:
|
||
"""Build decoded-token history for ``prefer_model_sampler`` models.
|
||
|
||
vLLM only populates ``sampling_metadata.output_token_ids`` when penalties
|
||
or logits processors require it. Models that opt into a custom sampler
|
||
(e.g. CosyVoice3 RAS sampler, HunyuanImage3 stage-transition sampler)
|
||
also depend on this history, so we reconstruct it directly from the
|
||
input batch. Shared by GPU and NPU AR runners.
|
||
"""
|
||
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]
|
||
|
||
return output_token_ids
|
||
|
||
def _sampling_metadata_for_model_sampler(self, sampling_metadata):
|
||
output_token_ids = self._build_model_sampler_output_token_ids()
|
||
if output_token_ids == sampling_metadata.output_token_ids:
|
||
return sampling_metadata
|
||
return replace(sampling_metadata, output_token_ids=output_token_ids)
|
||
|
||
def _init_mrope_positions(self, req_state: CachedRequestState):
|
||
"""Initialize M-RoPE positions for multimodal inputs.
|
||
|
||
Extracts multimodal feature metadata (image grids, video grids,
|
||
audio features) and computes M-RoPE positions for proper positional
|
||
encoding of multimodal tokens.
|
||
|
||
Args:
|
||
req_state: Cached request state containing multimodal features
|
||
|
||
Raises:
|
||
AssertionError: If the model does not support M-RoPE
|
||
|
||
Note:
|
||
Upstream vLLM (commit 470229c37) added a fallback for the case
|
||
where ``req_state.prompt_token_ids`` is None but
|
||
``req_state.prompt_embeds`` is available. Omni models always set
|
||
``prompt_token_ids``, so this fallback is deliberately omitted.
|
||
"""
|
||
image_grid_thw = []
|
||
video_grid_thw = []
|
||
second_per_grid_ts = []
|
||
audio_feature_lengths = []
|
||
use_audio_in_video = False
|
||
for mm_feature in req_state.mm_features:
|
||
mm_item = mm_feature.data
|
||
if mm_item is None:
|
||
continue
|
||
mm_input = mm_item.get_data()
|
||
if (t := mm_input.get("image_grid_thw")) is not None:
|
||
image_grid_thw.append(t.tolist())
|
||
if (t := mm_input.get("video_grid_thw")) is not None:
|
||
video_grid_thw.append(t.tolist())
|
||
if (t := mm_input.get("second_per_grid_ts")) is not None:
|
||
second_per_grid_ts.append(t)
|
||
if (t := mm_input.get("audio_feature_lengths")) is not None:
|
||
audio_feature_lengths.append(t)
|
||
# Check for use_audio_in_video
|
||
use_audio_in_video_value = mm_input.get("use_audio_in_video")
|
||
if use_audio_in_video_value is not None:
|
||
use_audio_in_video = bool(use_audio_in_video_value.item())
|
||
|
||
if supports_mrope(self.get_model()):
|
||
# Model implements SupportsMRoPE interface
|
||
# Pass all extracted metadata; models use what they need via **kwargs
|
||
sp_extra_args = getattr(req_state.sampling_params, "extra_args", {}) if req_state.sampling_params else {}
|
||
target_h = sp_extra_args.get("target_h") if isinstance(sp_extra_args, dict) else None
|
||
target_w = sp_extra_args.get("target_w") if isinstance(sp_extra_args, dict) else None
|
||
kwargs = dict(
|
||
mm_features=req_state.mm_features,
|
||
hf_config=self.model_config.hf_config,
|
||
image_grid_thw=image_grid_thw,
|
||
video_grid_thw=video_grid_thw,
|
||
second_per_grid_ts=second_per_grid_ts,
|
||
audio_feature_lengths=audio_feature_lengths,
|
||
use_audio_in_video=use_audio_in_video,
|
||
)
|
||
if target_h is not None:
|
||
kwargs["target_h"] = target_h
|
||
if target_w is not None:
|
||
kwargs["target_w"] = target_w
|
||
req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
|
||
req_state.prompt_token_ids,
|
||
**_filter_mrope_kwargs_for_model(self.model, kwargs),
|
||
)
|
||
else:
|
||
req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor(
|
||
req_state.prompt_token_ids,
|
||
hf_config=self.model_config.hf_config,
|
||
image_grid_thw=image_grid_thw,
|
||
video_grid_thw=video_grid_thw,
|
||
second_per_grid_ts=second_per_grid_ts,
|
||
audio_feature_lengths=audio_feature_lengths,
|
||
use_audio_in_video=use_audio_in_video,
|
||
)
|
||
|
||
def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"):
|
||
"""Calculate M-RoPE positions for scheduled tokens.
|
||
|
||
Delegates to the upstream implementation first, then applies a fixup
|
||
pass for models that pre-compute 2D spatial decode positions (e.g.
|
||
GLM-Image). This avoids duplicating the full upstream method while
|
||
still supporting non-linear decode position patterns.
|
||
|
||
Models opt-in by declaring ``precomputed_mrope_decode = True`` as a
|
||
class attribute. When set, ``get_mrope_input_positions`` is expected
|
||
to return positions covering **both** prefill and decode tokens.
|
||
"""
|
||
# Run upstream logic (handles prompt positions + linear decode fallback)
|
||
super()._calc_mrope_positions(scheduler_output)
|
||
|
||
# Only run the fixup if the model pre-computes decode M-RoPE positions
|
||
if not getattr(self.get_model(), "precomputed_mrope_decode", False):
|
||
return
|
||
|
||
self._fixup_precomputed_mrope_decode_positions(scheduler_output)
|
||
|
||
def _fixup_precomputed_mrope_decode_positions(self, scheduler_output: "SchedulerOutput") -> None:
|
||
"""Overwrite linear decode M-RoPE positions with pre-computed ones.
|
||
|
||
For image-generation models (like GLM-Image) that output tokens in 2D
|
||
grid order, ``get_mrope_input_positions`` returns positions for the
|
||
full sequence (prefill + decode). The upstream runner only uses the
|
||
prefill portion and falls back to linear increments for decode. This
|
||
method patches the decode slice with the correct pre-computed values.
|
||
"""
|
||
from vllm.utils import length_from_prompt_token_ids_or_embeds
|
||
|
||
mrope_pos_ptr = 0
|
||
for index, req_id in enumerate(self.input_batch.req_ids):
|
||
req = self.requests[req_id]
|
||
assert req.mrope_positions is not None
|
||
|
||
num_computed_tokens = self.input_batch.num_computed_tokens_cpu[index]
|
||
num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
|
||
num_prompt_tokens = length_from_prompt_token_ids_or_embeds(req.prompt_token_ids, req.prompt_embeds)
|
||
|
||
if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens:
|
||
prompt_part_len = max(0, num_prompt_tokens - num_computed_tokens)
|
||
completion_part_len = max(0, num_scheduled_tokens - prompt_part_len)
|
||
else:
|
||
prompt_part_len = num_scheduled_tokens
|
||
completion_part_len = 0
|
||
|
||
mrope_pos_ptr += prompt_part_len
|
||
|
||
if completion_part_len > 0:
|
||
dst_start = mrope_pos_ptr
|
||
decode_start = num_computed_tokens + prompt_part_len
|
||
decode_end = decode_start + completion_part_len
|
||
total_precomputed = req.mrope_positions.shape[1]
|
||
|
||
if decode_end <= total_precomputed:
|
||
# Overwrite the linear positions written by upstream with
|
||
# the correct pre-computed 2D spatial positions.
|
||
self.mrope_positions.cpu[:, dst_start : dst_start + completion_part_len] = req.mrope_positions[
|
||
:, decode_start:decode_end
|
||
]
|
||
|
||
mrope_pos_ptr += completion_part_len
|
||
|
||
def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None:
|
||
"""Update the cached states and the persistent batch with the scheduler
|
||
output.
|
||
|
||
The updated states are used by the `_prepare_inputs` function to create
|
||
the input GPU tensors for the model.
|
||
|
||
The SamplingMetadata is updated and copied to the GPU if there is a
|
||
new/resumed/paused/finished request in the batch.
|
||
"""
|
||
# Used for prefix cache
|
||
if self.omni_prefix_cache is not None:
|
||
self.omni_prefix_cache.reset_prefix_cached_new_req_ids()
|
||
|
||
# Remove finished requests from the cached states.
|
||
# cleanup_finished_request lives on OmniConnectorModelRunnerMixin and
|
||
# is only safe to call once init_omni_connectors() has finished
|
||
# populating mixin state (it sets ``_omni_connector_initialized = True``
|
||
# at the very end). Archs that inherit the method via MRO without
|
||
# running that init must be skipped, so gate on the explicit flag
|
||
# rather than probing private attribute names.
|
||
cleanup_finished_request = (
|
||
getattr(self, "cleanup_finished_request", None)
|
||
if getattr(self, "_omni_connector_initialized", False)
|
||
else None
|
||
)
|
||
for req_id in scheduler_output.finished_req_ids:
|
||
self.requests.pop(req_id, None)
|
||
self.model_intermediate_buffer.pop(req_id, None)
|
||
self.num_prompt_logprobs.pop(req_id, None)
|
||
if self.omni_prefix_cache is not None:
|
||
self.omni_prefix_cache.discard_deferred_mm_outputs(req_id)
|
||
if hasattr(self, "_downstream_payload_cache"):
|
||
self._downstream_payload_cache.pop(req_id, None)
|
||
if hasattr(self, "_talker_mtp_generators"):
|
||
self._talker_mtp_generators.pop(req_id, None)
|
||
if cleanup_finished_request is not None:
|
||
cleanup_finished_request(req_id)
|
||
|
||
self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids)
|
||
# Remove the finished requests from the persistent batch.
|
||
# NOTE(woosuk): There could be an edge case where finished_req_ids and
|
||
# scheduled_req_ids overlap. This happens when a request is aborted and
|
||
# then resubmitted with the same ID. In this case, we treat them as two
|
||
# distinct requests - clearing the cached states for the first request
|
||
# and handling the second as a new request.
|
||
for req_id in scheduler_output.finished_req_ids:
|
||
self.input_batch.remove_request(req_id)
|
||
|
||
# Zero GPU memory for freshly allocated cache blocks to prevent
|
||
# stale NaN/data from corrupting attention or SSM computation.
|
||
if scheduler_output.new_block_ids_to_zero:
|
||
self._zero_block_ids(scheduler_output.new_block_ids_to_zero)
|
||
|
||
# Free the cached encoder outputs.
|
||
for mm_hash in scheduler_output.free_encoder_mm_hashes:
|
||
self.encoder_cache.pop(mm_hash, None)
|
||
|
||
# Remove the unscheduled requests from the persistent batch.
|
||
# NOTE(woosuk): The unscheduled requests are either preempted requests
|
||
# or running requests that are not scheduled in this step. We remove
|
||
# them from the persistent batch but keep their cached states since
|
||
# they will be scheduled again sometime in the future.
|
||
scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
|
||
cached_req_ids = self.input_batch.req_id_to_index.keys()
|
||
resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
|
||
# NOTE(zhuohan): cached_req_ids and resumed_req_ids are usually disjoint,
|
||
# so `(scheduled_req_ids - resumed_req_ids) == scheduled_req_ids` holds
|
||
# apart from the forced-preemption case in reset_prefix_cache. And in
|
||
# that case we include the resumed_req_ids in the unscheduled set so
|
||
# that they get cleared from the persistent batch before being re-scheduled
|
||
# in the normal resumed request path.
|
||
unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids)
|
||
# NOTE(woosuk): The persistent batch optimization assumes that
|
||
# consecutive batches contain mostly the same requests. If batches
|
||
# have low request overlap (e.g., alternating between two distinct
|
||
# sets of requests), this optimization becomes very inefficient.
|
||
for req_id in unscheduled_req_ids:
|
||
self.input_batch.remove_request(req_id)
|
||
|
||
is_ngram_gpu = self.speculative_config is not None and self.speculative_config.use_ngram_gpu()
|
||
if is_ngram_gpu:
|
||
ngram_gpu_new_reqs: list[CachedRequestState] = []
|
||
|
||
reqs_to_add: list[CachedRequestState] = []
|
||
deferred_spec_decode_corrections = []
|
||
# Add new requests to the cached states.
|
||
for new_req_data in scheduler_output.scheduled_new_reqs:
|
||
req_id = new_req_data.req_id
|
||
if req_id in self.requests:
|
||
self._update_streaming_input_additional_info(new_req_data, req_id)
|
||
req_state = self._update_streaming_request(req_id, new_req_data)
|
||
reqs_to_add.append(req_state)
|
||
continue
|
||
|
||
# Since this is the first time the request has been scheduled,
|
||
# num_computed_tokens > 0 means that we have a hit in prefix
|
||
# caching; mark it so that we can manage the hidden states
|
||
# later on as needed.
|
||
if self.omni_prefix_cache is not None and new_req_data.num_computed_tokens > 0:
|
||
self.omni_prefix_cache.add_prefix_cached_new_req_id(req_id)
|
||
|
||
sampling_params = new_req_data.sampling_params
|
||
pooling_params = new_req_data.pooling_params
|
||
|
||
if sampling_params and sampling_params.sampling_type == SamplingType.RANDOM_SEED:
|
||
generator = torch.Generator(device=self.device)
|
||
generator.manual_seed(sampling_params.seed)
|
||
else:
|
||
generator = None
|
||
|
||
if self.is_pooling_model:
|
||
assert pooling_params is not None
|
||
task = pooling_params.task
|
||
assert task is not None, "You did not set `task` in the API"
|
||
|
||
model = cast(VllmModelForPooling, self.get_model())
|
||
to_update = model.pooler.get_pooling_updates(task)
|
||
to_update.apply(pooling_params)
|
||
|
||
req_state = CachedRequestState(
|
||
req_id=req_id,
|
||
prompt_token_ids=new_req_data.prompt_token_ids,
|
||
prompt_embeds=new_req_data.prompt_embeds,
|
||
mm_features=new_req_data.mm_features,
|
||
sampling_params=sampling_params,
|
||
pooling_params=pooling_params,
|
||
generator=generator,
|
||
block_ids=new_req_data.block_ids,
|
||
num_computed_tokens=new_req_data.num_computed_tokens,
|
||
output_token_ids=[],
|
||
lora_request=new_req_data.lora_request,
|
||
)
|
||
self.requests[req_id] = req_state
|
||
self.late_interaction_runner.register_request(req_id, pooling_params)
|
||
|
||
# If prompt embeddings are provided, decode and attach to inter_data
|
||
try:
|
||
if getattr(new_req_data, "prompt_embeds", None) is not None:
|
||
payload = new_req_data.prompt_embeds
|
||
dtype = getattr(np, payload.dtype)
|
||
arr = np.frombuffer(payload.data, dtype=dtype)
|
||
arr = arr.reshape(payload.shape)
|
||
pe_cpu = torch.from_numpy(arr)
|
||
setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
|
||
try:
|
||
new_req_data.prompt_embeds = pe_cpu # type: ignore[assignment]
|
||
except Exception:
|
||
pass
|
||
except Exception as e:
|
||
logger.error(f"Error decoding prompt embeds: {e}")
|
||
# Direct runner data-plane payloads populate
|
||
# model_intermediate_buffer without going through the deprecated
|
||
# additional_information request transport.
|
||
try:
|
||
model_buffer = getattr(new_req_data, "model_intermediate_buffer", None)
|
||
if isinstance(model_buffer, dict) and model_buffer:
|
||
self._update_intermediate_buffer(req_id, model_buffer)
|
||
except Exception as e:
|
||
logger.error(f"Error updating model intermediate buffer: {e}")
|
||
# Decode additional_information payloads (dictionary)
|
||
try:
|
||
if getattr(new_req_data, "additional_information", None) is not None:
|
||
info_dict = deserialize_additional_information(new_req_data.additional_information)
|
||
if info_dict:
|
||
self.model_intermediate_buffer[req_id] = info_dict
|
||
setattr(
|
||
self.requests[req_id],
|
||
"additional_information_cpu",
|
||
info_dict,
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Error decoding additional information: {e}")
|
||
|
||
if sampling_params and sampling_params.prompt_logprobs is not None:
|
||
self.num_prompt_logprobs[req_id] = (
|
||
self.input_batch.vocab_size
|
||
if sampling_params.prompt_logprobs == -1
|
||
else sampling_params.prompt_logprobs
|
||
)
|
||
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
|
||
if self.uses_mrope:
|
||
self._init_mrope_positions(req_state)
|
||
|
||
# Only relevant for models using XD-RoPE (e.g, HunYuan-VL)
|
||
if self.uses_xdrope_dim > 0:
|
||
self._init_xdrope_positions(req_state)
|
||
|
||
reqs_to_add.append(self.requests[req_id])
|
||
# Track new requests for ngram_gpu full tensor copy
|
||
if is_ngram_gpu:
|
||
ngram_gpu_new_reqs.append(req_state)
|
||
|
||
# Update the states of the running/resumed requests.
|
||
is_last_rank = get_pp_group().is_last_rank
|
||
req_data = scheduler_output.scheduled_cached_reqs
|
||
scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
|
||
|
||
# Save scheduler-allocated spec lengths before trimming so
|
||
# prev_num_draft_len keeps the optimistic count for rejection correction.
|
||
original_num_spec_per_req: dict[str, int] = {}
|
||
if self.speculative_config is not None and self.speculative_config.use_ngram_gpu():
|
||
for req_id, toks in scheduled_spec_tokens.items():
|
||
original_num_spec_per_req[req_id] = len(toks)
|
||
update_scheduler_for_invalid_drafts(
|
||
self._num_valid_draft_tokens_event,
|
||
self._num_valid_draft_tokens_cpu,
|
||
scheduler_output,
|
||
self.input_batch.req_id_to_index,
|
||
)
|
||
if self.use_async_spec_decode:
|
||
self.prev_num_draft_tokens.np.fill(0)
|
||
|
||
for i, req_id in enumerate(req_data.req_ids):
|
||
req_state = self.requests[req_id]
|
||
num_computed_tokens = req_data.num_computed_tokens[i]
|
||
new_block_ids = req_data.new_block_ids[i]
|
||
resumed_from_preemption = req_id in req_data.resumed_req_ids
|
||
num_output_tokens = req_data.num_output_tokens[i]
|
||
req_index = self.input_batch.req_id_to_index.get(req_id)
|
||
|
||
if req_state.prev_num_draft_len and self.use_async_scheduling:
|
||
# prev_num_draft_len is used in async scheduling mode with
|
||
# spec decode. it indicates if need to update num_computed_tokens
|
||
# of the request. for example:
|
||
# first step: num_computed_tokens = 0, spec_tokens = [],
|
||
# prev_num_draft_len = 0.
|
||
# second step: num_computed_tokens = 100(prompt length),
|
||
# spec_tokens = [a,b], prev_num_draft_len = 0.
|
||
# third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d],
|
||
# prev_num_draft_len = 2.
|
||
# num_computed_tokens in first step and second step doesn't contain
|
||
# the spec tokens length, but in third step it contains the
|
||
# spec tokens length. we only need to update num_computed_tokens
|
||
# when prev_num_draft_len > 0.
|
||
if req_index is None:
|
||
req_state.prev_num_draft_len = 0
|
||
else:
|
||
optimistic_num_accepted = req_state.prev_num_draft_len
|
||
req_state.output_token_ids.extend([-1] * optimistic_num_accepted)
|
||
|
||
deferred_spec_decode_corrections.append((req_id, optimistic_num_accepted, req_state))
|
||
|
||
prev_req_index = (
|
||
self.input_batch.prev_req_id_to_index.get(req_id)
|
||
if self.input_batch.prev_req_id_to_index
|
||
else None
|
||
)
|
||
if prev_req_index is not None:
|
||
self.prev_num_draft_tokens.np[prev_req_index] = optimistic_num_accepted
|
||
|
||
if is_ngram_gpu and optimistic_num_accepted > 0:
|
||
self.input_batch.num_tokens_no_spec[req_index] += optimistic_num_accepted
|
||
|
||
# Update the cached states.
|
||
req_state.num_computed_tokens = num_computed_tokens
|
||
|
||
if not is_last_rank:
|
||
if not req_data.new_token_ids:
|
||
new_token_ids: list[int] = []
|
||
else:
|
||
new_token_ids = req_data.new_token_ids[i]
|
||
num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens
|
||
if num_new_tokens == 1:
|
||
req_state.output_token_ids.append(new_token_ids[-1])
|
||
elif num_new_tokens > 0:
|
||
req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
|
||
elif num_output_tokens < len(req_state.output_token_ids):
|
||
# Some output tokens were discarded due to a sync-KV-load
|
||
# failure, or output_token_ids was inflated by the optimistic
|
||
# extend above (async spec decode). Align the cached state.
|
||
del req_state.output_token_ids[num_output_tokens:]
|
||
if req_index is not None:
|
||
end_idx = self.input_batch.num_prompt_tokens[req_index] + num_output_tokens
|
||
self.input_batch.num_tokens_no_spec[req_index] = end_idx
|
||
|
||
# Update the block IDs.
|
||
if not resumed_from_preemption:
|
||
if new_block_ids is not None:
|
||
# Append the new blocks to the existing block IDs.
|
||
for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
|
||
block_ids.extend(new_ids)
|
||
else:
|
||
assert req_index is None
|
||
assert new_block_ids is not None
|
||
# The request is resumed from preemption.
|
||
# Replace the existing block IDs with the new ones.
|
||
req_state.block_ids = new_block_ids
|
||
|
||
req_index = self.input_batch.req_id_to_index.get(req_id)
|
||
if req_index is None:
|
||
# The request is not in the persistent batch.
|
||
# The request was either preempted and resumed later, or was not
|
||
# scheduled in the previous step and needs to be added again.
|
||
|
||
if self.use_async_scheduling and num_output_tokens > 0:
|
||
# We must recover the output token ids for resumed requests in the
|
||
# async scheduling case, so that correct input_ids are obtained.
|
||
resumed_token_ids = req_data.all_token_ids[req_id]
|
||
req_state.output_token_ids = resumed_token_ids[-num_output_tokens:]
|
||
|
||
reqs_to_add.append(req_state)
|
||
# Track resumed requests for ngram_gpu full tensor copy
|
||
if is_ngram_gpu:
|
||
ngram_gpu_new_reqs.append(req_state)
|
||
continue
|
||
|
||
# Update the persistent batch.
|
||
self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
|
||
if new_block_ids is not None:
|
||
self.input_batch.block_table.append_row(new_block_ids, req_index)
|
||
|
||
# For the last rank, we don't need to update the token_ids_cpu
|
||
# because the sampled tokens are already cached.
|
||
if not is_last_rank:
|
||
# Add new_token_ids to token_ids_cpu.
|
||
start_token_index = num_computed_tokens
|
||
end_token_index = num_computed_tokens + len(new_token_ids)
|
||
self.input_batch.token_ids_cpu[req_index, start_token_index:end_token_index] = new_token_ids
|
||
self.input_batch.num_tokens_no_spec[req_index] = end_token_index
|
||
|
||
# Add spec_token_ids to token_ids_cpu.
|
||
self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens)
|
||
# Restore scheduler-side draft count after ngram trimming.
|
||
if original_num_spec_per_req:
|
||
orig = original_num_spec_per_req.get(req_id, 0)
|
||
if orig != req_state.prev_num_draft_len:
|
||
req_state.prev_num_draft_len = orig
|
||
|
||
# Add the new or resumed requests to the persistent batch.
|
||
# The smaller empty indices are filled first.
|
||
for request in reqs_to_add:
|
||
self.input_batch.add_request(request)
|
||
self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens)
|
||
|
||
# Condense the batched states if there are gaps left by removed requests
|
||
self.input_batch.condense()
|
||
# Allow attention backend to reorder the batch, potentially
|
||
self._may_reorder_batch(scheduler_output)
|
||
# Refresh batch metadata with any pending updates.
|
||
self.input_batch.refresh_metadata()
|
||
|
||
# Incrementally update ngram_gpu tensors after batch is stable
|
||
if is_ngram_gpu:
|
||
update_ngram_gpu_tensors_incremental(
|
||
self.input_batch,
|
||
self.token_ids_gpu_tensor,
|
||
self.num_tokens_no_spec_gpu,
|
||
ngram_gpu_new_reqs,
|
||
self.device,
|
||
_pinned_idx_buf=self._ngram_pinned_idx_buf,
|
||
_pinned_val_buf=self._ngram_pinned_val_buf,
|
||
)
|
||
|
||
if deferred_spec_decode_corrections:
|
||
|
||
def correct_spec_decode_token_counts():
|
||
valid_sampled_token_count = self._get_valid_sampled_token_count()
|
||
if not valid_sampled_token_count:
|
||
return
|
||
prev_req_id_to_index = self.input_batch.prev_req_id_to_index
|
||
if not prev_req_id_to_index:
|
||
return
|
||
for (
|
||
req_id,
|
||
optimistic_num_accepted,
|
||
req_state,
|
||
) in deferred_spec_decode_corrections:
|
||
prev_req_index = prev_req_id_to_index.get(req_id)
|
||
if prev_req_index is None:
|
||
continue
|
||
num_accepted = valid_sampled_token_count[prev_req_index] - 1
|
||
correction = optimistic_num_accepted - num_accepted
|
||
req_state.num_computed_tokens -= correction
|
||
cur_req_index = self.input_batch.req_id_to_index.get(req_id)
|
||
if cur_req_index is None:
|
||
continue
|
||
self.input_batch.num_computed_tokens_cpu[cur_req_index] -= correction
|
||
if is_ngram_gpu and correction > 0:
|
||
self.input_batch.num_tokens_no_spec[cur_req_index] -= correction
|
||
self.num_tokens_no_spec_gpu[cur_req_index] -= correction
|
||
|
||
return correct_spec_decode_token_counts
|
||
else:
|
||
return None
|
||
|
||
@torch.inference_mode()
|
||
def extract_multimodal_outputs(
|
||
self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput
|
||
) -> tuple[Any, Any]:
|
||
if (
|
||
hasattr(self.model, "have_multimodal_outputs")
|
||
and self.model.have_multimodal_outputs
|
||
and isinstance(hidden_states, OmniOutput)
|
||
):
|
||
text_hidden_states = hidden_states.text_hidden_states
|
||
multimodal_outputs = hidden_states.multimodal_outputs
|
||
|
||
elif isinstance(hidden_states, torch.Tensor):
|
||
text_hidden_states = hidden_states
|
||
multimodal_outputs = {}
|
||
elif isinstance(hidden_states, list) or isinstance(hidden_states, tuple):
|
||
text_hidden_states = hidden_states[0]
|
||
multimodal_outputs = {}
|
||
else:
|
||
raise ValueError(f"Invalid hidden states type: {type(hidden_states)}")
|
||
return text_hidden_states, multimodal_outputs
|
||
|
||
def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||
# Models loaded with load_format=dummy (e.g. MossTTSNano) may
|
||
# produce CPU hidden_states while the upstream sampler warmup
|
||
# expects GPU tensors. Skip the warmup in that case — a
|
||
# meaningful warmup requires real weights anyway.
|
||
if hidden_states.device != self.device:
|
||
return torch.tensor([])
|
||
return super()._dummy_sampler_run(hidden_states=hidden_states)
|
||
|
||
@torch.inference_mode()
|
||
def _dummy_run(
|
||
self,
|
||
num_tokens: int,
|
||
cudagraph_runtime_mode: CUDAGraphMode | None = None,
|
||
force_attention: bool = False,
|
||
uniform_decode: bool = False,
|
||
allow_microbatching: bool = True,
|
||
skip_eplb: bool = False,
|
||
is_profile: bool = False,
|
||
create_mixed_batch: bool = False,
|
||
remove_lora: bool = True,
|
||
is_graph_capturing: bool = False,
|
||
num_active_loras: int = 0,
|
||
profile_seq_lens: int | None = None,
|
||
randomize_inputs: bool = False,
|
||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
"""
|
||
Run a dummy forward pass to warm up/profile run or capture the
|
||
CUDA graph for the model.
|
||
|
||
Args:
|
||
num_tokens: Number of tokens to run the dummy forward pass.
|
||
cudagraph_runtime_mode: used to control the behavior.
|
||
- if not set will determine the cudagraph mode based on using
|
||
the self.cudagraph_dispatcher.
|
||
- CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
|
||
- CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
|
||
- CUDAGraphMode.FULL: Full cudagraph, attention metadata is
|
||
needed.
|
||
force_attention: If True, always create attention metadata. Used to
|
||
warm up attention backend when mode is NONE.
|
||
uniform_decode: If True, the batch is a uniform decode batch.
|
||
skip_eplb: If True, skip EPLB state update.
|
||
is_profile: If True, this is a profile run.
|
||
create_mixed_batch: If True, create a mixed batch with both decode
|
||
(1 token) and prefill (multiple tokens) requests.
|
||
remove_lora: If False, dummy LoRAs are not destroyed after the run
|
||
num_active_loras: Number of distinct active LoRAs to capture for.
|
||
LoRA is activated when num_active_loras > 0.
|
||
profile_seq_lens: If provided, use this value for seq_lens instead
|
||
of max_query_len. Used to profile attention workspace that
|
||
scales with context length.
|
||
randomize_inputs: If True, randomize dummy input ids to balance
|
||
expert selection. vLLM's kernel_warmup passes this when
|
||
autotuning flashinfer.
|
||
"""
|
||
mm_config = self.vllm_config.model_config.multimodal_config
|
||
if mm_config and mm_config.mm_encoder_only:
|
||
# The current dummy run only covers LM execution, so we can skip it.
|
||
# mm encoder dummy run may need to add in the future.
|
||
return torch.tensor([]), torch.tensor([])
|
||
|
||
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode()
|
||
|
||
# If cudagraph_mode.decode_mode() == FULL and
|
||
# cudagraph_mode.separate_routine(). This means that we are using
|
||
# different graphs and/or modes for mixed prefill-decode batches vs.
|
||
# uniform decode batches. A uniform decode batch means that all
|
||
# requests have identical query length, except a potential virtual
|
||
# request (shorter) in the batch account for padding.
|
||
# Uniform decode batch could either be common pure decode, where
|
||
# max_query_len == 1, or speculative decode, where
|
||
# max_query_len == 1 + num_spec_decode_tokens.
|
||
|
||
# When setting max_query_len = 1, we switch to and capture the optimized
|
||
# routine of FA2 for pure decode, i.e., Flashdecode + an optimization
|
||
# for GQA/MQA.
|
||
max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
|
||
|
||
# Set num_scheduled_tokens based on num_tokens and max_num_seqs
|
||
# for dummy run with LoRA so that the num_reqs collectively
|
||
# has num_tokens in total.
|
||
assert num_tokens <= self.max_num_tokens
|
||
max_num_reqs = self.scheduler_config.max_num_seqs
|
||
if create_mixed_batch:
|
||
assert not uniform_decode
|
||
# Create mixed batch:
|
||
# first half decode tokens, second half one prefill
|
||
num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2)
|
||
num_prefill_tokens = num_tokens - num_decode_tokens
|
||
num_reqs = num_decode_tokens + 1
|
||
|
||
# Create decode requests (1 token each) followed by prefill request
|
||
num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens]
|
||
# Note: Overriding max_query_len to be the prefill tokens
|
||
max_query_len = num_prefill_tokens
|
||
elif uniform_decode:
|
||
assert not create_mixed_batch
|
||
num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len))
|
||
num_scheduled_tokens_list = [max_query_len] * num_reqs
|
||
if num_tokens % max_query_len != 0:
|
||
num_scheduled_tokens_list[-1] = num_tokens % max_query_len
|
||
else:
|
||
num_reqs = min(num_tokens, max_num_reqs)
|
||
min_tokens_per_req = num_tokens // num_reqs
|
||
num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
|
||
num_scheduled_tokens_list[-1] += num_tokens % num_reqs
|
||
|
||
assert sum(num_scheduled_tokens_list) == num_tokens
|
||
assert len(num_scheduled_tokens_list) == num_reqs
|
||
num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
|
||
num_tokens_unpadded = int(num_scheduled_tokens.sum())
|
||
|
||
num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
|
||
|
||
_cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = (
|
||
self._determine_batch_execution_and_padding(
|
||
num_tokens=num_tokens_unpadded,
|
||
num_reqs=num_reqs,
|
||
num_scheduled_tokens_np=num_scheduled_tokens,
|
||
max_num_scheduled_tokens=max_query_len,
|
||
use_cascade_attn=False,
|
||
allow_microbatching=allow_microbatching,
|
||
force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE),
|
||
# `force_uniform_decode` is used for cudagraph capture; because for
|
||
# capturing mixed prefill-decode batches, we sometimes use
|
||
# num_tokens == num_reqs which looks like a uniform decode batch to the
|
||
# dispatcher; but we actually want to capture a piecewise cudagraph
|
||
force_uniform_decode=uniform_decode,
|
||
# `force_has_lora` is used for cudagraph capture; because LoRA is
|
||
# activated later in the context manager, but we need to know the
|
||
# LoRA state when determining the batch descriptor for capture
|
||
force_has_lora=num_active_loras > 0,
|
||
# `force_num_active_loras` is used for cudagraph capture; because we
|
||
# need to capture graphs for specific num_active_loras counts
|
||
force_num_active_loras=num_active_loras,
|
||
)
|
||
)
|
||
|
||
if cudagraph_runtime_mode is None:
|
||
cudagraph_runtime_mode = _cudagraph_mode
|
||
else:
|
||
assert cudagraph_runtime_mode == _cudagraph_mode, (
|
||
f"Cudagraph runtime mode mismatch in dummy_run. "
|
||
f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}."
|
||
)
|
||
|
||
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
|
||
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
|
||
should_ubatch,
|
||
num_scheduled_tokens,
|
||
num_tokens_padded,
|
||
num_reqs_padded,
|
||
self.vllm_config.parallel_config.num_ubatches,
|
||
)
|
||
logger.debug(
|
||
"ubatch_slices: %s, ubatch_slices_padded: %s",
|
||
ubatch_slices,
|
||
ubatch_slices_padded,
|
||
)
|
||
|
||
attn_metadata: PerLayerAttnMetadata | None = None
|
||
|
||
slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
|
||
num_tokens_padded=num_tokens_padded,
|
||
num_reqs_padded=num_reqs_padded,
|
||
num_tokens_unpadded=num_tokens_unpadded,
|
||
ubatch_slices=ubatch_slices_padded,
|
||
)
|
||
|
||
if slot_mappings_by_group is not None:
|
||
for sm in slot_mappings_by_group.values():
|
||
sm.fill_(-1)
|
||
|
||
with self.synchronize_input_prep():
|
||
# If force_attention is True, we always capture attention.
|
||
# Otherwise, it only happens for cudagraph_runtime_mode=FULL.
|
||
if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
|
||
if profile_seq_lens is not None:
|
||
seq_lens = profile_seq_lens # type: ignore[assignment]
|
||
elif create_mixed_batch:
|
||
seq_lens = torch.tensor( # type: ignore[assignment]
|
||
[1] * num_decode_tokens + [num_prefill_tokens + 1],
|
||
dtype=torch.int,
|
||
)
|
||
else:
|
||
seq_lens = max_query_len # type: ignore[assignment]
|
||
self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens
|
||
self.optimistic_seq_lens_cpu[num_reqs:].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, self.query_pos.np)
|
||
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)
|
||
|
||
pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
|
||
attn_metadata, _ = self._build_attention_metadata(
|
||
num_tokens=num_tokens_unpadded,
|
||
num_tokens_padded=num_tokens_padded if pad_attn else None,
|
||
num_reqs=num_reqs_padded,
|
||
max_query_len=max_query_len,
|
||
ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices),
|
||
for_cudagraph_capture=is_graph_capturing,
|
||
slot_mappings=slot_mappings_by_group,
|
||
use_spec_decode=self.speculative_config is not None,
|
||
)
|
||
self._maybe_attach_attention_metadata_extensions(
|
||
attn_metadata=attn_metadata,
|
||
num_reqs=num_reqs_padded,
|
||
num_reqs_padded=num_reqs_padded,
|
||
max_query_len=max_query_len,
|
||
pad_attn=True,
|
||
for_cudagraph_capture=is_graph_capturing,
|
||
)
|
||
|
||
with self.maybe_dummy_run_with_lora(
|
||
self.lora_config,
|
||
num_scheduled_tokens,
|
||
num_sampled_tokens,
|
||
remove_lora,
|
||
num_active_loras,
|
||
):
|
||
# Make sure padding doesn't exceed max_num_tokens
|
||
assert num_tokens_padded <= self.max_num_tokens
|
||
model_kwargs = self._init_model_kwargs()
|
||
if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
|
||
input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded)
|
||
|
||
model_kwargs = {
|
||
**model_kwargs,
|
||
**self._dummy_mm_kwargs(num_reqs),
|
||
}
|
||
elif self.enable_prompt_embeds:
|
||
input_ids = None
|
||
inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
|
||
model_kwargs = self._init_model_kwargs()
|
||
elif getattr(getattr(self, "model", None), "has_preprocess", False):
|
||
# Capture CUDA graph with inputs_embeds path so replay reads
|
||
# from the same buffer that _preprocess writes into.
|
||
input_ids = self.input_ids.gpu[:num_tokens_padded]
|
||
inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
|
||
else:
|
||
input_ids = self.input_ids.gpu[:num_tokens_padded]
|
||
inputs_embeds = None
|
||
|
||
if self.uses_mrope:
|
||
positions = self.mrope_positions.gpu[:, :num_tokens_padded]
|
||
elif self.uses_xdrope_dim > 0:
|
||
positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
|
||
else:
|
||
positions = self.positions[:num_tokens_padded]
|
||
|
||
if get_pp_group().is_first_rank:
|
||
intermediate_tensors = None
|
||
else:
|
||
if self.intermediate_tensors is None:
|
||
self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
|
||
batch_size=self.max_num_tokens,
|
||
dtype=self.model_config.dtype,
|
||
device=self.device,
|
||
)
|
||
|
||
intermediate_tensors = self.sync_and_gather_intermediate_tensors(num_tokens_padded, None, False)
|
||
|
||
if ubatch_slices_padded is not None:
|
||
# Adjust values to reflect a single ubatch.
|
||
# TODO(sage,lucas): this is cruft that should be addressed in
|
||
# the padding refactor.
|
||
num_tokens_padded = ubatch_slices_padded[0].num_tokens
|
||
if num_tokens_across_dp is not None:
|
||
num_tokens_across_dp[:] = num_tokens_padded
|
||
|
||
with (
|
||
self.maybe_randomize_inputs(input_ids, inputs_embeds, randomize_inputs=randomize_inputs),
|
||
set_forward_context(
|
||
attn_metadata,
|
||
self.vllm_config,
|
||
num_tokens=num_tokens_padded,
|
||
num_tokens_across_dp=num_tokens_across_dp,
|
||
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
||
batch_descriptor=batch_desc,
|
||
ubatch_slices=ubatch_slices_padded,
|
||
slot_mapping=slot_mappings,
|
||
),
|
||
):
|
||
if getattr(self.model, "talker", None) is not None and self.has_talker_mtp:
|
||
num_tokens_padded_talker_mtp = num_tokens_padded
|
||
if num_tokens_padded_talker_mtp == self.max_num_tokens:
|
||
num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
|
||
outputs = self.talker_mtp(
|
||
self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp],
|
||
self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp],
|
||
self.last_talker_hidden.gpu[:num_tokens_padded_talker_mtp],
|
||
self.text_step.gpu[:num_tokens_padded_talker_mtp],
|
||
)
|
||
self.compilation_config.cache_dir = None
|
||
outputs = self.model(
|
||
input_ids=input_ids,
|
||
positions=positions,
|
||
intermediate_tensors=intermediate_tensors,
|
||
inputs_embeds=inputs_embeds,
|
||
**model_kwargs,
|
||
)
|
||
|
||
if self.use_aux_hidden_state_outputs:
|
||
hidden_states, _ = outputs
|
||
else:
|
||
hidden_states = outputs
|
||
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
|
||
if self.speculative_config and (
|
||
self.speculative_config.use_eagle()
|
||
or self.speculative_config.uses_draft_model()
|
||
or self.speculative_config.uses_extract_hidden_states()
|
||
):
|
||
assert isinstance(
|
||
self.drafter,
|
||
EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer,
|
||
)
|
||
assert self.speculative_config is not None
|
||
# Eagle currently only supports PIECEWISE cudagraphs.
|
||
# Therefore only use cudagraphs if the main model uses PIECEWISE
|
||
# NOTE(lucas): this is a hack, need to clean up.
|
||
use_cudagraphs = (
|
||
(is_graph_capturing and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE)
|
||
or (not is_graph_capturing and cudagraph_runtime_mode != CUDAGraphMode.NONE)
|
||
) and not self.speculative_config.enforce_eager
|
||
|
||
# Note(gnovack) - We need to disable cudagraphs for one of the two
|
||
# lora cases when cudagraph_specialize_lora is enabled. This is a
|
||
# short term mitigation for issue mentioned in
|
||
# https://github.com/vllm-project/vllm/issues/28334
|
||
if self.compilation_config.cudagraph_specialize_lora and num_active_loras > 0:
|
||
use_cudagraphs = False
|
||
|
||
self.drafter.dummy_run(
|
||
num_tokens,
|
||
use_cudagraphs=use_cudagraphs,
|
||
is_graph_capturing=is_graph_capturing,
|
||
slot_mappings=slot_mappings,
|
||
)
|
||
|
||
# We register layerwise NVTX hooks here after the first dynamo tracing is
|
||
# done to avoid nvtx operations in hook functions being traced by
|
||
# torch dynamo and causing graph breaks.
|
||
# Note that for DYNAMO_ONCE and VLLM_COMPILE mode,
|
||
# compiled model's dynamo tracing is only done once and the compiled model's
|
||
# __call__ function is replaced by calling the compiled function.
|
||
# So it's safe to register hooks here. Hooks will be registered to
|
||
# both compiled and uncompiled models but they will never
|
||
# be called on the compiled model execution path.
|
||
self._register_layerwise_nvtx_hooks()
|
||
|
||
# This is necessary to avoid blocking DP.
|
||
# For dummy runs, we typically skip EPLB since we don't have any real
|
||
# requests to process.
|
||
# However, in DP settings, there may be cases when some DP ranks do
|
||
# not have any requests to process, so they're executing dummy batches.
|
||
# In such cases, we still have to trigger EPLB to make sure
|
||
# ranks execute the rearrangement in synchronization.
|
||
if not skip_eplb:
|
||
self.eplb_step(is_dummy=True, is_profile=is_profile)
|
||
|
||
logit_indices = np.cumsum(num_scheduled_tokens) - 1
|
||
logit_indices_device = torch.from_numpy(logit_indices).to(hidden_states.device, non_blocking=True)
|
||
return hidden_states, hidden_states[logit_indices_device]
|
||
|
||
# ------------------------------------------------------------------
|
||
# Payload decoding helpers (torch.Tensor passthrough + legacy
|
||
# PromptEmbedsPayload / AdditionalInformationPayload support)
|
||
# ------------------------------------------------------------------
|
||
|
||
@staticmethod
|
||
def _resolve_prompt_embeds_cpu(
|
||
pe: "torch.Tensor | object | None",
|
||
) -> torch.Tensor | None:
|
||
"""Convert *prompt_embeds* to a contiguous CPU tensor.
|
||
|
||
Accepts:
|
||
- ``torch.Tensor`` – moved to CPU as-is (the normal path after
|
||
upstream added ``prompt_embeds`` to ``EngineCoreRequest``).
|
||
- Legacy ``PromptEmbedsPayload`` (or any duck-typed object with
|
||
``.data``, ``.shape``, ``.dtype``) – decoded via numpy.
|
||
- ``None`` – returns ``None``.
|
||
"""
|
||
if pe is None:
|
||
return None
|
||
try:
|
||
if isinstance(pe, torch.Tensor):
|
||
return pe.detach().cpu().contiguous()
|
||
data = getattr(pe, "data", None)
|
||
shape = getattr(pe, "shape", None)
|
||
if data is not None and shape is not None:
|
||
dt = np.dtype(getattr(pe, "dtype", "float32"))
|
||
arr = np.frombuffer(data, dtype=dt).reshape(shape)
|
||
return torch.from_numpy(arr.copy())
|
||
except Exception:
|
||
logger.exception("Failed to decode prompt_embeds payload")
|
||
return None
|
||
|
||
def _decode_and_store_request_payloads(
|
||
self,
|
||
scheduler_output: "SchedulerOutput",
|
||
) -> None:
|
||
"""Decode per-request prompt_embeds and additional_information for
|
||
newly scheduled requests and store them on CPU in the request state.
|
||
"""
|
||
new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
|
||
if not new_reqs:
|
||
return
|
||
for nr in new_reqs:
|
||
req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
|
||
if req_id is None or req_id not in self.requests:
|
||
continue
|
||
pe_cpu = self._resolve_prompt_embeds_cpu(getattr(nr, "prompt_embeds", None))
|
||
if pe_cpu is not None:
|
||
setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
|
||
model_buffer = getattr(nr, "model_intermediate_buffer", None)
|
||
if isinstance(model_buffer, dict) and model_buffer:
|
||
self._update_intermediate_buffer(req_id, model_buffer)
|
||
info_payload = getattr(nr, "additional_information", None)
|
||
if info_payload is not None:
|
||
logger.warning_once(
|
||
"additional_information on request data is deprecated, use model_intermediate_buffer"
|
||
)
|
||
info_dict = deserialize_additional_information(info_payload)
|
||
if info_dict:
|
||
self.model_intermediate_buffer[req_id] = info_dict
|
||
setattr(self.requests[req_id], "additional_information_cpu", info_dict)
|
||
|
||
def _gather_runtime_additional_information(self) -> list[dict]:
|
||
"""Gather per-request model_intermediate_buffer in batch order."""
|
||
per_req_runtime_info = []
|
||
for req_id in self.input_batch.req_ids:
|
||
req_state = self.requests.get(req_id)
|
||
# MammothModa2 AR grid constraint: the model must emit a special
|
||
# end-of-line (EOL) token at the end of each image row. To determine
|
||
# whether the current decoding step falls on a row boundary, the
|
||
# constraint logic (see MammothModa2ARForConditionalGeneration.
|
||
# _apply_t2i_token_constraints) computes:
|
||
# column_id = generated_len % (ar_width + 1)
|
||
# and forces the EOL token when column_id == ar_width.
|
||
generated_len = len(req_state.output_token_ids) if req_state is not None else 0
|
||
info = self.model_intermediate_buffer.get(req_id, {})
|
||
if info:
|
||
info["generated_len"] = generated_len
|
||
per_req_runtime_info.append(info)
|
||
if "thinker_reply_part_per_request" in info:
|
||
q = info["thinker_reply_part_per_request"]
|
||
if hasattr(q, "shape"):
|
||
logger.debug(f"[OMNI] req={req_id} has thinker_reply_part_per_request queue shape: {q.shape}")
|
||
else:
|
||
per_req_runtime_info.append({})
|
||
return per_req_runtime_info
|
||
|
||
def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]:
|
||
"""Compute (start, end) token spans for each request within the flattened step sequence."""
|
||
req_token_spans: list[tuple[int, int]] = []
|
||
for req_index in range(len(self.input_batch.req_ids)):
|
||
start_offset = int(self.query_start_loc.cpu[req_index])
|
||
sched_tokens = int(num_scheduled_tokens_np[req_index])
|
||
req_token_spans.append((start_offset, start_offset + sched_tokens))
|
||
return req_token_spans
|
||
|
||
def _sync_local_stage_payloads(self) -> None:
|
||
"""Move received full-payload stage inputs into model_intermediate_buffer."""
|
||
cache = getattr(self, "_local_stage_payload_cache", None)
|
||
if not cache:
|
||
return
|
||
lock = getattr(self, "_lock", None)
|
||
ctx = lock if lock is not None else contextlib.nullcontext()
|
||
with ctx:
|
||
if not cache:
|
||
return
|
||
active_req_ids = set(getattr(self, "requests", {}))
|
||
pending = set(getattr(self, "_full_payload_pending_broadcast_req_ids", set()))
|
||
staged = {
|
||
req_id: payload
|
||
for req_id, payload in cache.items()
|
||
if req_id not in pending and req_id in active_req_ids and isinstance(payload, dict)
|
||
}
|
||
for req_id in staged:
|
||
cache.pop(req_id, None)
|
||
for req_id, payload in staged.items():
|
||
self._update_intermediate_buffer(req_id, payload)
|
||
|
||
def _build_model_kwargs_extra(self) -> dict:
|
||
"""Build extra keyword arguments passed to the model for this step."""
|
||
self._sync_local_stage_payloads()
|
||
model_kwargs_extra: dict[str, object] = {}
|
||
try:
|
||
buffer_map = self._gather_runtime_additional_information()
|
||
model_kwargs_extra["model_intermediate_buffer"] = buffer_map
|
||
# Backward compatible: also emit old name
|
||
model_kwargs_extra["runtime_additional_information"] = buffer_map
|
||
except Exception as e:
|
||
logger.error(f"[OMNI DEBUG] Error building model_kwargs_extra: {e}")
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
|
||
# Per-request (start, end) hidden-row spans so make_omni_output can map
|
||
# flat hidden rows to the right request in mixed prefill+decode steps,
|
||
# instead of assuming an equal rows-per-request split (which samples the
|
||
# wrong rows whenever per-request token counts differ).
|
||
nstp = self._omni_num_scheduled_tokens_np
|
||
if nstp is not None and len(nstp) == len(self.input_batch.req_ids):
|
||
try:
|
||
model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(nstp)
|
||
if getattr(self.model, "requires_request_sample_eligibility", False):
|
||
model_kwargs_extra["request_sample_eligible"] = [
|
||
bool(
|
||
(req := self.requests.get(req_id)) is not None
|
||
and int(req.num_computed_tokens) + int(nstp[req_index]) >= int(req.num_tokens)
|
||
)
|
||
for req_index, req_id in enumerate(self.input_batch.req_ids)
|
||
]
|
||
except Exception as e:
|
||
# Visible on purpose: the fallback is the equal rows-per-request
|
||
# split, which can re-introduce the cross-request corruption this
|
||
# plumbing fixes — a silent failure here must not pass unnoticed.
|
||
logger.warning("[OMNI] Failed to compute request_token_spans: %s", e)
|
||
|
||
if self._omni_query_start_loc_model_kwarg:
|
||
try:
|
||
num_reqs = len(self.input_batch.req_ids)
|
||
model_kwargs_extra["omni_query_start_loc"] = self.query_start_loc.gpu[: num_reqs + 1]
|
||
except Exception as e:
|
||
logger.debug("[OMNI] Failed to attach query_start_loc: %s", e)
|
||
|
||
if getattr(self.model_config, "has_sampling_extra_args", False):
|
||
extra_args_list: list[dict] = []
|
||
for req_id in self.input_batch.req_ids:
|
||
req = self.requests[req_id]
|
||
sp = req.sampling_params if req else None
|
||
extra_args_list.append(sp.extra_args if sp and sp.extra_args else {})
|
||
model_kwargs_extra["sampling_extra_args"] = extra_args_list
|
||
|
||
return model_kwargs_extra
|
||
|
||
def _process_additional_information_updates(
|
||
self,
|
||
hidden_states: torch.Tensor,
|
||
multimodal_outputs: Any,
|
||
num_scheduled_tokens_np: np.ndarray,
|
||
scheduler_output: "SchedulerOutput",
|
||
combined_hidden_states: dict[str, torch.Tensor] | None = None,
|
||
combined_multimodal_outputs: dict[str, Any] | None = None,
|
||
req_ids_filter: set[str] | None = None,
|
||
req_ids: list[str] | None = None,
|
||
query_start_loc_cpu: Any = None,
|
||
) -> None:
|
||
"""Process model-provided per-request updates and merge into model_intermediate_buffer."""
|
||
req_ids = req_ids if req_ids is not None else self.input_batch.req_ids
|
||
if query_start_loc_cpu is None:
|
||
query_start_loc_cpu = self.query_start_loc.cpu
|
||
if callable(query_start_loc_cpu):
|
||
query_start_loc_cpu = query_start_loc_cpu()
|
||
try:
|
||
# execute the custom postprocess function
|
||
# TODO(Peiqi): do we have a more elegant way to do this?
|
||
if hasattr(self.model, "has_postprocess") and self.model.has_postprocess:
|
||
postprocess_uses_hidden_states = getattr(self.model, "postprocess_uses_hidden_states", True)
|
||
postprocess_uses_multimodal_outputs = getattr(self.model, "postprocess_uses_multimodal_outputs", True)
|
||
postprocess_uses_req_infos = getattr(self.model, "postprocess_uses_req_infos", True)
|
||
for req_index, req_id in enumerate(req_ids):
|
||
if req_ids_filter is not None and req_id not in req_ids_filter:
|
||
continue
|
||
req_infos = self.model_intermediate_buffer.get(req_id, {}) if postprocess_uses_req_infos else {}
|
||
if postprocess_uses_hidden_states:
|
||
if combined_hidden_states:
|
||
# Combined hidden states contains all hidden states for every request
|
||
hidden_states_slice = combined_hidden_states[req_id]
|
||
else:
|
||
start_offset = int(query_start_loc_cpu[req_index])
|
||
sched_tokens = int(num_scheduled_tokens_np[req_index])
|
||
s, e = start_offset, start_offset + sched_tokens
|
||
# only consider to store data into update dict.
|
||
hidden_states_slice = hidden_states[s:e]
|
||
else:
|
||
hidden_states_slice = hidden_states
|
||
|
||
if not postprocess_uses_multimodal_outputs:
|
||
mm_out = None
|
||
elif combined_multimodal_outputs:
|
||
# NOTE this is a bit ugly, but the mm data is structured as a list of
|
||
# keys mapping to request IDs, and if enabled, we will always have all
|
||
# request IDs in every subdict, including for cache misses.
|
||
mm_out = {k: v[req_id] for k, v in combined_multimodal_outputs.items()}
|
||
else:
|
||
mm_out = multimodal_outputs
|
||
# Exclude 'hidden_states' from kwargs to avoid clash with
|
||
# the positional arg. The buffer entry must be preserved
|
||
# because preprocess reads hidden_states['last'] from it.
|
||
# TODO: pass req_infos as a single payload arg instead of **unpacking
|
||
# to avoid key collisions with positional args.
|
||
postprocess_kwargs = {k: v for k, v in req_infos.items() if k != "hidden_states"}
|
||
update_dict = self.model.postprocess(
|
||
hidden_states_slice,
|
||
multimodal_outputs=mm_out,
|
||
**postprocess_kwargs,
|
||
)
|
||
self._update_intermediate_buffer(req_id, update_dict)
|
||
except Exception as e:
|
||
logger.error(f"Error merging for requests:{req_ids} additional information update: {e}")
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
|
||
def _collect_additional_information_for_prefill(
|
||
self,
|
||
num_scheduled_tokens_np: np.ndarray,
|
||
) -> None:
|
||
"""Overlay per-request prompt_embeds for the prefill portion and collect
|
||
additional_information slices for this step."""
|
||
for req_index, req_id in enumerate(self.input_batch.req_ids):
|
||
req_state = self.requests[req_id]
|
||
pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
|
||
num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
|
||
prompt_len = len(req_state.prompt_token_ids)
|
||
prompt_remaining = max(0, prompt_len - num_computed_tokens)
|
||
sched_tokens = int(num_scheduled_tokens_np[req_index])
|
||
overlay_len = min(sched_tokens, prompt_remaining)
|
||
if overlay_len <= 0:
|
||
continue
|
||
if overlay_len > 0 and pe_cpu is not None:
|
||
src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to(
|
||
dtype=self.dtype, device=self.device, non_blocking=True
|
||
)
|
||
start_offset = int(self.query_start_loc.cpu[req_index])
|
||
self.inputs_embeds.gpu[start_offset : start_offset + overlay_len].copy_(src)
|
||
|
||
def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None:
|
||
replace = getattr(self.model, "replace_runtime_additional_information", False)
|
||
update_buffer = self._replace_intermediate_buffer if replace else self._update_intermediate_buffer
|
||
for new_req in scheduler_output.scheduled_new_reqs:
|
||
model_buffer = getattr(new_req, "model_intermediate_buffer", None)
|
||
if isinstance(model_buffer, dict) and model_buffer:
|
||
update_buffer(new_req.req_id, model_buffer)
|
||
if replace:
|
||
continue
|
||
payload_info = getattr(new_req, "additional_information", None)
|
||
decoded_info = deserialize_additional_information(payload_info)
|
||
if decoded_info:
|
||
update_buffer(new_req.req_id, decoded_info)
|
||
|
||
if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"):
|
||
cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {})
|
||
if isinstance(cached_infos, dict):
|
||
for req_id, req_infos in cached_infos.items():
|
||
decoded_info = deserialize_additional_information(req_infos)
|
||
if decoded_info:
|
||
update_buffer(req_id, decoded_info)
|
||
|
||
def _maybe_attach_mimo_audio_req_infos(
|
||
self,
|
||
req_state: CachedRequestState | None,
|
||
req_infos: dict | None,
|
||
req_id: str,
|
||
) -> dict | None:
|
||
"""Attach MiMoAudio-specific fields into req_infos if applicable.
|
||
|
||
This helper is intentionally small and self-contained so that it can be
|
||
unit-tested to prevent regressions when updating MiMoAudio handling.
|
||
"""
|
||
if req_state is None or self.model.__class__.__name__ != "MiMoAudioForConditionalGeneration":
|
||
return req_infos
|
||
|
||
# Always operate on a dict copy to avoid mutating shared instances.
|
||
req_infos = dict(req_infos) if isinstance(req_infos, dict) else {}
|
||
mm_features = getattr(req_state, "mm_features", None)
|
||
if mm_features and (not req_infos.get("mm_features")):
|
||
req_infos["mm_features"] = mm_features
|
||
req_infos["req_id"] = req_id
|
||
|
||
return req_infos
|
||
|
||
def _maybe_run_batch_preprocess(self, req_ids: list[str], device: torch.device) -> None:
|
||
"""Run an optional model-specific batch preprocess hook.
|
||
|
||
The generic runner only supplies current request ids and the runner-owned
|
||
intermediate buffer; model-specific code decides whether there is any
|
||
batchable work.
|
||
"""
|
||
preprocess_batch = getattr(self.model, "preprocess_batch", None)
|
||
if not callable(preprocess_batch):
|
||
return
|
||
preprocess_batch(
|
||
req_ids=req_ids,
|
||
model_intermediate_buffer=self.model_intermediate_buffer,
|
||
device=device,
|
||
)
|
||
|
||
def _preprocess(
|
||
self,
|
||
scheduler_output: "SchedulerOutput",
|
||
num_input_tokens: int,
|
||
intermediate_tensors: IntermediateTensors | None = None,
|
||
):
|
||
"""Align with v0.14.0 preprocess and omni's additional information handling.
|
||
|
||
Note:
|
||
Upstream vLLM (commit c621af169) added a conditional in the
|
||
``supports_mm_inputs`` path to handle precomputed ``prompt_embeds``
|
||
alongside multimodal inputs. Omni models that reach this override
|
||
always go through the omni-specific ``model.preprocess`` /
|
||
``has_preprocess`` code path below, so the upstream change is
|
||
deliberately not ported.
|
||
"""
|
||
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
|
||
is_first_rank = get_pp_group().is_first_rank
|
||
is_encoder_decoder = self.model_config.is_encoder_decoder
|
||
|
||
# _prepare_inputs may reorder the batch, so we must gather multi
|
||
# modal outputs after that to ensure the correct order
|
||
ec_connector_output = None
|
||
|
||
if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder:
|
||
# Run the multimodal encoder if any.
|
||
with self.maybe_get_ec_connector_output(
|
||
scheduler_output,
|
||
encoder_cache=self.encoder_cache,
|
||
) as ec_connector_output:
|
||
self._execute_mm_encoder(scheduler_output)
|
||
mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output)
|
||
|
||
# NOTE(woosuk): To unify token ids and soft tokens (vision
|
||
# embeddings), we always use embeddings (rather than token ids)
|
||
# as input to the multimodal model, even when the input is text.
|
||
inputs_embeds_scheduled = self.model.embed_input_ids(
|
||
self.input_ids.gpu[:num_scheduled_tokens],
|
||
multimodal_embeddings=mm_embeds,
|
||
is_multimodal=is_mm_embed,
|
||
)
|
||
|
||
# TODO(woosuk): Avoid the copy. Optimize.
|
||
self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
|
||
|
||
input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens)
|
||
model_kwargs = {
|
||
**self._init_model_kwargs(),
|
||
**self._extract_mm_kwargs(scheduler_output),
|
||
}
|
||
elif self.enable_prompt_embeds and is_first_rank:
|
||
# Get the input embeddings for the tokens that are not input embeds,
|
||
# then put them into the appropriate positions.
|
||
# TODO(qthequartermasterman): Since even when prompt embeds are
|
||
# enabled, (a) not all requests will use prompt embeds, and (b)
|
||
# after the initial prompt is processed, the rest of the generated
|
||
# tokens will be token ids, it is not desirable to have the
|
||
# embedding layer outside of the CUDA graph all the time. The v0
|
||
# engine avoids this by "double compiling" the CUDA graph, once
|
||
# with input_ids and again with inputs_embeds, for all num_tokens.
|
||
# If a batch only has token ids, then including the embedding layer
|
||
# in the CUDA graph will be more performant (like in the else case
|
||
# below).
|
||
token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
|
||
# Some tokens ids may need to become embeds
|
||
if token_ids_idx.numel() > 0:
|
||
token_ids = self.input_ids.gpu[token_ids_idx]
|
||
tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids)
|
||
self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
|
||
|
||
inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
|
||
model_kwargs = self._init_model_kwargs()
|
||
input_ids = None
|
||
elif getattr(self.model, "has_preprocess", False):
|
||
# Use pre-allocated buffer for CUDA graph compatibility.
|
||
input_ids = self.input_ids.gpu[:num_input_tokens]
|
||
inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
|
||
model_kwargs = self._init_model_kwargs()
|
||
else:
|
||
# For text-only models, we use token ids as input.
|
||
# While it is possible to use embeddings as input just like the
|
||
# multimodal models, it is not desirable for performance since
|
||
# then the embedding layer is not included in the CUDA graph.
|
||
input_ids = self.input_ids.gpu[:num_input_tokens]
|
||
inputs_embeds = None
|
||
model_kwargs = self._init_model_kwargs()
|
||
|
||
if self.uses_mrope:
|
||
positions = self.mrope_positions.gpu[:, :num_input_tokens]
|
||
elif self.uses_xdrope_dim > 0:
|
||
positions = self.xdrope_positions.gpu[:, :num_input_tokens]
|
||
else:
|
||
positions = self.positions[:num_input_tokens]
|
||
if num_input_tokens > num_scheduled_tokens:
|
||
self.positions[num_scheduled_tokens:num_input_tokens].zero_()
|
||
|
||
if is_first_rank:
|
||
intermediate_tensors = None
|
||
else:
|
||
assert intermediate_tensors is not None
|
||
intermediate_tensors = self.sync_and_gather_intermediate_tensors(
|
||
num_input_tokens, intermediate_tensors, True
|
||
)
|
||
|
||
if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
|
||
# Run the encoder, just like we do with other multimodal inputs.
|
||
# For an encoder-decoder model, our processing here is a bit
|
||
# simpler, because the outputs are just passed to the decoder.
|
||
# We are not doing any prompt replacement. We also will only
|
||
# ever have a single encoder input.
|
||
encoder_outputs = self._execute_mm_encoder(scheduler_output)
|
||
model_kwargs.update({"encoder_outputs": encoder_outputs})
|
||
|
||
req_ids = self.input_batch.req_ids
|
||
num_scheduled_tokens_np = np.array(
|
||
[scheduler_output.num_scheduled_tokens[rid] for rid in req_ids],
|
||
dtype=np.int32,
|
||
)
|
||
self._omni_num_scheduled_tokens_np = num_scheduled_tokens_np
|
||
|
||
# Note: only prefill need collect additional_information for now.
|
||
# Decode don't need per_req_additional_information anymore.
|
||
if inputs_embeds is not None:
|
||
# Prefill: overlay prompt_embeds and collect additional_information
|
||
self._collect_additional_information_for_prefill(num_scheduled_tokens_np)
|
||
|
||
# Keep per-request additional_information in sync for both new and
|
||
# cached requests. This is required for stages without preprocess
|
||
# (e.g., code2wav) so runtime_additional_information can be refreshed
|
||
# from scheduler cached infos on every step.
|
||
if hasattr(self.model, "has_preprocess") or hasattr(self.model, "enable_update_additional_information"):
|
||
self._update_additional_information(scheduler_output)
|
||
if not self.vllm_config.model_config.async_chunk:
|
||
# In full-payload (non-async-chunk) mode, connector-delivered
|
||
# stage payloads must override any earlier engine-level
|
||
# additional_information written by the legacy
|
||
# custom_process_input_func codec, so talker_preprocess reads
|
||
# the full thinker payload.
|
||
self._sync_local_stage_payloads()
|
||
|
||
if hasattr(self.model, "has_preprocess") and self.model.has_preprocess:
|
||
# Multimodal wrappers forward embeddings and therefore intentionally
|
||
# set ``input_ids`` to None. Model-specific preprocess hooks still
|
||
# need the scheduled token ids to build or replace those embeddings.
|
||
preprocess_input_ids = input_ids if input_ids is not None else self.input_ids.gpu[:num_input_tokens]
|
||
preprocess_device = preprocess_input_ids.device
|
||
self._maybe_run_batch_preprocess(self.input_batch.req_ids, preprocess_device)
|
||
|
||
# Overlay custom prompt_embeds per request for the prompt portion;
|
||
# collect additional_information (tensor/list) for prefill portion only
|
||
decode_req_ids: list[str] = []
|
||
decode_start_offsets: list[int] = []
|
||
decode_batch_items: list[tuple[str, int, dict[str, Any]]] = []
|
||
batch_decode_preprocess = getattr(self.model, "preprocess_decode_batch", None)
|
||
|
||
def flush_decode_batch() -> None:
|
||
nonlocal inputs_embeds
|
||
if not decode_batch_items:
|
||
return
|
||
|
||
assert callable(batch_decode_preprocess)
|
||
req_ids_b = [item[0] for item in decode_batch_items]
|
||
start_offsets_b = [item[1] for item in decode_batch_items]
|
||
req_infos_b = [item[2] for item in decode_batch_items]
|
||
ids_b = torch.stack(
|
||
[preprocess_input_ids[offset : offset + 1].reshape(-1)[0] for offset in start_offsets_b]
|
||
)
|
||
req_input_ids, req_embeds, last_talker_hidden, text_step, updates = batch_decode_preprocess(
|
||
input_ids=ids_b,
|
||
req_infos=req_infos_b,
|
||
)
|
||
if inputs_embeds is None:
|
||
inputs_embeds = torch.empty(
|
||
(preprocess_input_ids.shape[0], req_embeds.shape[-1]),
|
||
device=req_embeds.device,
|
||
dtype=req_embeds.dtype,
|
||
)
|
||
|
||
offsets_t = torch.tensor(start_offsets_b, device=req_embeds.device, dtype=torch.long)
|
||
inputs_embeds.index_copy_(0, offsets_t, req_embeds)
|
||
preprocess_input_ids.index_copy_(
|
||
0,
|
||
offsets_t,
|
||
req_input_ids.reshape(-1).to(dtype=preprocess_input_ids.dtype),
|
||
)
|
||
|
||
dst = slice(len(decode_req_ids), len(decode_req_ids) + len(req_ids_b))
|
||
self.talker_mtp_input_ids.gpu[dst].copy_(req_input_ids.reshape(-1))
|
||
self.talker_mtp_inputs_embeds.gpu[dst].copy_(req_embeds)
|
||
self.last_talker_hidden.gpu[dst].copy_(last_talker_hidden)
|
||
self.text_step.gpu[dst].copy_(text_step)
|
||
|
||
for req_id_b, update_dict_b in zip(req_ids_b, updates, strict=True):
|
||
self._update_intermediate_buffer(req_id_b, update_dict_b)
|
||
|
||
decode_req_ids.extend(req_ids_b)
|
||
decode_start_offsets.extend(start_offsets_b)
|
||
decode_batch_items.clear()
|
||
|
||
preprocess_input_ids = input_ids if input_ids is not None else self.input_ids.gpu[:num_input_tokens]
|
||
for req_index, req_id in enumerate(self.input_batch.req_ids):
|
||
req_infos = self.model_intermediate_buffer.get(req_id, {})
|
||
|
||
# mimo-audio check
|
||
req_state = self.requests.get(req_id)
|
||
req_infos = cast(dict[str, Any], self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id))
|
||
|
||
start_offset = int(self.query_start_loc.cpu[req_index])
|
||
sched_tokens = int(num_scheduled_tokens_np[req_index])
|
||
s, e = start_offset, start_offset + sched_tokens
|
||
span_len = int(e) - int(s)
|
||
|
||
# call the custom process function
|
||
req_infos["request_id"] = req_id
|
||
req_infos["duplex_token_offset"] = int(self.input_batch.num_computed_tokens_cpu[req_index])
|
||
req_infos["duplex_prompt_len"] = len(req_state.prompt_token_ids) if req_state is not None else None
|
||
prompt_token_ids = getattr(req_state, "prompt_token_ids", ()) if req_state is not None else ()
|
||
prompt_len = len(prompt_token_ids or ())
|
||
num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
|
||
is_prefill = num_computed_tokens < prompt_len
|
||
req_infos["_omni_prompt_len"] = prompt_len
|
||
req_infos["_omni_num_computed_tokens"] = num_computed_tokens
|
||
req_infos["_omni_is_prefill"] = is_prefill
|
||
# The has_preprocess branch offers a reusable embedding buffer
|
||
# that is not initialized for ordinary text requests. Models
|
||
# may reuse its contents only when this runner filled them.
|
||
req_infos["_omni_input_embeds_precomputed"] = bool(
|
||
is_first_rank
|
||
and (
|
||
(self.supports_mm_inputs and not is_encoder_decoder)
|
||
or self.enable_prompt_embeds
|
||
)
|
||
)
|
||
# Output-token cap, so a model that must ship a payload on the
|
||
# request's final step can tell which step that is. A finished
|
||
# request drops out of req_ids_output_copy, and downstream
|
||
# routing keys off that list, so a payload produced after its
|
||
# last step never reaches the caller.
|
||
sampling_params = getattr(req_state, "sampling_params", None)
|
||
req_infos["_omni_max_tokens"] = getattr(sampling_params, "max_tokens", None)
|
||
# Seed, so a model that samples inside forward() can be
|
||
# reproducible: vLLM's own sampler seeding does not reach it.
|
||
req_infos["_omni_seed"] = getattr(sampling_params, "seed", None)
|
||
if callable(batch_decode_preprocess) and self.has_talker_mtp and span_len == 1 and not is_prefill:
|
||
decode_batch_items.append((req_id, s, req_infos))
|
||
continue
|
||
|
||
flush_decode_batch()
|
||
|
||
embed_slice = inputs_embeds[s:e] if inputs_embeds is not None else None
|
||
req_input_ids, req_embeds, update_dict = self.model.preprocess(
|
||
input_ids=preprocess_input_ids[s:e],
|
||
input_embeds=embed_slice,
|
||
**req_infos,
|
||
)
|
||
if inputs_embeds is None:
|
||
inputs_embeds = torch.empty(
|
||
(preprocess_input_ids.shape[0], req_embeds.shape[-1]),
|
||
device=req_embeds.device,
|
||
dtype=req_embeds.dtype,
|
||
)
|
||
input_ids = preprocess_input_ids
|
||
|
||
if self.has_talker_mtp and span_len == 1 and not is_prefill:
|
||
last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
|
||
decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1)
|
||
self.talker_mtp_input_ids.gpu[decode_slice].copy_(req_input_ids)
|
||
self.talker_mtp_inputs_embeds.gpu[decode_slice].copy_(req_embeds)
|
||
self.last_talker_hidden.gpu[decode_slice].copy_(last_talker_hidden)
|
||
self.text_step.gpu[decode_slice].copy_(text_step)
|
||
decode_req_ids.append(req_id)
|
||
decode_start_offsets.append(s)
|
||
|
||
# TODO(Peiqi): the merge stage could move out from the critical path
|
||
self._update_intermediate_buffer(req_id, update_dict)
|
||
|
||
# update the inputs_embeds and input_ids
|
||
seg_len = min(span_len, req_embeds.shape[0])
|
||
inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
|
||
if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
|
||
preprocess_input_ids[s : s + seg_len] = req_input_ids
|
||
if input_ids is None:
|
||
input_ids = preprocess_input_ids
|
||
|
||
flush_decode_batch()
|
||
|
||
# run talker mtp decode
|
||
if self.has_talker_mtp:
|
||
self._talker_mtp_forward(decode_req_ids, inputs_embeds, decode_start_offsets)
|
||
|
||
return (
|
||
input_ids,
|
||
inputs_embeds,
|
||
positions,
|
||
intermediate_tensors,
|
||
model_kwargs,
|
||
ec_connector_output,
|
||
)
|
||
|
||
def _talker_mtp_forward(
|
||
self,
|
||
decode_req_ids: list[str],
|
||
inputs_embeds: torch.Tensor,
|
||
start_offsets: list[int] | None = None,
|
||
) -> None:
|
||
decode_batch_size = len(decode_req_ids)
|
||
if decode_batch_size == 0:
|
||
return
|
||
_cudagraph_mode, batch_desc, _, _, _ = self._determine_batch_execution_and_padding(
|
||
num_tokens=decode_batch_size,
|
||
num_reqs=decode_batch_size,
|
||
num_scheduled_tokens_np=np.ones(decode_batch_size, dtype=np.int32),
|
||
max_num_scheduled_tokens=1,
|
||
use_cascade_attn=False,
|
||
)
|
||
# Force eager for unwrapped code predictors (AR loops / multinomial).
|
||
# When talker_mtp is not wrapped by the platform's full-graph wrapper,
|
||
# it manages its own device graphs internally (code_predictor has its
|
||
# own bucket sizes).
|
||
if not isinstance(self.talker_mtp, current_omni_platform.get_graph_wrapper_cls()):
|
||
_cudagraph_mode = CUDAGraphMode.NONE
|
||
num_tokens_padded = decode_batch_size
|
||
else:
|
||
num_tokens_padded = batch_desc.num_tokens
|
||
req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded]
|
||
req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded]
|
||
last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded]
|
||
text_step = self.text_step.gpu[:num_tokens_padded]
|
||
subtalker_params = getattr(self.vllm_config.model_config, "subtalker_sampling_params", None)
|
||
if not isinstance(subtalker_params, dict):
|
||
subtalker_params = {}
|
||
|
||
def _explicit_talker_seed(req_id: str) -> int | None:
|
||
sampling_params = getattr(self.requests[req_id], "sampling_params", None)
|
||
extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
|
||
seed = None
|
||
if isinstance(extra_args, dict):
|
||
seed = extra_args.get("tts_local_seed")
|
||
return int(seed) if seed is not None else None
|
||
|
||
def _row_generator(req_id: str) -> torch.Generator | None:
|
||
seed = _explicit_talker_seed(req_id)
|
||
if seed is None:
|
||
return None
|
||
cache = getattr(self, "_talker_mtp_generators", None)
|
||
if cache is None:
|
||
cache = {}
|
||
self._talker_mtp_generators = cache
|
||
generator = cache.get(req_id)
|
||
if generator is None or generator.device != req_input_ids.device:
|
||
generator = torch.Generator(device=req_input_ids.device)
|
||
generator.manual_seed(seed)
|
||
cache[req_id] = generator
|
||
return generator
|
||
|
||
row_generators = [_row_generator(req_id) for req_id in decode_req_ids]
|
||
cache = getattr(self, "_talker_mtp_generators", None)
|
||
if cache:
|
||
# Generators live as long as their request; drop finished ones.
|
||
for stale_id in [rid for rid in cache if rid not in self.requests]:
|
||
del cache[stale_id]
|
||
|
||
if (
|
||
decode_batch_size > 1
|
||
and any(generator is not None for generator in row_generators)
|
||
and not getattr(self.model, "talker_mtp_accepts_per_row_generators", False)
|
||
):
|
||
# A torch.Generator is a single stream. Using one generator for a
|
||
# multi-row batch would make explicitly-seeded requests depend on
|
||
# other rows in the same scheduler step, so keep that path scalar.
|
||
saved_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size].clone()
|
||
saved_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].clone()
|
||
saved_hidden = self.last_talker_hidden.gpu[:decode_batch_size].clone()
|
||
saved_text = self.text_step.gpu[:decode_batch_size].clone()
|
||
try:
|
||
for row, req_id in enumerate(decode_req_ids):
|
||
self.talker_mtp_input_ids.gpu[:1].copy_(saved_input_ids[row : row + 1])
|
||
self.talker_mtp_inputs_embeds.gpu[:1].copy_(saved_embeds[row : row + 1])
|
||
self.last_talker_hidden.gpu[:1].copy_(saved_hidden[row : row + 1])
|
||
self.text_step.gpu[:1].copy_(saved_text[row : row + 1])
|
||
row_offsets = None if start_offsets is None else [start_offsets[row]]
|
||
self._talker_mtp_forward([req_id], inputs_embeds, row_offsets)
|
||
finally:
|
||
self.talker_mtp_input_ids.gpu[:decode_batch_size].copy_(saved_input_ids)
|
||
self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].copy_(saved_embeds)
|
||
self.last_talker_hidden.gpu[:decode_batch_size].copy_(saved_hidden)
|
||
self.text_step.gpu[:decode_batch_size].copy_(saved_text)
|
||
return
|
||
|
||
talker_kwargs = {
|
||
"do_sample": subtalker_params.get("do_sample"),
|
||
"temperature": subtalker_params.get("temperature"),
|
||
"top_k": subtalker_params.get("top_k"),
|
||
"top_p": subtalker_params.get("top_p"),
|
||
}
|
||
if decode_batch_size == 1:
|
||
if row_generators[0] is not None:
|
||
talker_kwargs["generator"] = row_generators[0]
|
||
elif any(generator is not None for generator in row_generators):
|
||
talker_kwargs["generators"] = row_generators
|
||
if getattr(self.model, "talker_mtp_accepts_req_infos", False):
|
||
talker_kwargs["req_ids"] = decode_req_ids
|
||
talker_kwargs["req_infos"] = [
|
||
self.model_intermediate_buffer.setdefault(req_id, {}) for req_id in decode_req_ids
|
||
]
|
||
from vllm_omni.edge.step_stats import StepStats
|
||
|
||
with (
|
||
current_omni_platform.set_forward_context(
|
||
None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
|
||
),
|
||
StepStats.get().timed("model.mtp_graph_ms", sync_device=getattr(self, "device", None)),
|
||
):
|
||
req_embeds, code_predictor_codes = self.talker_mtp(
|
||
req_input_ids,
|
||
req_embeds,
|
||
last_talker_hidden,
|
||
text_step,
|
||
**talker_kwargs,
|
||
)
|
||
# update the inputs_embeds and code_predictor_codes
|
||
out_key = getattr(self.model, "talker_mtp_output_key", ("codes", "audio"))
|
||
if not isinstance(out_key, tuple) or len(out_key) != 2:
|
||
raise TypeError(f"talker_mtp_output_key must be a 2-tuple, got {type(out_key).__name__}: {out_key!r}")
|
||
if start_offsets is None:
|
||
id_to_index = self.input_batch.req_id_to_index
|
||
start_offsets = [int(self.query_start_loc.cpu[id_to_index[req_id]]) for req_id in decode_req_ids]
|
||
if start_offsets == list(range(decode_batch_size)):
|
||
inputs_embeds[:decode_batch_size].copy_(req_embeds[:decode_batch_size])
|
||
else:
|
||
# A device tensor constructor waits for the preceding MTP graph.
|
||
# Enqueue the small host index transfer without synchronizing it.
|
||
offsets = torch.tensor(start_offsets, device="cpu", dtype=torch.long).to(
|
||
inputs_embeds.device, non_blocking=True
|
||
)
|
||
inputs_embeds.index_copy_(0, offsets, req_embeds[:decode_batch_size])
|
||
if code_predictor_codes is not None:
|
||
if out_key in getattr(self.model, "gpu_resident_buffer_keys", set()):
|
||
# One owned batch snapshot protects all rows from graph reuse.
|
||
owned_codes = code_predictor_codes[:decode_batch_size].detach().clone()
|
||
for req_id, row in zip(decode_req_ids, owned_codes.split(1), strict=True):
|
||
req_state = self.requests.get(req_id)
|
||
if req_state is not None:
|
||
existing = self.model_intermediate_buffer.setdefault(req_id, {})
|
||
existing.setdefault(out_key[0], {})[out_key[1]] = row
|
||
req_state.additional_information_cpu = existing
|
||
else:
|
||
for idx, req_id in enumerate(decode_req_ids):
|
||
update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
|
||
self._update_intermediate_buffer(req_id, update_dict)
|
||
|
||
def _accepted_forward_kwargs(self) -> frozenset[str] | None:
|
||
"""Names the stage model's ``forward`` will accept, or ``None`` for any.
|
||
|
||
Omni-native stage models take ``**kwargs`` and pick out what they need,
|
||
so the AR runner hands every stage ``sampling_metadata``,
|
||
``logits_index`` and ``sampler`` unconditionally. A model that vLLM owns
|
||
does not: ``LlamaForCausalLM.forward`` is
|
||
``(input_ids, positions, intermediate_tensors, inputs_embeds)`` and
|
||
nothing else, so passing the omni extras is an immediate ``TypeError``
|
||
on the first forward -- which is what "the model is in the registry" not
|
||
meaning "the model runs under Omni" looks like in practice (proposal
|
||
2.4, item 12).
|
||
|
||
Dropping them is the correct behaviour rather than a workaround: those
|
||
arguments exist for stages that sample inside the model, and a
|
||
vLLM-native causal LM samples through vLLM's own sampler afterwards.
|
||
Computed once per runner -- ``inspect.signature`` on every decode step
|
||
would be absurd.
|
||
"""
|
||
cached = getattr(self, "_omni_forward_kwargs_cache", "unset")
|
||
if cached != "unset":
|
||
return cached
|
||
|
||
import inspect
|
||
|
||
accepted: frozenset[str] | None
|
||
try:
|
||
# CUDA graph wrappers delegate attributes to their wrapped model;
|
||
# their class does not define ``forward``. Inspect the resolved
|
||
# callable so eager and captured native models use the same filter.
|
||
forward = getattr(self.model, "forward", self.model)
|
||
parameters = inspect.signature(forward).parameters
|
||
except (TypeError, ValueError):
|
||
accepted = None # not introspectable; assume it copes, as before
|
||
else:
|
||
if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()):
|
||
accepted = None
|
||
else:
|
||
accepted = frozenset(parameters) - {"self"}
|
||
self._omni_forward_kwargs_cache = accepted
|
||
return accepted
|
||
|
||
def _model_forward(
|
||
self,
|
||
input_ids: torch.Tensor | None = None,
|
||
positions: torch.Tensor | None = None,
|
||
intermediate_tensors: IntermediateTensors | None = None,
|
||
inputs_embeds: torch.Tensor | None = None,
|
||
**model_kwargs: dict[str, Any],
|
||
):
|
||
"""Inject omni-specific kwargs into forward and cache model output"""
|
||
model_kwargs_extra = self._build_model_kwargs_extra()
|
||
accepted = self._accepted_forward_kwargs()
|
||
if accepted is not None:
|
||
dropped = sorted((set(model_kwargs) | set(model_kwargs_extra)) - accepted)
|
||
if dropped and not getattr(self, "_omni_logged_dropped_kwargs", False):
|
||
self._omni_logged_dropped_kwargs = True
|
||
logger.info(
|
||
"[Omni] %s.forward does not accept %s; passing the rest. This is "
|
||
"the expected path for a vLLM-native model used as an omni AR "
|
||
"stage -- it samples through vLLM's sampler, not its own.",
|
||
type(self.model).__name__,
|
||
dropped,
|
||
)
|
||
model_kwargs = {k: v for k, v in model_kwargs.items() if k in accepted}
|
||
model_kwargs_extra = {k: v for k, v in model_kwargs_extra.items() if k in accepted}
|
||
update_decode_metadata = getattr(self.model, "update_decode_step_metadata", None)
|
||
if getattr(self.model, "supports_omni_decode_step_metadata", False) and callable(update_decode_metadata):
|
||
update_decode_metadata(
|
||
input_ids=input_ids,
|
||
positions=positions,
|
||
inputs_embeds=inputs_embeds,
|
||
omni_query_start_loc=model_kwargs_extra.get("omni_query_start_loc"),
|
||
req_ids=self.input_batch.req_ids,
|
||
)
|
||
|
||
model_output = super()._model_forward(
|
||
input_ids=input_ids,
|
||
positions=positions,
|
||
intermediate_tensors=intermediate_tensors,
|
||
inputs_embeds=inputs_embeds,
|
||
**model_kwargs,
|
||
**model_kwargs_extra,
|
||
)
|
||
if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
|
||
model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra)
|
||
# Cache model output so later sample_tokens can consume multimodal results.
|
||
self._omni_last_model_output = model_output
|
||
return model_output
|
||
|
||
def _store_value(self, dest: dict, key: str, value: Any, gpu_keys: set) -> None:
|
||
if isinstance(value, torch.Tensor):
|
||
if key in gpu_keys:
|
||
dest[key] = value.detach().clone()
|
||
else:
|
||
t = value.detach()
|
||
if t.is_cuda:
|
||
dest[key] = t.to("cpu").contiguous()
|
||
else:
|
||
# If the tensor is already on the CPU, there is no need to unload it to the CPU.
|
||
dest[key] = t.contiguous()
|
||
elif isinstance(value, list):
|
||
dest[key] = [
|
||
(item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in value
|
||
]
|
||
else:
|
||
dest[key] = value
|
||
|
||
def _update_intermediate_buffer(self, req_id: str, upd: dict) -> None:
|
||
if not isinstance(upd, dict) or not upd:
|
||
return
|
||
req_state = self.requests.get(req_id)
|
||
if req_state is None:
|
||
return
|
||
# Check if the model declares keys that should stay on GPU (tuples of (type_key, qualifier))
|
||
gpu_keys: set[tuple[str, str]] = set()
|
||
if hasattr(self, "model") and hasattr(self.model, "gpu_resident_buffer_keys"):
|
||
gpu_keys = self.model.gpu_resident_buffer_keys
|
||
existing = self.model_intermediate_buffer.setdefault(req_id, {})
|
||
for k, v in upd.items():
|
||
if isinstance(v, dict):
|
||
existing_sub = existing.setdefault(k, {})
|
||
for qual, val in v.items():
|
||
self._store_value(existing_sub, qual, val, {q for tk, q in gpu_keys if tk == k})
|
||
else:
|
||
self._store_value(existing, k, v, set())
|
||
# Backward compatible: mirror to old setattr location
|
||
setattr(req_state, "additional_information_cpu", existing)
|
||
|
||
def _replace_intermediate_buffer(self, req_id: str, upd: dict) -> None:
|
||
"""Replace one request's runtime payload with the current chunk."""
|
||
if not isinstance(upd, dict):
|
||
return
|
||
req_state = self.requests.get(req_id)
|
||
if req_state is None:
|
||
return
|
||
snapshot = dict(upd)
|
||
previous = self.model_intermediate_buffer.get(req_id)
|
||
previous_meta = previous.get("meta") if isinstance(previous, dict) else None
|
||
incoming_meta = snapshot.get("meta")
|
||
meta = dict(incoming_meta) if isinstance(incoming_meta, dict) else {}
|
||
if isinstance(previous_meta, dict):
|
||
for key in ("num_processed_tokens", "resumable"):
|
||
if key not in meta and key in previous_meta:
|
||
meta[key] = previous_meta[key]
|
||
if meta or "meta" in snapshot:
|
||
snapshot["meta"] = meta
|
||
self.model_intermediate_buffer[req_id] = snapshot
|
||
setattr(req_state, "additional_information_cpu", snapshot)
|
||
|
||
def _update_streaming_input_additional_info(self, new_req_data, req_id):
|
||
# For streaming input prefill case only. Update buffer from last segment input.
|
||
cached_additional_info = self.model_intermediate_buffer.get(req_id, {})
|
||
inc_info = getattr(new_req_data, "model_intermediate_buffer", None)
|
||
if not isinstance(inc_info, dict) or not inc_info:
|
||
payload_info = getattr(new_req_data, "additional_information", None)
|
||
inc_info = deserialize_additional_information(payload_info)
|
||
if getattr(self.model, "replace_runtime_additional_information", False):
|
||
snapshot = dict(inc_info) if isinstance(inc_info, dict) else {}
|
||
meta = snapshot.get("meta")
|
||
meta = dict(meta) if isinstance(meta, dict) else {}
|
||
meta["num_processed_tokens"] = 0
|
||
meta["resumable"] = True
|
||
snapshot["meta"] = meta
|
||
self._replace_intermediate_buffer(req_id, snapshot)
|
||
return
|
||
if not isinstance(inc_info, dict) or not inc_info:
|
||
return
|
||
merged_info = dict(cached_additional_info) if isinstance(cached_additional_info, dict) else {}
|
||
for key, value in inc_info.items():
|
||
if isinstance(value, dict):
|
||
existing_sub = merged_info.get(key)
|
||
merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {}
|
||
for sk, sv in value.items():
|
||
merged_sub[sk] = sv
|
||
merged_info[key] = merged_sub
|
||
else:
|
||
merged_info[key] = value
|
||
merged_info.setdefault("meta", {})["num_processed_tokens"] = 0
|
||
merged_info.setdefault("meta", {})["resumable"] = True
|
||
self.model_intermediate_buffer[req_id] = merged_info
|
||
setattr(self.requests[req_id], "additional_information_cpu", merged_info)
|