OmniRequestOutput inherits RequestOutput

Signed-off-by: Boao Shi <aoibosh@connect.hku.hk>

docs/serving/audio_generate_api.md

@@ -66,7 +66,7 @@ Content-Type: application/json
 |-----------|------|---------|-------------|
 | `input` | string | **required** | Text prompt describing the audio to generate |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
 
 #### Diffusion Parameters
@@ -90,7 +90,6 @@ Returns binary audio data with the appropriate `Content-Type` header:
 | `mp3` | `audio/mpeg` |
 | `flac` | `audio/flac` |
 | `pcm` | `audio/pcm` |
-| `aac` | `audio/aac` |
 | `opus` | `audio/opus` |
 
 ## Examples
@@ -285,7 +284,7 @@ Pydantic validation failure (e.g. invalid `response_format`, `speed` out of rang
     "detail": [
         {
             "type": "literal_error",
-            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
+            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3' or 'opus'",
             ...
         }
     ]

docs/serving/speech_api.md

@@ -109,7 +109,7 @@ Content-Type: application/json
 | `input` | string | **required** | The text to synthesize into speech |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
 | `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25-4.0) |
 
 #### vLLM-Omni Extension Parameters

docs/user_guide/examples/online_serving/text_to_audio.md

@@ -139,7 +139,7 @@ Content-Type: application/json
 |-----------|------|---------|-------------|
 | `input` | string | **required** | Text prompt describing the audio to generate |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
 | `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
 | `audio_start` | float | 0.0 | Audio start time in seconds |

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -7,6 +7,7 @@ E2E Online tests for Qwen3-Omni model.
 import os
 
 import pytest
+from openai import BadRequestError
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -610,3 +611,21 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
     text = responses[0].text_content if responses else ""
     word_count = len(text.split())
     assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True)
+def test_invalid_audio_format_rejected(omni_server, openai_client) -> None:
+    """Ensure invalid audio format is rejected with 400 before streaming starts."""
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text=get_prompt(),
+    )
+    with pytest.raises(BadRequestError, match="audio format"):
+        openai_client.client.chat.completions.create(
+            model=omni_server.model,
+            messages=messages,
+            modalities=["text", "audio"],
+            audio={"voice": "alloy", "format": "aac"},
+            stream=True,
+        )

tests/entrypoints/openai_api/test_audio_format.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for audio output format handling in chat completions.
+
+Covers:
+- #4716: audio.format parameter must be respected, not hardcoded to WAV
+- Format validation rejects unsupported formats
+- pcm16 is mapped to pcm for soundfile compatibility
+- Default format is WAV when not specified
+- create_audio encodes correctly for each supported format
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    DEFAULT_AUDIO_FORMAT,
+    SUPPORTED_AUDIO_FORMATS,
+    SUPPORTED_CHAT_AUDIO_FORMATS,
+    CreateAudio,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestAudioFormatConstants:
+    def test_default_format_is_wav(self):
+        assert DEFAULT_AUDIO_FORMAT == "wav"
+
+    def test_supported_formats_no_aac(self):
+        assert "aac" not in SUPPORTED_AUDIO_FORMATS
+
+    def test_chat_formats_include_pcm16(self):
+        assert "pcm16" in SUPPORTED_CHAT_AUDIO_FORMATS
+
+    def test_chat_formats_superset_of_audio_formats(self):
+        assert SUPPORTED_AUDIO_FORMATS <= SUPPORTED_CHAT_AUDIO_FORMATS
+
+
+class TestCreateAudio:
+    @pytest.fixture
+    def mixin(self):
+        return AudioMixin()
+
+    @pytest.fixture
+    def audio_tensor(self):
+        return np.sin(np.linspace(0, 2 * np.pi, 24000)).astype(np.float32)
+
+    @pytest.mark.parametrize("fmt", ["wav", "mp3", "flac", "opus", "pcm"])
+    def test_create_audio_supported_formats(self, mixin, audio_tensor, fmt):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format=fmt,
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert len(response.audio_data) > 0
+
+    def test_wav_magic_bytes(self, mixin, audio_tensor):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="wav",
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert response.audio_data[:4] == b"RIFF"
+        assert response.media_type == "audio/wav"
+
+    def test_mp3_encoding(self, mixin, audio_tensor):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="mp3",
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert response.media_type == "audio/mpeg"
+        assert response.audio_data[:4] != b"RIFF"
+
+    def test_flac_magic_bytes(self, mixin, audio_tensor):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="flac",
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert response.audio_data[:4] == b"fLaC"
+        assert response.media_type == "audio/flac"
+
+    def test_unsupported_format_falls_back_to_default(self, mixin, audio_tensor):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="aac",
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert response.audio_data[:4] == b"RIFF"
+        assert response.media_type == "audio/wav"
+
+    def test_base64_encoding(self, mixin, audio_tensor):
+        import base64
+
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="wav",
+            speed=1.0,
+            base64_encode=True,
+        )
+        response = mixin.create_audio(audio_obj)
+        decoded = base64.b64decode(response.audio_data)
+        assert decoded[:4] == b"RIFF"
+
+
+class TestResolveAudioFormat:
+    """Test _resolve_audio_format via the serving chat class."""
+
+    @pytest.fixture
+    def serving_chat(self):
+        from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+        return object.__new__(OmniOpenAIServingChat)
+
+    def _make_request(self, audio_params=None):
+        from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+
+        req = ChatCompletionRequest(
+            model="test-model",
+            messages=[{"role": "user", "content": "hello"}],
+        )
+        if audio_params is not None:
+            req.audio = audio_params
+        return req
+
+    def test_default_format_when_no_audio_params(self, serving_chat):
+        request = self._make_request()
+        result = serving_chat._resolve_audio_format(request)
+        assert result == "wav"
+
+    def test_extracts_mp3_format(self, serving_chat):
+        request = self._make_request({"format": "mp3", "voice": "alloy"})
+        result = serving_chat._resolve_audio_format(request)
+        assert result == "mp3"
+
+    def test_pcm16_mapped_to_pcm(self, serving_chat):
+        request = self._make_request({"format": "pcm16", "voice": "alloy"})
+        result = serving_chat._resolve_audio_format(request)
+        assert result == "pcm"
+
+    def test_invalid_format_returns_error(self, serving_chat):
+        from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+
+        request = self._make_request({"format": "aac", "voice": "alloy"})
+        result = serving_chat._resolve_audio_format(request)
+        assert isinstance(result, ErrorResponse)
+        assert "aac" in result.error.message
+
+    def test_all_supported_formats_accepted(self, serving_chat):
+        from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+
+        for fmt in SUPPORTED_CHAT_AUDIO_FORMATS:
+            request = self._make_request({"format": fmt, "voice": "alloy"})
+            result = serving_chat._resolve_audio_format(request)
+            assert not isinstance(result, ErrorResponse), f"Format {fmt} should be accepted"

