[RFC]: Decouple Multimodal Output Channel & Simplify Output Processor

Summary

The current Output Processor in vLLM-Omni has a fundamental architectural problem: multimodal outputs (image/audio/latent) repurpose vLLM's pooler_output/pooling_output channel, causing type mismatches, semantic conflicts, and making it impossible for diffusion models to use vLLM engines as encoders. Additionally, there are numerous unnecessary per-modality methods (all with identical logic), multimodal tensors attached via setattr, divergent sync/async wrapping, and overly broad exception handling.

Core Insights:

1. Multimodal outputs need a dedicated data channel and must not occupy pooling_output. OmniModelRunnerOutput already has a multimodal_outputs field but it's not used in the output pipeline — it should be threaded through to EngineCoreOutput and the Output Processor.
2. The current 6 _process_xxx_output methods all do exactly the same thing. The Output Processor does not need to distinguish between modalities.

This RFC proposes:

1. Separate the multimodal output channel — Add a multimodal_output field through ModelRunnerOutput → EngineCoreOutput → OutputProcessor, stop repurposing pooling_output
2. Eliminate the false modality routing — Delete 6 _process_xxx_output methods and the string if/elif chain, replace with unified logic
3. Introduce OutputModality enum — Provide type safety at the configuration layer, replacing free-text strings
4. Introduce MultimodalCompletionOutput — Replace setattr dynamic attachment
5. Unify OmniRequestOutput — Eliminate sync/async divergence and diffusion double-wrapping
6. Explicit tensor consolidation strategies — Replace audio special-casing and bare except

Motivation

Comparison with Upstream vLLM

vLLM's OutputProcessor design is clean and well-separated:

Component	Responsibility	Extension Method
OutputProcessor	Orchestration: stats → detokenize → generate RequestOutput	Single internal loop, single entry point
RequestState	Per-request state management (detokenizer, logprobs)	from_new_request factory
IncrementalDetokenizer	Incremental detokenize + stop string detection	Fast/Slow implementations
RequestOutput / CompletionOutput	API-facing output data structures	Fixed fields, PoolingOutput uses separate branch

vLLM has only two clean output paths:

• Generation path: new_token_ids → detokenize → CompletionOutput → RequestOutput
• Pooling path: pooling_output → PoolingOutput → PoolingRequestOutput

Routing between them is simply pooling_output is not None.

Current Problems in vLLM-Omni

vLLM-Omni introduced multimodal output needs (image, audio, latent, etc.) on top of this, but the extension approach lacks systematic design:

Problem 0 (Fundamental): Multimodal Outputs Repurpose pooling_output Channel

vLLM's pooler_output/pooling_output was designed for embedding/pooling tasks:

# vLLM ModelRunnerOutput: pooler_output: list[torch.Tensor | None] | None
# vLLM EngineCoreOutput: pooling_output: torch.Tensor | None
# vLLM Scheduler: if request.pooling_params and pooler_output is not None → FINISHED_STOPPED

vLLM-Omni stuffs multimodal outputs (dict[str, torch.Tensor]) into this channel:

# AR Model Runner: multimodal dict into pooler_output
pooler_output: list[dict[str, object]] = []
payload = {"hidden": hidden_slice}
payload.update(mm_payload)  # ← multimodal mixed in

# OmniEngineCoreOutput: type override
class OmniEngineCoreOutput(EngineCoreOutput):
    pooling_output: dict[str, torch.Tensor] | None = None  # was torch.Tensor | None

This creates 4 severe conflicts:

Conflict	Impact
Type mismatch	vLLM expects torch.Tensor, Omni passes dict. torch.as_tensor(dict) fails
Semantic conflict	Scheduler's pooling_params and pooler_output completion check triggered by multimodal outputs
Diffusion encoder blocked	If diffusion model needs vLLM engine as encoder, pooler_output is already occupied
Output Processor hack	Must set eco.pooling_output = None to undo the occupation

Ironically, OmniModelRunnerOutput already defines multimodal_outputs: dict[str, torch.Tensor] | None but it's not used in the output pipeline.

Problem 1: False Modality Routing — 6 Methods Doing the Same Thing

# output_processor.py L381-396
if output_type == "image":
    self._process_image_output(eco)
