anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Refactor][OutputProcessor 1/8]: clean up multimodal payload accumulation and move payload logic into MultimodalPayload & directory refactor (#4980)

Browse Source 
Signed-off-by: Boao Shi <aoibosh@connect.hku.hk>
Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: bjf-frz <frz123db@gmail.com>
This commit is contained in:
Shi Boao Bowieandbjf-frz authored and GitHub committed 2026-07-15 15:50:19 +08:00
1 parent 795e4e0160
commit a432b8a325
13 files changed
+1080 -1018
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.gitignore
+1
View File
@@ -248,6 +248,7 @@ datasets/
# Output files
outputs/
!vllm_omni/outputs/
results/
generated/
output_*/
docs/api/README.md
+6 -6
View File
@@ -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
tests/engine/test_output_modality.py
+9 -5
View File
@@ -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
tests/engine/test_output_processor.py
+2 -2
View File
@@ -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]
vllm_omni/engine/arg_utils.py
+1 -1
View File
@@ -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
vllm_omni/engine/mm_outputs.py
+7 -125
View File
@@ -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,
)
vllm_omni/engine/output_modality.py
+7 -118
View File
@@ -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,
)
vllm_omni/engine/output_processor.py
+8 -760
View File
@@ -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
)
vllm_omni/engine/stage_init_utils.py
+1 -1
View File
@@ -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
vllm_omni/outputs/mm_outputs.py
+261
View File
@@ -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})"
vllm_omni/outputs/output_modality.py
+120
View File
@@ -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
vllm_omni/outputs/output_processor.py
+633
View File
@@ -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
)
vllm_omni/outputs/utils.py
+24
View File
@@ -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)