tests/model_tests/conftest.py

@@ -22,7 +22,7 @@ def tiny_model_paths(request, run_level):
         if not hasattr(item, "callspec"):
             continue
         model_name = item.callspec.params.get("model_name")
-        if model_name is None:
+        if model_name is None or model_name not in DIFFUSION_TEST_SETTINGS:
             continue
         if model_name not in model_paths:
             settings = DIFFUSION_TEST_SETTINGS[model_name]

vllm_omni/benchmarks/data_modules/seed_tts_eval.py

@@ -613,7 +613,12 @@ def compute_seed_tts_wer_metrics(
                 e,
             )
 
-        wav_16k = _pcm_s16le_to_f32_16k(pcm)
+        # Request functions normalize ``tts_output_pcm_bytes`` to Seed-TTS WER
+        # format before it reaches this evaluator: 24 kHz mono int16 PCM.
+        # Do not apply the streamed-response env format here again; for MOSS
+        # local that would reinterpret already-normalized 24 kHz mono audio as
+        # 48 kHz stereo and corrupt ASR/WER.
+        wav_16k = _pcm_s16le_to_f32_16k(pcm, pcm_sample_rate=24000, channels=1)
         if len(wav_16k) == 0:
             asr_failed += 1
             if include_per_item:

vllm_omni/distributed/omni_connectors/utils/serialization.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Mapping
-from dataclasses import asdict, is_dataclass
+from dataclasses import asdict, fields, is_dataclass
 from typing import Any
 
 import msgspec
@@ -59,7 +59,10 @@ class OmniMsgpackEncoder:
         if isinstance(obj, Image.Image):
             return self._encode_pil_image(obj)
 
-        # RequestOutput (not a dataclass, needs special handling)
+        # RequestOutput (not a dataclass, needs special handling).
+        # Note: OmniRequestOutput subclasses RequestOutput but is a dataclass,
+        # so msgspec encodes it natively (field map) and it never reaches
+        # this hook.
         if isinstance(obj, RequestOutput):
             return self._encode_request_output(obj)
 
@@ -252,23 +255,30 @@ class OmniMsgpackDecoder:
     def _decode_omni_request_output(self, obj: dict[str, Any]) -> Any:
         """Decode dict to OmniRequestOutput.
 
-        OmniRequestOutput is a dataclass, so we can use msgspec.convert
-        or construct it directly.
+        By the time this runs, ``_post_process`` has already reconstructed
+        nested values bottom-up (CompletionOutput, tensors, images, ...), so
+        the object is built directly from the known dataclass fields.
         """
         from vllm_omni.outputs import OmniRequestOutput
 
+        data = dict(obj)
+        # Legacy wire format (pre-inheritance): the stage output was nested
+        # under "request_output". Merge its content via the compat setter.
+        legacy_inner = data.pop("request_output", None)
+        multi_modal_placeholders = data.pop("multi_modal_placeholders", None)
+        field_names = {f.name for f in fields(OmniRequestOutput)}
+        known = {k: v for k, v in data.items() if k in field_names}
         try:
-            # Use msgspec.convert for dataclass reconstruction
-            return msgspec.convert(obj, OmniRequestOutput)
+            out = OmniRequestOutput(**known)
         except Exception:
-            try:
-                # Fallback: construct directly if msgspec.convert fails
-                # (e.g., if some fields are missing or have wrong types)
-                return OmniRequestOutput(**obj)
-            except Exception:
-                # If both attempts fail, return dict as-is (defensive fallback)
-                # This should rarely happen if _is_omni_request_output is correct
-                return obj
+            # Defensive fallback; should rarely happen if
+            # _is_omni_request_output is correct.
+            return obj
+        if legacy_inner is not None and not isinstance(legacy_inner, dict):
+            out.request_output = legacy_inner
+        if multi_modal_placeholders is not None:
+            setattr(out, "multi_modal_placeholders", multi_modal_placeholders)
+        return out
 
     def _decode_tensor(self, obj: dict[str, Any]) -> torch.Tensor:
         """Decode dict to torch.Tensor."""

vllm_omni/entrypoints/omni_base.py

@@ -577,7 +577,6 @@ class OmniBase(PDDisaggregationMixin):
         self.prom_metrics.set_running(running)
         self.prom_metrics.set_waiting(max(0, total - running))
 
-        images = getattr(engine_outputs, "images", []) if output_type == "image" else []
         response_metrics: dict[str, Any] = {}
         stage_metrics: dict[str, dict[str, Any]] = {}
         rid_key = str(req_id)
@@ -602,17 +601,14 @@ class OmniBase(PDDisaggregationMixin):
                 response_metrics["final_output_type"] = current_stage_metrics["final_output_type"]
                 response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"]
                 response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"]
+        # Generation content (outputs, prompt, images, trajectory_*, ...) is
+        # copied from engine_outputs onto the returned object by
+        # OmniRequestOutput.__post_init__.
         return OmniRequestOutput(
             request_id=req_id or "",
             stage_id=stage_id,
             final_output_type=output_type,
             request_output=engine_outputs,
-            images=images,
-            trajectory_latents=getattr(engine_outputs, "trajectory_latents", None),
-            trajectory_timesteps=getattr(engine_outputs, "trajectory_timesteps", None),
-            trajectory_log_probs=getattr(engine_outputs, "trajectory_log_probs", None),
-            trajectory_decoded=getattr(engine_outputs, "trajectory_decoded", None),
-            _custom_output=getattr(engine_outputs, "_custom_output", {}),
             metrics=response_metrics,
             stage_durations=stage_durations,
             peak_memory_mb=peak_memory_mb,

vllm_omni/entrypoints/openai/audio_utils_mixin.py

@@ -5,7 +5,7 @@ import torch
 import torchaudio
 from vllm.logger import init_logger
 
-from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
+from vllm_omni.entrypoints.openai.protocol.audio import DEFAULT_AUDIO_FORMAT, AudioResponse, CreateAudio
 
 try:
     import soundfile
@@ -49,13 +49,12 @@ class AudioMixin:
             "pcm": ("RAW", "audio/pcm", {"subtype": "PCM_16"}),
             "flac": ("FLAC", "audio/flac", {}),
             "mp3": ("MP3", "audio/mpeg", {}),
-            "aac": ("AAC", "audio/aac", {}),
             "opus": ("OGG", "audio/ogg", {"subtype": "OPUS"}),
         }
 
         if response_format not in supported_formats:
-            logger.warning(f"Unsupported response format '{response_format}', defaulting to 'wav'.")
-            response_format = "wav"
+            logger.warning(f"Unsupported response format '{response_format}', defaulting to '{DEFAULT_AUDIO_FORMAT}'.")
+            response_format = DEFAULT_AUDIO_FORMAT
 
         soundfile_format, media_type, kwargs = supported_formats[response_format]
 

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -6,6 +6,10 @@ from pydantic import AliasChoices, BaseModel, Field, field_validator, model_vali
 
 _MAX_EMBEDDING_DIM = 8192
 
+SUPPORTED_AUDIO_FORMATS: frozenset[str] = frozenset({"wav", "pcm", "flac", "mp3", "opus"})
+SUPPORTED_CHAT_AUDIO_FORMATS: frozenset[str] = SUPPORTED_AUDIO_FORMATS | {"pcm16"}
+DEFAULT_AUDIO_FORMAT: str = "wav"
+
 
 def _normalize_ref_audio_value(value):
     if value is None:
@@ -59,7 +63,7 @@ class OpenAICreateSpeechRequest(BaseModel):
         default=None,
         description="Instructions for voice style/emotion (maps to 'instruct' for Qwen3-TTS)",
     )
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(
         default=1.0,
         ge=0.25,
@@ -325,7 +329,7 @@ class OpenAICreateAudioGenerateRequest(BaseModel):
         description="Text prompt describing the audio to generate",
     )
     model: str | None = None
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(
         default=1.0,
         ge=0.25,
@@ -391,7 +395,7 @@ class SpeechBatchItem(BaseModel):
     input: str
     voice: str | None = Field(default=None, validation_alias=AliasChoices("voice", "speaker"))
     instructions: str | None = None
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] | None = None
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] | None = None
     speed: float | None = Field(default=None, ge=0.25, le=4.0)
     task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None
     language: str | None = None
@@ -411,7 +415,7 @@ class BatchSpeechRequest(BaseModel):
     items: list[SpeechBatchItem] = Field(..., min_length=1)
     voice: str | None = Field(default=None, validation_alias=AliasChoices("voice", "speaker"))
     instructions: str | None = None
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
     task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None
     language: str | None = None
@@ -500,7 +504,7 @@ class StreamingSpeechSessionConfig(BaseModel):
     task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None
     language: str | None = None
     instructions: str | None = None
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
     max_new_tokens: int | None = Field(default=None, ge=1)
     initial_codec_chunk_frames: int | None = Field(

vllm_omni/entrypoints/openai/serving_chat.py

@@ -97,7 +97,12 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
 from vllm_omni.entrypoints.openai.image_api_utils import encode_image_base64_with_compression, validate_layered_layers
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
-from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    DEFAULT_AUDIO_FORMAT,
+    SUPPORTED_CHAT_AUDIO_FORMATS,
+    AudioResponse,
+    CreateAudio,
+)
 from vllm_omni.entrypoints.openai.protocol.images import (
     ImageData,
     ImageEditARDeltaChunk,
@@ -450,6 +455,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         output_modalities = getattr(request, "modalities", engine_output_modalities)
         request.modalities = output_modalities if output_modalities is not None else engine_output_modalities
 
+        if request.modalities and "audio" in request.modalities:
+            audio_format_check = self._resolve_audio_format(request)
+            if isinstance(audio_format_check, ErrorResponse):
+                return audio_format_check
+
         num_inference_steps = None
         extra_body: dict[str, Any] = {}
         # Omni multistage image generation: Stage-0 (AR) should receive a clean
@@ -2561,10 +2571,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         else:
             sample_rate = int(sr_raw)
 
+        audio_format = self._resolve_audio_format(request)
+        if isinstance(audio_format, ErrorResponse):
+            return audio_format
+
         audio_obj = CreateAudio(
             audio_tensor=audio_tensor,
             sample_rate=sample_rate,
-            response_format="wav",
+            response_format=audio_format,
             speed=1.0,
             base64_encode=True,
         )
@@ -3099,7 +3113,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 request_id=request_id,
             )
 
-        images = getattr(result.request_output, "images", [])
+        images = getattr(result, "images", [])
         stage_durations = result.stage_durations
         peak_memory_mb = result.peak_memory_mb
         cot_output = None
@@ -3170,7 +3184,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         )
                         yield f"data: {chunk.model_dump_json()}\n\n"
                 elif final_output_type == "image":