elif output_type in ("text+image", "text,image", "image+text"):
    self._process_text_image_output(eco)
elif output_type in ("latents", "latent"):
    self._process_latents_output(eco)
elif output_type in ("audio", "speech"):
    self._process_audio_output(eco)
elif output_type == "text":
    self._process_text_output(eco)

This looks like per-modality dispatch, but each method's implementation is identical:

def _process_image_output(self, eco):
    if eco.pooling_output is None:
        tensor = self._extract_from_multimodal_outputs(eco, keys=("image", "images", ...))
        if tensor is not None:
            eco.pooling_output = tensor

def _process_audio_output(self, eco):
    if eco.pooling_output is None:
        tensor = self._extract_from_multimodal_outputs(eco, keys=("audio", "audios", ...))
        if tensor is not None:
            eco.pooling_output = tensor

def _process_latents_output(self, eco):
    if eco.pooling_output is None:
        tensor = self._extract_from_multimodal_outputs(eco, keys=("latent", "latents", ...))
        if tensor is not None:
            eco.pooling_output = tensor

Moreover, _extract_from_multimodal_outputs has a fallback that grabs the first tensor in the dict:

# Try the first tensor in the dict as a fallback
for v in mm.values():
    if isinstance(v, torch.Tensor):
        return v

So the per-modality key lists are effectively useless. The entire routing + 6 methods = ~90 lines of code can be replaced with 3 lines.

Additionally, the string routing itself has problems:

• Same semantics have multiple aliases ("latents" vs "latent", "audio" vs "speech", "text+image" vs "text,image" vs "image+text")
• No compile-time type safety

Problem 2: Inconsistent Multimodal Tensor Attachment

# OmniRequestState._new_completion_output L230-240
if not hasattr(base_output, "multimodal_output"):
    setattr(base_output, "multimodal_output", {})
mm_out = getattr(base_output, "multimodal_output")

multimodal_output is dynamically attached to CompletionOutput via setattr, while CompletionOutput's dataclass has no such field. Consumers must use getattr to access it. OmniRequestOutput has yet another multi-level nested access pattern:

# outputs.py L130-141
for req_out in request_outputs:
    for output in getattr(req_out, "outputs", []):
        if mm := getattr(output, "multimodal_output", None):
            return mm
    if mm := getattr(req_out, "multimodal_output", None):
        return mm

Problem 3: Type Annotations Don't Match Actual Returns

OmniRequestState.make_request_output declares return type OmniRequestOutput | PoolingRequestOutput | None, but actually returns vLLM's RequestOutput via self._new_request_output().

Problem 4: Divergent Sync/Async Orchestration Wrapping

Sync path (omni.py):

output_to_yield = OmniRequestOutput(
    stage_id=stage_id, final_output_type=stage.final_output_type,
    request_output=engine_outputs,
)

Async path (async_omni.py):

if stage.final_output_type == "image":
    output_to_yield = OmniRequestOutput(
        ..., images=images, finished=finished,
    )
else:
    output_to_yield = OmniRequestOutput(
        ..., finished=finished,
    )

Sync doesn't pass images or finished; async branches by final_output_type. Same request produces different OmniRequestOutput fields across paths.

Problem 5: Diffusion Output Double-Wrapping

Diffusion engine directly produces OmniRequestOutput, then orchestration layer wraps again with OmniRequestOutput(request_output=existing_OmniRequestOutput), creating nesting.

Problem 6: Overly Broad Exception Handling

except Exception:
    logger.exception("Error accumulating multimodal tensor")

All exceptions silently swallowed. Multimodal data may be silently lost.

Current Code Quantification

File	Problematic Lines	Main Issues
engine/output_processor.py	~470 lines	False modality routing, setattr, bare except, type mismatch
outputs.py	~256 lines	OmniRequestOutput dual pipeline/diffusion mode
entrypoints/omni.py	~20 lines	Sync wrapping
entrypoints/async_omni.py	~20 lines	Async wrapping (diverges from sync)
distributed/.../serialization.py	~50 lines	multimodal_output serialization (relies on getattr)

Total: ~816 lines with structural problems.

Design

Core Design Principles

