diff --git a/.gitignore b/.gitignore index d9e8bd68b..224dac867 100644 --- a/.gitignore +++ b/.gitignore @@ -248,6 +248,7 @@ datasets/ # Output files outputs/ +!vllm_omni/outputs/ results/ generated/ output_*/ diff --git a/docs/api/README.md b/docs/api/README.md index d4c580629..cfea6d175 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -46,17 +46,17 @@ Engine classes for offline and online inference. - [vllm_omni.engine.PromptEmbedsPayload][] - [vllm_omni.engine.arg_utils.OmniEngineArgs][] - [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][] -- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][] -- [vllm_omni.engine.mm_outputs.MultimodalPayload][] - [vllm_omni.engine.orchestrator.Orchestrator][] - [vllm_omni.engine.orchestrator.OrchestratorRequestState][] -- [vllm_omni.engine.output_modality.OutputModality][] -- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][] -- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][] -- [vllm_omni.engine.output_processor.OmniRequestState][] - [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][] - [vllm_omni.engine.stage_init_utils.StageMetadata][] - [vllm_omni.engine.stage_init_utils.StartedLlmStage][] +- [vllm_omni.outputs.mm_outputs.MultimodalCompletionOutput][] +- [vllm_omni.outputs.mm_outputs.MultimodalPayload][] +- [vllm_omni.outputs.output_modality.OutputModality][] +- [vllm_omni.outputs.output_modality.TensorAccumulationStrategy][] +- [vllm_omni.outputs.output_processor.MultimodalOutputProcessor][] +- [vllm_omni.outputs.output_processor.OmniRequestState][] ## Core diff --git a/tests/engine/test_output_modality.py b/tests/engine/test_output_modality.py index 0abfe6ea0..ff15e5824 100644 --- a/tests/engine/test_output_modality.py +++ b/tests/engine/test_output_modality.py @@ -14,7 +14,7 @@ import torch # ── Load modules without triggering vllm_omni.__init__ ───────────── pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -_ENGINE_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "engine" +_OUTPUTS_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "outputs" def _load_module(name: str, filepath: Path): @@ -26,12 +26,16 @@ def _load_module(name: str, filepath: Path): _om_mod = _load_module( - "vllm_omni.engine.output_modality", - _ENGINE_DIR / "output_modality.py", + "vllm_omni.outputs.output_modality", + _OUTPUTS_DIR / "output_modality.py", +) +_load_module( + "vllm_omni.outputs.utils", + _OUTPUTS_DIR / "utils.py", ) _mm_mod = _load_module( - "vllm_omni.engine.mm_outputs", - _ENGINE_DIR / "mm_outputs.py", + "vllm_omni.outputs.mm_outputs", + _OUTPUTS_DIR / "mm_outputs.py", ) OutputModality = _om_mod.OutputModality diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index f9b08ca15..86e71fcbc 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -11,8 +11,8 @@ from vllm.outputs import PoolingRequestOutput from vllm.sampling_params import RequestOutputKind from vllm.v1.engine import FinishReason -from vllm_omni.engine.output_modality import OutputModalityNames -from vllm_omni.engine.output_processor import OmniRequestState +from vllm_omni.outputs.output_modality import OutputModalityNames +from vllm_omni.outputs.output_processor import OmniRequestState pytestmark = [pytest.mark.core_model, pytest.mark.cpu] diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 2fcfd43f7..bdf43711c 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -9,7 +9,7 @@ from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.logger import init_logger from vllm_omni.config import OmniModelConfig -from vllm_omni.engine.output_modality import OutputModality +from vllm_omni.outputs.output_modality import OutputModality from vllm_omni.platforms import current_omni_platform from vllm_omni.plugins import load_omni_general_plugins diff --git a/vllm_omni/engine/mm_outputs.py b/vllm_omni/engine/mm_outputs.py index 3c9da2496..336bf7d9a 100644 --- a/vllm_omni/engine/mm_outputs.py +++ b/vllm_omni/engine/mm_outputs.py @@ -1,127 +1,9 @@ -"""Multimodal output data structures for vLLM-Omni. +import warnings -This module defines structured types for multimodal outputs. +from vllm_omni.outputs.mm_outputs import * # noqa: F401,F403 -""" - -from __future__ import annotations - -from collections.abc import Iterator, Mapping -from dataclasses import dataclass, field -from typing import Any - -import torch -from vllm.outputs import CompletionOutput - - -@dataclass(eq=False) -class MultimodalPayload(Mapping): - """Structured multimodal output payload. - - Implements ``collections.abc.Mapping`` so that ``isinstance(payload, dict)`` - style checks in downstream code can be replaced with duck-typing, and - ``payload.get(key)``, ``payload[key]``, ``key in payload``, ``len(payload)`` - all work seamlessly for both tensors and metadata. - - Attributes: - tensors: Dictionary mapping modality/key names to their tensors. - metadata: Optional dictionary for non-tensor metadata - (e.g., sample rate for audio, image dimensions). - """ - - tensors: dict[str, torch.Tensor] = field(default_factory=dict) - metadata: dict[str, Any] = field(default_factory=dict) - - @property - def primary_tensor(self) -> torch.Tensor | None: - """Return the first tensor in the payload, or None if empty.""" - if self.tensors: - return next(iter(self.tensors.values())) - return None - - @property - def is_empty(self) -> bool: - """Return True if the payload has no tensors and no metadata.""" - return len(self.tensors) == 0 and len(self.metadata) == 0 - - def get(self, key: str, default: Any = None) -> Any: - """Get a value by key, searching tensors first then metadata.""" - if key in self.tensors: - return self.tensors[key] - return self.metadata.get(key, default) - - def __contains__(self, key: str) -> bool: - return key in self.tensors or key in self.metadata - - def __getitem__(self, key: str) -> Any: - """Dict-like indexing: search tensors first, then metadata.""" - if key in self.tensors: - return self.tensors[key] - if key in self.metadata: - return self.metadata[key] - raise KeyError(key) - - def __len__(self) -> int: - return len(self.tensors) + len(self.metadata) - - def __iter__(self) -> Iterator[str]: - yield from self.tensors - yield from self.metadata - - def __bool__(self) -> bool: - return bool(self.tensors) or bool(self.metadata) - - def __eq__(self, other: object) -> bool: - """Support equality with plain dicts and other Mappings.""" - if isinstance(other, MultimodalPayload): - return self.tensors == other.tensors and self.metadata == other.metadata - if isinstance(other, Mapping): - return dict(self) == dict(other) - return NotImplemented - - def to_dict(self) -> dict[str, Any]: - """Convert back to a plain dict (tensors + metadata merged).""" - result: dict[str, Any] = dict(self.tensors) - result.update(self.metadata) - return result - - @classmethod - def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None: - """Create a MultimodalPayload from a raw dictionary. - - Separates torch.Tensor values into tensors and everything - else into metadata. - """ - if not data: - return None - tensors: dict[str, torch.Tensor] = {} - metadata: dict[str, Any] = {} - for k, v in data.items(): - if isinstance(v, torch.Tensor): - tensors[k] = v - else: - metadata[k] = v - if not tensors and not metadata: - return None - return cls(tensors=tensors, metadata=metadata) - - -@dataclass -class MultimodalCompletionOutput(CompletionOutput): - """CompletionOutput with multimodal support. - - Inherits all CompletionOutput fields and adds multimodal_output. - As a CompletionOutput subclass, compatible with all existing vLLM consumers. - """ - - def __init__( - self, - multimodal_output: MultimodalPayload | None = None, - **kwargs: Any, - ): - super().__init__(**kwargs) - self.multimodal_output = multimodal_output - - def __repr__(self) -> str: - base = super().__repr__() - return f"{base[:-1]}, multimodal_output={self.multimodal_output!r})" +warnings.warn( + "Importing from 'vllm_omni.engine.mm_outputs' is deprecated. Use 'vllm_omni.outputs.mm_outputs' instead.", + DeprecationWarning, + stacklevel=2, +) diff --git a/vllm_omni/engine/output_modality.py b/vllm_omni/engine/output_modality.py index 2536a8ade..364519bed 100644 --- a/vllm_omni/engine/output_modality.py +++ b/vllm_omni/engine/output_modality.py @@ -1,120 +1,9 @@ -"""Output modality types for vLLM-Omni. +import warnings -This module defines the OutputModality enum and TensorAccumulationStrategy -for type-safe multimodal output routing and tensor merging. +from vllm_omni.outputs.output_modality import * # noqa: F401,F403 -""" - -from __future__ import annotations - -import re -from enum import Enum, Flag, StrEnum, auto - -_MODALITY_ALIASES: dict[str, str] = { - "speech": "audio", - "images": "image", - "latents": "latent", - "wav": "audio", - "waveform": "audio", - "pixel_values": "image", - "pixels": "image", - "token_ids": "text", - "tokens": "text", -} - - -class OutputModalityNames(StrEnum): - """Keys for output modalities. - - TODO: (Alex) Integrate this with the big-flag enum below + throughout the code - for better type safety (currently only used for output processor). - """ - - TEXT = "text" - IMAGE = "image" - AUDIO = "audio" - LATENT = "latent" - - -# Specify which output modalities may be drained when handling delta messages. -# For some types, e.g., latents, we need to be careful to ensure the full context -# is passed as the stream yields due to assumptions in the I/O processing and model -# when async chunk isn't enabled. -NON_DRAINABLE_MODALITIES = {OutputModalityNames.TEXT, OutputModalityNames.LATENT} -DRAINABLE_MODALITIES = {mod for mod in OutputModalityNames if mod not in NON_DRAINABLE_MODALITIES} - - -class OutputModality(Flag): - """Bit-flag enum for output modalities. - - Compose freely with ``|`` — no need to enumerate every combination. - - Single: ``OutputModality.TEXT``, ``OutputModality.IMAGE``, ... - Compound: ``OutputModality.TEXT | OutputModality.IMAGE`` (text+image) - - Note: POOLING is intentionally excluded. Pooling/embedding is vLLM's - native path (pooling_output → PoolingRequestOutput), handled entirely - by the base OutputProcessor. vLLM-Omni's layer does not participate. - """ - - TEXT = auto() - IMAGE = auto() - AUDIO = auto() - LATENT = auto() - - @classmethod - def from_string(cls, s: str | None) -> OutputModality: - """Parse a free-text modality string into an OutputModality flag. - - Handles common aliases and compound strings separated by + or ,. - - Examples:: - - OutputModality.from_string("text+image") - # → OutputModality.TEXT | OutputModality.IMAGE - """ - if not s or not s.strip(): - return cls.TEXT - - parts = [p.strip().lower() for p in re.split(r"[+,]", s.strip())] - result = cls(0) - for p in parts: - p = _MODALITY_ALIASES.get(p, p) - try: - result |= cls[p.upper()] - except KeyError: - raise ValueError(f"Unknown modality: {p!r}. Supported: {[m.name.lower() for m in cls]}") - return result - - @property - def has_text(self) -> bool: - return OutputModality.TEXT in self - - @property - def has_multimodal(self) -> bool: - return bool(self & ~OutputModality.TEXT) - - -class TensorAccumulationStrategy(Enum): - """Strategy for merging incremental multimodal tensors.""" - - CONCAT_DIM0 = "concat_dim0" - """Concatenate along dimension 0. Used for image/latent tensors.""" - - CONCAT_LAST = "concat_last" - """Concatenate along the last dimension. Used for audio waveforms.""" - - APPEND_LIST = "append_list" - """Append to a list (no tensor concatenation).""" - - REPLACE = "replace" - """Replace previous tensor entirely with the latest one.""" - - -def get_accumulation_strategy(modality: OutputModality) -> TensorAccumulationStrategy: - """Determine tensor merge strategy from the multimodal flags.""" - if OutputModality.AUDIO in modality: - return TensorAccumulationStrategy.CONCAT_LAST - if OutputModality.IMAGE in modality or OutputModality.LATENT in modality: - return TensorAccumulationStrategy.CONCAT_DIM0 - return TensorAccumulationStrategy.CONCAT_DIM0 # default +warnings.warn( + "Importing from 'vllm_omni.engine.output_modality' is deprecated. Use 'vllm_omni.outputs.output_modality' instead.", + DeprecationWarning, + stacklevel=2, +) diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 71303269f..0d87583a8 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -1,762 +1,10 @@ -from dataclasses import fields as dataclass_fields -from typing import Any +import warnings -import torch -from vllm.logger import init_logger -from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput -from vllm.sampling_params import RequestOutputKind -from vllm.tokenizers import TokenizerLike -from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason -from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor -from vllm.v1.engine.output_processor import ( - OutputProcessorOutput, - RequestOutputCollector, - RequestState, +from vllm_omni.outputs.output_processor import * # noqa: F401,F403 + +warnings.warn( + "Importing from 'vllm_omni.engine.output_processor' is deprecated. " + "Use 'vllm_omni.outputs.output_processor' instead.", + DeprecationWarning, + stacklevel=2, ) -from vllm.v1.engine.parallel_sampling import ParentRequest -from vllm.v1.metrics.stats import IterationStats, RequestStateStats - -from vllm_omni.data_entry_keys import unflatten_payload -from vllm_omni.engine.mm_outputs import MultimodalCompletionOutput, MultimodalPayload -from vllm_omni.engine.output_modality import ( - DRAINABLE_MODALITIES, - OutputModality, - TensorAccumulationStrategy, - get_accumulation_strategy, -) -from vllm_omni.outputs import OmniRequestOutput - -logger = init_logger(__name__) - -# Keys whose values may arrive as scalar tensors (e.g. sample rate) but -# represent metadata, not concatenatable data. When these appear in -# ``MultimodalPayload.tensors`` they must be relocated to ``.metadata`` -# before consolidation so the tensor concatenation path does not attempt -# a dimensioned cat on a 0-d tensor. -_METADATA_TENSOR_KEYS: frozenset[str] = frozenset({"sr", "sample_rate", "audio_sample_rate"}) - -# --------------------------------------------------------------------------- -# Module-level helpers -# --------------------------------------------------------------------------- - - -def _to_cpu(x: Any) -> Any: - """Move a tensor to CPU; pass through non-tensors unchanged.""" - if isinstance(x, torch.Tensor): - try: - return x.detach().to("cpu", non_blocking=True).contiguous() - except (RuntimeError, AttributeError): - return x - if isinstance(x, dict): - return {str(sk): _to_cpu(sv) for sk, sv in x.items()} - return x - - -def _merge_payload(target: MultimodalPayload, incoming: MultimodalPayload) -> None: - """Merge *incoming* into *target* using list-based deferred concatenation.""" - for k, v in incoming.tensors.items(): - if k not in target.tensors: - target.tensors[k] = v - else: - existing = target.tensors[k] - if isinstance(existing, list): - existing.append(v) - elif isinstance(existing, torch.Tensor) and isinstance(v, torch.Tensor): - target.tensors[k] = [existing, v] - else: - target.tensors[k] = v - - for k, v in incoming.metadata.items(): - if k not in target.metadata: - target.metadata[k] = v - else: - existing = target.metadata[k] - if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor): - target.metadata[k] = [existing, v] - elif isinstance(v, torch.Tensor) and isinstance(existing, list): - existing.append(v) - else: - # Metadata: replace with latest value - target.metadata[k] = v - - -def _cat_tensors( - tensors: list[torch.Tensor], - strategy: TensorAccumulationStrategy, -) -> torch.Tensor: - """Concatenate a list of tensors according to *strategy*.""" - if strategy == TensorAccumulationStrategy.CONCAT_LAST: - return torch.cat(tensors, dim=-1) - if strategy == TensorAccumulationStrategy.REPLACE: - return tensors[-1] - # CONCAT_DIM0 / APPEND_LIST / default - return torch.cat(tensors, dim=0) - - -def _modality_to_type_string(modality: OutputModality) -> str: - """Convert an OutputModality flag to a lowercase type string.""" - try: - if OutputModality.AUDIO in modality: - return "audio" - if OutputModality.IMAGE in modality: - return "image" - if OutputModality.LATENT in modality: - return "latent" - except TypeError: - # Flag identity mismatch (e.g. after module reload in tests). - name = getattr(modality, "name", "") or "" - lowered = name.lower() - if "audio" in lowered: - return "audio" - if "image" in lowered: - return "image" - if "latent" in lowered: - return "latent" - return "text" - - -class OmniRequestState(RequestState): - """Request state for omni models, tracking multimodal outputs. - - Extends the base RequestState with support for accumulating - multimodal tensor outputs (e.g., images, audio, latents) that - are produced incrementally during generation. - """ - - def __init__( - self, - *args, - **kwargs, - ): - # arrival_time is always passed as a keyword argument by the sole caller - # (from_new_request). The *args fallback via positional index was removed - # because upstream RequestState.__init__ reorders parameters between - # releases, making args[12] fragile. - arrival_time = kwargs.get("arrival_time") - super().__init__(*args, **kwargs) - self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0)) - # Omni-specific: multimodal output accumulation - # TODO: mm_type is per-request, not per-key. If a model ever produces - # both audio and latent outputs, the modality type would flip on each - # add_multimodal_tensor() call. Consider tracking per-key modality - # types (e.g. dict[str, str]) for future multi-output models. - self.mm_type: str | None = None - self.mm_accumulated: MultimodalPayload = MultimodalPayload() - - def apply_streaming_update(self, update) -> None: - super().apply_streaming_update(update) - self.native_text_stats.arrival_time = update.arrival_time - - @staticmethod - def _to_cpu(x): - """Try to convert to CPU tensor if needed.""" - if isinstance(x, torch.Tensor): - try: - return x.detach().to("cpu", non_blocking=True).contiguous() - except Exception: - return x - return x - - def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None: - """Accumulate a multimodal tensor payload into the request state. - - Normalizes incoming payloads (dict or raw tensor) into a - MultimodalPayload and merges with any previously accumulated data. - Uses list-based deferred concatenation to avoid O(n²) repeated - torch.cat calls. - """ - if payload is None: - return - try: - if mm_type: - self.mm_type = (mm_type or "").lower() - - target_key = self.mm_type or "hidden" - - # Build incoming MultimodalPayload - if isinstance(payload, dict): - remapped: dict[str, Any] = {} - for k, v in payload.items(): - # Normalize producer keys to the modality name. - # AR runners produce {"hidden": ...} and generation - # runners produce {"model_outputs": ...}; remap both - # to the semantic modality key (e.g. "audio", "latent"). - if k == "model_outputs": - k = target_key - elif k == "hidden" and target_key != "hidden": - k = target_key - remapped[k] = _to_cpu(v) - incoming = MultimodalPayload.from_dict(remapped) - elif isinstance(payload, MultimodalPayload): - incoming = payload - else: - incoming = MultimodalPayload.from_dict({target_key: _to_cpu(payload)}) - - if incoming is None: - return - - if self.mm_accumulated.is_empty: - self.mm_accumulated = incoming - else: - _merge_payload(self.mm_accumulated, incoming) - except (ValueError, TypeError, RuntimeError): - logger.exception("Error accumulating multimodal tensor") - - def _consolidate_multimodal_tensors(self) -> None: - """Consolidate accumulated tensor lists into single tensors. - - Uses TensorAccumulationStrategy derived from the output modality - to determine concatenation behavior. Metadata values always use - REPLACE (keep latest). - """ - if self.mm_accumulated.is_empty: - return - - try: - modality = OutputModality.from_string(self.mm_type) - except (ValueError, KeyError): - modality = OutputModality.TEXT - strategy = get_accumulation_strategy(modality) - - # Move known metadata scalar tensors (e.g. sample rate) from - # tensors to metadata so they are handled by the REPLACE path - # below instead of failing tensor concatenation. - for key in _METADATA_TENSOR_KEYS: - v = self.mm_accumulated.tensors.pop(key, None) - if v is not None: - if isinstance(v, list): - # Keep the last value (REPLACE semantics) - self.mm_accumulated.metadata[key] = v[-1] - else: - self.mm_accumulated.metadata[key] = v - - try: - for k, v in list(self.mm_accumulated.tensors.items()): - if isinstance(v, list) and v and isinstance(v[0], torch.Tensor): - try: - self.mm_accumulated.tensors[k] = _cat_tensors(v, strategy) - except RuntimeError: - if k == "audio": - try: - self.mm_accumulated.tensors[k] = torch.cat(v, dim=-1) - except RuntimeError: - self.mm_accumulated.tensors[k] = torch.cat([t.reshape(-1) for t in v], dim=0) - else: - logger.warning( - "Error concatenating tensor for key %s; keeping last tensor", - k, - ) - self.mm_accumulated.tensors[k] = v[-1] - - # Metadata: consolidate any deferred tensor lists (REPLACE semantics) - for k, v in list(self.mm_accumulated.metadata.items()): - if isinstance(v, list) and v and isinstance(v[0], torch.Tensor): - self.mm_accumulated.metadata[k] = v[-1] - elif isinstance(v, dict): - for sk, sv in list(v.items()): - if isinstance(sv, list) and sv and isinstance(sv[0], torch.Tensor): - v[sk] = sv[-1] - except (RuntimeError, TypeError, KeyError): - logger.exception("Error consolidating multimodal tensors") - - # Override: do not route to pooling-only path; always create completion - # outputs, and attach pooling_result into the CompletionOutput. - def make_request_output( - self, - new_token_ids: list[int], - pooling_output: torch.Tensor | None, - finish_reason: FinishReason | None, - stop_reason: int | str | None, - kv_transfer_params: dict[str, Any] | None = None, - routed_experts: Any = None, - ) -> OmniRequestOutput | PoolingRequestOutput | None: - """Create a request output from generation results. - - Creates a RequestOutput or PoolingRequestOutput from the generated - tokens and accumulated multimodal outputs. Attaches multimodal - tensors to the completion output if available. - - Args: - new_token_ids: List of newly generated token IDs - pooling_output: Optional pooling output tensor - finish_reason: Optional finish reason indicating why generation stopped - stop_reason: Optional stop reason (token ID or stop string) - kv_transfer_params: Optional KV cache transfer parameters - routed_experts: Optional MoE routed-expert ids for this step, - attached to the completion output for generation stages - - Returns: - OmniRequestOutput or PoolingRequestOutput if output should be - emitted (based on finish status and output kind), None otherwise - """ - # Pooling-only requests should follow base behavior. - if self.detokenizer is None and pooling_output is not None: - return super().make_request_output( - new_token_ids, - pooling_output, - finish_reason, - stop_reason, - kv_transfer_params, - ) - - finished = finish_reason is not None - final_only = self.output_kind == RequestOutputKind.FINAL_ONLY - is_delta = self.output_kind == RequestOutputKind.DELTA - - if not finished and final_only: - return None - - # Consolidate accumulated tensors when finishing, or for every - # CUMULATIVE step so the consumer always sees a single tensor. - if finished or not is_delta: - self._consolidate_multimodal_tensors() - - if self.stream_interval > 1 and self.detokenizer is not None: - # Send output request only when - # 1. It has finished, or - # 2. It is the first token, or - # 3. It has reached the stream interval number of tokens - if not ( - finished - or self.sent_tokens_offset == 0 - or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval - ): - return None - - if self.output_kind == RequestOutputKind.DELTA: - # Send tokens from the offset in DELTA mode, otherwise all - # tokens are sent. - new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :] - self.sent_tokens_offset = self.detokenizer.num_output_tokens() - - external_req_id = self.external_req_id - - output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts) - - if self.parent_req is None: - outputs = [output] - else: - outputs, finished = self.parent_req.get_outputs(self.request_id, output) - if not outputs: - return None - external_req_id = self.parent_req.external_req_id - - return self._new_request_output( - external_req_id, - outputs, - finished, - kv_transfer_params, - ) - - def _new_completion_output( - self, - token_ids: list[int], - finish_reason: FinishReason | None, - stop_reason: int | str | None, - routed_experts: Any = None, - ) -> MultimodalCompletionOutput | CompletionOutput: - """Create a completion output with multimodal data attached. - - Returns a MultimodalCompletionOutput when multimodal data has been - accumulated, otherwise returns the base CompletionOutput. Snapshots - the accumulated payload before draining DELTA-mode modality keys so - that callers receive an immutable view. - """ - # When there is no detokenizer (generation stages), build a minimal - # CompletionOutput directly since upstream asserts detokenizer != None. - if self.detokenizer is None: - finished = finish_reason is not None - base_output = CompletionOutput( - index=self.request_index, - text="", - token_ids=token_ids, - logprobs=None, - cumulative_logprob=None, - finish_reason=str(finish_reason) if finished else None, - stop_reason=stop_reason if finished else None, - routed_experts=routed_experts, - ) - else: - base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason) - - # Always provide cumulative token IDs for inter-stage processors. - if self.detokenizer is not None: - base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids) - else: - base_output.cumulative_token_ids = list(token_ids) - - # Attach cumulative_text only at the final step for inter-stage use. - if finish_reason is not None and hasattr(self.detokenizer, "output_text"): - base_output.cumulative_text = self.detokenizer.output_text - - try: - if self.mm_accumulated and not self.mm_accumulated.is_empty: - # Snapshot: copy current tensors/metadata so drain doesn't - # mutate the payload already handed to the caller. - # Unflatten dotted keys (e.g. "hidden_states.layer_0") back - # to nested dicts so downstream consumers (thinker2talker etc.) - # can access with .get("hidden_states", {}).get("layers", {}). - merged = {**self.mm_accumulated.tensors, **self.mm_accumulated.metadata} - nested = unflatten_payload(merged) - snapshot = MultimodalPayload.from_dict(nested) or MultimodalPayload() - kwargs = {f.name: getattr(base_output, f.name) for f in dataclass_fields(CompletionOutput)} - output = MultimodalCompletionOutput( - multimodal_output=snapshot, - **kwargs, - ) - output.cumulative_token_ids = base_output.cumulative_token_ids - if hasattr(base_output, "cumulative_text"): - output.cumulative_text = base_output.cumulative_text - - # DELTA mode: drain modality keys (e.g. audio) so the next - # step only sees freshly accumulated data for those keys. - if self.output_kind == RequestOutputKind.DELTA: - for modality_key in DRAINABLE_MODALITIES: - self.mm_accumulated.tensors.pop(modality_key, None) - - return output - except (RuntimeError, TypeError, AttributeError): - logger.exception("Error creating MultimodalCompletionOutput") - return base_output - - def _new_request_output( - self, - external_req_id: str, - outputs: list, - finished: bool, - kv_transfer_params: dict[str, Any] | None = None, - ) -> RequestOutput | PoolingRequestOutput: - """Create request output, handling no-detokenizer generation stages. - - Upstream asserts ``self.logprobs_processor is not None`` which fails - for generation stages (talker, code2wav) that have no tokenizer. - When the logprobs processor is absent we build the RequestOutput - directly with ``prompt_logprobs=None``. - """ - if self.logprobs_processor is not None: - return super()._new_request_output( - external_req_id, - outputs, - finished, - kv_transfer_params, - ) - - # No-detokenizer path: build RequestOutput directly. - prompt_token_ids = self.prompt_token_ids - if prompt_token_ids is None and self.prompt_embeds is not None: - prompt_token_ids = [0] * len(self.prompt_embeds) - if prompt_token_ids is None: - prompt_token_ids = [] - - return RequestOutput( - request_id=external_req_id, - lora_request=self.lora_request, - prompt=self.prompt, - prompt_token_ids=prompt_token_ids, - prompt_logprobs=None, - outputs=outputs, - finished=finished, - kv_transfer_params=kv_transfer_params, - num_cached_tokens=self.num_cached_tokens, - metrics=self.stats, - ) - - -class MultimodalOutputProcessor(VLLMOutputProcessor): - """Handles multimodal output processing. - - Captures multimodal outputs from OmniEngineCoreOutput and accumulates - them as MultimodalPayload in OmniRequestState, before delegating to - the base vLLM OutputProcessor for text handling. - - The data flow is: - 1. For each EngineCoreOutput with multimodal_output: - - Capture into OmniRequestState.add_multimodal_tensor() - 2. Base vLLM OutputProcessor handles text detokenization - 3. On finish, _consolidate_multimodal_tensors() concatenates accumulated - tensors using strategy-based dispatch - 4. _new_completion_output() returns MultimodalCompletionOutput - """ - - def __init__( - self, - tokenizer: TokenizerLike | None, - *, - log_stats: bool, - stream_interval: int = 1, - tracing_enabled: bool = False, - engine_core_output_type: str | None = None, - output_modality: OutputModality = OutputModality.TEXT, - ): - """Initialize the multimodal output processor. - - Args: - tokenizer: Tokenizer for detokenizing text outputs - log_stats: Whether to log statistics - stream_interval: Stream interval for output generation - engine_core_output_type: Optional output type string (e.g., - "image", "audio", "latent"). Converted to OutputModality - internally. Kept for backward compatibility with - stage_init_utils. - output_modality: Type-safe output modality flag. Used to tag - multimodal outputs with the correct modality key when - per-output type info is unavailable. - """ - super().__init__( - tokenizer=tokenizer, - log_stats=log_stats, - stream_interval=stream_interval, - tracing_enabled=tracing_enabled, - ) - # Convert string-based engine_core_output_type to OutputModality - if engine_core_output_type is not None: - self.output_modality = OutputModality.from_string(engine_core_output_type) - else: - self.output_modality = output_modality - self.engine_core_output_type = engine_core_output_type - self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {} - - def _native_text_metric_record(self, request_id: str) -> dict[str, Any]: - return self._native_text_metrics_by_request.setdefault( - request_id, - { - "vllm_ttft_ms": 0.0, - "vllm_tpot_ms": 0.0, - "vllm_itl_ms": 0.0, - "vllm_itls_ms": [], - }, - ) - - def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: - return self._native_text_metrics_by_request.pop(request_id, {}) - - def abort_requests(self, request_ids, internal: bool) -> list[str]: - request_ids = list(request_ids) - for request_id in request_ids: - if internal: - req_state = self.request_states.get(request_id) - if req_state is not None: - self._native_text_metrics_by_request.pop(req_state.external_req_id, None) - else: - self._native_text_metrics_by_request.pop(request_id, None) - return super().abort_requests(request_ids, internal) - - def add_request( - self, - request: EngineCoreRequest, - prompt: str | None, - parent_req: ParentRequest | None = None, - request_index: int = 0, - queue: RequestOutputCollector | None = None, - ) -> None: - """Add a new request to be processed. - - Creates an OmniRequestState for the request and registers it - for output processing. - - Args: - request: Engine core request to add - prompt: Optional prompt string for the request - parent_req: Optional parent request for parallel sampling - request_index: Index of the request in the batch - queue: Optional queue for collecting outputs - - Raises: - ValueError: If the request ID is already registered - """ - request_id = request.request_id - req_state = self.request_states.get(request_id) - if req_state is not None: - self._update_streaming_request_state(req_state, request, prompt) - return - - req_state = OmniRequestState.from_new_request( - tokenizer=self.tokenizer, - request=request, - prompt=prompt, - parent_req=parent_req, - request_index=request_index, - queue=queue, - log_stats=self.log_stats, - stream_interval=self.stream_interval, - ) - self.request_states[request_id] = req_state - if parent_req: - self.parent_requests[parent_req.request_id] = parent_req - self.external_req_ids[req_state.external_req_id].append(request_id) - - def remove_request(self, request_id: str) -> None: - """Rollback one previously registered request if it was never submitted.""" - req_state = self.request_states.pop(request_id, None) - if req_state is None: - return - - external_req_id = getattr(req_state, "external_req_id", None) - if external_req_id is not None: - self._native_text_metrics_by_request.pop(external_req_id, None) - request_ids = self.external_req_ids.get(external_req_id) - if request_ids is not None: - self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id] - if not self.external_req_ids[external_req_id]: - self.external_req_ids.pop(external_req_id, None) - - parent_req = getattr(req_state, "parent_req", None) - if parent_req is not None: - self.parent_requests.pop(parent_req.request_id, None) - - def process_outputs( - self, - engine_core_outputs: list[EngineCoreOutput], - engine_core_timestamp: float | None = None, - iteration_stats: IterationStats | None = None, - ) -> OutputProcessorOutput: - default_mm_type = _modality_to_type_string(self.output_modality) - - # Separate outputs that upstream can handle (has detokenizer or - # pooling) from multimodal-only outputs (no detokenizer, no pooling) - # that would trigger upstream's `assert detokenizer is not None`. - upstream_outputs: list[EngineCoreOutput] = [] - mm_only_outputs: list[EngineCoreOutput] = [] - - for eco in engine_core_outputs: - req_state = self.request_states.get(eco.request_id) - if req_state is None: - continue - - # Accumulate multimodal tensors regardless of path. - if isinstance(req_state, OmniRequestState): - mm_output = getattr(eco, "multimodal_output", None) - if mm_output is not None: - mm_type = getattr(eco, "output_type", None) or default_mm_type - req_state.add_multimodal_tensor(mm_output, mm_type) - - # Route: if no detokenizer and no pooling output, handle locally - # to avoid upstream's assert on detokenizer. - if req_state.detokenizer is None and eco.pooling_output is None: - mm_only_outputs.append(eco) - else: - upstream_outputs.append(eco) - - # Handle multimodal-only outputs (generation stages) locally. - self._process_mm_only_outputs(mm_only_outputs) - - # Delegate text/pooling outputs to upstream. - return super().process_outputs( - upstream_outputs, - engine_core_timestamp=engine_core_timestamp, - iteration_stats=iteration_stats, - ) - - def _process_mm_only_outputs( - self, - engine_core_outputs: list[EngineCoreOutput], - ) -> None: - """Handle outputs from generation stages that have no detokenizer. - - These cannot go through upstream process_outputs because it asserts - detokenizer is not None when pooling_output is None. - """ - for eco in engine_core_outputs: - req_state = self.request_states.get(eco.request_id) - if req_state is None or not isinstance(req_state, OmniRequestState): - continue - - new_token_ids = eco.new_token_ids - finish_reason = eco.finish_reason - stop_reason = eco.stop_reason - kv_transfer_params = eco.kv_transfer_params - routed_experts = eco.routed_experts - req_state.num_cached_tokens = eco.num_cached_tokens - req_state.is_prefilling = False - - if request_output := req_state.make_request_output( - new_token_ids, - None, # pooling_output - finish_reason, - stop_reason, - kv_transfer_params, - routed_experts, - ): - if req_state.queue is not None: - req_state.queue.put(request_output) - - if finish_reason is not None: - self._finish_request(req_state) - - def _update_stats_from_output( - self, - req_state: RequestState, - engine_core_output: EngineCoreOutput, - engine_core_timestamp: float | None, - iteration_stats: IterationStats | None, - ): - was_prefilling = req_state.is_prefilling - native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None - previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0 - - # NOTE: We pass ``None`` for *iteration_stats* to the parent so that - # the upstream's ``_update_stats_from_output`` logs stats via its own - # ``RequestStateStats`` (req_state.stats) without touching the omni - # ``IterationStats`` object. Below we re-apply the same update using - # *native_stats* (the omni-specific ``RequestStateStats`` attached to - # ``OmniRequestState``), which correctly feeds the omni metrics path. - # If the upstream adds new logic in ``_update_stats_from_output`` (e.g. - # additional stat fields), this override must be kept in sync. - super()._update_stats_from_output( - req_state, - engine_core_output, - engine_core_timestamp, - None, - ) - - if iteration_stats is None or engine_core_timestamp is None or native_stats is None: - return - - iteration_stats.update_from_output( - engine_core_output, - engine_core_timestamp, - was_prefilling, - native_stats, - self.lora_states, - req_state.lora_name, - ) - record = self._native_text_metric_record(req_state.external_req_id) - if was_prefilling: - record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0) - return - - if previous_last_token_ts > 0: - itl_ms = max((engine_core_timestamp - previous_last_token_ts) * 1000.0, 0.0) - itls_ms = record.setdefault("vllm_itls_ms", []) - itls_ms.append(itl_ms) - record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms)) - - def _update_stats_from_finished( - self, - req_state: RequestState, - finish_reason: FinishReason | None, - iteration_stats: IterationStats | None, - ): - super()._update_stats_from_finished(req_state, finish_reason, None) - - native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None - if iteration_stats is None or native_stats is None or finish_reason is None: - return - - iteration_stats.update_from_finished_request( - finish_reason=finish_reason, - request_id=req_state.external_req_id, - num_prompt_tokens=req_state.prompt_len, - max_tokens_param=req_state.max_tokens_param, - req_stats=native_stats, - num_cached_tokens=req_state.num_cached_tokens, - ) - if not iteration_stats.finished_requests: - return - - finished_request = iteration_stats.finished_requests[-1] - if finished_request.request_id != req_state.external_req_id: - return - self._native_text_metric_record(req_state.external_req_id)["vllm_tpot_ms"] = ( - float(finished_request.mean_time_per_output_token) * 1000.0 - ) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index c3cf11b72..d43160895 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -27,11 +27,11 @@ from vllm.v1.executor import Executor from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs -from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams from vllm_omni.inputs.preprocess import OmniInputPreprocessor +from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig diff --git a/vllm_omni/outputs/mm_outputs.py b/vllm_omni/outputs/mm_outputs.py new file mode 100644 index 000000000..8c5fe6442 --- /dev/null +++ b/vllm_omni/outputs/mm_outputs.py @@ -0,0 +1,261 @@ +"""Multimodal output data structures for vLLM-Omni. + +This module defines structured types for multimodal outputs. + +""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from dataclasses import dataclass, field +from typing import Any + +import torch +from vllm.logger import init_logger +from vllm.outputs import CompletionOutput + +from vllm_omni.outputs.output_modality import TensorAccumulationStrategy +from vllm_omni.outputs.utils import _is_tensor_list, _to_cpu + +logger = init_logger(__name__) + +# Keys whose values are metadata scalars (e.g. audio sample rate) but may +# arrive as 0-d torch.Tensors — from_dict routes all tensors into .tensors, +# so we relocate them to .metadata before consolidation to avoid a bogus +# torch.cat attempt and its warn-and-keep-last fallback. +_METADATA_TENSOR_KEYS: frozenset[str] = frozenset({"sr", "sample_rate", "audio_sample_rate"}) + + +def _cat_tensors( + tensors: list[torch.Tensor], + strategy: TensorAccumulationStrategy, +) -> torch.Tensor: + """Concatenate a list of tensors according to *strategy*.""" + if strategy == TensorAccumulationStrategy.CONCAT_LAST: + return torch.cat(tensors, dim=-1) + if strategy == TensorAccumulationStrategy.REPLACE: + return tensors[-1] + # CONCAT_DIM0 / APPEND_LIST / default + return torch.cat(tensors, dim=0) + + +def _consolidate_tensor_list( + key: str, + tensor_list: list[torch.Tensor], + strategy: TensorAccumulationStrategy, +) -> torch.Tensor: + """Concatenate a deferred tensor list, with fallbacks on shape mismatch.""" + try: + return _cat_tensors(tensor_list, strategy) + except RuntimeError: + # [TODO: this part is for async chunk, not for audio only. can we come up with a design + # that identify this by whether using async chunk instead of modality?] + if key != "audio": + logger.warning("Error concatenating tensor for key %s; keeping last tensor", key) + return tensor_list[-1] + # Audio chunks may have mismatched shapes; retry along the last dim, + # then fall back to flattening each chunk. + try: + return torch.cat(tensor_list, dim=-1) + except RuntimeError: + return torch.cat([chunk.reshape(-1) for chunk in tensor_list], dim=0) + + +def _append_entries(store: dict[str, Any], incoming: dict[str, Any]) -> None: + """Append *incoming* entries into *store*. + + Tensor values are collected into lists (deferred concatenation, avoiding + O(n²) repeated torch.cat); non-tensor values replace the existing entry. + """ + for key, new_value in incoming.items(): + existing = store.get(key) + if isinstance(existing, list) and isinstance(new_value, torch.Tensor): + existing.append(new_value) + elif isinstance(existing, torch.Tensor) and isinstance(new_value, torch.Tensor): + store[key] = [existing, new_value] + else: + store[key] = new_value + + +@dataclass(eq=False) +class MultimodalPayload(Mapping): + """Structured multimodal output payload. + + Implements ``collections.abc.Mapping`` so that ``isinstance(payload, dict)`` + style checks in downstream code can be replaced with duck-typing, and + ``payload.get(key)``, ``payload[key]``, ``key in payload``, ``len(payload)`` + all work seamlessly for both tensors and metadata. + + Attributes: + tensors: Dictionary mapping modality/key names to their tensors. + metadata: Optional dictionary for non-tensor metadata + (e.g., sample rate for audio, image dimensions). + """ + + tensors: dict[str, torch.Tensor] = field(default_factory=dict) + metadata: dict[str, Any] = field(default_factory=dict) + + @property + def primary_tensor(self) -> torch.Tensor | None: + """Return the first tensor in the payload, or None if empty.""" + if self.tensors: + return next(iter(self.tensors.values())) + return None + + @property + def is_empty(self) -> bool: + """Return True if the payload has no tensors and no metadata.""" + return len(self.tensors) == 0 and len(self.metadata) == 0 + + def get(self, key: str, default: Any = None) -> Any: + """Get a value by key, searching tensors first then metadata.""" + if key in self.tensors: + return self.tensors[key] + return self.metadata.get(key, default) + + def __contains__(self, key: str) -> bool: + return key in self.tensors or key in self.metadata + + def __getitem__(self, key: str) -> Any: + """Dict-like indexing: search tensors first, then metadata.""" + if key in self.tensors: + return self.tensors[key] + if key in self.metadata: + return self.metadata[key] + raise KeyError(key) + + def __len__(self) -> int: + return len(self.tensors) + len(self.metadata) + + def __iter__(self) -> Iterator[str]: + yield from self.tensors + yield from self.metadata + + def __bool__(self) -> bool: + return bool(self.tensors) or bool(self.metadata) + + def __eq__(self, other: object) -> bool: + """Support equality with plain dicts and other Mappings.""" + if isinstance(other, MultimodalPayload): + return self.tensors == other.tensors and self.metadata == other.metadata + if isinstance(other, Mapping): + return dict(self) == dict(other) + return NotImplemented + + def merged_with(self, incoming: MultimodalPayload) -> MultimodalPayload: + """Merge *incoming* onto this payload and return the result. + + Tensor values accumulate into lists for deferred concatenation; + non-tensor values are replaced with the latest. When this payload + is empty, *incoming* is returned as-is, so callers should use the + return value: ``accumulated = accumulated.merged_with(incoming)``. + """ + if self.is_empty: + return incoming + _append_entries(self.tensors, incoming.tensors) + _append_entries(self.metadata, incoming.metadata) + return self + + def consolidate_tensors(self, strategy: TensorAccumulationStrategy) -> None: + """Concatenate deferred tensor lists into single tensors. + + Tensors are generated content accumulated as chunks, so lists are + concatenated according to *strategy* (e.g. audio chunks along the + time dimension, latent frames along the batch dimension). + """ + # Relocate scalar metadata tensors (e.g. sample rate) that from_dict + # routed into .tensors, so they take the REPLACE path via .metadata + # instead of failing torch.cat as 0-d tensors. + for key in _METADATA_TENSOR_KEYS: + value = self.tensors.pop(key, None) + if value is None: + continue + self.metadata[key] = value[-1] if isinstance(value, list) else value + + for key, value in list(self.tensors.items()): + if _is_tensor_list(value): + self.tensors[key] = _consolidate_tensor_list(key, value, strategy) + + def consolidate_metadata(self) -> None: + """Resolve deferred tensor lists in metadata by keeping the latest value. + + Metadata values are per-step snapshots (e.g. sample rate), not + content deltas, so the latest value supersedes earlier ones. + Nested dicts (unflattened payloads) are resolved one level down. + """ + for key, value in list(self.metadata.items()): + if _is_tensor_list(value): + self.metadata[key] = value[-1] + elif isinstance(value, dict): + for sub_key, sub_value in list(value.items()): + if _is_tensor_list(sub_value): + value[sub_key] = sub_value[-1] + + def to_dict(self) -> dict[str, Any]: + """Convert back to a plain dict (tensors + metadata merged).""" + result: dict[str, Any] = dict(self.tensors) + result.update(self.metadata) + return result + + @classmethod + def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None: + """Create a MultimodalPayload from a raw dictionary. + + Separates torch.Tensor values into tensors and everything + else into metadata. + """ + if not data: + return None + tensors: dict[str, torch.Tensor] = {} + metadata: dict[str, Any] = {} + for k, v in data.items(): + if isinstance(v, torch.Tensor): + tensors[k] = v + else: + metadata[k] = v + if not tensors and not metadata: + return None + return cls(tensors=tensors, metadata=metadata) + + @classmethod + def from_raw(cls, payload: Any, modality_key: str) -> MultimodalPayload | None: + """Create a MultimodalPayload from a raw producer payload. + + Accepts a MultimodalPayload (returned as-is), a dict, or a bare + tensor (stored under *modality_key*). Tensors are moved to CPU. + Producer-specific dict keys are remapped to the semantic modality + key (e.g. "audio", "latent"): AR runners produce {"hidden": ...} + and generation runners produce {"model_outputs": ...}. + """ + if isinstance(payload, MultimodalPayload): + return payload + + if not isinstance(payload, dict): + return cls.from_dict({modality_key: _to_cpu(payload)}) + + remapped: dict[str, Any] = {} + for key, value in payload.items(): + is_producer_key = key == "model_outputs" or (key == "hidden" and modality_key != "hidden") + remapped[modality_key if is_producer_key else key] = _to_cpu(value) + return cls.from_dict(remapped) + + +@dataclass +class MultimodalCompletionOutput(CompletionOutput): + """CompletionOutput with multimodal support. + + Inherits all CompletionOutput fields and adds multimodal_output. + As a CompletionOutput subclass, compatible with all existing vLLM consumers. + """ + + def __init__( + self, + multimodal_output: MultimodalPayload | None = None, + **kwargs: Any, + ): + super().__init__(**kwargs) + self.multimodal_output = multimodal_output + + def __repr__(self) -> str: + base = super().__repr__() + return f"{base[:-1]}, multimodal_output={self.multimodal_output!r})" diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py new file mode 100644 index 000000000..2536a8ade --- /dev/null +++ b/vllm_omni/outputs/output_modality.py @@ -0,0 +1,120 @@ +"""Output modality types for vLLM-Omni. + +This module defines the OutputModality enum and TensorAccumulationStrategy +for type-safe multimodal output routing and tensor merging. + +""" + +from __future__ import annotations + +import re +from enum import Enum, Flag, StrEnum, auto + +_MODALITY_ALIASES: dict[str, str] = { + "speech": "audio", + "images": "image", + "latents": "latent", + "wav": "audio", + "waveform": "audio", + "pixel_values": "image", + "pixels": "image", + "token_ids": "text", + "tokens": "text", +} + + +class OutputModalityNames(StrEnum): + """Keys for output modalities. + + TODO: (Alex) Integrate this with the big-flag enum below + throughout the code + for better type safety (currently only used for output processor). + """ + + TEXT = "text" + IMAGE = "image" + AUDIO = "audio" + LATENT = "latent" + + +# Specify which output modalities may be drained when handling delta messages. +# For some types, e.g., latents, we need to be careful to ensure the full context +# is passed as the stream yields due to assumptions in the I/O processing and model +# when async chunk isn't enabled. +NON_DRAINABLE_MODALITIES = {OutputModalityNames.TEXT, OutputModalityNames.LATENT} +DRAINABLE_MODALITIES = {mod for mod in OutputModalityNames if mod not in NON_DRAINABLE_MODALITIES} + + +class OutputModality(Flag): + """Bit-flag enum for output modalities. + + Compose freely with ``|`` — no need to enumerate every combination. + + Single: ``OutputModality.TEXT``, ``OutputModality.IMAGE``, ... + Compound: ``OutputModality.TEXT | OutputModality.IMAGE`` (text+image) + + Note: POOLING is intentionally excluded. Pooling/embedding is vLLM's + native path (pooling_output → PoolingRequestOutput), handled entirely + by the base OutputProcessor. vLLM-Omni's layer does not participate. + """ + + TEXT = auto() + IMAGE = auto() + AUDIO = auto() + LATENT = auto() + + @classmethod + def from_string(cls, s: str | None) -> OutputModality: + """Parse a free-text modality string into an OutputModality flag. + + Handles common aliases and compound strings separated by + or ,. + + Examples:: + + OutputModality.from_string("text+image") + # → OutputModality.TEXT | OutputModality.IMAGE + """ + if not s or not s.strip(): + return cls.TEXT + + parts = [p.strip().lower() for p in re.split(r"[+,]", s.strip())] + result = cls(0) + for p in parts: + p = _MODALITY_ALIASES.get(p, p) + try: + result |= cls[p.upper()] + except KeyError: + raise ValueError(f"Unknown modality: {p!r}. Supported: {[m.name.lower() for m in cls]}") + return result + + @property + def has_text(self) -> bool: + return OutputModality.TEXT in self + + @property + def has_multimodal(self) -> bool: + return bool(self & ~OutputModality.TEXT) + + +class TensorAccumulationStrategy(Enum): + """Strategy for merging incremental multimodal tensors.""" + + CONCAT_DIM0 = "concat_dim0" + """Concatenate along dimension 0. Used for image/latent tensors.""" + + CONCAT_LAST = "concat_last" + """Concatenate along the last dimension. Used for audio waveforms.""" + + APPEND_LIST = "append_list" + """Append to a list (no tensor concatenation).""" + + REPLACE = "replace" + """Replace previous tensor entirely with the latest one.""" + + +def get_accumulation_strategy(modality: OutputModality) -> TensorAccumulationStrategy: + """Determine tensor merge strategy from the multimodal flags.""" + if OutputModality.AUDIO in modality: + return TensorAccumulationStrategy.CONCAT_LAST + if OutputModality.IMAGE in modality or OutputModality.LATENT in modality: + return TensorAccumulationStrategy.CONCAT_DIM0 + return TensorAccumulationStrategy.CONCAT_DIM0 # default diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py new file mode 100644 index 000000000..e39b995b0 --- /dev/null +++ b/vllm_omni/outputs/output_processor.py @@ -0,0 +1,633 @@ +from dataclasses import fields as dataclass_fields +from typing import Any + +import torch +from vllm.logger import init_logger +from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput +from vllm.sampling_params import RequestOutputKind +from vllm.tokenizers import TokenizerLike +from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason +from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor +from vllm.v1.engine.output_processor import ( + OutputProcessorOutput, + RequestOutputCollector, + RequestState, +) +from vllm.v1.engine.parallel_sampling import ParentRequest +from vllm.v1.metrics.stats import IterationStats, RequestStateStats + +from vllm_omni.data_entry_keys import unflatten_payload +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.mm_outputs import MultimodalCompletionOutput, MultimodalPayload +from vllm_omni.outputs.output_modality import ( + DRAINABLE_MODALITIES, + OutputModality, + get_accumulation_strategy, +) + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Module-level helpers +# --------------------------------------------------------------------------- + + +def _modality_to_type_string(modality: OutputModality) -> str: + """Convert an OutputModality flag to a lowercase type string.""" + try: + if OutputModality.AUDIO in modality: + return "audio" + if OutputModality.IMAGE in modality: + return "image" + if OutputModality.LATENT in modality: + return "latent" + except TypeError: + # Flag identity mismatch (e.g. after module reload in tests). + name = getattr(modality, "name", "") or "" + lowered = name.lower() + if "audio" in lowered: + return "audio" + if "image" in lowered: + return "image" + if "latent" in lowered: + return "latent" + return "text" + + +class OmniRequestState(RequestState): + """Request state for omni models, tracking multimodal outputs. + + Extends the base RequestState with support for accumulating + multimodal tensor outputs (e.g., images, audio, latents) that + are produced incrementally during generation. + """ + + def __init__( + self, + *args, + **kwargs, + ): + # arrival_time is always passed as a keyword argument by the sole caller + # (from_new_request). The *args fallback via positional index was removed + # because upstream RequestState.__init__ reorders parameters between + # releases, making args[12] fragile. + arrival_time = kwargs.get("arrival_time") + super().__init__(*args, **kwargs) + self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0)) + # Omni-specific: multimodal output accumulation + # TODO: mm_type is per-request, not per-key. If a model ever produces + # both audio and latent outputs, the modality type would flip on each + # add_multimodal_tensor() call. Consider tracking per-key modality + # types (e.g. dict[str, str]) for future multi-output models. + self.mm_type: str | None = None + self.mm_accumulated: MultimodalPayload = MultimodalPayload() + + def apply_streaming_update(self, update) -> None: + super().apply_streaming_update(update) + self.native_text_stats.arrival_time = update.arrival_time + + def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None: + """Accumulate a multimodal tensor payload into the request state. + + Normalizes incoming payloads (dict or raw tensor) into a + MultimodalPayload and merges with any previously accumulated data. + Uses list-based deferred concatenation to avoid O(n²) repeated + torch.cat calls. + """ + if payload is None: + return + try: + if mm_type: + self.mm_type = mm_type.lower() + modality_key = self.mm_type or "hidden" + + incoming = MultimodalPayload.from_raw(payload, modality_key) + if incoming is not None: + self.mm_accumulated = self.mm_accumulated.merged_with(incoming) + except (ValueError, TypeError, RuntimeError): + logger.exception("Error accumulating multimodal tensor") + + def _consolidate_multimodal_tensors(self) -> None: + """Consolidate accumulated tensor lists into single tensors. + + Uses TensorAccumulationStrategy derived from the output modality + to determine concatenation behavior. Metadata values always use + REPLACE (keep latest). + """ + if self.mm_accumulated.is_empty: + return + + try: + modality = OutputModality.from_string(self.mm_type) + except (ValueError, KeyError): + modality = OutputModality.TEXT + strategy = get_accumulation_strategy(modality) + + try: + self.mm_accumulated.consolidate_tensors(strategy) + self.mm_accumulated.consolidate_metadata() + except (RuntimeError, TypeError, KeyError): + logger.exception("Error consolidating multimodal tensors") + + # Override: do not route to pooling-only path; always create completion + # outputs, and attach pooling_result into the CompletionOutput. + def make_request_output( + self, + new_token_ids: list[int], + pooling_output: torch.Tensor | None, + finish_reason: FinishReason | None, + stop_reason: int | str | None, + kv_transfer_params: dict[str, Any] | None = None, + routed_experts: Any = None, + ) -> OmniRequestOutput | PoolingRequestOutput | None: + """Create a request output from generation results. + + Creates a RequestOutput or PoolingRequestOutput from the generated + tokens and accumulated multimodal outputs. Attaches multimodal + tensors to the completion output if available. + + Args: + new_token_ids: List of newly generated token IDs + pooling_output: Optional pooling output tensor + finish_reason: Optional finish reason indicating why generation stopped + stop_reason: Optional stop reason (token ID or stop string) + kv_transfer_params: Optional KV cache transfer parameters + routed_experts: Optional MoE routed-expert ids for this step, + attached to the completion output for generation stages + + Returns: + OmniRequestOutput or PoolingRequestOutput if output should be + emitted (based on finish status and output kind), None otherwise + """ + # Pooling-only requests should follow base behavior. + if self.detokenizer is None and pooling_output is not None: + return super().make_request_output( + new_token_ids, + pooling_output, + finish_reason, + stop_reason, + kv_transfer_params, + ) + + finished = finish_reason is not None + final_only = self.output_kind == RequestOutputKind.FINAL_ONLY + is_delta = self.output_kind == RequestOutputKind.DELTA + + if not finished and final_only: + return None + + # Consolidate accumulated tensors when finishing, or for every + # CUMULATIVE step so the consumer always sees a single tensor. + if finished or not is_delta: + self._consolidate_multimodal_tensors() + + if self.stream_interval > 1 and self.detokenizer is not None: + # Send output request only when + # 1. It has finished, or + # 2. It is the first token, or + # 3. It has reached the stream interval number of tokens + if not ( + finished + or self.sent_tokens_offset == 0 + or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval + ): + return None + + if self.output_kind == RequestOutputKind.DELTA: + # Send tokens from the offset in DELTA mode, otherwise all + # tokens are sent. + new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :] + self.sent_tokens_offset = self.detokenizer.num_output_tokens() + + external_req_id = self.external_req_id + + output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts) + + if self.parent_req is None: + outputs = [output] + else: + outputs, finished = self.parent_req.get_outputs(self.request_id, output) + if not outputs: + return None + external_req_id = self.parent_req.external_req_id + + return self._new_request_output( + external_req_id, + outputs, + finished, + kv_transfer_params, + ) + + def _new_completion_output( + self, + token_ids: list[int], + finish_reason: FinishReason | None, + stop_reason: int | str | None, + routed_experts: Any = None, + ) -> MultimodalCompletionOutput | CompletionOutput: + """Create a completion output with multimodal data attached. + + Returns a MultimodalCompletionOutput when multimodal data has been + accumulated, otherwise returns the base CompletionOutput. Snapshots + the accumulated payload before draining DELTA-mode modality keys so + that callers receive an immutable view. + """ + # When there is no detokenizer (generation stages), build a minimal + # CompletionOutput directly since upstream asserts detokenizer != None. + if self.detokenizer is None: + finished = finish_reason is not None + base_output = CompletionOutput( + index=self.request_index, + text="", + token_ids=token_ids, + logprobs=None, + cumulative_logprob=None, + finish_reason=str(finish_reason) if finished else None, + stop_reason=stop_reason if finished else None, + routed_experts=routed_experts, + ) + else: + base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason) + + # Always provide cumulative token IDs for inter-stage processors. + if self.detokenizer is not None: + base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids) + else: + base_output.cumulative_token_ids = list(token_ids) + + # Attach cumulative_text only at the final step for inter-stage use. + if finish_reason is not None and hasattr(self.detokenizer, "output_text"): + base_output.cumulative_text = self.detokenizer.output_text + + try: + if self.mm_accumulated and not self.mm_accumulated.is_empty: + # Snapshot: copy current tensors/metadata so drain doesn't + # mutate the payload already handed to the caller. + # Unflatten dotted keys (e.g. "hidden_states.layer_0") back + # to nested dicts so downstream consumers (thinker2talker etc.) + # can access with .get("hidden_states", {}).get("layers", {}). + merged = {**self.mm_accumulated.tensors, **self.mm_accumulated.metadata} + nested = unflatten_payload(merged) + snapshot = MultimodalPayload.from_dict(nested) or MultimodalPayload() + kwargs = {f.name: getattr(base_output, f.name) for f in dataclass_fields(CompletionOutput)} + output = MultimodalCompletionOutput( + multimodal_output=snapshot, + **kwargs, + ) + output.cumulative_token_ids = base_output.cumulative_token_ids + if hasattr(base_output, "cumulative_text"): + output.cumulative_text = base_output.cumulative_text + + # DELTA mode: drain modality keys (e.g. audio) so the next + # step only sees freshly accumulated data for those keys. + if self.output_kind == RequestOutputKind.DELTA: + for modality_key in DRAINABLE_MODALITIES: + self.mm_accumulated.tensors.pop(modality_key, None) + + return output + except (RuntimeError, TypeError, AttributeError): + logger.exception("Error creating MultimodalCompletionOutput") + return base_output + + def _new_request_output( + self, + external_req_id: str, + outputs: list, + finished: bool, + kv_transfer_params: dict[str, Any] | None = None, + ) -> RequestOutput | PoolingRequestOutput: + """Create request output, handling no-detokenizer generation stages. + + Upstream asserts ``self.logprobs_processor is not None`` which fails + for generation stages (talker, code2wav) that have no tokenizer. + When the logprobs processor is absent we build the RequestOutput + directly with ``prompt_logprobs=None``. + """ + if self.logprobs_processor is not None: + return super()._new_request_output( + external_req_id, + outputs, + finished, + kv_transfer_params, + ) + + # No-detokenizer path: build RequestOutput directly. + prompt_token_ids = self.prompt_token_ids + if prompt_token_ids is None and self.prompt_embeds is not None: + prompt_token_ids = [0] * len(self.prompt_embeds) + if prompt_token_ids is None: + prompt_token_ids = [] + + return RequestOutput( + request_id=external_req_id, + lora_request=self.lora_request, + prompt=self.prompt, + prompt_token_ids=prompt_token_ids, + prompt_logprobs=None, + outputs=outputs, + finished=finished, + kv_transfer_params=kv_transfer_params, + num_cached_tokens=self.num_cached_tokens, + metrics=self.stats, + ) + + +class MultimodalOutputProcessor(VLLMOutputProcessor): + """Handles multimodal output processing. + + Captures multimodal outputs from OmniEngineCoreOutput and accumulates + them as MultimodalPayload in OmniRequestState, before delegating to + the base vLLM OutputProcessor for text handling. + + The data flow is: + 1. For each EngineCoreOutput with multimodal_output: + - Capture into OmniRequestState.add_multimodal_tensor() + 2. Base vLLM OutputProcessor handles text detokenization + 3. On finish, _consolidate_multimodal_tensors() concatenates accumulated + tensors using strategy-based dispatch + 4. _new_completion_output() returns MultimodalCompletionOutput + """ + + def __init__( + self, + tokenizer: TokenizerLike | None, + *, + log_stats: bool, + stream_interval: int = 1, + tracing_enabled: bool = False, + engine_core_output_type: str | None = None, + output_modality: OutputModality = OutputModality.TEXT, + ): + """Initialize the multimodal output processor. + + Args: + tokenizer: Tokenizer for detokenizing text outputs + log_stats: Whether to log statistics + stream_interval: Stream interval for output generation + engine_core_output_type: Optional output type string (e.g., + "image", "audio", "latent"). Converted to OutputModality + internally. Kept for backward compatibility with + stage_init_utils. + output_modality: Type-safe output modality flag. Used to tag + multimodal outputs with the correct modality key when + per-output type info is unavailable. + """ + super().__init__( + tokenizer=tokenizer, + log_stats=log_stats, + stream_interval=stream_interval, + tracing_enabled=tracing_enabled, + ) + # Convert string-based engine_core_output_type to OutputModality + if engine_core_output_type is not None: + self.output_modality = OutputModality.from_string(engine_core_output_type) + else: + self.output_modality = output_modality + self.engine_core_output_type = engine_core_output_type + self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {} + + def _native_text_metric_record(self, request_id: str) -> dict[str, Any]: + return self._native_text_metrics_by_request.setdefault( + request_id, + { + "vllm_ttft_ms": 0.0, + "vllm_tpot_ms": 0.0, + "vllm_itl_ms": 0.0, + "vllm_itls_ms": [], + }, + ) + + def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: + return self._native_text_metrics_by_request.pop(request_id, {}) + + def abort_requests(self, request_ids, internal: bool) -> list[str]: + request_ids = list(request_ids) + for request_id in request_ids: + if internal: + req_state = self.request_states.get(request_id) + if req_state is not None: + self._native_text_metrics_by_request.pop(req_state.external_req_id, None) + else: + self._native_text_metrics_by_request.pop(request_id, None) + return super().abort_requests(request_ids, internal) + + def add_request( + self, + request: EngineCoreRequest, + prompt: str | None, + parent_req: ParentRequest | None = None, + request_index: int = 0, + queue: RequestOutputCollector | None = None, + ) -> None: + """Add a new request to be processed. + + Creates an OmniRequestState for the request and registers it + for output processing. + + Args: + request: Engine core request to add + prompt: Optional prompt string for the request + parent_req: Optional parent request for parallel sampling + request_index: Index of the request in the batch + queue: Optional queue for collecting outputs + + Raises: + ValueError: If the request ID is already registered + """ + request_id = request.request_id + req_state = self.request_states.get(request_id) + if req_state is not None: + self._update_streaming_request_state(req_state, request, prompt) + return + + req_state = OmniRequestState.from_new_request( + tokenizer=self.tokenizer, + request=request, + prompt=prompt, + parent_req=parent_req, + request_index=request_index, + queue=queue, + log_stats=self.log_stats, + stream_interval=self.stream_interval, + ) + self.request_states[request_id] = req_state + if parent_req: + self.parent_requests[parent_req.request_id] = parent_req + self.external_req_ids[req_state.external_req_id].append(request_id) + + def remove_request(self, request_id: str) -> None: + """Rollback one previously registered request if it was never submitted.""" + req_state = self.request_states.pop(request_id, None) + if req_state is None: + return + + external_req_id = getattr(req_state, "external_req_id", None) + if external_req_id is not None: + self._native_text_metrics_by_request.pop(external_req_id, None) + request_ids = self.external_req_ids.get(external_req_id) + if request_ids is not None: + self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id] + if not self.external_req_ids[external_req_id]: + self.external_req_ids.pop(external_req_id, None) + + parent_req = getattr(req_state, "parent_req", None) + if parent_req is not None: + self.parent_requests.pop(parent_req.request_id, None) + + def process_outputs( + self, + engine_core_outputs: list[EngineCoreOutput], + engine_core_timestamp: float | None = None, + iteration_stats: IterationStats | None = None, + ) -> OutputProcessorOutput: + default_mm_type = _modality_to_type_string(self.output_modality) + + # Separate outputs that upstream can handle (has detokenizer or + # pooling) from multimodal-only outputs (no detokenizer, no pooling) + # that would trigger upstream's `assert detokenizer is not None`. + upstream_outputs: list[EngineCoreOutput] = [] + mm_only_outputs: list[EngineCoreOutput] = [] + + for eco in engine_core_outputs: + req_state = self.request_states.get(eco.request_id) + if req_state is None: + continue + + # Accumulate multimodal tensors regardless of path. + if isinstance(req_state, OmniRequestState): + mm_output = getattr(eco, "multimodal_output", None) + if mm_output is not None: + mm_type = getattr(eco, "output_type", None) or default_mm_type + req_state.add_multimodal_tensor(mm_output, mm_type) + + # Route: if no detokenizer and no pooling output, handle locally + # to avoid upstream's assert on detokenizer. + if req_state.detokenizer is None and eco.pooling_output is None: + mm_only_outputs.append(eco) + else: + upstream_outputs.append(eco) + + # Handle multimodal-only outputs (generation stages) locally. + self._process_mm_only_outputs(mm_only_outputs) + + # Delegate text/pooling outputs to upstream. + return super().process_outputs( + upstream_outputs, + engine_core_timestamp=engine_core_timestamp, + iteration_stats=iteration_stats, + ) + + def _process_mm_only_outputs( + self, + engine_core_outputs: list[EngineCoreOutput], + ) -> None: + """Handle outputs from generation stages that have no detokenizer. + + These cannot go through upstream process_outputs because it asserts + detokenizer is not None when pooling_output is None. + """ + for eco in engine_core_outputs: + req_state = self.request_states.get(eco.request_id) + if req_state is None or not isinstance(req_state, OmniRequestState): + continue + + new_token_ids = eco.new_token_ids + finish_reason = eco.finish_reason + stop_reason = eco.stop_reason + kv_transfer_params = eco.kv_transfer_params + routed_experts = eco.routed_experts + req_state.num_cached_tokens = eco.num_cached_tokens + req_state.is_prefilling = False + + if request_output := req_state.make_request_output( + new_token_ids, + None, # pooling_output + finish_reason, + stop_reason, + kv_transfer_params, + routed_experts, + ): + if req_state.queue is not None: + req_state.queue.put(request_output) + + if finish_reason is not None: + self._finish_request(req_state) + + def _update_stats_from_output( + self, + req_state: RequestState, + engine_core_output: EngineCoreOutput, + engine_core_timestamp: float | None, + iteration_stats: IterationStats | None, + ): + was_prefilling = req_state.is_prefilling + native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None + previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0 + + # NOTE: We pass ``None`` for *iteration_stats* to the parent so that + # the upstream's ``_update_stats_from_output`` logs stats via its own + # ``RequestStateStats`` (req_state.stats) without touching the omni + # ``IterationStats`` object. Below we re-apply the same update using + # *native_stats* (the omni-specific ``RequestStateStats`` attached to + # ``OmniRequestState``), which correctly feeds the omni metrics path. + # If the upstream adds new logic in ``_update_stats_from_output`` (e.g. + # additional stat fields), this override must be kept in sync. + super()._update_stats_from_output( + req_state, + engine_core_output, + engine_core_timestamp, + None, + ) + + if iteration_stats is None or engine_core_timestamp is None or native_stats is None: + return + + iteration_stats.update_from_output( + engine_core_output, + engine_core_timestamp, + was_prefilling, + native_stats, + self.lora_states, + req_state.lora_name, + ) + record = self._native_text_metric_record(req_state.external_req_id) + if was_prefilling: + record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0) + return + + if previous_last_token_ts > 0: + itl_ms = max((engine_core_timestamp - previous_last_token_ts) * 1000.0, 0.0) + itls_ms = record.setdefault("vllm_itls_ms", []) + itls_ms.append(itl_ms) + record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms)) + + def _update_stats_from_finished( + self, + req_state: RequestState, + finish_reason: FinishReason | None, + iteration_stats: IterationStats | None, + ): + super()._update_stats_from_finished(req_state, finish_reason, None) + + native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None + if iteration_stats is None or native_stats is None or finish_reason is None: + return + + iteration_stats.update_from_finished_request( + finish_reason=finish_reason, + request_id=req_state.external_req_id, + num_prompt_tokens=req_state.prompt_len, + max_tokens_param=req_state.max_tokens_param, + req_stats=native_stats, + num_cached_tokens=req_state.num_cached_tokens, + ) + if not iteration_stats.finished_requests: + return + + finished_request = iteration_stats.finished_requests[-1] + if finished_request.request_id != req_state.external_req_id: + return + self._native_text_metric_record(req_state.external_req_id)["vllm_tpot_ms"] = ( + float(finished_request.mean_time_per_output_token) * 1000.0 + ) diff --git a/vllm_omni/outputs/utils.py b/vllm_omni/outputs/utils.py new file mode 100644 index 000000000..88fc11fb5 --- /dev/null +++ b/vllm_omni/outputs/utils.py @@ -0,0 +1,24 @@ +"""Shared helpers for multimodal output handling.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def _to_cpu(x: Any) -> Any: + """Move a tensor to CPU; pass through non-tensors unchanged.""" + if isinstance(x, torch.Tensor): + try: + return x.detach().to("cpu", non_blocking=True).contiguous() + except (RuntimeError, AttributeError): + return x + if isinstance(x, dict): + return {str(sub_key): _to_cpu(sub_value) for sub_key, sub_value in x.items()} + return x + + +def _is_tensor_list(value: Any) -> bool: + """Return True if *value* is a non-empty list of tensors (deferred concat).""" + return isinstance(value, list) and bool(value) and isinstance(value[0], torch.Tensor)