-                    images = self._flatten_diffusion_images(getattr(output.request_output, "images", []))
+                    images = self._flatten_diffusion_images(getattr(output, "images", []))
                     if not images:
                         raise RuntimeError("Streaming image edit produced an empty final image output.")
                     image_data = [
@@ -3474,8 +3488,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             # Image output path (text2img / img2img)
             final_output_type = getattr(result, "final_output_type", "image")
-            # Handle nested OmniRequestOutput structure where images might be in request_output
-            images = getattr(result.request_output, "images", [])
+            images = getattr(result, "images", [])
             multimodal_output = getattr(result, "multimodal_output", {}) or {}
             stage_durations = result.stage_durations
             peak_memory_mb = result.peak_memory_mb
@@ -3515,10 +3528,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     raise ValueError(f"Unexpected audio tensor rank {audio_tensor.ndim}; expected 1-3 dims.")
                 audio_array = audio_tensor.numpy()
 
+                audio_format = self._resolve_audio_format(request)
+                if isinstance(audio_format, ErrorResponse):
+                    return audio_format
+
                 audio_obj = CreateAudio(
                     audio_tensor=audio_array,
                     sample_rate=sample_rate,
-                    response_format="wav",
+                    response_format=audio_format,
                     speed=1.0,
                     base64_encode=True,
                 )
@@ -3734,6 +3751,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         return height, width
 
+    def _resolve_audio_format(self, request: ChatCompletionRequest) -> str | ErrorResponse:
+        """Extract and validate the audio output format from a chat completion request."""
+        audio_params = getattr(request, "audio", None)
+        if isinstance(audio_params, dict):
+            audio_format = audio_params.get("format", DEFAULT_AUDIO_FORMAT)
+        else:
+            audio_format = DEFAULT_AUDIO_FORMAT
+        if audio_format not in SUPPORTED_CHAT_AUDIO_FORMATS:
+            return self._create_error_response(
+                f"Invalid audio format '{audio_format}'. Supported formats: {sorted(SUPPORTED_CHAT_AUDIO_FORMATS)}",
+            )
+        if audio_format == "pcm16":
+            audio_format = "pcm"
+        return audio_format
+
     def _create_error_response(
         self,
         message: str,

vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py

@@ -0,0 +1,164 @@
+# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+"""CUDA graph wrapper for MOSS-TTS Local streaming codec decode."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.cuda import CUDAGraph
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _CapturedStreamingDecodeGraph:
+    graph: CUDAGraph
+    static_codes: torch.Tensor
+    static_lengths: torch.Tensor
+    static_audio: torch.Tensor
+    static_audio_lengths: torch.Tensor
+
+
+class CUDAGraphStreamingDecoderWrapper:
+    """Replay fixed-size streaming ``_decode_frame`` steps with CUDA graphs.
+
+    The wrapped codec must already be inside ``codec.streaming(batch_size)``.
+    Graphs are keyed by the streaming step frame count ``T``. The batch width is
+    fixed to the codec stream slot count.
+    """
+
+    def __init__(
+        self,
+        codec: nn.Module,
+        *,
+        batch_size: int,
+        num_quantizers: int,
+        reset_streaming_state: Callable[[], None] | None = None,
+    ) -> None:
+        self.codec = codec
+        self.batch_size = int(batch_size)
+        self.num_quantizers = int(num_quantizers)
+        self.reset_streaming_state = reset_streaming_state
+        self.graphs: dict[int, _CapturedStreamingDecodeGraph] = {}
+        self._pool = None
+        self._warmed_up = False
+
+    @property
+    def is_ready(self) -> bool:
+        return bool(self.graphs)
+
+    @property
+    def capture_sizes(self) -> list[int]:
+        return sorted(self.graphs)
+
+    @torch.no_grad()
+    def warmup(self, device: torch.device, capture_sizes: list[int]) -> None:
+        if self._warmed_up:
+            return
+        self._warmed_up = True
+        if device.type != "cuda":
+            return
+        sizes = sorted({int(size) for size in capture_sizes if int(size) > 0}, reverse=True)
+        if not sizes:
+            return
+
+        logger.info(
+            "MOSS-TTS streaming decoder CUDA graph warmup: batch_size=%d n_vq=%d capture_sizes=%s",
+            self.batch_size,
+            self.num_quantizers,
+            list(reversed(sizes)),
+        )
+        start_s = time.perf_counter()
+        with torch.cuda.device(device):
+            for size in sizes:
+                try:
+                    self._capture(size, device)
+                    logger.info("  Captured MOSS-TTS streaming decoder CUDA graph for size=%d", size)
+                except Exception:
+                    self.graphs.pop(size, None)
+                    logger.warning(
+                        "  Failed to capture MOSS-TTS streaming decoder CUDA graph for size=%d; "
+                        "this size will use eager decode",
+                        size,
+                        exc_info=True,
+                    )
+        logger.info(
+            "MOSS-TTS streaming decoder CUDA graph warmup complete: %d/%d captured in %.1f ms",
+            len(self.graphs),
+            len(sizes),
+            (time.perf_counter() - start_s) * 1000.0,
+        )
+
+    @torch.no_grad()
+    def _capture(self, size: int, device: torch.device) -> None:
+        codes = torch.zeros(
+            self.num_quantizers,
+            self.batch_size,
+            size,
+            dtype=torch.long,
+            device=device,
+        )
+        lengths = torch.full((self.batch_size,), size, dtype=torch.long, device=device)
+        exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=device)
+
+        self.codec._set_streaming_exec_mask(exec_mask)
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            _ = self.codec._decode_frame(codes, lengths)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.accelerator.synchronize(device)
+
+        if self.reset_streaming_state is not None:
+            self.reset_streaming_state()
+        self.codec._set_streaming_exec_mask(exec_mask)
+        if self._pool is None:
+            self._pool = current_platform.get_global_graph_pool()
+        graph = CUDAGraph()
+        with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"):
+            output = self.codec._decode_frame(codes, lengths)
+
+        if output.audio is None or output.audio_lengths is None:
+            raise RuntimeError("MOSS-TTS streaming decoder CUDA graph capture produced empty audio.")
+
+        self.graphs[size] = _CapturedStreamingDecodeGraph(
+            graph=graph,
+            static_codes=codes,
+            static_lengths=lengths,
+            static_audio=output.audio,
+            static_audio_lengths=output.audio_lengths,
+        )
+
+    @torch.no_grad()
+    def decode(
+        self,
+        codes_step: torch.Tensor,
+        exec_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor] | None:
+        if not codes_step.is_cuda or torch.cuda.is_current_stream_capturing():
+            return None
+        n_vq, batch_size, step_t = codes_step.shape
+        if int(n_vq) != self.num_quantizers or int(batch_size) != self.batch_size:
+            return None
+        # The streaming codec advances its internal KV offsets by tensor T, so
+        # bucket-padding a smaller step into a larger captured graph would write
+        # padded frames into state. Capture exact step sizes instead.
+        entry = self.graphs.get(int(step_t))
+        if entry is None:
+            return None
+
+        self.codec._set_streaming_exec_mask(exec_mask)
+        entry.static_codes.copy_(codes_step)
+        entry.graph.replay()
+        return entry.static_audio, entry.static_audio_lengths
+
+
+__all__ = ["CUDAGraphStreamingDecoderWrapper"]

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -26,6 +26,9 @@ from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
 from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import (
     MossAudioTokenizerConfig as MossAudioTokenizerV2Config,
 )
+from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import (
+    CUDAGraphStreamingDecoderWrapper,
+)
 from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
     MossTTSCUDAGraphCodecWrapper,
 )
@@ -43,6 +46,7 @@ class _MossCodecStreamSession:
         *,
         stream_slots: int,
         n_vq: int,
+        cudagraph_capture_sizes: list[int] | None = None,
     ) -> None:
         self._codec = codec
         self._stream_slots = int(stream_slots)
@@ -54,6 +58,19 @@ class _MossCodecStreamSession:
         self._closed = False
         with torch.no_grad():
             self._exit_stack.enter_context(codec.streaming(self._batch_size))
+        self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None
+        capture_sizes = sorted({int(size) for size in (cudagraph_capture_sizes or []) if int(size) > 0})
+        if capture_sizes and self._device.type == "cuda":
+            self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper(
+                codec,
+                batch_size=self._batch_size,
+                num_quantizers=self._n_vq,
+                reset_streaming_state=lambda: self.reset_slots(list(range(self._batch_size))),
+            )
+            self._cudagraph_wrapper.warmup(self._device, capture_sizes)
+            self.reset_slots(list(range(self._batch_size)))
+            if not self._cudagraph_wrapper.is_ready:
+                self._cudagraph_wrapper = None
 
     def acquire(self) -> int | None:
         if not self._free_stream_slots:
@@ -109,12 +126,23 @@ class _MossCodecStreamSession:
             codes_lengths[slot] = int(codes.shape[1])
             exec_mask[slot] = True
 