1. Multimodal outputs must have a dedicated data channel — they must not repurpose pooler_output/pooling_output. pooling_output returns to its original vLLM semantics (encoder output for embedding/pooling tasks).
2. The Output Processor does not need to distinguish between modalities. All multimodal output processing logic is uniform: read from multimodal_output field → accumulate in request state. Modality information is only meaningful at the configuration layer and tensor consolidation layer.

Overall Architecture

┌──────────────────────────────────────────────────────────────────────────┐
│  ModelRunner                                                              │
│    pooler_output  ─────→  Scheduler ─→ EngineCoreOutput.pooling_output    │
│    (torch.Tensor,          (vLLM semantics: embedding/pooling)             │
│     pooling tasks only)                                                    │
│                                                                          │
│    multimodal_output ──→  Scheduler ─→ EngineCoreOutput.multimodal_output │
│    (dict[str, Any],        (new: dedicated channel for image/audio/latent) │
│     image/audio/latent)                                                   │
└──────────────────────────────────────────────────────────────────────────┘
                                  │
                                  ▼
┌──────────────────────────────────────────────────────────────────────────┐
│                     MultimodalOutputProcessor                             │
│                                                                          │
│  Unified processing (modality-agnostic):                                  │
│  1. Read from eco.multimodal_output (don't touch pooling_output)          │
│  2. Accumulate into OmniRequestState.mm_accumulated                       │
│  3. Base processor handles text / pooling paths normally                   │
│                                                                          │
│  ┌──────────────────────────────────────────┐                           │
│  │ OmniRequestState → MultimodalCompletionOutput                        │
│  └──────────────────────────────────────────┘                           │
│  ┌──────────────────────────────────────────┐                           │
│  │ OmniRequestOutput (unified, non-nestable)                             │
│  └──────────────────────────────────────────┘                           │
└──────────────────────────────────────────────────────────────────────────┘

---

I. Separate Multimodal Output Channel — Stop Repurposing pooling_output

Problem: Multimodal outputs (dict[str, torch.Tensor]) are stuffed into pooler_output/pooling_output (designed for torch.Tensor), causing type conflicts, semantic confusion, and blocking diffusion encoder scenarios.

Design: Add a dedicated multimodal_output field at every layer of the pipeline: Model Runner → Scheduler → EngineCoreOutput → Output Processor. pooler_output/pooling_output returns to vLLM's original semantics.

1.1 ModelRunnerOutput Layer

OmniModelRunnerOutput already has a multimodal_outputs field, but it's currently unused in the output pipeline. Use it directly after refactoring:

class OmniModelRunnerOutput(ModelRunnerOutput):
    multimodal_outputs: dict[str, torch.Tensor] | None = None

Model Runner changes:

# Before (gpu_ar_model_runner.py): multimodal data stuffed into pooler_output
pooler_output: list[dict[str, object]] = []
payload = {"hidden": hidden_slice}
payload.update(mm_payload)   # ← multimodal data mixed into pooler_output
pooler_output.append(payload)
output = OmniModelRunnerOutput(pooler_output=pooler_output, ...)

# After: pooler_output only carries hidden states (or None), multimodal uses dedicated field
pooler_output = [hidden_slice for ...] if needs_hidden else None
output = OmniModelRunnerOutput(
    pooler_output=pooler_output,                    # torch.Tensor | None, vLLM semantics
    multimodal_outputs=per_request_mm_outputs,       # dict[str, list[dict]] | None, dedicated channel
    ...
)

1.2 EngineCoreOutput Layer

# Before
class OmniEngineCoreOutput(EngineCoreOutput):
    pooling_output: dict[str, torch.Tensor] | None = None  # type override, breaks vLLM contract

# After
class OmniEngineCoreOutput(EngineCoreOutput):
    # pooling_output inherits base class, stays torch.Tensor | None
    multimodal_output: dict[str, Any] | None = None         # new dedicated field

1.3 Scheduler Layer

# Before (omni_ar_scheduler.py): multimodal dict passed through pooling_output
pooler_output = pooler_outputs[req_index]  # actually a dict
outputs.append(EngineCoreOutput(pooling_output=pooler_output, ...))

# After: two separate channels
pooler_output = ...  # torch.Tensor | None (pooling tasks only)
mm_output = ...      # dict[str, Any] | None (multimodal only)
outputs.append(OmniEngineCoreOutput(
    pooling_output=pooler_output,    # vLLM semantics
    multimodal_output=mm_output,     # dedicated channel
    ...
))

# Scheduler's pooling completion check no longer mis-triggered by multimodal
if request.pooling_params and pooler_output is not None:
    request.status = RequestStatus.FINISHED_STOPPED  # ← only real pooling triggers this

1.4 ChunkTransferAdapter Layer

# Before: type annotation says Tensor but actually receives dict
def save_async(self, pooling_output: torch.Tensor | None, request):
    ...

# After: type matches reality
def save_async(self, multimodal_output: dict[str, Any] | None, request):
    ...

1.5 Output Processor Layer

# Before: extract multimodal from pooling_output, then set None to hack
if eco.pooling_output is not None and req_state.detokenizer is not None:
    req_state.add_multimodal_tensor(eco.pooling_output)
    eco.pooling_output = None  # ← hack: undo occupation, force text path

# After: read from dedicated field, no hack needed
if eco.multimodal_output is not None:
    req_state.add_multimodal_tensor(eco.multimodal_output)
# pooling_output untouched, base processor handles pooling/text paths normally

Impact:

• pooling_output returns to vLLM original semantics, type restored to torch.Tensor | None
• Diffusion models can use vLLM engine as encoder (pooler_output not occupied)
• Scheduler's pooling_params and pooler_output completion check no longer triggered by multimodal outputs
• Output Processor no longer needs eco.pooling_output = None hack
• ChunkTransferAdapter type annotations match reality
• Multimodal outputs and pooling outputs can coexist — same request can have text + multimodal + pooling output simultaneously

---

II. Eliminate False Modality Routing — Unified Output Processor

Problem: 6 _process_xxx_output methods do the exact same thing; the string if/elif routing is accidental complexity.

Design: Delete all per-modality methods and routing logic. Output Processor reads from eco.multimodal_output (the new dedicated field) instead of eco.pooling_output.

Before (~90 lines):

def process_outputs(self, engine_core_outputs, ...):
    for eco in engine_core_outputs:
        self._route_and_normalize(eco)                        # ← string routing
        req_state = self.request_states.get(eco.request_id)
        ...
        if eco.pooling_output is not None and req_state.detokenizer is not None:
            req_state.add_multimodal_tensor(eco.pooling_output, ...)
            eco.pooling_output = None                         # ← hack
    return super().process_outputs(...)

def _route_and_normalize(self, eco):
    output_type = (getattr(eco, "output_type", ...) or "").lower()
    if output_type == "image":                            # ← if/elif chain
        self._process_image_output(eco)
    elif output_type in ("text+image", "text,image", "image+text"):
        self._process_text_image_output(eco)
    ...  # 6 branches + fallback

# + 6 methods, ~8 lines each, identical logic

After (~12 lines):

class MultimodalOutputProcessor(VLLMOutputProcessor):

    def __init__(self, tokenizer, log_stats, stream_interval=1,
                 output_modality: OutputModality = OutputModality.TEXT):
        super().__init__(tokenizer=tokenizer, log_stats=log_stats,
                         stream_interval=stream_interval)
        self._output_modality = output_modality

    def process_outputs(self, engine_core_outputs, ...):
        for eco in engine_core_outputs:
            req_state = self.request_states.get(eco.request_id)
            if not isinstance(req_state, OmniRequestState):
                continue

            # Read from dedicated multimodal_output field — don't touch pooling_output
            mm_output = getattr(eco, "multimodal_output", None)
            if mm_output is not None:
                req_state.add_multimodal_tensor(mm_output)

        return super().process_outputs(engine_core_outputs, ...)

Impact:

• Delete 8 methods (~90 lines), core logic ~12 lines
• No more eco.pooling_output = None hack — pooling_output is untouched
• No per-modality logic needed — all modalities processed identically

---

III. OutputModality Enum — Configuration-Layer Type Safety

Problem: engine_output_type uses free-text strings with multiple aliases for the same semantics.

Design: Introduce OutputModality enum, used only at the configuration layer and final output wrapping. The Output Processor's internal processing logic does not depend on it.

# vllm_omni/engine/output_modality.py

from enum import Flag, auto

class OutputModality(Flag):
    """Bit-flag enum — compose freely with `|`, no need to enumerate combinations.

    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()

    _ALIASES: ClassVar[dict[str, str]] = {
        "speech": "audio",
        "images": "image",
        "latents": "latent",
    }

    @classmethod
    def from_string(cls, s: str | None) -> "OutputModality":
        """Parse free-text.

        "text+image" / "image,text" / "image+text" → TEXT | IMAGE
        "speech" → AUDIO
        """
        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 = cls._ALIASES.get(p, p)
            try:
                result |= cls[p.upper()]
            except KeyError:
                raise ValueError(f"Unknown modality: {p!r}. "
                                 f"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)

Usage examples:

m = OutputModality.from_string("text+image")
assert m == OutputModality.TEXT | OutputModality.IMAGE
assert m.has_text and m.has_multimodal

# Any combination works — no pre-defined composite members needed
m2 = OutputModality.TEXT | OutputModality.AUDIO | OutputModality.IMAGE
assert m2.has_text and m2.has_multimodal

Where the enum is used (configuration and final output only):

Location	Purpose
OmniEngineArgs.output_modality	Normalize at stage YAML parsing time
MultimodalOutputProcessor._output_modality	Check has_text to decide whether to preserve text path
OmniRequestState._consolidation_strategy	Determine tensor merge strategy
OmniRequestOutput.output_modality	Inform API layer of output type

Explicitly NOT used for:

• Modality routing inside Output Processor (eliminated)
• Per-modality processing method selection (eliminated)

---

IV. MultimodalCompletionOutput — Replacing setattr Attachment

Problem: multimodal_output is dynamically attached to CompletionOutput via setattr; the dataclass has no such field.

Design: Introduce MultimodalCompletionOutput inheriting CompletionOutput with multimodal_output as a first-class field.

# vllm_omni/engine/outputs.py

from dataclasses import dataclass, field
from typing import Any
import torch
from vllm.outputs import CompletionOutput


@dataclass
class MultimodalPayload:
    """Structured multimodal output payload.

    Replaces dict[str, Any] with type-safe access.
    """
    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."""
        if self.tensors:
            return next(iter(self.tensors.values()))
        return None


@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.
    """
    multimodal_output: MultimodalPayload | None = None

In OmniRequestState:

class OmniRequestState(RequestState):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.mm_accumulated: MultimodalPayload | None = None

    def add_multimodal_tensor(self, payload: Any | None) -> None:
        """Unified tensor accumulation — modality-agnostic."""
        if payload is None:
            return
        incoming = self._normalize_to_dict(payload)
        if self.mm_accumulated is None:
            self.mm_accumulated = MultimodalPayload(tensors=incoming)
        else:
            self._merge_tensors(incoming)

    def _new_completion_output(self, token_ids, finish_reason, stop_reason, routed_experts=None):
        base = super()._new_completion_output(token_ids, finish_reason, stop_reason, routed_experts)
        if self.mm_accumulated is None:
            return base
        return MultimodalCompletionOutput(
            index=base.index,
            text=base.text,
            token_ids=base.token_ids,
            cumulative_logprob=base.cumulative_logprob,
            logprobs=base.logprobs,
            routed_experts=base.routed_experts,
            finish_reason=base.finish_reason,
            stop_reason=base.stop_reason,
            multimodal_output=self.mm_accumulated,
        )

Impact:

• Eliminates all setattr / getattr dynamic attribute operations
• multimodal_output becomes a formal field — IDE completion + type checking
• add_multimodal_tensor no longer takes mm_type parameter — processing logic is the same for all modalities

---

V. Unified OmniRequestOutput — Eliminating Double-Wrapping

Problem: OmniRequestOutput serves both pipeline and diffusion modes with mixed fields; diffusion outputs get double-wrapped; sync/async construction logic diverges.

Design: Refactor OmniRequestOutput into a non-nestable unified output container, and provide a _wrap_output unified constructor function shared by sync/async orchestration layers.

@dataclass
class OmniRequestOutput:
    """Unified request output. Non-nestable."""

    request_id: str
    finished: bool = True
    output_modality: OutputModality = OutputModality.TEXT

    # Base output from vLLM
    request_output: RequestOutput | PoolingRequestOutput | None = None

    # Multimodal output
    multimodal: MultimodalPayload | None = None

    # Images (Diffusion)
    images: list[Image.Image] = field(default_factory=list)

    # Metadata
    stage_id: int | None = None
    prompt: OmniPromptType | None = None
    metrics: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_engine_output(cls, request_output, stage_id, output_modality):
        """Construct from engine's RequestOutput. Shared by sync/async."""
        multimodal = None
        for comp_out in getattr(request_output, "outputs", []):
            if isinstance(comp_out, MultimodalCompletionOutput) and comp_out.multimodal_output:
                multimodal = comp_out.multimodal_output
                break
        return cls(
            request_id=request_output.request_id,
            finished=request_output.finished,
            output_modality=output_modality,
            request_output=request_output,
            multimodal=multimodal,
            stage_id=stage_id,
        )

    @classmethod
    def from_diffusion(cls, request_id, images, ...):
        """Construct from Diffusion engine. No double-wrapping."""
        return cls(request_id=request_id, finished=True,
                   output_modality=OutputModality.IMAGE, images=images, ...)

Orchestration layer unification:

# Shared by omni.py and async_omni.py
def _wrap_output(engine_output, stage_id, output_modality):
    if isinstance(engine_output, OmniRequestOutput):
        engine_output.stage_id = stage_id
        return engine_output
    return OmniRequestOutput.from_engine_output(engine_output, stage_id, output_modality)

---

VI. Tensor Consolidation Strategies

Problem: _consolidate_multimodal_tensors has audio special-case skip and bare except.

Design: Bind strategy to OutputModality — this is where modality distinction is actually meaningful, since different modalities have different tensor shape semantics.

class TensorAccumulationStrategy(Enum):
    CONCAT_DIM0 = "concat_dim0"      # image / latent
    CONCAT_LAST = "concat_last"       # audio
    APPEND_LIST = "append_list"
    REPLACE = "replace"

def get_accumulation_strategy(modality: OutputModality) -> TensorAccumulationStrategy:
    """Determine tensor merge strategy from the multimodal flags.

    Uses Flag bit checks — no need to enumerate combinations.
    """
    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

---

VII. Exception Handling — Explicit Propagation

Design: Distinguish recoverable from non-recoverable exceptions.

class MultimodalOutputError(Exception):
    """Base class for multimodal output processing errors."""
    pass

class TensorAccumulationError(MultimodalOutputError):
    """Tensor accumulation failed (recoverable: skip this step's tensor)."""
    pass

class TensorConsolidationError(MultimodalOutputError):
    """Tensor consolidation failed (non-recoverable: request should be marked as failed)."""
    pass

Exception Type	Handling
TensorAccumulationError	Log warning + skip this step's tensor
TensorConsolidationError	Log error + set mm_accumulated = None + record in metrics
Other	Propagate upstream

---

Data Flow Comparison

Before:

OmniModelRunnerOutput (multimodal_outputs field unused)
  pooler_output = [{"hidden": tensor, "audio": tensor}, ...]  ← mm dict stuffed in
  → Scheduler
    pooler_output → EngineCoreOutput.pooling_output            ← type: Tensor → dict
    pooling_params and pooler_output → false completion         ← semantic conflict
  → MultimodalOutputProcessor._route_and_normalize()           ← string if/elif
    → _process_audio_output / _process_image_output / ...      ← 6 methods (same logic)
    → req_state.add_multimodal_tensor(eco.pooling_output)
    → eco.pooling_output = None                                ← hack: undo occupation
  → super().process_outputs()
    → setattr(base_output, "multimodal_output", ...)           ← dynamic attachment
  → Orchestrator → OmniRequestOutput(request_output=...)       ← sync/async diverge
  → API → getattr(output, "multimodal_output", None)           ← duck typing

After:

OmniModelRunnerOutput
  pooler_output = [tensor, ...]                 ← vLLM semantics (Tensor | None)
  multimodal_outputs = {"audio": tensor, ...}   ← dedicated channel
  → Scheduler
    pooler_output → EngineCoreOutput.pooling_output       ← correct type: Tensor | None
    mm_output → OmniEngineCoreOutput.multimodal_output    ← dedicated field
    (pooling_params completion check unaffected)
  → MultimodalOutputProcessor.process_outputs()
    → req_state.add_multimodal_tensor(eco.multimodal_output) ← read from dedicated field
    (pooling_output untouched, no hack needed)
  → super().process_outputs()
    → MultimodalCompletionOutput(multimodal_output=...)      ← subclass, formal field
  → Orchestrator → _wrap_output() → OmniRequestOutput        ← unified function
  → API → output.multimodal.primary_tensor                    ← type-safe access

---

Architecture Diagram

classDiagram
    class OmniEngineCoreOutput {
        +pooling_output: Tensor | None
        +multimodal_output: dict | None
    }

    class OutputModality {
        <<Flag>>
        TEXT
        IMAGE
        AUDIO
        LATENT
        +from_string(s) OutputModality
        +has_text bool
        +has_multimodal bool
    }

    class TensorAccumulationStrategy {
        <<enum>>
        CONCAT_DIM0
        CONCAT_LAST
        APPEND_LIST
        REPLACE
    }

    class MultimodalPayload {
        +tensors: dict
        +metadata: dict
        +primary_tensor Tensor
    }

    class MultimodalCompletionOutput {
        +multimodal_output: MultimodalPayload
    }

    class OmniRequestState {
        +mm_accumulated: MultimodalPayload
        +add_multimodal_tensor(payload)
        +_consolidate_multimodal_tensors()
    }

    class MultimodalOutputProcessor {
        -_output_modality: OutputModality
        +process_outputs()
    }

    class OmniRequestOutput {
        +output_modality: OutputModality
        +multimodal: MultimodalPayload
        +request_output: RequestOutput
        +from_engine_output()
        +from_diffusion()
    }

    EngineCoreOutput <|-- OmniEngineCoreOutput
    OmniEngineCoreOutput ..> MultimodalOutputProcessor : multimodal_output
    MultimodalOutputProcessor --> OutputModality
    MultimodalOutputProcessor --> OmniRequestState
    OmniRequestState --> MultimodalPayload
    OmniRequestState --> TensorAccumulationStrategy
    MultimodalCompletionOutput --> MultimodalPayload
    OmniRequestOutput --> MultimodalPayload
    OmniRequestOutput --> OutputModality
    CompletionOutput <|-- MultimodalCompletionOutput
    RequestState <|-- OmniRequestState
    VLLMOutputProcessor <|-- MultimodalOutputProcessor

Migration Plan

Phase 1: Foundation Types (0.5 day, lowest risk)

• Create vllm_omni/engine/output_modality.py — OutputModality enum + TensorAccumulationStrategy
• Create vllm_omni/engine/outputs.py — MultimodalPayload + MultimodalCompletionOutput
• Add output_modality property to OmniEngineArgs
• Introduce only, no old logic replaced — zero regression

Phase 2: Separate Multimodal Output Channel (1.5 days, most critical)

• Modify OmniEngineCoreOutput: add multimodal_output: dict[str, Any] | None, restore pooling_output to inherit base class torch.Tensor | None
• Modify gpu_ar_model_runner.py: multimodal data no longer stuffed into pooler_output, use OmniModelRunnerOutput.multimodal_outputs instead; pooler_output only holds hidden states tensor or None
• Modify gpu_generation_model_runner.py: same
• Modify omni_ar_scheduler.py: extract pooler_output (→ pooling_output) and multimodal_outputs (→ multimodal_output) separately from OmniModelRunnerOutput
• Modify omni_generation_scheduler.py: same
• Modify ChunkTransferAdapter.save_async: parameter from pooling_output: torch.Tensor to multimodal_output: dict[str, Any]
• End-to-end validation: Qwen2.5-Omni, Qwen3-Omni, GLM-Image, MiMoAudio

Phase 3: Simplify Output Processor + OmniRequestState Refactoring (1 day)

• MultimodalOutputProcessor.__init__ accepts OutputModality
• Rewrite process_outputs: read from eco.multimodal_output, don't touch pooling_output
• Delete _route_and_normalize, 6 _process_xxx_output methods, etc. (~90 lines)
• mm_accumulated type → MultimodalPayload | None
• _new_completion_output returns MultimodalCompletionOutput
• Replace bare except Exception with precise exception types

Phase 4: OmniRequestOutput Unification + Orchestration (1 day)

• Refactor OmniRequestOutput: add output_modality, multimodal fields; implement factory methods
• Implement _wrap_output unified wrapping function
• Update omni.py and async_omni.py to use _wrap_output
• Update serialization code for MultimodalPayload

Phase 5: Cleanup (0.5 day)

• Delete old code and compatibility layers
• Update stage YAML documentation

Directory Structure

vllm_omni/engine/
├── __init__.py                   # OmniEngineCoreOutput
├── output_modality.py            # OutputModality enum + TensorAccumulationStrategy  [NEW]
├── outputs.py                    # MultimodalPayload + MultimodalCompletionOutput     [NEW]
├── output_processor.py           # MultimodalOutputProcessor (greatly slimmed) + OmniRequestState
├── arg_utils.py                  # OmniEngineArgs (add output_modality property)
└── ...

vllm_omni/
├── outputs.py                    # OmniRequestOutput (refactored) + OmniModelRunnerOutput
└── ...

Testing Strategy

Layer	Test Method
OutputModality.from_string	Unit: all known aliases normalize correctly, unknown strings raise ValueError
MultimodalPayload	Unit: tensor store/retrieve, primary_tensor access
MultimodalCompletionOutput	Unit: inherits all CompletionOutput fields + multimodal_output
MultimodalOutputProcessor	Unit: unified processing behaves identically for image/audio/latent
OmniRequestState	Unit: add_multimodal_tensor accumulation; each TensorAccumulationStrategy
OmniRequestOutput	Unit: from_engine_output, from_diffusion; compatibility properties
End-to-end	Qwen2.5-Omni (audio), Qwen3-Omni (audio), GLM-Image (image), MiMoAudio (audio)
Regression	All existing tests preserved, full CI after each Phase

Risks and Mitigations

Risk	Mitigation
Phase 2 touches Model Runner + Scheduler + Transfer Adapter full chain	Validate per-model: fix AR runner first (Qwen2.5-Omni), verify, then Generation runner
After pooler_output separation, downstream stage input processors depending on hidden key need adaptation	Check all custom_process_next_stage_input_func implementations, unify to read from multimodal_output
New field on OmniEngineCoreOutput affects msgspec serialization	multimodal_output set with omit_defaults=True, only serialized when present
MultimodalCompletionOutput incompatible with vLLM serving type checks	Is CompletionOutput subclass, isinstance passes
OutputModality.from_string misses non-standard strings	Phase 1 scans entire codebase for all engine_output_type values
Serialization compatibility	Phase 4 updates serialization.py simultaneously

Out of Scope

Area	Rationale
Streaming multimodal output	Needs separate RFC
IOProcessor plugin system	vLLM pooling post-processing, defer to later
Multimodal OpenAI API compatibility	Serving layer concern
OmniOutput model internal template	Model executor layer (RFC-003)
additional_information protocol	Independent infrastructure concern

Open Questions

1. After separation, should the AR Model Runner's pooler_output still carry hidden states? Or should hidden states also go through multimodal_output? Currently pooler_output = [{"hidden": tensor, ...}] blurs the line between inter-stage data and pooling semantics.
2. Should TensorAccumulationStrategy be declared by the model side (via stage config) rather than bound to OutputModality in the output processor layer?
3. Does MultimodalPayload.tensors need stronger types, or is dict[str, torch.Tensor] flexibility sufficient?
4. Should diffusion pipelines use MultimodalPayload internally for full unification?
5. Do we need to support simultaneous multi-modality combinations (e.g., image + audio) at the OmniRequestOutput level?
6. Should the Scheduler have completion logic for multimodal_output (i.e., multimodal_output is not None → mark finished), similar to pooling, or should completion be driven entirely by finish_reason / new_token_ids?

Feedback Period.

No response

CC List.

@hsliuustc0106 @ZJY0516 @princepride @Gaohan123

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.