1172 lines
55 KiB
Python
1172 lines
55 KiB
Python
"""AR GPU Model Runner for vLLM-Omni.
|
|
|
|
Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
|
|
and also outputs sampled tokens.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import nullcontext
|
|
from copy import copy
|
|
from dataclasses import replace
|
|
from typing import Any, NamedTuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from vllm.config import CUDAGraphMode
|
|
from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
|
|
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
|
|
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
|
|
from vllm.forward_context import set_forward_context
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
|
|
extract_routed_experts_for_current_batch,
|
|
get_global_experts_capturer,
|
|
issue_routing_d2h_copy,
|
|
)
|
|
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
|
|
from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
|
|
from vllm.v1.spec_decode.dflash import DFlashProposer
|
|
from vllm.v1.spec_decode.draft_model import DraftModelProposer
|
|
from vllm.v1.spec_decode.eagle import EagleProposer
|
|
from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
|
|
from vllm.v1.spec_decode.gemma4 import Gemma4Proposer
|
|
from vllm.v1.structured_output.utils import apply_grammar_bitmask
|
|
from vllm.v1.utils import record_function_or_nullcontext
|
|
from vllm.v1.worker.gpu_model_runner import (
|
|
EMPTY_MODEL_RUNNER_OUTPUT,
|
|
AsyncGPUModelRunnerOutput,
|
|
IntermediateTensors,
|
|
)
|
|
from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
|
|
from vllm.v1.worker.utils import is_residual_scattered_for_sp
|
|
|
|
from vllm_omni.data_entry_keys import flatten_payload
|
|
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
|
|
from vllm_omni.outputs import OmniModelRunnerOutput
|
|
from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element
|
|
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
|
|
from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class ExecuteModelState(NamedTuple):
|
|
scheduler_output: SchedulerOutput
|
|
logits: torch.Tensor | None
|
|
spec_decode_metadata: Any
|
|
spec_decode_common_attn_metadata: Any
|
|
hidden_states: torch.Tensor
|
|
sample_hidden_states: torch.Tensor
|
|
aux_hidden_states: list[torch.Tensor] | None
|
|
ec_connector_output: Any
|
|
cudagraph_stats: Any
|
|
# OMNI: multimodal_outputs field for omni-specific multimodal handling
|
|
multimodal_outputs: Any
|
|
# slot_mappings for attention/drafter (aligned with upstream v1 API)
|
|
slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None
|
|
|
|
|
|
class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
|
|
"""Autoregressive GPU model runner that returns hidden states per request.
|
|
|
|
Follows the v0.12 two-phase execute/sample flow from GPUModelRunner, and
|
|
reuses Omni hooks for additional_information / multimodal outputs. This
|
|
class only overrides sample_tokens to expose hidden states + multimodal
|
|
outputs per request while keeping Async output semantics.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
|
|
# each model stage has their own hidden size
|
|
self.hidden_size = self.model_config.hf_text_config.hidden_size
|
|
self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
|
|
# Initialize KV cache manager (preserve vllm_config fallback behavior)
|
|
self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
|
|
# Worker-connector init is gated by a per-`model_arch` allowlist
|
|
# (covers both producer-side and consumer-side runners for the
|
|
# arches below). Consumer-wait stages must be registered
|
|
# separately as `(model_arch, model_stage)` tuples in
|
|
# `omni_scheduling_coordinator._FULL_PAYLOAD_INPUT_STAGES`;
|
|
# forgetting that produces a Stage-1 hang on the consumer.
|
|
_OMNI_CONNECTOR_INIT_ARCHS = {
|
|
"Qwen3OmniMoeForConditionalGeneration",
|
|
"Qwen2_5OmniForConditionalGeneration",
|
|
"CovoAudioForConditionalGeneration",
|
|
"MiMoAudioModel",
|
|
"Qwen3TTSTalkerForConditionalGeneration",
|
|
"Qwen3TTSCode2Wav",
|
|
"CosyVoice3Model",
|
|
"DyninOmniForConditionalGeneration",
|
|
}
|
|
if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
|
|
self.init_omni_connectors(
|
|
vllm_config=self.vllm_config,
|
|
model_config=self.model_config,
|
|
kv_transfer_manager=self.kv_transfer_manager,
|
|
)
|
|
self._downstream_payload_cache: dict[str, bool] = {}
|
|
|
|
def _make_buffer(self, *size, dtype, numpy=True):
|
|
# Prevent ray from pinning the buffer due to large size
|
|
from vllm_omni.distributed.ray_utils.utils import (
|
|
calculate_total_bytes,
|
|
maybe_disable_pin_memory_for_ray,
|
|
)
|
|
|
|
total_bytes = calculate_total_bytes(size, dtype)
|
|
|
|
# Use the context manager to temporarily disable pinning if needed
|
|
with maybe_disable_pin_memory_for_ray(self, total_bytes):
|
|
return super()._make_buffer(*size, dtype=dtype, numpy=numpy)
|
|
|
|
def _build_model_sampler_output_token_ids(self) -> list[list[int]]:
|
|
"""Build decoded-token history for custom model samplers.
|
|
|
|
vLLM only populates sampling_metadata.output_token_ids when penalties or
|
|
logits processors require it. CosyVoice3's custom RAS sampler also
|
|
depends on this history, so we reconstruct it directly from the input
|
|
batch for prefer_model_sampler models.
|
|
"""
|
|
req_output_token_ids = getattr(self.input_batch, "req_output_token_ids", [])
|
|
req_ids = list(getattr(self.input_batch, "req_ids", []))
|
|
output_token_ids = [list(req_output_token_ids[idx] or []) for idx in range(len(req_ids))]
|
|
|
|
sampled_token_ids_cpu = getattr(self.input_batch, "sampled_token_ids_cpu", None)
|
|
async_copy_ready_event = getattr(self.input_batch, "async_copy_ready_event", None)
|
|
prev_req_id_to_index = getattr(self.input_batch, "prev_req_id_to_index", None)
|
|
if sampled_token_ids_cpu is None or not output_token_ids or prev_req_id_to_index is None:
|
|
return output_token_ids
|
|
|
|
sampled_token_ids: list[list[int]] | None = None
|
|
for index, req_id in enumerate(req_ids):
|
|
prev_index = prev_req_id_to_index.get(req_id)
|
|
if prev_index is None:
|
|
continue
|
|
req_history = output_token_ids[index]
|
|
if not req_history or req_history[-1] != -1:
|
|
continue
|
|
if sampled_token_ids is None:
|
|
assert async_copy_ready_event is not None
|
|
async_copy_ready_event.synchronize()
|
|
sampled_token_ids = sampled_token_ids_cpu.tolist()
|
|
new_ids = list(sampled_token_ids[prev_index])
|
|
if not new_ids:
|
|
continue
|
|
num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1)
|
|
first_placeholder = req_history.index(-1)
|
|
num_placeholders = len(req_history) - first_placeholder
|
|
num_to_replace = min(num_sampled_ids, num_placeholders)
|
|
req_history[first_placeholder : first_placeholder + num_to_replace] = new_ids[:num_to_replace]
|
|
|
|
for index, req_history in enumerate(output_token_ids):
|
|
if -1 in req_history:
|
|
output_token_ids[index] = req_history[: req_history.index(-1)]
|
|
|
|
return output_token_ids
|
|
|
|
def _sampling_metadata_for_model_sampler(self, sampling_metadata):
|
|
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 _request_final_stage_id(self, req_id: str) -> int | None:
|
|
info = self.model_intermediate_buffer.get(req_id)
|
|
if not isinstance(info, dict):
|
|
req_state = self.requests.get(req_id)
|
|
info = getattr(req_state, "additional_information_cpu", None)
|
|
if not isinstance(info, dict):
|
|
return None
|
|
val = info.get("omni_final_stage_id")
|
|
try:
|
|
return int(val)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
def _request_needs_downstream_stage_payload(self, req_id: str) -> bool:
|
|
cached = self._downstream_payload_cache.get(req_id)
|
|
if cached is not None:
|
|
return cached
|
|
# Conservative default: keep payload if marker is missing.
|
|
final_stage_id = self._request_final_stage_id(req_id)
|
|
needs_payload = final_stage_id is None or final_stage_id > 0
|
|
self._downstream_payload_cache[req_id] = needs_payload
|
|
return needs_payload
|
|
|
|
def _resolve_pooler_payload_req_ids(self, req_ids_output_copy: list[str]) -> tuple[str, list[str]]:
|
|
downstream_req_ids = [rid for rid in req_ids_output_copy if self._request_needs_downstream_stage_payload(rid)]
|
|
engine_output_type = (self.vllm_config.model_config.engine_output_type or "").lower()
|
|
# Single-stage AR TTS models (e.g. VoxCPM2) finish on this stage but still
|
|
# need multimodal payloads for final audio postprocess/output.
|
|
if engine_output_type == "audio" and not downstream_req_ids:
|
|
downstream_req_ids = req_ids_output_copy
|
|
return engine_output_type, downstream_req_ids
|
|
|
|
def capture_model(self) -> int:
|
|
result = super().capture_model()
|
|
self._capture_talker_mtp_graphs()
|
|
return result
|
|
|
|
def _capture_talker_mtp_graphs(self) -> None:
|
|
from vllm_omni.worker.gpu_model_runner import CUDAGraphWrapper
|
|
|
|
if not self.has_talker_mtp or not isinstance(self.talker_mtp, CUDAGraphWrapper):
|
|
return
|
|
|
|
from vllm.compilation.monitor import set_cudagraph_capturing_enabled
|
|
from vllm.distributed.parallel_state import graph_capture
|
|
|
|
capture_sizes = self.compilation_config.cudagraph_capture_sizes
|
|
num_warmups = self.compilation_config.cudagraph_num_of_warmups
|
|
capture_sizes = sorted(capture_sizes, reverse=True)
|
|
logger.info("Capturing talker_mtp graphs for sizes %s", capture_sizes)
|
|
|
|
set_cudagraph_capturing_enabled(True)
|
|
try:
|
|
with torch.inference_mode(), graph_capture(device=self.device):
|
|
for bsz in capture_sizes:
|
|
_, batch_desc, _, _, _ = self._determine_batch_execution_and_padding(
|
|
num_tokens=bsz,
|
|
num_reqs=bsz,
|
|
num_scheduled_tokens_np=np.ones(bsz, dtype=np.int32),
|
|
max_num_scheduled_tokens=1,
|
|
use_cascade_attn=False,
|
|
)
|
|
n = batch_desc.num_tokens
|
|
ids = self.talker_mtp_input_ids.gpu[:n]
|
|
emb = self.talker_mtp_inputs_embeds.gpu[:n]
|
|
hid = self.last_talker_hidden.gpu[:n]
|
|
ts = self.text_step.gpu[:n]
|
|
|
|
for _ in range(num_warmups):
|
|
with set_forward_context(
|
|
None,
|
|
self.vllm_config,
|
|
cudagraph_runtime_mode=CUDAGraphMode.NONE,
|
|
batch_descriptor=batch_desc,
|
|
):
|
|
self.talker_mtp(ids, emb, hid, ts)
|
|
|
|
with set_forward_context(
|
|
None,
|
|
self.vllm_config,
|
|
cudagraph_runtime_mode=CUDAGraphMode.FULL,
|
|
batch_descriptor=batch_desc,
|
|
):
|
|
self.talker_mtp(ids, emb, hid, ts)
|
|
torch.accelerator.synchronize()
|
|
|
|
logger.info("Captured talker_mtp graphs for %d sizes", len(capture_sizes))
|
|
except RuntimeError as e:
|
|
raise RuntimeError(
|
|
f"talker_mtp graph capture failed for a model that declared talker_mtp_graph_safe=True: {e}"
|
|
) from e
|
|
finally:
|
|
set_cudagraph_capturing_enabled(False)
|
|
|
|
def _model_needs_full_prefix_hidden_states(self) -> bool:
|
|
"""Opt-out hook for models whose postprocess only consumes the tail.
|
|
|
|
When False, we skip both the per-step GPU->CPU hidden-state write into
|
|
OmniTensorPrefixCache and the merged-tensor reconstruction on hits;
|
|
postprocess receives the normal scheduled-token slice instead. Models
|
|
that need the full cached_prefix + new_tail span (default) are not
|
|
affected.
|
|
"""
|
|
model = getattr(self, "model", None)
|
|
return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True))
|
|
|
|
def _deferred_prefix_cache_mm_keys(self) -> set[str]:
|
|
"""Model-declared multimodal keys whose prefix-cache writes are deferred."""
|
|
model = getattr(self, "model", None)
|
|
keys = getattr(model, "deferred_prefix_cache_mm_keys", ())
|
|
return set(keys or ())
|
|
|
|
def _maybe_update_prefix_cache(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
multimodal_outputs: dict,
|
|
num_tokens_unpadded: int,
|
|
num_tokens_padded: int,
|
|
):
|
|
"""If prefix caching is enabled and it's the last pipeline parallelism rank,
|
|
retrieve the hidden states & multimodal outputs from the prefix cache based
|
|
on our batch slot mappings.
|
|
"""
|
|
# Cache hidden states if we've enabled hidden state prefix caching
|
|
# unless this isn't the last pipeline parallelism rank.
|
|
if self.omni_prefix_cache is not None and get_pp_group().is_last_rank:
|
|
# If this happens, it generally means the model is not following the correct
|
|
# interface yet and is therefore currently not compatible with prefix cache.
|
|
if multimodal_outputs is not None and not isinstance(multimodal_outputs, dict):
|
|
logger.warning_once(
|
|
"prefix caching expects mm outputs to be a dict, but got %s",
|
|
type(multimodal_outputs),
|
|
)
|
|
|
|
hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
|
|
# FIX: The .cpu attribute of slot_mapping is stale (not updated by the Triton
|
|
# _compute_slot_mapping_kernel which only writes to .gpu). We must use .gpu and
|
|
# sync back to CPU to get the correctly computed slot mapping.
|
|
slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu
|
|
slot_mapping_cpu = slot_mapping_gpu[:num_tokens_padded].cpu()
|
|
self.omni_prefix_cache.update_omni_tensor_prefix_cache(
|
|
hidden_states=hs_for_cache,
|
|
multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs,
|
|
num_tokens_unpadded=num_tokens_unpadded,
|
|
slot_mapping=slot_mapping_cpu,
|
|
num_tokens_padded=num_tokens_padded,
|
|
skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(),
|
|
)
|
|
|
|
def _maybe_get_combined_prefix_cache_tensors(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
multimodal_outputs: dict,
|
|
num_scheduled_tokens: dict[str, int],
|
|
) -> tuple[dict[str, torch.Tensor] | None, dict | None]:
|
|
"""If prefix caching is enabled, extract the merged hidden states and multimodal outputs for
|
|
all requests in the batch (including those that aren't a hit on Prefix cache).
|
|
"""
|
|
# Prior to applying the post-processing func, extract
|
|
# the prefix cached hidden states and multimodal states.
|
|
combined_hidden_states, combined_multimodal_outputs = None, None
|
|
if self.omni_prefix_cache is not None:
|
|
if (
|
|
not self._model_needs_full_prefix_hidden_states()
|
|
and not self.omni_prefix_cache.has_prefix_cached_new_req_ids()
|
|
):
|
|
return None, None
|
|
if self._model_needs_full_prefix_hidden_states():
|
|
combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
|
|
query_start_loc=self.query_start_loc.cpu,
|
|
input_batch=self.input_batch,
|
|
hidden_states=hidden_states,
|
|
num_scheduled_tokens=num_scheduled_tokens,
|
|
)
|
|
combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states(
|
|
query_start_loc=self.query_start_loc.cpu,
|
|
input_batch=self.input_batch,
|
|
multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs,
|
|
num_scheduled_tokens=num_scheduled_tokens,
|
|
)
|
|
return combined_hidden_states, combined_multimodal_outputs
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None:
|
|
if self.execute_model_state is not None:
|
|
raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
|
|
|
|
if not getattr(self, "_warmup_state_cleared", False):
|
|
self._warmup_state_cleared = True
|
|
if hasattr(self.model, "_clear_warmup_state"):
|
|
self.model._clear_warmup_state()
|
|
|
|
# [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
|
|
finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
|
|
if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"):
|
|
for req_id, data in finished_reqs.items():
|
|
try:
|
|
# NOTE: seq_len is the same as num_computed_tokens_cpu in current
|
|
# async scheduling, since both exclude async placeholders. We use
|
|
# seq_len since we control it, just in case upstream async scheduler
|
|
# semantics change in the future.
|
|
num_computed = data.get("seq_len")
|
|
|
|
model_meta = self.model.get_kv_transfer_metadata(
|
|
req_id,
|
|
num_computed_tokens=num_computed,
|
|
)
|
|
if model_meta:
|
|
existing = data.get("custom_metadata") or {}
|
|
existing.update(model_meta)
|
|
data["custom_metadata"] = existing
|
|
except Exception as e:
|
|
logger.warning(f"Failed to get custom metadata from model for {req_id}: {e}")
|
|
self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
|
|
finished_reqs=finished_reqs,
|
|
kv_caches=self.kv_caches,
|
|
block_size=self.cache_config.block_size,
|
|
cache_dtype=str(self.cache_config.cache_dtype),
|
|
request_id_resolver=self._resolve_global_request_id,
|
|
)
|
|
|
|
if hasattr(self, "_omni_connector"):
|
|
for request in getattr(scheduler_output, "pending_input_registrations", []):
|
|
self.register_chunk_recv(request)
|
|
self.recv_full_payload_inputs(scheduler_output)
|
|
if self._pending_full_payload_send:
|
|
flush_ids = set(getattr(scheduler_output, "finished_req_ids", set()))
|
|
flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests})
|
|
if flush_ids:
|
|
self.flush_full_payload_outputs(flush_ids)
|
|
|
|
if self.omni_prefix_cache is not None and scheduler_output.finished_req_ids:
|
|
self.omni_prefix_cache.commit_deferred_mm_outputs(
|
|
set(scheduler_output.finished_req_ids),
|
|
self.input_batch,
|
|
)
|
|
|
|
if self.routed_experts_initialized:
|
|
capturer = get_global_experts_capturer()
|
|
if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
|
|
capturer.finalize_pending_copy()
|
|
|
|
# If ngram_gpu is used, we need to copy the scheduler_output to avoid
|
|
# the modification has influence on the scheduler_output in engine core process.
|
|
# The replace is much faster than deepcopy.
|
|
if self.speculative_config is not None and self.speculative_config.use_ngram_gpu():
|
|
num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy()
|
|
spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy()
|
|
scheduler_output = replace(
|
|
scheduler_output,
|
|
num_scheduled_tokens=num_scheduled_tokens_copy,
|
|
scheduled_spec_decode_tokens=spec_decode_tokens_copy,
|
|
)
|
|
|
|
if has_kv_transfer_group():
|
|
kv_connector_metadata = scheduler_output.kv_connector_metadata
|
|
if kv_connector_metadata is not None:
|
|
get_kv_transfer_group().handle_preemptions(kv_connector_metadata)
|
|
|
|
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
|
|
with (
|
|
record_function_or_nullcontext("gpu_model_runner: preprocess"),
|
|
self.synchronize_input_prep(),
|
|
):
|
|
# Update persistent batch states.
|
|
deferred_state_corrections_fn = self._update_states(scheduler_output)
|
|
|
|
# Notify model of finished requests for state cleanup
|
|
if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
|
|
self.model.on_requests_finished(scheduler_output.finished_req_ids)
|
|
|
|
if has_ec_transfer() and not get_ec_transfer().is_consumer:
|
|
with self.maybe_get_ec_connector_output(
|
|
scheduler_output,
|
|
encoder_cache=self.encoder_cache,
|
|
) as ec_connector_output:
|
|
self._execute_mm_encoder(scheduler_output)
|
|
|
|
kv_ids = self.kv_extracted_req_ids
|
|
self.kv_extracted_req_ids = None
|
|
|
|
output = make_empty_encoder_model_runner_output(scheduler_output)
|
|
if kv_ids:
|
|
output = copy(output)
|
|
output.kv_extracted_req_ids = kv_ids
|
|
return self.attach_omni_connector_output(output)
|
|
|
|
if not num_scheduled_tokens:
|
|
if (
|
|
self.parallel_config.distributed_executor_backend == "external_launcher"
|
|
and self.parallel_config.data_parallel_size > 1
|
|
):
|
|
self._dummy_run(1)
|
|
|
|
# Capture KV extraction results before early return;
|
|
# sample_tokens() is skipped on this path so the IDs
|
|
# would otherwise be silently overwritten next step.
|
|
kv_ids = self.kv_extracted_req_ids
|
|
self.kv_extracted_req_ids = None
|
|
|
|
if not has_kv_transfer_group():
|
|
output = EMPTY_MODEL_RUNNER_OUTPUT
|
|
else:
|
|
output = self.kv_connector_no_forward(scheduler_output, self.vllm_config)
|
|
|
|
if kv_ids:
|
|
output = copy(output)
|
|
output.kv_extracted_req_ids = kv_ids
|
|
|
|
return self.attach_omni_connector_output(output)
|
|
|
|
if self.cache_config.kv_sharing_fast_prefill:
|
|
assert not self.num_prompt_logprobs, (
|
|
"--kv-sharing-fast-prefill produces incorrect "
|
|
"logprobs for prompt tokens, tokens, please disable "
|
|
"it when the requests need prompt logprobs"
|
|
)
|
|
|
|
num_reqs = self.input_batch.num_reqs
|
|
req_ids = self.input_batch.req_ids
|
|
tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
|
|
num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
|
|
max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
|
|
num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
|
|
|
|
logits_indices, spec_decode_metadata = self._prepare_inputs(
|
|
scheduler_output,
|
|
num_scheduled_tokens_np,
|
|
)
|
|
|
|
cascade_attn_prefix_lens = None
|
|
# Disable cascade attention when using microbatching (DBO)
|
|
if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
|
|
# Pre-compute cascade attention prefix lengths
|
|
cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
|
|
num_scheduled_tokens_np,
|
|
self.input_batch.num_computed_tokens_cpu[:num_reqs],
|
|
scheduler_output.num_common_prefix_blocks,
|
|
)
|
|
|
|
(
|
|
cudagraph_mode,
|
|
batch_desc,
|
|
should_ubatch,
|
|
num_tokens_across_dp,
|
|
cudagraph_stats,
|
|
) = self._determine_batch_execution_and_padding(
|
|
num_tokens=num_tokens_unpadded,
|
|
num_reqs=num_reqs,
|
|
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
|
max_num_scheduled_tokens=max_num_scheduled_tokens,
|
|
use_cascade_attn=cascade_attn_prefix_lens is not None,
|
|
num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
|
|
)
|
|
|
|
num_tokens_padded = batch_desc.num_tokens
|
|
num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
|
|
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
|
|
should_ubatch,
|
|
num_scheduled_tokens_np,
|
|
num_tokens_padded,
|
|
num_reqs_padded,
|
|
self.parallel_config.num_ubatches,
|
|
)
|
|
|
|
pad_attn = cudagraph_mode == CUDAGraphMode.FULL
|
|
|
|
use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
|
|
ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
|
|
|
|
# True if any attention backend handles KV cache update separately
|
|
# from forward() (i.e., forward_includes_kv_cache_update=False). When true,
|
|
# slot_mappings must use padded dimensions to match the key/value tensors.
|
|
from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec
|
|
|
|
has_separate_kv_update = not all(
|
|
all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id])
|
|
for id, spec in enumerate(self.kv_cache_config.kv_cache_groups)
|
|
if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec)
|
|
)
|
|
|
|
slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
|
|
num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded,
|
|
num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs),
|
|
num_tokens_unpadded=num_tokens_unpadded,
|
|
ubatch_slices=ubatch_slices_padded,
|
|
)
|
|
|
|
attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
|
|
num_tokens=num_tokens_unpadded,
|
|
num_tokens_padded=num_tokens_padded if pad_attn else None,
|
|
num_reqs=num_reqs,
|
|
num_reqs_padded=num_reqs_padded if pad_attn else None,
|
|
max_query_len=max_num_scheduled_tokens,
|
|
ubatch_slices=ubatch_slices_attn,
|
|
logits_indices=logits_indices,
|
|
use_spec_decode=use_spec_decode,
|
|
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
|
|
cascade_attn_prefix_lens=cascade_attn_prefix_lens,
|
|
slot_mappings=slot_mappings_by_group,
|
|
)
|
|
self._maybe_attach_attention_metadata_extensions(
|
|
attn_metadata=attn_metadata,
|
|
num_reqs=num_reqs,
|
|
num_reqs_padded=num_reqs_padded,
|
|
max_query_len=max_num_scheduled_tokens,
|
|
pad_attn=pad_attn,
|
|
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
|
)
|
|
|
|
(
|
|
input_ids,
|
|
inputs_embeds,
|
|
positions,
|
|
intermediate_tensors,
|
|
model_kwargs,
|
|
ec_connector_output,
|
|
) = self._preprocess(scheduler_output, num_tokens_padded, intermediate_tensors)
|
|
|
|
# Let the model adjust inputs before forward (e.g. restore input_ids
|
|
# for multimodal position detection, fix decode position offsets).
|
|
if hasattr(self.model, "prepare_runner_inputs"):
|
|
input_ids, positions = self.model.prepare_runner_inputs(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
inputs_embeds=inputs_embeds,
|
|
req_ids=req_ids[:num_reqs],
|
|
num_computed_tokens=[int(self.input_batch.num_computed_tokens_cpu[i]) for i in range(num_reqs)],
|
|
num_scheduled_tokens=[int(num_scheduled_tokens_np[i]) for i in range(num_reqs)],
|
|
input_ids_buffer=self.input_ids.gpu[:num_tokens_padded],
|
|
)
|
|
|
|
# Set cudagraph mode to none if calc_kv_scales is true.
|
|
# KV scales calculation involves dynamic operations that are incompatible
|
|
# with CUDA graph capture.
|
|
if self.calculate_kv_scales:
|
|
cudagraph_mode = CUDAGraphMode.NONE
|
|
# Mark KV scales as calculated after the first forward pass
|
|
self.calculate_kv_scales = False
|
|
|
|
# Run the model.
|
|
# Use persistent buffers for CUDA graphs.
|
|
# When spec decode is enabled, defer connector finalization
|
|
# (wait_for_save + clear metadata) until after draft model runs.
|
|
defer_kv_connector_finalize = self.speculative_config is not None
|
|
with (
|
|
nullcontext(),
|
|
set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_tokens_padded,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
cudagraph_runtime_mode=cudagraph_mode,
|
|
batch_descriptor=batch_desc,
|
|
ubatch_slices=ubatch_slices_padded,
|
|
slot_mapping=slot_mappings, # OMNI: required for KV cache operations
|
|
),
|
|
record_function_or_nullcontext("gpu_model_runner: forward"),
|
|
self.maybe_get_kv_connector_output(
|
|
scheduler_output,
|
|
defer_finalize=defer_kv_connector_finalize,
|
|
) as kv_connector_output,
|
|
):
|
|
model_output = self._model_forward(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
**model_kwargs,
|
|
sampling_metadata=self.input_batch.sampling_metadata,
|
|
logits_index=logits_indices,
|
|
sampler=self.sampler,
|
|
)
|
|
|
|
# [Omni] Map pending ropes metadata to req_ids.
|
|
if hasattr(self.model, "flush_pending_metadata"):
|
|
self.model.flush_pending_metadata(list(req_ids))
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: postprocess"):
|
|
if self.use_aux_hidden_state_outputs:
|
|
# True when EAGLE 3 is used.
|
|
hidden_states, aux_hidden_states = model_output
|
|
else:
|
|
# Common case.
|
|
hidden_states = model_output
|
|
aux_hidden_states = None
|
|
|
|
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output)
|
|
|
|
# Cache hidden states & multimodal outputs if we've enabled hidden state
|
|
# prefix caching unless this isn't the last pipeline parallelism rank.
|
|
self._maybe_update_prefix_cache(
|
|
hidden_states=hidden_states,
|
|
multimodal_outputs=multimodal_outputs,
|
|
num_tokens_unpadded=num_tokens_unpadded,
|
|
num_tokens_padded=num_tokens_padded,
|
|
)
|
|
|
|
if not self.broadcast_pp_output:
|
|
# Common case.
|
|
if not get_pp_group().is_last_rank:
|
|
# Return the intermediate tensors.
|
|
assert isinstance(hidden_states, IntermediateTensors)
|
|
hidden_states.kv_connector_output = kv_connector_output
|
|
self.kv_connector_output = kv_connector_output
|
|
return hidden_states
|
|
|
|
if self.is_pooling_model:
|
|
# Return the pooling output.
|
|
return self._pool(
|
|
hidden_states,
|
|
num_scheduled_tokens,
|
|
num_scheduled_tokens_np,
|
|
kv_connector_output,
|
|
)
|
|
|
|
sample_hidden_states = hidden_states[logits_indices]
|
|
# Try with sampling_metadata first; fall back to without for models that don't support it
|
|
try:
|
|
logits = self.model.compute_logits(
|
|
sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
|
|
)
|
|
except TypeError:
|
|
logits = self.model.compute_logits(sample_hidden_states)
|
|
else:
|
|
# Rare case.
|
|
assert not self.is_pooling_model
|
|
|
|
sample_hidden_states = hidden_states[logits_indices]
|
|
if not get_pp_group().is_last_rank:
|
|
all_gather_tensors = {
|
|
"residual": not is_residual_scattered_for_sp(self.vllm_config, num_tokens_padded)
|
|
}
|
|
get_pp_group().send_tensor_dict(
|
|
hidden_states.tensors,
|
|
all_gather_group=get_tp_group(),
|
|
all_gather_tensors=all_gather_tensors,
|
|
)
|
|
logits = None
|
|
else:
|
|
# Try with sampling_metadata first; fall back to without for models that don't support it
|
|
try:
|
|
logits = self.model.compute_logits(
|
|
sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
|
|
)
|
|
except TypeError:
|
|
logits = self.model.compute_logits(sample_hidden_states)
|
|
|
|
model_output_broadcast_data: dict[str, Any] = {}
|
|
if logits is not None:
|
|
model_output_broadcast_data["logits"] = logits.contiguous()
|
|
|
|
broadcasted = get_pp_group().broadcast_tensor_dict(
|
|
model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
|
|
)
|
|
assert broadcasted is not None
|
|
logits = broadcasted["logits"]
|
|
|
|
self.execute_model_state = ExecuteModelState(
|
|
scheduler_output,
|
|
logits,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
hidden_states,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
ec_connector_output,
|
|
cudagraph_stats,
|
|
multimodal_outputs,
|
|
slot_mappings, # OMNI: pass slot_mappings for drafter
|
|
)
|
|
self.kv_connector_output = kv_connector_output
|
|
|
|
if deferred_state_corrections_fn:
|
|
deferred_state_corrections_fn()
|
|
|
|
if self.routed_experts_initialized and hasattr(self, "_positions_cpu"):
|
|
issue_routing_d2h_copy(
|
|
input_batch_req_ids=self.input_batch.req_ids,
|
|
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
|
|
positions=self.positions,
|
|
positions_cpu=self._positions_cpu,
|
|
)
|
|
|
|
return None
|
|
|
|
def _sample(
|
|
self,
|
|
logits: torch.Tensor | None,
|
|
spec_decode_metadata: Any,
|
|
):
|
|
sampling_metadata = self.input_batch.sampling_metadata
|
|
if spec_decode_metadata is None:
|
|
model_sample = getattr(self.model, "sample", None)
|
|
self.input_batch.update_async_output_token_ids()
|
|
if logits is not None and callable(model_sample) and getattr(self.model, "prefer_model_sampler", False):
|
|
# Apply logit bias (min_tokens, allowed_token_ids) before
|
|
# the custom model sampler — the standard GPU sampler does
|
|
# this internally, but prefer_model_sampler bypasses it.
|
|
if hasattr(self.sampler, "logit_bias_state"):
|
|
self.sampler.logit_bias_state.apply_logit_bias(
|
|
logits,
|
|
self.input_batch.expanded_idx_mapping,
|
|
self.input_batch.idx_mapping_np,
|
|
self.input_batch.positions[self.input_batch.logits_indices],
|
|
)
|
|
sampler_output = model_sample(
|
|
logits,
|
|
self._sampling_metadata_for_model_sampler(sampling_metadata),
|
|
)
|
|
if sampler_output is not None:
|
|
return sampler_output
|
|
return self.sampler(
|
|
logits=logits,
|
|
sampling_metadata=sampling_metadata,
|
|
)
|
|
|
|
return super()._sample(logits, spec_decode_metadata)
|
|
|
|
@staticmethod
|
|
def _resolve_req_hidden_states(
|
|
hidden_states_cpu: torch.Tensor,
|
|
combined_hidden_states: dict[str, torch.Tensor] | None,
|
|
rid: str,
|
|
start: int,
|
|
end: int,
|
|
):
|
|
if combined_hidden_states is not None:
|
|
# We always have all request IDs for prefix cache, even for
|
|
# partial cache misses, so this should never happen.
|
|
if rid not in combined_hidden_states:
|
|
raise RuntimeError("Request IDs in the batch are missing from the merged states!")
|
|
return combined_hidden_states[rid]
|
|
# Prefix caching is disabled
|
|
return hidden_states_cpu[start:end]
|
|
|
|
@torch.inference_mode()
|
|
def sample_tokens(
|
|
self,
|
|
grammar_output: GrammarOutput | None,
|
|
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
|
|
kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None)
|
|
self.kv_extracted_req_ids = None
|
|
|
|
# Used for prefix cache
|
|
combined_hidden_states = None
|
|
combined_multimodal_outputs = None
|
|
# Used when we don't use prefix cache; prefix cache builds the payloads
|
|
# internally since it already needs to do this for the cached tensors
|
|
mm_cpu = {}
|
|
|
|
if self.execute_model_state is None:
|
|
kv_connector_output = self.kv_connector_output
|
|
self.kv_connector_output = None
|
|
# Nothing to do (PP non-final rank case), output isn't used.
|
|
if not kv_connector_output:
|
|
return None # type: ignore[return-value]
|
|
|
|
# In case of PP with kv transfer, we need to pass through the
|
|
# kv_connector_output
|
|
if kv_connector_output.is_empty():
|
|
return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
|
|
|
|
output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
|
|
output.kv_connector_output = kv_connector_output
|
|
return self.attach_omni_connector_output(output)
|
|
|
|
# Unpack ephemeral state.
|
|
(
|
|
scheduler_output,
|
|
logits,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
hidden_states,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
ec_connector_output,
|
|
cudagraph_stats,
|
|
multimodal_outputs,
|
|
slot_mappings, # OMNI: unpack slot_mappings for drafter
|
|
) = self.execute_model_state
|
|
self.execute_model_state = None
|
|
seq_len = hidden_states.shape[0]
|
|
|
|
# Apply structured output bitmasks if present.
|
|
if grammar_output is not None:
|
|
apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits)
|
|
|
|
# Correct padding values of prompt_token_ids to match the logits vocabulary size
|
|
if logits is not None and not self.input_batch.sampling_metadata.no_penalties:
|
|
smd = self.input_batch.sampling_metadata
|
|
if smd.prompt_token_ids is not None:
|
|
logits_vocab = logits.shape[-1]
|
|
if self.input_batch.vocab_size > logits_vocab:
|
|
smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab)
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: sample"):
|
|
sampler_output = self._sample(logits, spec_decode_metadata)
|
|
|
|
self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output)
|
|
|
|
self._draft_token_ids = None
|
|
self._draft_token_req_ids = None
|
|
self.valid_sampled_token_count_gpu = None
|
|
self.input_batch.prev_sampled_token_ids = None
|
|
|
|
def propose_draft_token_ids(sampled_token_ids):
|
|
assert spec_decode_common_attn_metadata is not None
|
|
with record_function_or_nullcontext("gpu_model_runner: draft"):
|
|
self._draft_token_ids = self.propose_draft_token_ids(
|
|
scheduler_output,
|
|
sampled_token_ids,
|
|
self.input_batch.sampling_metadata,
|
|
hidden_states,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
slot_mappings, # OMNI: pass slot_mappings to drafter (upstream v1 API)
|
|
)
|
|
self._copy_draft_token_ids_to_cpu(scheduler_output)
|
|
|
|
spec_config = self.speculative_config
|
|
propose_drafts_after_bookkeeping = False
|
|
if spec_config is not None:
|
|
input_fits_in_drafter = spec_decode_common_attn_metadata is not None and (
|
|
spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
|
|
<= self.effective_drafter_max_model_len
|
|
)
|
|
use_gpu_toks = (
|
|
spec_config.use_eagle() or spec_config.uses_draft_model() or spec_config.uses_extract_hidden_states()
|
|
) and not spec_config.disable_padded_drafter_batch
|
|
if use_gpu_toks:
|
|
assert isinstance(
|
|
self.drafter,
|
|
EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer,
|
|
)
|
|
sampled_token_ids = sampler_output.sampled_token_ids
|
|
if input_fits_in_drafter:
|
|
propose_draft_token_ids(sampled_token_ids)
|
|
elif self.valid_sampled_token_count_event is not None:
|
|
assert spec_decode_common_attn_metadata is not None
|
|
next_token_ids, valid_sampled_tokens_count = self.drafter.prepare_next_token_ids_padded(
|
|
self.optimistic_seq_lens_cpu,
|
|
sampled_token_ids,
|
|
self.requests,
|
|
self.input_batch,
|
|
self.discard_request_mask.gpu,
|
|
)
|
|
self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
|
|
# Since we couldn't run the drafter,
|
|
# just use zeros for the draft tokens.
|
|
self._draft_token_ids = torch.zeros(1, device=self.device, dtype=torch.int32).expand(
|
|
len(self.input_batch.req_ids), self.num_spec_tokens
|
|
)
|
|
self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
|
|
else:
|
|
propose_drafts_after_bookkeeping = input_fits_in_drafter
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
|
|
(
|
|
num_nans_in_logits,
|
|
logprobs_lists,
|
|
valid_sampled_token_ids,
|
|
prompt_logprobs_dict,
|
|
req_ids_output_copy,
|
|
req_id_to_index_output_copy,
|
|
invalid_req_indices,
|
|
) = self._bookkeeping_sync(
|
|
scheduler_output,
|
|
sampler_output,
|
|
logits,
|
|
hidden_states,
|
|
scheduler_output.total_num_scheduled_tokens,
|
|
)
|
|
|
|
if propose_drafts_after_bookkeeping:
|
|
# ngram and other speculative decoding methods use the sampled
|
|
# tokens on the CPU, so they are run after bookkeeping.
|
|
propose_draft_token_ids(valid_sampled_token_ids)
|
|
|
|
# Finalize KV connector (wait_for_save + clear metadata) after
|
|
# draft model runs. Deferred from target model forward to allow
|
|
# draft model to also save its KV cache.
|
|
if self.speculative_config is not None:
|
|
self.finalize_kv_connector()
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: eplb"):
|
|
self.eplb_step()
|
|
|
|
# kv_connector_output may be modified during drafting
|
|
kv_connector_output = self.kv_connector_output
|
|
self.kv_connector_output = None
|
|
|
|
engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
|
|
needs_pooler_payload = len(downstream_req_ids) > 0
|
|
downstream_req_id_set = set(downstream_req_ids)
|
|
hidden_states_cpu = None
|
|
req_hidden_states_cpu: dict[str, torch.Tensor] | None = None
|
|
if needs_pooler_payload:
|
|
num_valid_tokens = min(
|
|
int(scheduler_output.total_num_scheduled_tokens),
|
|
int(hidden_states.shape[0]),
|
|
)
|
|
if len(downstream_req_ids) == len(req_ids_output_copy):
|
|
hidden_states_cpu = hidden_states[:num_valid_tokens].detach().to("cpu").contiguous()
|
|
else:
|
|
req_hidden_states_cpu = {}
|
|
num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
|
|
if num_scheduled_tokens_np is None:
|
|
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,
|
|
)
|
|
query_start_loc_cpu = self.query_start_loc.cpu
|
|
if self.omni_prefix_cache is not None:
|
|
deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys()
|
|
if deferred_mm_cache_keys:
|
|
self.omni_prefix_cache.stage_deferred_mm_outputs(
|
|
query_start_loc=query_start_loc_cpu,
|
|
input_batch=self.input_batch,
|
|
multimodal_outputs=(
|
|
flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs
|
|
),
|
|
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
|
|
deferred_mm_cache_keys=deferred_mm_cache_keys,
|
|
)
|
|
|
|
pooler_output: list[dict[str, object]] | None = None
|
|
if needs_pooler_payload:
|
|
mm_cpu = None
|
|
# Prior to applying the post-processing func, extract
|
|
# the prefix cached hidden states and multimodal states.
|
|
if self.omni_prefix_cache is not None:
|
|
(
|
|
combined_hidden_states,
|
|
combined_multimodal_outputs,
|
|
) = self._maybe_get_combined_prefix_cache_tensors(
|
|
hidden_states,
|
|
multimodal_outputs,
|
|
scheduler_output.num_scheduled_tokens,
|
|
)
|
|
# Otherwise we don't have the mm CPU data yet, so we still need to build it
|
|
if self.omni_prefix_cache is None or combined_multimodal_outputs is None:
|
|
mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs)
|
|
|
|
self._process_additional_information_updates(
|
|
hidden_states,
|
|
multimodal_outputs,
|
|
num_scheduled_tokens_np,
|
|
scheduler_output,
|
|
combined_hidden_states,
|
|
combined_multimodal_outputs,
|
|
req_ids_filter=downstream_req_id_set,
|
|
)
|
|
|
|
if req_hidden_states_cpu is not None and combined_hidden_states is None:
|
|
for rid in downstream_req_ids:
|
|
idx = req_id_to_index_output_copy[rid]
|
|
start = int(query_start_loc_cpu[idx])
|
|
sched = int(num_scheduled_tokens_np[idx])
|
|
end = start + sched
|
|
req_hidden_states_cpu[rid] = hidden_states[start:end].detach().to("cpu").contiguous()
|
|
|
|
pooler_output = []
|
|
for out_idx, rid in enumerate(req_ids_output_copy):
|
|
if rid not in downstream_req_id_set:
|
|
pooler_output.append({})
|
|
continue
|
|
idx = req_id_to_index_output_copy[rid]
|
|
start = int(query_start_loc_cpu[idx])
|
|
sched = int(num_scheduled_tokens_np[idx])
|
|
end = start + sched
|
|
# If prefix cache is enabled, we have already split everything
|
|
# by request and converted the states to CPU tensors
|
|
if req_hidden_states_cpu is not None and combined_hidden_states is None:
|
|
req_hidden_states = req_hidden_states_cpu[rid]
|
|
else:
|
|
req_hidden_states = self._resolve_req_hidden_states(
|
|
hidden_states_cpu,
|
|
combined_hidden_states,
|
|
rid,
|
|
start,
|
|
end,
|
|
)
|
|
payload: dict[str, object] = {"hidden": req_hidden_states}
|
|
|
|
mm_payload: dict[str, object] = {}
|
|
if combined_multimodal_outputs or mm_cpu:
|
|
if combined_multimodal_outputs:
|
|
# Prefix cache enabled; all items have already been processed
|
|
# and split apart for each request as needed, and all tensors
|
|
# have already been detached to the CPU. Lists are kept as
|
|
# passthrough data for consistent behavior in postprocess.
|
|
# Recurse into nested dicts so list-valued sub-keys (e.g.
|
|
# embed.tts_bos = [tensor]) are unwrapped to bare tensors
|
|
# at the leaves; downstream flatten_payload then yields a
|
|
# wire-clean dict[str, torch.Tensor].
|
|
def _unwrap_lists(v):
|
|
if isinstance(v, list):
|
|
return v[idx] if idx < len(v) else v[0]
|
|
if isinstance(v, dict):
|
|
return {k: _unwrap_lists(sv) for k, sv in v.items()}
|
|
return v
|
|
|
|
for mm_key in combined_multimodal_outputs.keys():
|
|
mm_payload[mm_key] = _unwrap_lists(combined_multimodal_outputs[mm_key][rid])
|
|
|
|
else:
|
|
# Prefix cache disabled; we still need to process the data
|
|
for mm_key, mm_val in mm_cpu.items():
|
|
mm_payload[mm_key] = to_payload_element(
|
|
element=mm_val,
|
|
idx=idx,
|
|
start=start,
|
|
end=end,
|
|
pass_lists_through=False,
|
|
seq_len=seq_len,
|
|
)
|
|
payload.update(mm_payload)
|
|
# Flatten nested dicts to dotted keys so pooling_output
|
|
# stays dict[str, torch.Tensor] for msgspec serialization.
|
|
pooler_output.append(flatten_payload(payload))
|
|
|
|
if pooler_output and self._should_accumulate_full_payload_output():
|
|
for i, rid in enumerate(req_ids_output_copy):
|
|
req_state = self.requests.get(rid)
|
|
if req_state is not None and pooler_output[i]:
|
|
self.accumulate_full_payload_output(rid, pooler_output[i], req_state)
|
|
|
|
with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
|
|
routed_experts_dict = None
|
|
if self.routed_experts_initialized:
|
|
routed_experts_dict = extract_routed_experts_for_current_batch(
|
|
req_ids=req_ids_output_copy,
|
|
requests=self.requests,
|
|
req_id_to_index=self.input_batch.req_id_to_index,
|
|
num_tokens_no_spec=self.input_batch.num_tokens_no_spec,
|
|
max_model_len=self.max_model_len,
|
|
)
|
|
output = OmniModelRunnerOutput(
|
|
req_ids=req_ids_output_copy,
|
|
req_id_to_index=req_id_to_index_output_copy,
|
|
sampled_token_ids=valid_sampled_token_ids,
|
|
logprobs=logprobs_lists,
|
|
prompt_logprobs_dict=prompt_logprobs_dict,
|
|
pooler_output=(pooler_output if engine_output_type != "text" and needs_pooler_payload else None),
|
|
kv_connector_output=kv_connector_output,
|
|
ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
|
|
num_nans_in_logits=num_nans_in_logits,
|
|
cudagraph_stats=cudagraph_stats,
|
|
routed_experts_dict=routed_experts_dict,
|
|
)
|
|
output.kv_extracted_req_ids = kv_extracted_req_ids
|
|
output.omni_connector_output = self.get_omni_connector_output()
|
|
|
|
if not self.use_async_scheduling:
|
|
return output
|
|
with record_function_or_nullcontext("gpu_model_runner: AsyncGPUModelRunnerOutput"):
|
|
async_output = AsyncGPUModelRunnerOutput(
|
|
model_runner_output=output,
|
|
sampled_token_ids=sampler_output.sampled_token_ids,
|
|
logprobs_tensors=sampler_output.logprobs_tensors,
|
|
invalid_req_indices=invalid_req_indices,
|
|
async_output_copy_stream=self.async_output_copy_stream,
|
|
vocab_size=self.input_batch.vocab_size,
|
|
)
|
|
with record_function_or_nullcontext("gpu_model_runner: set_async_sampled_token_ids"):
|
|
# Save ref of sampled_token_ids CPU tensor if the batch contains
|
|
# any requests with sampling params that require output ids.
|
|
self.input_batch.set_async_sampled_token_ids(
|
|
async_output.sampled_token_ids_cpu,
|
|
async_output.async_copy_ready_event,
|
|
)
|
|
|
|
return async_output
|
|
|
|
def _resolve_global_request_id(self, req_id: str) -> str:
|
|
"""Resolve global request ID from request state."""
|
|
req_state = self.requests.get(req_id)
|
|
if not req_state:
|
|
return req_id
|
|
|
|
add_info = self.model_intermediate_buffer.get(req_id, {})
|
|
global_id = add_info.get("global_request_id")
|
|
if global_id:
|
|
if isinstance(global_id, list) and global_id:
|
|
global_id = global_id[0]
|
|
if isinstance(global_id, bytes):
|
|
return global_id.decode("utf-8")
|
|
return str(global_id)
|
|
return req_id
|