-        self._codec._set_streaming_exec_mask(exec_mask)
-        result = self._codec._decode_frame(codes_step, codes_lengths)
-        if result.audio is None:
-            return {}
-        audio = result.audio.detach().to("cpu", torch.float32)
-        lengths = result.audio_lengths.detach().to("cpu") if result.audio_lengths is not None else None
+        graph_output: tuple[torch.Tensor, torch.Tensor] | None = None
+        if self._cudagraph_wrapper is not None:
+            graph_output = self._cudagraph_wrapper.decode(codes_step, exec_mask)
+
+        used_cudagraph = graph_output is not None
+        if used_cudagraph:
+            audio_tensor, lengths_tensor = graph_output
+        else:
+            self._codec._set_streaming_exec_mask(exec_mask)
+            result = self._codec._decode_frame(codes_step, codes_lengths)
+            if result.audio is None:
+                return {}
+            audio_tensor = result.audio
+            lengths_tensor = result.audio_lengths
+
+        audio = audio_tensor.detach().to("cpu", torch.float32)
+        lengths = lengths_tensor.detach().to("cpu") if lengths_tensor is not None else None
         out: dict[int, torch.Tensor] = {}
         for slot in slot_codes:
             wav = audio[slot]
@@ -174,10 +202,17 @@ class MossTTSCodecDecoder(nn.Module):
         self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32)
         self._stream_session: _MossCodecStreamSession | None = None
         self._stream_slots: int = self._connector_int("codec_stream_slots", default=0)
-        self._stream_max_step_frames: int = self._connector_int("codec_max_step_frames", default=100)
+        self._stream_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0)
+        self._stream_max_step_frames: int = self._stream_chunk_frames or 100
         self._stream_req_slots: dict[str, int] = {}
         self._stream_pending_codes: dict[str, list[torch.Tensor]] = {}
         self._stream_starved_reqs: set[str] = set()
+        self._codec_streaming: bool = self._connector_bool("codec_streaming", default=False)
+        default_graph_capture_max = self._stream_chunk_frames or min(self._stream_max_step_frames, 32)
+        default_graph_capture_sizes = list(range(1, max(1, default_graph_capture_max) + 1))
+        self._streaming_cudagraph_capture_sizes = self._streaming_cudagraph_capture_sizes_from_compilation_config(
+            default_graph_capture_sizes
+        )
 
     # ------------------------------------------------------------------
     # vLLM-Omni stubs (codec has no AR loop)
@@ -240,6 +275,15 @@ class MossTTSCodecDecoder(nn.Module):
                 },
             )
 
+        if self._codec_streaming and runtime_additional_information is None:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={
+                    "model_outputs": [empty] * num_req,
+                    "sr": [sr_tensor] * num_req,
+                },
+            )
+
         audios: list[torch.Tensor] = [empty] * num_req
         srs: list[torch.Tensor] = [sr_tensor] * num_req
         device = next(self._codec.parameters()).device
@@ -293,7 +337,7 @@ class MossTTSCodecDecoder(nn.Module):
                 continue
             meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {}
             finished = bool(meta.get("stream_finished", meta.get("finished", False)))
-            streaming_enabled = bool(meta.get("codec_streaming", False))
+            streaming_enabled = bool(meta.get("codec_streaming", self._codec_streaming))
             code_flat_numel = meta.get("code_flat_numel")
             if streaming_enabled and finished and code_flat_numel is not None and int(code_flat_numel) == 0:
                 for _, wav in self._finish_empty_streaming_requests([info]).items():
@@ -382,7 +426,7 @@ class MossTTSCodecDecoder(nn.Module):
             if not isinstance(info, dict):
                 continue
             meta = (info.get("meta", {}) or {}) if isinstance(info.get("meta", {}), dict) else {}
-            if not bool(meta.get("codec_streaming", False)):
+            if not bool(meta.get("codec_streaming", self._codec_streaming)):
                 continue
             finished = bool(meta.get("stream_finished", meta.get("finished", False)))
             if not finished:
@@ -446,6 +490,7 @@ class MossTTSCodecDecoder(nn.Module):
             self._codec,
             stream_slots=max(1, slots),
             n_vq=self._n_vq,
+            cudagraph_capture_sizes=self._streaming_cudagraph_capture_sizes,
         )
         return self._stream_session
 
@@ -597,6 +642,30 @@ class MossTTSCodecDecoder(nn.Module):
             return int(extra_cfg[name])
         return default
 
+    def _connector_bool(self, name: str, default: bool = False) -> bool:
+        model_cfg = getattr(self.vllm_config, "model_config", None)
+        connector_cfg = getattr(model_cfg, "stage_connector_config", None)
+        if isinstance(connector_cfg, dict):
+            extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg)
+        else:
+            extra_cfg = getattr(connector_cfg, "extra", None)
+        if isinstance(extra_cfg, dict) and name in extra_cfg:
+            return bool(extra_cfg[name])
+        return default
+
+    def _streaming_cudagraph_capture_sizes_from_compilation_config(
+        self,
+        default: list[int],
+    ) -> list[int]:
+        if getattr(self.vllm_config.model_config, "enforce_eager", True):
+            return []
+        compilation_config = getattr(self.vllm_config, "compilation_config", None)
+        capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None)
+        if capture_sizes:
+            max_step_frames = max(1, int(self._stream_max_step_frames))
+            return sorted({int(size) for size in capture_sizes if 0 < int(size) <= max_step_frames})
+        return list(default)
+
     # ------------------------------------------------------------------
     # Weight loading
     # ------------------------------------------------------------------
@@ -716,7 +785,9 @@ class MossTTSCodecDecoder(nn.Module):
             self._n_channels,
         )
 
-        self._maybe_enable_decoder_cudagraph(device)
+        self._configure_decoder_cudagraph(device)
+        if self._codec_streaming and self._streaming_cudagraph_capture_sizes:
+            self._ensure_stream_session()
 
         # vLLM's track_weights_loading() compares the returned set against
         # ``self.named_parameters()``. After ``self._codec = codec`` above,
@@ -739,12 +810,30 @@ class MossTTSCodecDecoder(nn.Module):
         codec = MossAudioTokenizerModel(codec_cfg)
         return codec_cfg, codec
 
-    def _maybe_enable_decoder_cudagraph(self, device: torch.device) -> None:
-        """Capture CUDA Graphs for the codec decoder if enforce_eager is False."""
+    def _configure_decoder_cudagraph(self, device: torch.device) -> None:
+        """Select the codec CUDA Graph path.
+
+        ``enforce_eager`` is the single graph on/off switch. If graphing is
+        enabled, ``codec_streaming`` decides whether decode uses the persistent
+        streaming-state wrapper or the offline full-chunk wrapper.
+        """
         if getattr(self.vllm_config.model_config, "enforce_eager", True):
+            self._streaming_cudagraph_capture_sizes = []
             return
         if self._codec is None:
             return
+        if self._codec_streaming:
+            logger.info(
+                "MOSS-TTS codec CUDA Graph selected streaming wrapper: capture_sizes=%s",
+                self._streaming_cudagraph_capture_sizes,
+            )
+            return
+
+        self._enable_non_streaming_decoder_cudagraph(device)
+
+    def _enable_non_streaming_decoder_cudagraph(self, device: torch.device) -> None:
+        if self._codec is None:
+            return
 
         # Read capture sizes from the connector's extra config (same convention
         # as Qwen3-TTS), falling back to a sensible default covering common

vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py

@@ -1,6 +1,6 @@
 """CUDA Graph acceleration for the MOSS Audio Tokenizer codec decoder.
 
-Captures the codec decode method (_decode_frame, or _decode on the vendored codec) for a set of fixed frame-count
+Captures MossAudioTokenizerModel._decode_frame for a set of fixed frame-count
 bucket sizes, then replays the captured graph at inference time to eliminate
 kernel-launch overhead.  Inputs that exceed all captured sizes fall back to
 eager execution transparently.
@@ -24,7 +24,7 @@ logger = init_logger(__name__)
 
 
 class MossTTSCUDAGraphCodecWrapper:
-    """CUDA Graph wrapper for the codec decode method (_decode_frame, or _decode on the vendored codec).
+    """CUDA Graph wrapper for MossAudioTokenizerModel._decode_frame.
 
     Graphs are keyed by padded_T (int).  On each call the actual T is
     bucket-matched to the smallest pre-captured size >= T.  The static code
@@ -51,15 +51,6 @@ class MossTTSCUDAGraphCodecWrapper:
         enabled: bool = True,
     ) -> None:
         self.model = model
-        # Upstream remote-code tokenizer exposes _decode_frame; the vendored
-        # MossAudioTokenizerModel still exposes _decode. Resolve once here so
-        # both codec implementations work.
-        decode_fn = getattr(model, "_decode_frame", None) or getattr(model, "_decode", None)
-        if decode_fn is None:
-            raise AttributeError(
-                f"{type(model).__module__}.{type(model).__name__} exposes neither _decode_frame nor _decode"
-            )
-        self._decode_fn = decode_fn
         self.capture_sizes: list[int] = sorted(capture_sizes)
         self.num_quantizers = num_quantizers
         self.enabled = enabled
@@ -108,7 +99,7 @@ class MossTTSCUDAGraphCodecWrapper:
             dummy_codes = torch.zeros(nq, 1, size, dtype=torch.long, device=device)
             dummy_lengths = torch.tensor([size], dtype=torch.long, device=device)
             with torch.no_grad():
-                _ = self._decode_fn(dummy_codes, dummy_lengths)
+                _ = self.model._decode_frame(dummy_codes, dummy_lengths)
 
         torch.accelerator.synchronize(device)
 
@@ -141,13 +132,13 @@ class MossTTSCUDAGraphCodecWrapper:
 
         # Extra eager warmup inside capture to ensure all kernels are compiled.
         with torch.no_grad():
-            _ = self._decode_fn(static_codes, static_lengths)
+            _ = self.model._decode_frame(static_codes, static_lengths)
         torch.accelerator.synchronize(device)
 
         graph = CUDAGraph()
         with torch.no_grad():
             with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
-                static_out = self._decode_fn(static_codes, static_lengths)
+                static_out = self.model._decode_frame(static_codes, static_lengths)
 
         self.graphs[size] = graph
         self.static_codes[size] = static_codes

vllm_omni/outputs/__init__.py

@@ -3,7 +3,7 @@ from typing import Any
 
 import torch
 from PIL import Image
-from vllm.outputs import RequestOutput
+from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.v1.outputs import ModelRunnerOutput
 
 from vllm_omni.inputs.data import OmniPromptType
@@ -59,12 +59,58 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
     omni_connector_output: OmniConnectorOutput | None = None
 
 
+# RequestOutput attributes that hold generation content.  When wrapping a
+# stage output, these are copied onto the OmniRequestOutput itself so that
+# the object IS the RequestOutput instead of holding one.
+_REQUEST_OUTPUT_CONTENT_ATTRS = (
+    "prompt",
+    "prompt_token_ids",
+    "prompt_logprobs",
+    "outputs",
+    "finished",
+    "lora_request",
+    "encoder_prompt",
+    "encoder_prompt_token_ids",
+    "num_cached_tokens",
+    "kv_transfer_params",
+)
+
+# Omni-specific content copied when the wrapped stage output is itself an
+# OmniRequestOutput (e.g. a diffusion stage inside a multi-stage pipeline).
+_OMNI_CONTENT_ATTRS = (
+    "images",
+    "latents",
+    "trajectory_latents",
+    "trajectory_timesteps",
+    "trajectory_log_probs",
+    "trajectory_decoded",
+    "_multimodal_output",
+    "_custom_output",
+)
+
+# Dynamic attributes that must not be copied verbatim onto the wrapper:
+# they collide with properties defined on OmniRequestOutput.
+_COPY_SKIP_ATTRS = frozenset({"request_output", "multimodal_output", "custom_output"})
+
+
 @dataclass
-class OmniRequestOutput:
+class OmniRequestOutput(RequestOutput):
     """Unified request output for both pipeline stages and diffusion models.
 
+    Extends vLLM's ``RequestOutput`` so that omni outputs can flow directly
+    through vLLM serving codepaths (which expect ``prompt_token_ids``,
+    ``outputs``, etc. as real attributes).  The inherited fields store the
+    LLM generation content; omni-specific fields store pipeline/diffusion
+    extras.
+
+    Note: ``RequestOutput`` is a plain class (not a dataclass), so all of its
+    attributes are redeclared below as dataclass fields with defaults — the
+    dataclass-generated ``__init__`` replaces ``RequestOutput.__init__`` and
+    must set them itself.
+
     This class handles outputs from:
-    1. Multi-stage LLM pipelines (with stage_id, final_output_type, request_output)
+    1. Multi-stage LLM pipelines (with stage_id, final_output_type, and the
+       inherited RequestOutput fields carrying the stage's generation content)
     2. Diffusion models (with images, prompt, metrics)
 
     Attributes:
@@ -72,30 +118,38 @@ class OmniRequestOutput:
         finished: Whether generation is complete
         stage_id: Identifier of the stage that produced this output (pipeline mode)
         final_output_type: Type of output ("text", "image", "audio", "latents")
-        request_output: The underlying RequestOutput from the stage (pipeline mode)
         images: List of generated PIL images (diffusion mode)
-        prompt: The prompt used for generation (diffusion mode)
+        prompt: The prompt used for generation
         latents: Optional tensor of latent representations (diffusion mode)
-        metrics: Optional dictionary of generation metrics
+        metrics: Generation metrics. A plain dict for omni outputs; may carry
+            vLLM's request stats object when copied from a raw RequestOutput.
     """
 
+    # --- Inherited RequestOutput attributes (redeclared as fields) ---
     request_id: str = ""
+    prompt: OmniPromptType | None = None
+    prompt_token_ids: list[int] | None = None
+    prompt_logprobs: Any = None
+    outputs: list[CompletionOutput] = field(default_factory=list)
     finished: bool = True
+    metrics: Any = field(default_factory=dict)
+    lora_request: Any = None
+    encoder_prompt: str | None = None
+    encoder_prompt_token_ids: list[int] | None = None
+    num_cached_tokens: int | None = None
+    kv_transfer_params: dict[str, Any] | None = None
 
-    # Pipeline stage fields
+    # --- Pipeline stage fields ---
     stage_id: int | None = None
     final_output_type: str = "text"
-    request_output: RequestOutput | None = None
 
-    # Diffusion model fields
+    # --- Diffusion model fields ---
     images: list[Image.Image] = field(default_factory=list)
-    prompt: OmniPromptType | None = None
     latents: torch.Tensor | None = None
     trajectory_latents: torch.Tensor | None = None
     trajectory_timesteps: torch.Tensor | None = None
     trajectory_log_probs: torch.Tensor | None = None
     trajectory_decoded: list | None = None
-    metrics: dict[str, Any] = field(default_factory=dict)
     _multimodal_output: dict[str, Any] = field(default_factory=dict)
     _custom_output: dict[str, Any] = field(default_factory=dict)
 
@@ -110,6 +164,44 @@ class OmniRequestOutput:
     error_status_code: int | None = None
     error_type: str | None = None
 
+    def _copy_content_from(self, source: Any) -> None:
+        """Copy generation content from a stage output into this object.
+
+        ``source`` may be a vLLM ``RequestOutput``, another
+        ``OmniRequestOutput``, or a duck-typed stand-in (tests use mocks and
+        SimpleNamespace). Attribute presence is checked against the source's
+        instance ``__dict__`` so that auto-created Mock attributes are not
+        copied. Control fields set by the wrapper (stage_id,
+        final_output_type, metrics, stage_durations, ...) are intentionally
+        not copied.
+        """
+        source_vars = getattr(source, "__dict__", None)
+
+        def _has(name: str) -> bool:
+            if source_vars is not None:
+                return name in source_vars
+            return hasattr(source, name)
+
+        for name in _REQUEST_OUTPUT_CONTENT_ATTRS + _OMNI_CONTENT_ATTRS:
+            if _has(name):
+                setattr(self, name, getattr(source, name))
+        if not self.request_id:
+            self.request_id = getattr(source, "request_id", "") or ""
+        # Dynamic attribute that vLLM code attaches to RequestOutput; it is a
+        # read-only property on this class, so store it in _multimodal_output.
+        if not isinstance(source, OmniRequestOutput) and _has("multimodal_output"):
+            mm = source.multimodal_output
+            if mm is not None:
+                self._multimodal_output = mm
+        # Preserve remaining dynamic attributes (e.g. dynamically annotated
+        # stage outputs) so they stay reachable through `request_output`.
+        if source_vars:
+            known = type(self).__dataclass_fields__.keys()
+            for name, value in source_vars.items():
+                if name.startswith("_") or name in known or name in _COPY_SKIP_ATTRS:
+                    continue
+                setattr(self, name, value)
+
     @classmethod
     def from_error(
         cls,
@@ -143,7 +235,10 @@ class OmniRequestOutput:
         final_output_type: str,
         request_output: RequestOutput,
     ) -> "OmniRequestOutput":
-        """Create output from pipeline stage.
+        """Create output from a pipeline stage.
+
+        The stage output's content is copied into the new object's own
+        (inherited) fields.
 
         Args:
             stage_id: Stage identifier
@@ -153,13 +248,13 @@ class OmniRequestOutput:
         Returns:
             OmniRequestOutput configured for pipeline mode
         """
-        return cls(
-            request_id=getattr(request_output, "request_id", ""),
+        out = cls(
             stage_id=stage_id,
             final_output_type=final_output_type,
             request_output=request_output,
-            finished=True,
         )
+        out.finished = True
+        return out
 
     @classmethod
     def from_diffusion(
@@ -220,35 +315,21 @@ class OmniRequestOutput:
 
     @property
     def multimodal_output(self) -> Any:
-        """Return multimodal output from the underlying request output or local field.
+        """Return the multimodal output payload.
 
-        For pipeline outputs, this checks completion outputs first, then request_output.
-        For diffusion outputs, this returns the local _multimodal_output field.
+        Checks completion outputs first (where multimodal_output is attached
+        by AR stages), then the local _multimodal_output field.
 
         Returns either a MultimodalPayload (Phase 3+) or a plain dict (legacy).
         """
-        if self.request_output is None:
-            return self._multimodal_output
-
-        # Check completion outputs first (where multimodal_output is attached)
-        for output in getattr(self.request_output, "outputs", []):
+        for output in self.outputs:
             if mm := getattr(output, "multimodal_output", None):
                 return mm
-        if mm := getattr(self.request_output, "multimodal_output", None):
-            return mm
         return self._multimodal_output
 
     @property
     def custom_output(self) -> dict[str, Any]:
-        """Return custom output data from diffusion pipelines.
-
-        For diffusion outputs, returns the local _custom_output field.
-        For pipeline outputs with an inner OmniRequestOutput, forwards
-        the custom_output from the inner request output.
-        """
-        if self.request_output is not None:
-            if isinstance(self.request_output, OmniRequestOutput):
-                return self.request_output._custom_output
+        """Return custom output data from diffusion pipelines."""
         return self._custom_output
 
     @custom_output.setter
@@ -260,58 +341,6 @@ class OmniRequestOutput:
         """Return the number of generated images."""
         return len(self.images)
 
-    # Pass-through properties keep vLLM serving codepaths compatible with
-    # OmniRequestOutput for pipeline outputs (Issue #345).
-    @property
-    def prompt_token_ids(self) -> list[int] | None:
-        """Return prompt token IDs from the underlying request output.
-
-        This property is required for compatibility with vLLM's streaming
-        chat completion generator which checks res.prompt_token_ids.
-        """
-        if self.request_output is not None:
-            return getattr(self.request_output, "prompt_token_ids", None)
-        return None
-
-    @property
-    def outputs(self) -> list[Any]:
-        """Return outputs from the underlying request output.
-
-        This property is required for compatibility with vLLM's streaming
-        and non-streaming chat completion generators.
-        """
-        if self.request_output is not None:
-            return getattr(self.request_output, "outputs", [])
-        return []
-
-    @property
-    def encoder_prompt_token_ids(self) -> list[int] | None:
-        """Return encoder prompt token IDs from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "encoder_prompt_token_ids", None)
-        return None
-
-    @property
-    def prompt_logprobs(self) -> Any:
-        """Return prompt logprobs from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "prompt_logprobs", None)
-        return None
-
-    @property
-    def num_cached_tokens(self) -> int | None:
-        """Return number of cached tokens from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "num_cached_tokens", None)
-        return None
-
-    @property
-    def kv_transfer_params(self) -> Any:
-        """Return KV transfer params from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "kv_transfer_params", None)
-        return None
-
     @property
     def is_diffusion_output(self) -> bool:
         """Check if this is a diffusion model output."""
@@ -320,60 +349,33 @@ class OmniRequestOutput:
     @property
     def is_pipeline_output(self) -> bool:
         """Check if this is a pipeline stage output."""
-        return self.stage_id is not None and self.request_output is not None
+        return self.stage_id is not None
 
     def unwrap(self) -> "OmniRequestOutput":
-        """Unwrap nested OmniRequestOutput to get the innermost result.
+        """Return the output itself.
 
-        This helper handles the common pattern where pipeline outputs may wrap
-        other OmniRequestOutput instances. It recursively unwraps until it reaches
-        the final output with actual content (images, text, etc.).
-
-        Returns:
-            The innermost OmniRequestOutput containing the actual generation results.
-
-        Example:
-            ```python
-            result = omni.generate(...)
-            output = OmniRequestOutput.unwrap_result(result)
-            if output.images:
-                # Access images directly
-                video_frames = output.images
-            ```
+        Kept for backward compatibility: outputs used to nest a stage's
+        OmniRequestOutput inside ``request_output``; content is now flattened
+        onto the object itself, so there is nothing to unwrap.
         """
-        current = self
-        # Unwrap nested pipeline outputs
-        while current.is_pipeline_output and current.request_output is not None:
-            if isinstance(current.request_output, OmniRequestOutput):
-                current = current.request_output
-            else:
-                break
-        return current
+        return self
 
     @staticmethod
     def unwrap_result(result: Any) -> "OmniRequestOutput":
-        """Unwrap result from omni.generate() to get the final OmniRequestOutput.
+        """Extract the final OmniRequestOutput from omni.generate() results.
 
-        This static helper handles the full unwrapping pattern including:
+        Handles:
         1. Extracting from list if needed
         2. Type validation
-        3. Recursive unwrapping of nested pipeline outputs
 
         Args:
             result: The result from omni.generate() - may be a list or OmniRequestOutput
 
         Returns:
-            The innermost OmniRequestOutput with actual content
+            The OmniRequestOutput with actual content
 
         Raises:
             ValueError: If result is not an OmniRequestOutput or list containing one
-
-        Example:
-            ```python
-            result = omni.generate(...)
-            output = OmniRequestOutput.unwrap_result(result)
-            # output is guaranteed to be the final OmniRequestOutput
-            ```
         """
         # Handle list wrapper
         if isinstance(result, list):
@@ -385,7 +387,6 @@ class OmniRequestOutput:
         if not isinstance(result, OmniRequestOutput):
             raise ValueError(f"Expected OmniRequestOutput, got {type(result)}")
 
-        # Unwrap nested outputs
         return result.unwrap()
 
     def to_dict(self) -> dict[str, Any]:
@@ -425,7 +426,8 @@ class OmniRequestOutput:
             f"finished={self.finished}",
             f"stage_id={self.stage_id}",
             f"final_output_type={self.final_output_type!r}",
-            f"request_output={self.request_output}",
+            f"prompt_token_ids={self.prompt_token_ids}",
+            f"outputs={self.outputs}",
             f"images={images_repr}",
             f"prompt={self.prompt!r}",
             f"latents={self.latents}",
@@ -437,3 +439,45 @@ class OmniRequestOutput:
         ]
 
         return f"OmniRequestOutput({', '.join(parts)})"
+
+
+def _request_output_get(self: OmniRequestOutput) -> RequestOutput | None:
+    """Backward-compat view of the wrapped stage output.
+
+    The omni output now IS the RequestOutput, so pipeline/LLM outputs return
+    ``self``; pure diffusion or error outputs return None (matching the old
+    composition semantics where ``request_output`` was unset for them).
+    """
+    if self.stage_id is not None or self.outputs:
+        return self
+    return None
+
+
+def _request_output_set(self: OmniRequestOutput, value: Any) -> None:
+    if value is not None:
+        self._copy_content_from(value)
+
+
+# `request_output` is attached after class creation, NOT declared as a field
+# or InitVar: msgspec serializes dataclasses natively via __dataclass_fields__
+# (which includes InitVar pseudo-fields), so a request_output entry would be
+# read through the property during encoding — returning `self` and recursing
+# forever. As a plain property it is invisible to both dataclasses and
+# msgspec.
+OmniRequestOutput.request_output = property(  # type: ignore[assignment]
+    _request_output_get, _request_output_set
+)
+
+# Keep `request_output=` working as a constructor keyword for
+# backward compatibility: the wrapped stage output's content is copied onto
+# the new object (see _copy_content_from).
+_dataclass_init = OmniRequestOutput.__init__
+
+
+def _omni_request_output_init(self: OmniRequestOutput, *args: Any, request_output: Any = None, **kwargs: Any) -> None:
+    _dataclass_init(self, *args, **kwargs)
+    if request_output is not None:
+        self._copy_content_from(request_output)
+
+
+OmniRequestOutput.__init__ = _omni_request_output_init  # type: ignore[method-assign]

vllm_omni/worker/gpu_generation_model_runner.py

@@ -76,10 +76,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
 
     def _update_request_states(self, scheduler_output: SchedulerOutput):
         # remove requests
-        # Some stateful vocoder model may need to clean the state
-        # to avoid the leak of slots when the requests have been aborted.
-        if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
-            self.model.on_requests_finished(scheduler_output.finished_req_ids)
         for req_id in scheduler_output.finished_req_ids:
             self.input_batch.remove_request(req_id)
         scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
@@ -157,6 +153,13 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             if self.model_config.async_chunk and num_scheduled_tokens:
                 self._update_request_states(scheduler_output)
             deferred_state_corrections_fn = self._update_states(scheduler_output)
+
+            # Notify stateful models of finished requests before the
+            # zero-token early return. Request cleanup is a model lifecycle
+            # concern and does not depend on the inter-stage transfer mode.
+            if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
+                self.model.on_requests_finished(scheduler_output.finished_req_ids)
+
             if not scheduler_output.total_num_scheduled_tokens:
                 return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)