[Refactor] Output Processor Phase 2: separate multimodal output channel (#1601) (#2744)

Signed-off-by: Megha Agarwal <agarwalmegha1308@gmail.com>
Signed-off-by: meghaagr13 <divyanshsinghvi@gmail.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: meghaagr13 <divyanshsinghvi@gmail.com>
Co-authored-by: Boao Shi <aoibosh@connect.hku.hk>
Co-authored-by: Taichang Zhou <tzhouam@connect.ust.hk>
Co-authored-by: Cursor <cursoragent@cursor.com>

pyproject.toml

@@ -281,6 +281,7 @@ extend-ignore-identifiers-re = [
     ".*[Oo]no_[Aa]nna.*",
     ".*cann.*",
     ".*[Nn][Dd]3.*",
+    ".*updator.*",
     "byt5.*",
 ]
 extend-ignore-re = [

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -200,9 +200,9 @@ def test_save_async(build_adapter):
     request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
 
     adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": False}
-    adapter.save_async(pooling_output=None, request=request)
+    adapter.save_async(multimodal_output=None, request=request)
     adapter.custom_process_next_stage_input_func = lambda **kwargs: {}
-    adapter.save_async(pooling_output=None, request=request)
+    adapter.save_async(multimodal_output=None, request=request)
 
     task = adapter._pending_save_reqs.popleft()
     assert task["is_finished"] is False
@@ -214,7 +214,7 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad
     request.num_computed_tokens = 10
     request.num_output_placeholders = 2
 
-    adapter.save_async(pooling_output=None, request=request)
+    adapter.save_async(multimodal_output=None, request=request)
 
     assert adapter.requests_num_chunks_sent["external-async"] == 8
     assert len(adapter._pending_save_reqs) == 1
@@ -236,7 +236,7 @@ def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, m
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
 
     adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
     )
 
     assert cleanup_calls == []  # no terminal cleanup; meta.finished is false
@@ -256,7 +256,7 @@ def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatc
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
 
     adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
     )
 
     assert connector.put.called
@@ -274,7 +274,7 @@ def test_send_single_request_struct_preserves_segment_finished(build_adapter, mo
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
 
     adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
+        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
     )
 
     sent_payload = connector.put.call_args.kwargs["data"]
@@ -289,13 +289,13 @@ def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapt
     request.num_computed_tokens = 0
     adapter.requests_num_chunks_sent["ext-stream"] = 111
 
-    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
 
     assert len(adapter._pending_save_reqs) == 0
     assert adapter.requests_num_chunks_sent["ext-stream"] == 111
 
     adapter.requests_num_chunks_sent.pop("ext-stream")
-    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
 
     assert len(adapter._pending_save_reqs) == 1
     assert adapter.requests_num_chunks_sent["ext-stream"] == 0
@@ -312,7 +312,7 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
 
     adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
+        {"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
     )
 
     assert len(cleanup_calls) == 1

tests/e2e/offline_inference/test_voxcpm2_tts.py

@@ -1,6 +1,7 @@
 """E2E test for VoxCPM2 native AR offline inference."""
 
 import os
+from collections.abc import Mapping
 
 import pytest
 import torch
@@ -23,7 +24,7 @@ def voxcpm2_engine():
 
 def _extract_audio(multimodal_output: dict) -> torch.Tensor:
     """Extract the final complete audio tensor from multimodal output."""
-    assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}"
+    assert isinstance(multimodal_output, (dict, Mapping)), f"Expected dict/Mapping, got {type(multimodal_output)}"
 
     # Output processor accumulates per-step audio chunks under "audio".
     audio = multimodal_output.get("audio")

tests/engine/test_ensure_tensor_values.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for _ensure_tensor_values boundary enforcement."""
+
+import pytest
+import torch
+
+from vllm_omni.worker.gpu_ar_model_runner import _ensure_tensor_values
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_tensors_pass_through():
+    """Tensor values pass through unchanged."""
+    t = torch.randn(10)
+    result = _ensure_tensor_values({"audio": t})
+    assert torch.equal(result["audio"], t)
+
+
+def test_scalar_int_wrapped():
+    """Integer scalars are wrapped in torch.tensor."""
+    result = _ensure_tensor_values({"sr": 24000})
+    assert isinstance(result["sr"], torch.Tensor)
+    assert result["sr"].item() == 24000
+
+
+def test_scalar_float_wrapped():
+    """Float scalars are wrapped in torch.tensor."""
+    result = _ensure_tensor_values({"gain": 0.5})
+    assert isinstance(result["gain"], torch.Tensor)
+    assert result["gain"].item() == pytest.approx(0.5)
+
+
+def test_list_converted():
+    """Lists of numbers are converted to tensors."""
+    result = _ensure_tensor_values({"ids": [1, 2, 3]})
+    assert isinstance(result["ids"], torch.Tensor)
+    assert list(result["ids"]) == [1, 2, 3]
+
+
+def test_non_tensorizable_dropped():
+    """Values that cannot be tensorized are dropped with warning."""
+    result = _ensure_tensor_values({"ok": torch.randn(5), "bad": {"nested": "dict"}})
+    assert "ok" in result
+    assert "bad" not in result
+
+
+def test_mixed_payload():
+    """Mixed tensor + scalar + list payload is properly sanitized."""
+    payload = {
+        "audio": torch.randn(100),
+        "sr": 24000,
+        "codes": [1, 2, 3],
+        "hidden": torch.randn(10, 64),
+    }
+    result = _ensure_tensor_values(payload)
+    assert len(result) == 4
+    for v in result.values():
+        assert isinstance(v, torch.Tensor)

tests/engine/test_output_modality.py

@@ -81,7 +81,7 @@ def test_multimodal_payload_and_completion_output():
     assert p is not None
     assert "waveform" in p.tensors and torch.equal(p.primary_tensor, data["waveform"])
     assert p.metadata["sample_rate"] == 16000
-    assert not p.is_empty and len(p) == 1
+    assert not p.is_empty and len(p) == 2  # 1 tensor + 1 metadata
 
     # None/empty returns None
     assert MultimodalPayload.from_dict(None) is None
@@ -125,7 +125,7 @@ def test_output_modality_printed_examples(capsys):
     print(f"    tensors keys : {list(p.tensors.keys())}")
     print(f"    primary_tensor: shape={p.primary_tensor.shape}, dtype={p.primary_tensor.dtype}")
     print(f"    metadata      : {p.metadata}")
-    print(f"    is_empty={p.is_empty}, len={len(p)}")
+    print(f"    is_empty={p.is_empty}, len={len(p)}")  # len = tensors + metadata
     print(f"  from_dict(None) -> {MultimodalPayload.from_dict(None)}")
     print(f"  from_dict({{}})   -> {MultimodalPayload.from_dict({})}")
 

tests/engine/test_output_processor.py

@@ -209,3 +209,64 @@ def test_cumulative_token_ids_is_a_copy():
     s = _make_state(RequestOutputKind.DELTA)
     out = s._new_completion_output([42], None, None)
     assert out.cumulative_token_ids is not _DETOK.output_token_ids
+
+
+# ---------------------------------------------------------------------------
+# No-detokenizer regression tests (Concern 2)
+# ---------------------------------------------------------------------------
+
+_NO_DETOK_STATE_KWARGS = {
+    **_DEFAULT_STATE_KWARGS,
+    "detokenizer": None,
+    "logprobs_processor": None,
+}
+
+
+def _make_no_detok_state(output_kind: RequestOutputKind):
+    return OmniRequestState(**_NO_DETOK_STATE_KWARGS, output_kind=output_kind)
+
+
+def test_no_detokenizer_completion_output():
+    """_new_completion_output works when detokenizer is None (generation stages)."""
+    s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
+    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
+    out = s._new_completion_output([1, 2], None, None)
+    # Should still produce output with empty text and the token_ids as-is
+    assert out.text == ""
+    assert list(out.token_ids) == [1, 2]
+    assert hasattr(out, "cumulative_token_ids")
+    assert list(out.cumulative_token_ids) == [1, 2]
+    # Multimodal data should be attached
+    assert AUDIO in out.multimodal_output
+
+
+def test_no_detokenizer_make_request_output():
+    """make_request_output works without detokenizer when multimodal data is present."""
+    s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
+    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
+    result = s.make_request_output([], None, FinishReason.STOP, None)
+    assert result is not None
+    assert not isinstance(result, PoolingRequestOutput)
+    assert AUDIO in result.outputs[0].multimodal_output
+
+
+def test_no_detokenizer_stream_interval_skipped():
+    """stream_interval > 1 does not crash when detokenizer is None."""
+    kwargs = {**_NO_DETOK_STATE_KWARGS, "stream_interval": 5}
+    s = OmniRequestState(**kwargs, output_kind=RequestOutputKind.DELTA)
+    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
+    # Should not assert; stream_interval logic is simply skipped
+    result = s.make_request_output([], None, FinishReason.STOP, None)
+    assert result is not None
+
+
+def test_no_detokenizer_final_only():
+    """FINAL_ONLY with no detokenizer returns None until finish."""
+    s = _make_no_detok_state(RequestOutputKind.FINAL_ONLY)
+    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
+    # Non-finish step should return None
+    assert s.make_request_output([], None, None, None) is None
+    # Finish step should return output
+    result = s.make_request_output([], None, FinishReason.STOP, None)
+    assert result is not None
+    assert AUDIO in result.outputs[0].multimodal_output

tests/engine/test_wire_multimodal_output.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Msgpack roundtrip tests for OmniEngineCoreOutputs with multimodal_output.
+
+Validates that tensor-only payloads survive msgspec encode/decode and that
+non-tensor values fail decoding (enforcing the wire invariant).
+"""
+
+import pytest
+import torch
+from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
+
+from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _roundtrip(outputs: OmniEngineCoreOutputs) -> OmniEngineCoreOutputs:
+    encoder = MsgpackEncoder()
+    encoded = encoder.encode(outputs)
+    decoder = MsgpackDecoder(OmniEngineCoreOutputs)
+    return decoder.decode(encoded)
+
+
+def test_tensor_only_roundtrip():
+    """Tensor-only multimodal_output survives msgpack roundtrip."""
+    audio = torch.randn(16000)
+    sr = torch.tensor(24000)
+    eco = OmniEngineCoreOutput(
+        request_id="req-1",
+        new_token_ids=[1, 2, 3],
+        finish_reason=None,
+        multimodal_output={"audio": audio, "sr": sr},
+    )
+    decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
+    assert len(decoded.outputs) == 1
+    out = decoded.outputs[0]
+    assert out.request_id == "req-1"
+    assert isinstance(out.multimodal_output, dict)
+    assert torch.allclose(out.multimodal_output["audio"], audio)
+    assert out.multimodal_output["sr"].item() == 24000
+
+
+def test_empty_multimodal_roundtrip():
+    """None multimodal_output roundtrips correctly."""
+    eco = OmniEngineCoreOutput(
+        request_id="req-2",
+        new_token_ids=[4],
+        finish_reason=None,
+        multimodal_output=None,
+    )
+    decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
+    assert decoded.outputs[0].multimodal_output is None
+
+
+def test_multiple_tensor_keys_roundtrip():
+    """Multiple tensor keys survive roundtrip."""
+    hidden = torch.randn(10, 4096)
+    codes = torch.randint(0, 1024, (1, 100))
+    eco = OmniEngineCoreOutput(
+        request_id="req-3",
+        new_token_ids=[],
+        finish_reason=None,
+        multimodal_output={"hidden": hidden, "codes": codes},
+    )
+    decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
+    out = decoded.outputs[0]
+    assert torch.allclose(out.multimodal_output["hidden"], hidden)
+    assert torch.equal(out.multimodal_output["codes"], codes)
+
+
+def test_non_tensor_value_fails_decode():
+    """Non-tensor values in multimodal_output cause decode failure.
+
+    This validates that the wire field must be tensor-only and that
+    the _ensure_tensor_values boundary in model runners is necessary.
+    """
+    eco = OmniEngineCoreOutput(
+        request_id="req-bad",
+        new_token_ids=[],
+        finish_reason=None,
+        multimodal_output={"audio": torch.randn(10), "sr": 24000},  # scalar!
+    )
+    encoder = MsgpackEncoder()
+    encoded = encoder.encode(OmniEngineCoreOutputs(outputs=[eco]))
+    decoder = MsgpackDecoder(OmniEngineCoreOutputs)
+    with pytest.raises(Exception):
+        decoder.decode(encoded)

tests/helpers/assertions.py

@@ -481,12 +481,32 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
             if "text" in modalities:
                 transcript = (response.audio_content or "").strip()
                 text_output = (response.text_content or "").strip()
-                similarity = cosine_similarity_text(
-                    transcript.lower(),
-                    text_output.lower(),
-                )
-                print(f"similarity is: {similarity}")
-                assert similarity > similarity_threshold, "The audio content is not same as the text"
+                # For very short outputs (e.g. one-word answers), n-gram cosine
+                # similarity with length penalty is unreliable because Whisper
+                # may hallucinate extra context around the short utterance.  Use
+                # a containment check instead: the shorter text must appear in
+                # the longer one (after preprocessing removes punctuation).
+                _SHORT_TEXT_THRESHOLD = 15
+                if len(text_output) <= _SHORT_TEXT_THRESHOLD or len(transcript) <= _SHORT_TEXT_THRESHOLD:
+                    shorter = text_output.lower() if len(text_output) <= len(transcript) else transcript.lower()
+                    longer = transcript.lower() if len(text_output) <= len(transcript) else text_output.lower()
+                    import re as _re
+
+                    shorter_clean = _re.sub(r"[^\w\s]", "", shorter).strip()
+                    longer_clean = _re.sub(r"[^\w\s]", "", longer).strip()
+                    assert shorter_clean and (shorter_clean in longer_clean), (
+                        f"The audio content is not same as the text "
+                        f"(short-text containment check failed: "
+                        f"text={text_output!r}, transcript={transcript!r})"
+                    )
+                    print(f"short-text containment check passed: {shorter_clean!r} in {longer_clean!r}")
+                else:
+                    similarity = cosine_similarity_text(
+                        transcript.lower(),
+                        text_output.lower(),
+                    )
+                    print(f"similarity is: {similarity}")
+                    assert similarity > similarity_threshold, "The audio content is not same as the text"
             if audio_ref_text:
                 audio_similarity = cosine_similarity_text(
                     response.audio_content.lower(),

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -140,7 +140,7 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=True,
     )
@@ -167,7 +167,7 @@ def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=True,
     )
@@ -188,13 +188,13 @@ def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
 
     payload1 = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=False,
     )
     payload2 = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=False,
     )
@@ -216,14 +216,14 @@ def test_talker2code2wav_async_chunk_waits_for_prelookahead_and_emits_cumulative
 
     payload_pending = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=False,
     )
     request.output_token_ids = [1, 2, 3]
     payload_ready = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=False,
     )
@@ -246,14 +246,14 @@ def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
 
     payload_stream = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=False,
     )
     request.output_token_ids = [3, 4, 5, 6]
     payload_final = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=True,
     )
@@ -281,14 +281,14 @@ def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
 
     payload_pending = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=False,
     )
     request.output_token_ids = [8, 9, 10, 11]
     payload_ready = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=False,
     )
@@ -310,13 +310,13 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
 
     payload_stream = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=False,
     )
     payload_final = talker2code2wav_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
         is_finished=True,
     )

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -56,7 +56,7 @@ def _call(tm, rid, *, n_frames, finished=False, req_ic=None):
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)]
     return talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,))}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,))}},
         request=_req(rid, finished=finished, initial_codec_chunk_frames=req_ic),
         is_finished=finished,
     )
@@ -66,7 +66,7 @@ def test_empty_returns_none():
     tm = _tm()
     p = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,))}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,))}},
         request=_req("r", finished=False),
     )
     assert p is None
@@ -76,7 +76,7 @@ def test_eof_marker_when_finished_empty():
     tm = _tm()
     p = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output=None,
+        multimodal_output=None,
         request=_req("r", finished=True),
         is_finished=True,
     )
@@ -89,7 +89,7 @@ def test_flush_on_finish():
     tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(24)]
     p = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output=None,
+        multimodal_output=None,
         request=_req("r", finished=True),
         is_finished=True,
     )
@@ -262,7 +262,7 @@ def test_first_streaming_chunk_prepends_ref_code_context():
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
@@ -286,7 +286,7 @@ def test_followup_ref_code_context_is_sent_as_metadata_handle():
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
@@ -319,7 +319,7 @@ def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
         request=_req(rid, finished=False),
         is_finished=False,
     )
@@ -338,7 +338,7 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
 
     first_payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]]), "ref": ref_code}},
+        multimodal_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]]), "ref": ref_code}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
@@ -348,14 +348,14 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
     for _ in range(8):
         talker2code2wav_async_chunk(
             transfer_manager=tm,
-            pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
+            multimodal_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
             request=_req(rid, finished=False, initial_codec_chunk_frames=10),
             is_finished=False,
         )
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
+        multimodal_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )

tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py

@@ -92,7 +92,7 @@ def test_empty_chunk_when_not_finished():
 
     payload = generator2tokenizer_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output={"audio": torch.zeros((0,))},
+        multimodal_output={"audio": torch.zeros((0,))},
         request=request,
     )
 
@@ -109,7 +109,7 @@ def test_flush_tail_when_finished():
 
     payload = generator2tokenizer_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,  # e.g. EOS step with no audio
+        multimodal_output=None,  # e.g. EOS step with no audio
         request=request,
     )
 
@@ -133,7 +133,7 @@ def test_eof_marker_when_finished_with_no_frames():
 
     payload = generator2tokenizer_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
     )
 
@@ -153,10 +153,10 @@ def test_normal_chunk_emission():
 
     # Feed 25 frames one by one
     for i in range(25):
-        pooling_output = {"audio": torch.tensor([float(i)] * 4)}
+        multimodal_output = {"audio": torch.tensor([float(i)] * 4)}
         payload = generator2tokenizer_async_chunk(
             transfer_manager=transfer_manager,
-            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
             request=request,
         )
 
@@ -181,10 +181,10 @@ def test_small_initial_chunks():
 
     # Feed 5 frames (should trigger emission because codec_chunk_frames_at_begin=5)
     for i in range(5):
-        pooling_output = {"audio": torch.tensor([float(i + 1)] * 3)}
+        multimodal_output = {"audio": torch.tensor([float(i + 1)] * 3)}
         payload = generator2tokenizer_async_chunk(
             transfer_manager=transfer_manager,
-            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
             request=request,
         )
 
@@ -208,10 +208,10 @@ def test_no_emission_between_boundaries():
 
     # Feed 3 frames (3 % 5 != 0, should not emit)
     for i in range(3):
-        pooling_output = {"audio": torch.tensor([float(i)] * 4)}
+        multimodal_output = {"audio": torch.tensor([float(i)] * 4)}
         payload = generator2tokenizer_async_chunk(
             transfer_manager=transfer_manager,
-            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
             request=request,
         )
 
@@ -230,10 +230,10 @@ def test_context_handling_format():
 
     # Feed 5 frames to trigger a chunk (chunk_size_at_begin=5)
     for i in range(5):
-        pooling_output = {"audio": torch.tensor([float(i + 10)] * 2)}  # codebook_dim=2
+        multimodal_output = {"audio": torch.tensor([float(i + 10)] * 2)}  # codebook_dim=2
         payload = generator2tokenizer_async_chunk(
             transfer_manager=transfer_manager,
-            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
             request=request,
         )
 
@@ -250,14 +250,14 @@ def test_context_handling_format():
     assert flat_codes.numel() == total_window_frames * 2  # codebook_dim=2
 
 
-def test_none_pooling_output_not_finished_returns_none():
-    """None pooling_output when not finished returns None."""
+def test_none_multimodal_output_not_finished_returns_none():
+    """None multimodal_output when not finished returns None."""
     transfer_manager = _make_transfer_manager()
     request = _req("rid-none", finished=False)
 
     payload = generator2tokenizer_async_chunk(
         transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
         request=request,
     )
 

tests/unit/higgs_audio_v3/test_higgs_audio_v3.py

@@ -8,7 +8,8 @@ without requiring the actual checkpoint or GPU.
 
 import asyncio
 import time
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -1156,6 +1157,31 @@ class TestStageInputProcessor:
         assert result[0, 1].item() == 200
         assert result[0, 2].item() == 300
 
+    def test_async_chunk_accepts_multimodal_output_keyword(self):
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            talker2code2wav_async_chunk,
+        )
+
+        request_id = "req-higgs"
+        transfer_manager = SimpleNamespace(
+            code_prompt_token_ids=defaultdict(list),
+            connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 1, "codec_right_holdback_frames": 0}}),
+        )
+        request = SimpleNamespace(external_req_id=request_id, is_finished=lambda: False)
+
+        result = None
+        for step in range(8):
+            result = talker2code2wav_async_chunk(
+                transfer_manager=transfer_manager,
+                multimodal_output={"codes": {"audio": torch.full((1, 8), step, dtype=torch.long)}},
+                request=request,
+                is_finished=False,
+            )
+
+        assert result is not None
+        assert result.codes is not None
+        assert result.codes.audio.numel() > 0
+
 
 # ---- AC-10: Registry ----
 

tests/worker/test_gpu_ar_model_runner.py

@@ -92,6 +92,9 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
     runner.supports_mm_inputs = False
     runner.use_async_scheduling = False
     runner._omni_num_scheduled_tokens_np = None
+    runner.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(engine_output_type="audio"),
+    )
 
     monkeypatch.setattr(
         GPUARModelRunner, "_sample", lambda self, logits, spec_decode_metadata: SimpleNamespace(sampled_token_ids=[])
@@ -125,8 +128,8 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
 
     output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)
 
-    assert torch.equal(output.pooler_output[0]["hidden"], torch.tensor([[1.0, 10.0]]))
+    assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]]))
     assert torch.equal(
-        output.pooler_output[1]["hidden"],
+        output.multimodal_outputs[1]["hidden"],
         torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
     )

tests/worker/test_gpu_generation_model_runner.py

@@ -49,8 +49,8 @@ def test_sample_tokens_tensor_output():
 
     output = GPUGenerationModelRunner.sample_tokens(runner)
 
-    assert len(output.pooler_output) == 1
-    assert output.pooler_output[0]["model_outputs"].shape == (2, 3)
+    assert len(output.multimodal_outputs) == 1
+    assert output.multimodal_outputs[0]["model_outputs"].shape == (2, 3)
 
 
 def test_sample_tokens_list_output():
@@ -59,8 +59,8 @@ def test_sample_tokens_list_output():
 
     output = GPUGenerationModelRunner.sample_tokens(runner)
 
-    assert len(output.pooler_output) == 1
-    assert output.pooler_output[0]["model_outputs"].shape == (2, 1)
+    assert len(output.multimodal_outputs) == 1
+    assert output.multimodal_outputs[0]["model_outputs"].shape == (2, 1)
 
 
 def test_sample_tokens_list_allows_none_output():
@@ -69,8 +69,8 @@ def test_sample_tokens_list_allows_none_output():
 
     output = GPUGenerationModelRunner.sample_tokens(runner)
 
-    assert len(output.pooler_output) == 1
-    assert output.pooler_output[0]["model_outputs"] is None
+    assert len(output.multimodal_outputs) == 1
+    assert output.multimodal_outputs[0]["model_outputs"] is None
 
 
 def test_sample_tokens_dict_output():
@@ -79,7 +79,7 @@ def test_sample_tokens_dict_output():
 
     output = GPUGenerationModelRunner.sample_tokens(runner)
 
-    assert len(output.pooler_output) == 1
-    assert "audio" in output.pooler_output[0]
-    assert "unused" not in output.pooler_output[0]
-    assert output.pooler_output[0]["audio"].shape == (1, 4)
+    assert len(output.multimodal_outputs) == 1
+    assert "audio" in output.multimodal_outputs[0]
+    assert "unused" not in output.multimodal_outputs[0]
+    assert output.multimodal_outputs[0]["audio"].shape == (1, 4)

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -288,6 +288,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
         pooler_outputs = model_runner_output.pooler_output
+        mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
         cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
@@ -389,6 +390,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             new_logprobs = None
             new_token_ids = generated_token_ids
             pooler_output = pooler_outputs[req_index] if pooler_outputs else None
+            mm_output = mm_outputs[req_index] if mm_outputs else None
             kv_transfer_params = None
             status_before_stop = request.status
             finish_reason = None
@@ -463,7 +465,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
+            if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
                     OmniEngineCoreOutput(
@@ -473,6 +475,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                         new_logprobs=new_logprobs,
                         new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                         pooling_output=pooler_output,
+                        multimodal_output=mm_output,
                         stop_reason=request.stop_reason,
                         events=request.take_events(),
                         prefill_stats=request.take_prefill_stats(),
@@ -485,7 +488,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     )
                 )
                 if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.save_async(pooler_output, request, is_segment_finished)
+                    self.chunk_transfer_adapter.save_async(mm_output, request, is_segment_finished)
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -420,6 +420,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
         pooler_outputs = model_runner_output.pooler_output
+        mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
 
@@ -498,6 +499,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             new_token_ids = generated_token_ids
             kv_transfer_params = None
             pooler_output = pooler_outputs[req_index] if pooler_outputs else None
+            mm_output = mm_outputs[req_index] if mm_outputs else None
             status_before_stop = request.status
             finish_reason = None
             routed_experts = None
@@ -558,7 +560,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
+            if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
+                # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
                     OmniEngineCoreOutput(
                         request_id=req_id,
@@ -567,6 +570,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                         new_logprobs=new_logprobs,
                         new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                         pooling_output=pooler_output,
+                        multimodal_output=mm_output,
                         stop_reason=request.stop_reason,
                         events=request.take_events(),
                         prefill_stats=request.take_prefill_stats(),

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -3,7 +3,7 @@
 
 import importlib
 from collections import defaultdict, deque
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import Any
 
 import torch
@@ -146,7 +146,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
     def save_async(
         self,
-        pooling_output: torch.Tensor | None = None,
+        multimodal_output: dict[str, Any] | None = None,
         request: Request | None = None,
         is_segment_finished: bool = False,
     ):
@@ -163,7 +163,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         as ``request.is_finished()``.
 
         Args:
-            pooling_output: Partial pooling output dictionary
+            multimodal_output: Per-request multimodal output dictionary
             request: Request object
             is_segment_finished: whether the segment of request is finished
         """
@@ -183,7 +183,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens
         task = {
-            "pooling_output": pooling_output,
+            "multimodal_output": multimodal_output,
             "request": request,
             "is_finished": is_finished,
             "is_segment_finished": is_segment_finished,
@@ -292,8 +292,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         return False
 
     def _send_single_request(self, task: dict):
-        raw_po = task["pooling_output"]
-        pooling_output = unflatten_payload(raw_po) if isinstance(raw_po, dict) else raw_po
+        raw_mm = task["multimodal_output"]
+        multimodal_output = unflatten_payload(raw_mm) if isinstance(raw_mm, Mapping) else raw_mm
         request = task["request"]
         is_finished = task["is_finished"]
         is_segment_finished = task["is_segment_finished"]
@@ -308,7 +308,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             try:
                 payload_data = self.custom_process_next_stage_input_func(
                     transfer_manager=self,
-                    pooling_output=pooling_output,
+                    multimodal_output=multimodal_output,
                     request=request,
                     # Existing processors use is_finished as a flush signal.
                     is_finished=is_segment_finished,

vllm_omni/distributed/omni_connectors/utils/serialization.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Mapping
 from dataclasses import asdict, is_dataclass
 from typing import Any
 
@@ -150,10 +151,13 @@ class OmniMsgpackEncoder:
             "multi_modal_placeholders": getattr(obj, "multi_modal_placeholders", None),
             "kv_transfer_params": obj.kv_transfer_params,
         }
-        # Handle dynamically added multimodal_output attribute
+        # Handle multimodal_output attribute (MultimodalPayload or dict)
         mm_output = getattr(obj, "multimodal_output", None)
         if mm_output is not None:
-            result["multimodal_output"] = mm_output
+            if isinstance(mm_output, Mapping):
+                result["multimodal_output"] = dict(mm_output)
+            else:
+                result["multimodal_output"] = mm_output
         return result
 
     def _encode_completion_output(self, obj: CompletionOutput) -> dict[str, Any]:
@@ -161,7 +165,11 @@ class OmniMsgpackEncoder:
         result = asdict(obj)
         mm_output = getattr(obj, "multimodal_output", None)
         if mm_output is not None:
-            result["multimodal_output"] = mm_output
+            # Convert MultimodalPayload to plain dict for wire format
+            if isinstance(mm_output, Mapping):
+                result["multimodal_output"] = dict(mm_output)
+            else:
+                result["multimodal_output"] = mm_output
         return result
 
 

vllm_omni/engine/__init__.py

@@ -117,7 +117,10 @@ class OmniEngineCoreRequest(EngineCoreRequest):
 
 
 class OmniEngineCoreOutput(EngineCoreOutput):
-    pooling_output: dict[str, torch.Tensor] | None = None
+    # Dedicated channel for multimodal outputs (image/audio/latent).
+    # pooling_output is inherited from EngineCoreOutput as torch.Tensor | None
+    # and retains its original vLLM semantics for pooling/embedding tasks.
+    multimodal_output: dict[str, torch.Tensor] | None = None
     # Finished flag for streaming input segment
     is_segment_finished: bool | None = False
     # Streaming update prompt length

vllm_omni/engine/async_omni_engine.py

@@ -1480,7 +1480,7 @@ class AsyncOmniEngine:
             return
 
         mm_data = prompt.get("multi_modal_data")
-        if not isinstance(mm_data, dict) or not mm_data:
+        if not isinstance(mm_data, Mapping) or not mm_data:
             return
 
         from vllm.multimodal.hasher import MultiModalHasher

vllm_omni/engine/mm_outputs.py

@@ -6,6 +6,7 @@ This module defines structured types for multimodal outputs.
 
 from __future__ import annotations
 
+from collections.abc import Iterator, Mapping
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -13,10 +14,15 @@ import torch
 from vllm.outputs import CompletionOutput
 
 
-@dataclass
-class MultimodalPayload:
+@dataclass(eq=False)
+class MultimodalPayload(Mapping):
     """Structured multimodal output payload.
 
+    Implements ``collections.abc.Mapping`` so that ``isinstance(payload, dict)``
+    style checks in downstream code can be replaced with duck-typing, and
+    ``payload.get(key)``, ``payload[key]``, ``key in payload``, ``len(payload)``
+    all work seamlessly for both tensors and metadata.
+
     Attributes:
         tensors: Dictionary mapping modality/key names to their tensors.
         metadata: Optional dictionary for non-tensor metadata
@@ -35,18 +41,49 @@ class MultimodalPayload:
 
     @property
     def is_empty(self) -> bool:
-        """Return True if the payload has no tensors."""
-        return len(self.tensors) == 0
+        """Return True if the payload has no tensors and no metadata."""
+        return len(self.tensors) == 0 and len(self.metadata) == 0
 
-    def get(self, key: str) -> torch.Tensor | None:
-        """Get a tensor by key, returning None if not found."""
-        return self.tensors.get(key)
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get a value by key, searching tensors first then metadata."""
+        if key in self.tensors:
+            return self.tensors[key]
+        return self.metadata.get(key, default)
 
     def __contains__(self, key: str) -> bool:
-        return key in self.tensors
+        return key in self.tensors or key in self.metadata
+
+    def __getitem__(self, key: str) -> Any:
+        """Dict-like indexing: search tensors first, then metadata."""
+        if key in self.tensors:
+            return self.tensors[key]
+        if key in self.metadata:
+            return self.metadata[key]
+        raise KeyError(key)
 
     def __len__(self) -> int:
-        return len(self.tensors)
+        return len(self.tensors) + len(self.metadata)
+
+    def __iter__(self) -> Iterator[str]:
+        yield from self.tensors
+        yield from self.metadata
+
+    def __bool__(self) -> bool:
+        return bool(self.tensors) or bool(self.metadata)
+
+    def __eq__(self, other: object) -> bool:
+        """Support equality with plain dicts and other Mappings."""
+        if isinstance(other, MultimodalPayload):
+            return self.tensors == other.tensors and self.metadata == other.metadata
+        if isinstance(other, Mapping):
+            return dict(self) == dict(other)
+        return NotImplemented
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert back to a plain dict (tensors + metadata merged)."""
+        result: dict[str, Any] = dict(self.tensors)
+        result.update(self.metadata)
+        return result
 
     @classmethod
     def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None:

vllm_omni/engine/output_modality.py

@@ -21,6 +21,8 @@ _MODALITY_ALIASES: dict[str, str] = {
     "waveform": "audio",
     "pixel_values": "image",
     "pixels": "image",
+    "token_ids": "text",
+    "tokens": "text",
 }
 
 

vllm_omni/engine/output_processor.py

@@ -1,8 +1,9 @@
+from dataclasses import fields as dataclass_fields
 from typing import Any
 
 import torch
 from vllm.logger import init_logger
-from vllm.outputs import PoolingRequestOutput
+from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput
 from vllm.sampling_params import RequestOutputKind
 from vllm.tokenizers import TokenizerLike
 from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
@@ -16,12 +17,98 @@ from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.metrics.stats import IterationStats, RequestStateStats
 
 from vllm_omni.data_entry_keys import unflatten_payload
-from vllm_omni.engine.output_modality import DRAINABLE_MODALITIES
+from vllm_omni.engine.mm_outputs import MultimodalCompletionOutput, MultimodalPayload
+from vllm_omni.engine.output_modality import (
+    DRAINABLE_MODALITIES,
+    OutputModality,
+    TensorAccumulationStrategy,
+    get_accumulation_strategy,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
 
 
+# ---------------------------------------------------------------------------
+# Module-level helpers
+# ---------------------------------------------------------------------------
+
+
+def _to_cpu(x: Any) -> Any:
+    """Move a tensor to CPU; pass through non-tensors unchanged."""
+    if isinstance(x, torch.Tensor):
+        try:
+            return x.detach().to("cpu", non_blocking=True).contiguous()
+        except (RuntimeError, AttributeError):
+            return x
+    if isinstance(x, dict):
+        return {str(sk): _to_cpu(sv) for sk, sv in x.items()}
+    return x
+
+
+def _merge_payload(target: MultimodalPayload, incoming: MultimodalPayload) -> None:
+    """Merge *incoming* into *target* using list-based deferred concatenation."""
+    for k, v in incoming.tensors.items():
+        if k not in target.tensors:
+            target.tensors[k] = v
+        else:
+            existing = target.tensors[k]
+            if isinstance(existing, list):
+                existing.append(v)
+            elif isinstance(existing, torch.Tensor) and isinstance(v, torch.Tensor):
+                target.tensors[k] = [existing, v]
+            else:
+                target.tensors[k] = v
+
+    for k, v in incoming.metadata.items():
+        if k not in target.metadata:
+            target.metadata[k] = v
+        else:
+            existing = target.metadata[k]
+            if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
+                target.metadata[k] = [existing, v]
+            elif isinstance(v, torch.Tensor) and isinstance(existing, list):
+                existing.append(v)
+            else:
+                # Metadata: replace with latest value
+                target.metadata[k] = v
+
+
+def _cat_tensors(
+    tensors: list[torch.Tensor],
+    strategy: TensorAccumulationStrategy,
+) -> torch.Tensor:
+    """Concatenate a list of tensors according to *strategy*."""
+    if strategy == TensorAccumulationStrategy.CONCAT_LAST:
+        return torch.cat(tensors, dim=-1)
+    if strategy == TensorAccumulationStrategy.REPLACE:
+        return tensors[-1]
+    # CONCAT_DIM0 / APPEND_LIST / default
+    return torch.cat(tensors, dim=0)
+
+
+def _modality_to_type_string(modality: OutputModality) -> str:
+    """Convert an OutputModality flag to a lowercase type string."""
+    try:
+        if OutputModality.AUDIO in modality:
+            return "audio"
+        if OutputModality.IMAGE in modality:
+            return "image"
+        if OutputModality.LATENT in modality:
+            return "latent"
+    except TypeError:
+        # Flag identity mismatch (e.g. after module reload in tests).
+        name = getattr(modality, "name", "") or ""
+        lowered = name.lower()
+        if "audio" in lowered:
+            return "audio"
+        if "image" in lowered:
+            return "image"
+        if "latent" in lowered:
+            return "latent"
+    return "text"
+
+
 class OmniRequestState(RequestState):
     """Request state for omni models, tracking multimodal outputs.
 
@@ -41,10 +128,12 @@ class OmniRequestState(RequestState):
         super().__init__(*args, **kwargs)
         self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0))
         # Omni-specific: multimodal output accumulation
-        # NOTE: Keys in mm_accumulated matter, because they dictate which
-        # outputs are drained (i.e., only drain modality keys, don't drain
-        # hidden states).
-        self.mm_accumulated: dict[str, Any] = {}
+        # TODO: mm_type is per-request, not per-key. If a model ever produces
+        # both audio and latent outputs, the modality type would flip on each
+        # add_multimodal_tensor() call. Consider tracking per-key modality
+        # types (e.g. dict[str, str]) for future multi-output models.
+        self.mm_type: str | None = None
+        self.mm_accumulated: MultimodalPayload = MultimodalPayload()
 
     def apply_streaming_update(self, update) -> None:
         super().apply_streaming_update(update)
@@ -53,8 +142,6 @@ class OmniRequestState(RequestState):
     @staticmethod
     def _to_cpu(x):
         """Try to convert to CPU tensor if needed."""
-        # TODO: Make this more robust and unify with other payload
-        # building utils, we do this in multiple places.
         if isinstance(x, torch.Tensor):
             try:
                 return x.detach().to("cpu", non_blocking=True).contiguous()
@@ -63,109 +150,94 @@ class OmniRequestState(RequestState):
         return x
 
     def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None:
+        """Accumulate a multimodal tensor payload into the request state.
+
+        Normalizes incoming payloads (dict or raw tensor) into a
+        MultimodalPayload and merges with any previously accumulated data.
+        Uses list-based deferred concatenation to avoid O(n²) repeated
+        torch.cat calls.
+        """
         if payload is None:
             return
-
-        mm_type = (mm_type or "").lower()
         try:
+            if mm_type:
+                self.mm_type = (mm_type or "").lower()
+
+            target_key = self.mm_type or "hidden"
+
+            # Build incoming MultimodalPayload
             if isinstance(payload, dict):
-                # Keep payload flat (dotted keys like "hidden_states.layer_0")
-                # during accumulation so that all values are tensors/scalars and
-                # the merge logic below works correctly.  Unflatten happens
-                # later in _consolidate_multimodal_tensors after concatenation.
-
-                incoming: dict[str, Any] = {}
-                # TODO (Alex): Clean up and simplify key management
-                target_key = mm_type or "hidden"
-
+                remapped: dict[str, Any] = {}
                 for k, v in payload.items():
+                    # Normalize producer keys to the modality name.
+                    # AR runners produce {"hidden": ...} and generation
+                    # runners produce {"model_outputs": ...}; remap both
+                    # to the semantic modality key (e.g. "audio", "latent").
                     if k == "model_outputs":
                         k = target_key
                     elif k == "hidden" and target_key != "hidden":
                         k = target_key
-
-                    incoming[k] = self._to_cpu(v)
+                    remapped[k] = _to_cpu(v)
+                incoming = MultimodalPayload.from_dict(remapped)
+            elif isinstance(payload, MultimodalPayload):
+                incoming = payload
             else:
-                key = mm_type or "hidden"
-                incoming = {key: self._to_cpu(payload)}
+                incoming = MultimodalPayload.from_dict({target_key: _to_cpu(payload)})
 
-            if not self.mm_accumulated:
+            if incoming is None:
+                return
+
+            if self.mm_accumulated.is_empty:
                 self.mm_accumulated = incoming
             else:
-                for k, v in incoming.items():
-                    if k not in self.mm_accumulated:
-                        self.mm_accumulated[k] = v
-                    else:
-                        existing = self.mm_accumulated[k]
-                        if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
-                            self.mm_accumulated[k] = [existing, v]
-                        elif isinstance(v, torch.Tensor) and isinstance(existing, list):
-                            existing.append(v)
-                        elif isinstance(v, dict) and isinstance(existing, dict):
-                            for sk, sv in v.items():
-                                if sk not in existing:
-                                    existing[sk] = sv
-                                elif isinstance(sv, torch.Tensor) and isinstance(existing[sk], torch.Tensor):
-                                    existing[sk] = [existing[sk], sv]
-                                elif isinstance(sv, torch.Tensor) and isinstance(existing[sk], list):
-                                    existing[sk].append(sv)
-                                else:
-                                    existing[sk] = sv
-                        else:
-                            self.mm_accumulated[k] = v
-        except Exception:
-            # Log and continue without crashing the output pipeline
+                _merge_payload(self.mm_accumulated, incoming)
+        except (ValueError, TypeError, RuntimeError):
             logger.exception("Error accumulating multimodal tensor")
 
     def _consolidate_multimodal_tensors(self) -> None:
-        """Consolidate accumulated tensor lists into single tensors via concatenation.
+        """Consolidate accumulated tensor lists into single tensors.
 
-        Only DELTA drains modality keys per-step, so they will never be lists here and
-        can be skipped.  For CUMULATIVE and FINAL_ONLY, modality keys accumulate across
-        steps and need consolidation.
+        Uses TensorAccumulationStrategy derived from the output modality
+        to determine concatenation behavior. Metadata values always use
+        REPLACE (keep latest).
         """
-        if not self.mm_accumulated:
+        if self.mm_accumulated.is_empty:
             return
 
-        skip_modality = self.output_kind == RequestOutputKind.DELTA
         try:
-            for k, v in self.mm_accumulated.items():
-                if skip_modality and k in DRAINABLE_MODALITIES:
-                    continue
+            modality = OutputModality.from_string(self.mm_type)
+        except (ValueError, KeyError):
+            modality = OutputModality.TEXT
+        strategy = get_accumulation_strategy(modality)
+
+        try:
+            for k, v in list(self.mm_accumulated.tensors.items()):
                 if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
                     try:
+                        self.mm_accumulated.tensors[k] = _cat_tensors(v, strategy)
+                    except RuntimeError:
                         if k == "audio":
-                            # Preserve channel dimension when chunks are compatible;
-                            # fall back to 1-D waveform concatenation for uneven chunks.
                             try:
-                                self.mm_accumulated[k] = torch.cat(v, dim=-1)
+                                self.mm_accumulated.tensors[k] = torch.cat(v, dim=-1)
                             except RuntimeError:
-                                self.mm_accumulated[k] = torch.cat([t.reshape(-1) for t in v], dim=0)
-                        elif k == "sr":
-                            # Sample rate is a constant scalar, keep last value.
-                            self.mm_accumulated[k] = v[-1]
+                                self.mm_accumulated.tensors[k] = torch.cat([t.reshape(-1) for t in v], dim=0)
                         else:
-                            self.mm_accumulated[k] = torch.cat(v, dim=0)
-                    except Exception:
-                        # Keep last tensor on failure
-                        logger.warning(f"Error concatenating tensor for key {k}; keeping last tensor")
-                        self.mm_accumulated[k] = v[-1]
-                elif isinstance(v, dict):
-                    for sk, sv in v.items():
-                        if isinstance(sv, list) and sv and isinstance(sv[0], torch.Tensor):
-                            try:
-                                v[sk] = torch.cat(sv, dim=0)
-                            except Exception:
-                                v[sk] = sv[-1]
-        except Exception:
-            logger.exception("Error consolidating multimodal tensors")
+                            logger.warning(
+                                "Error concatenating tensor for key %s; keeping last tensor",
+                                k,
+                            )
+                            self.mm_accumulated.tensors[k] = v[-1]
 
-        # Restore nested structure from flat dotted keys now that all tensor
-        # lists have been concatenated into single tensors.
-        try:
-            self.mm_accumulated = unflatten_payload(self.mm_accumulated)
-        except Exception:
-            logger.exception("Error unflattening consolidated multimodal tensors")
+            # Metadata: consolidate any deferred tensor lists (REPLACE semantics)
+            for k, v in list(self.mm_accumulated.metadata.items()):
+                if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
+                    self.mm_accumulated.metadata[k] = v[-1]
+                elif isinstance(v, dict):
+                    for sk, sv in list(v.items()):
+                        if isinstance(sv, list) and sv and isinstance(sv[0], torch.Tensor):
+                            v[sk] = sv[-1]
+        except (RuntimeError, TypeError, KeyError):
+            logger.exception("Error consolidating multimodal tensors")
 
     # Override: do not route to pooling-only path; always create completion
     # outputs, and attach pooling_result into the CompletionOutput.
@@ -205,18 +277,18 @@ class OmniRequestState(RequestState):
             )
 
         finished = finish_reason is not None
-        is_final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
+        final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
         is_delta = self.output_kind == RequestOutputKind.DELTA
 
-        if not finished and is_final_only:
+        if not finished and final_only:
             return None
 
+        # Consolidate accumulated tensors when finishing, or for every
+        # CUMULATIVE step so the consumer always sees a single tensor.
         if finished or not is_delta:
             self._consolidate_multimodal_tensors()
 
-        if self.stream_interval > 1:
-            assert self.detokenizer is not None
-
+        if self.stream_interval > 1 and self.detokenizer is not None:
             # Send output request only when
             # 1. It has finished, or
             # 2. It is the first token, or
@@ -228,7 +300,7 @@ class OmniRequestState(RequestState):
             ):
                 return None
 
-            if is_delta:
+            if self.output_kind == RequestOutputKind.DELTA:
                 # Send tokens from the offset in DELTA mode, otherwise all
                 # tokens are sent.
                 new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :]
@@ -258,50 +330,127 @@ class OmniRequestState(RequestState):
         token_ids: list[int],
         finish_reason: FinishReason | None,
         stop_reason: int | str | None,
-    ) -> Any:
-        # Reuse base text/logprobs logic, then annotate with pooling_result.
-        # Note: upstream _new_completion_output no longer accepts routed_experts
-        # as a parameter; it reads from ``self.routed_experts_chunks`` internally.
-        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
+    ) -> MultimodalCompletionOutput | CompletionOutput:
+        """Create a completion output with multimodal data attached.
 
-        # Inter-stage processors need the full cumulative token sequence.
-        # In DELTA mode, base_output.token_ids only has the latest step's
-        # tokens, so we always store a cumulative copy here.
-        base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids)
+        Returns a MultimodalCompletionOutput when multimodal data has been
+        accumulated, otherwise returns the base CompletionOutput.  Snapshots
+        the accumulated payload before draining DELTA-mode modality keys so
+        that callers receive an immutable view.
+        """
+        # When there is no detokenizer (generation stages), build a minimal
+        # CompletionOutput directly since upstream asserts detokenizer != None.
+        if self.detokenizer is None:
+            finished = finish_reason is not None
+            base_output = CompletionOutput(
+                index=self.request_index,
+                text="",
+                token_ids=token_ids,
+                logprobs=None,
+                cumulative_logprob=None,
+                finish_reason=str(finish_reason) if finished else None,
+                stop_reason=stop_reason if finished else None,
+            )
+        else:
+            base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
+
+        # Always provide cumulative token IDs for inter-stage processors.
+        if self.detokenizer is not None:
+            base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids)
+        else:
+            base_output.cumulative_token_ids = list(token_ids)
 
         # Attach cumulative_text only at the final step for inter-stage use.
         if finish_reason is not None and hasattr(self.detokenizer, "output_text"):
             base_output.cumulative_text = self.detokenizer.output_text
 
-        if not hasattr(base_output, "multimodal_output"):
-            setattr(base_output, "multimodal_output", {})
-        if self.mm_accumulated:
-            mm_out = getattr(base_output, "multimodal_output")
-            if isinstance(mm_out, dict):
-                for k, v in self.mm_accumulated.items():
-                    mm_out[k] = v
-            else:
-                setattr(base_output, "multimodal_output", self.mm_accumulated)
+        try:
+            if self.mm_accumulated and not self.mm_accumulated.is_empty:
+                # Snapshot: copy current tensors/metadata so drain doesn't
+                # mutate the payload already handed to the caller.
+                # Unflatten dotted keys (e.g. "hidden_states.layer_0") back
+                # to nested dicts so downstream consumers (thinker2talker etc.)
+                # can access with .get("hidden_states", {}).get("layers", {}).
+                merged = {**self.mm_accumulated.tensors, **self.mm_accumulated.metadata}
+                nested = unflatten_payload(merged)
+                snapshot = MultimodalPayload.from_dict(nested) or MultimodalPayload()
+                kwargs = {f.name: getattr(base_output, f.name) for f in dataclass_fields(CompletionOutput)}
+                output = MultimodalCompletionOutput(
+                    multimodal_output=snapshot,
+                    **kwargs,
+                )
+                output.cumulative_token_ids = base_output.cumulative_token_ids
+                if hasattr(base_output, "cumulative_text"):
+                    output.cumulative_text = base_output.cumulative_text
 
-        if self.output_kind == RequestOutputKind.DELTA:
-            for modality_key in DRAINABLE_MODALITIES:
-                self.mm_accumulated.pop(modality_key, None)
+                # DELTA mode: drain modality keys (e.g. audio) so the next
+                # step only sees freshly accumulated data for those keys.
+                if self.output_kind == RequestOutputKind.DELTA:
+                    for modality_key in DRAINABLE_MODALITIES:
+                        self.mm_accumulated.tensors.pop(modality_key, None)
 
+                return output
+        except (RuntimeError, TypeError, AttributeError):
+            logger.exception("Error creating MultimodalCompletionOutput")
         return base_output
 
+    def _new_request_output(
+        self,
+        external_req_id: str,
+        outputs: list,
+        finished: bool,
+        kv_transfer_params: dict[str, Any] | None = None,
+    ) -> RequestOutput | PoolingRequestOutput:
+        """Create request output, handling no-detokenizer generation stages.
+
+        Upstream asserts ``self.logprobs_processor is not None`` which fails
+        for generation stages (talker, code2wav) that have no tokenizer.
+        When the logprobs processor is absent we build the RequestOutput
+        directly with ``prompt_logprobs=None``.
+        """
+        if self.logprobs_processor is not None:
+            return super()._new_request_output(
+                external_req_id,
+                outputs,
+                finished,
+                kv_transfer_params,
+            )
+
+        # No-detokenizer path: build RequestOutput directly.
+        prompt_token_ids = self.prompt_token_ids
+        if prompt_token_ids is None and self.prompt_embeds is not None:
+            prompt_token_ids = [0] * len(self.prompt_embeds)
+        if prompt_token_ids is None:
+            prompt_token_ids = []
+
+        return RequestOutput(
+            request_id=external_req_id,
+            lora_request=self.lora_request,
+            prompt=self.prompt,
+            prompt_token_ids=prompt_token_ids,
+            prompt_logprobs=None,
+            outputs=outputs,
+            finished=finished,
+            kv_transfer_params=kv_transfer_params,
+            num_cached_tokens=self.num_cached_tokens,
+            metrics=self.stats,
+        )
+
 
 class MultimodalOutputProcessor(VLLMOutputProcessor):
-    """Handles multimodal output processing by capturing pooling_output
-    from EngineCoreOutput and accumulating it as multimodal tensors,
-    before delegating to the base vLLM OutputProcessor for text handling.
+    """Handles multimodal output processing.
 
-    The actual data flow is:
-    1. For each EngineCoreOutput with pooling_output and a detokenizer:
-       - Capture pooling_output into OmniRequestState.add_multimodal_tensor()
-       - Clear eco.pooling_output to force text path in base processor
+    Captures multimodal outputs from OmniEngineCoreOutput and accumulates
+    them as MultimodalPayload in OmniRequestState, before delegating to
+    the base vLLM OutputProcessor for text handling.
+
+    The data flow is:
+    1. For each EngineCoreOutput with multimodal_output:
+       - Capture into OmniRequestState.add_multimodal_tensor()
     2. Base vLLM OutputProcessor handles text detokenization
-    3. On finish, _consolidate_multimodal_tensors() concatenates accumulated tensors
-    4. _new_completion_output() attaches mm_accumulated to CompletionOutput
+    3. On finish, _consolidate_multimodal_tensors() concatenates accumulated
+       tensors using strategy-based dispatch
+    4. _new_completion_output() returns MultimodalCompletionOutput
     """
 
     def __init__(
@@ -312,6 +461,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         stream_interval: int = 1,
         tracing_enabled: bool = False,
         engine_core_output_type: str | None = None,
+        output_modality: OutputModality = OutputModality.TEXT,
     ):
         """Initialize the multimodal output processor.
 
@@ -319,9 +469,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             tokenizer: Tokenizer for detokenizing text outputs
             log_stats: Whether to log statistics
             stream_interval: Stream interval for output generation
-            engine_core_output_type: Optional output type specification
-                (e.g., "image", "audio", "latent"). Used to tag multimodal
-                outputs with the correct modality key.
+            engine_core_output_type: Optional output type string (e.g.,
+                "image", "audio", "latent"). Converted to OutputModality
+                internally. Kept for backward compatibility with
+                stage_init_utils.
+            output_modality: Type-safe output modality flag. Used to tag
+                multimodal outputs with the correct modality key when
+                per-output type info is unavailable.
         """
         super().__init__(
             tokenizer=tokenizer,
@@ -329,6 +483,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             stream_interval=stream_interval,
             tracing_enabled=tracing_enabled,
         )
+        # Convert string-based engine_core_output_type to OutputModality
+        if engine_core_output_type is not None:
+            self.output_modality = OutputModality.from_string(engine_core_output_type)
+        else:
+            self.output_modality = output_modality
         self.engine_core_output_type = engine_core_output_type
         self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {}
 
@@ -426,21 +585,78 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         engine_core_timestamp: float | None = None,
         iteration_stats: IterationStats | None = None,
     ) -> OutputProcessorOutput:
+        default_mm_type = _modality_to_type_string(self.output_modality)
+
+        # Separate outputs that upstream can handle (has detokenizer or
+        # pooling) from multimodal-only outputs (no detokenizer, no pooling)
+        # that would trigger upstream's `assert detokenizer is not None`.
+        upstream_outputs: list[EngineCoreOutput] = []
+        mm_only_outputs: list[EngineCoreOutput] = []
+
+        for eco in engine_core_outputs:
+            req_state = self.request_states.get(eco.request_id)
+            if req_state is None:
+                continue
+
+            # Accumulate multimodal tensors regardless of path.
+            if isinstance(req_state, OmniRequestState):
+                mm_output = getattr(eco, "multimodal_output", None)
+                if mm_output is not None:
+                    mm_type = getattr(eco, "output_type", None) or default_mm_type
+                    req_state.add_multimodal_tensor(mm_output, mm_type)
+
+            # Route: if no detokenizer and no pooling output, handle locally
+            # to avoid upstream's assert on detokenizer.
+            if req_state.detokenizer is None and eco.pooling_output is None:
+                mm_only_outputs.append(eco)
+            else:
+                upstream_outputs.append(eco)
+
+        # Handle multimodal-only outputs (generation stages) locally.
+        self._process_mm_only_outputs(mm_only_outputs)
+
+        # Delegate text/pooling outputs to upstream.
+        return super().process_outputs(
+            upstream_outputs,
+            engine_core_timestamp=engine_core_timestamp,
+            iteration_stats=iteration_stats,
+        )
+
+    def _process_mm_only_outputs(
+        self,
+        engine_core_outputs: list[EngineCoreOutput],
+    ) -> None:
+        """Handle outputs from generation stages that have no detokenizer.
+
+        These cannot go through upstream process_outputs because it asserts
+        detokenizer is not None when pooling_output is None.
+        """
         for eco in engine_core_outputs:
             req_state = self.request_states.get(eco.request_id)
             if req_state is None or not isinstance(req_state, OmniRequestState):
                 continue
-            if eco.pooling_output is not None and req_state.detokenizer is not None:
-                mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
-                req_state.add_multimodal_tensor(eco.pooling_output, mm_type)
-                # Force text path in base processor for multimodal outputs.
-                eco.pooling_output = None
 
-        return super().process_outputs(
-            engine_core_outputs,
-            engine_core_timestamp=engine_core_timestamp,
-            iteration_stats=iteration_stats,
-        )
+            new_token_ids = eco.new_token_ids
+            finish_reason = eco.finish_reason
+            stop_reason = eco.stop_reason
+            kv_transfer_params = eco.kv_transfer_params
+            routed_experts = eco.routed_experts
+            req_state.num_cached_tokens = eco.num_cached_tokens
+            req_state.is_prefilling = False
+
+            if request_output := req_state.make_request_output(
+                new_token_ids,
+                None,  # pooling_output
+                finish_reason,
+                stop_reason,
+                kv_transfer_params,
+                routed_experts,
+            ):
+                if req_state.queue is not None:
+                    req_state.queue.put(request_output)
+
+            if finish_reason is not None:
+                self._finish_request(req_state)
 
     def _update_stats_from_output(
         self,

vllm_omni/engine/stage_engine_core_client.py

@@ -18,6 +18,8 @@ from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
 
 import psutil
+import vllm.v1.engine as _vllm_engine_module
+import vllm.v1.engine.core_client as _vllm_core_client_module
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
@@ -30,6 +32,7 @@ from vllm_omni.distributed.omni_connectors.utils.initialization import (
     KV_TRANSFER_PORT_OFFSET,
 )
 from vllm_omni.distributed.omni_connectors.utils.kv_utils import kv_zmq_port
+from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
 from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata
 
@@ -174,6 +177,21 @@ class StageEngineCoreClientBase(StageClientBase):
             self.stage_id,
             self.replica_id,
         )
+
+        # Patch the output decoder type so the client decodes
+        # OmniEngineCoreOutputs (which carries multimodal_output per
+        # EngineCoreOutput) instead of the base EngineCoreOutputs.
+        # Must happen BEFORE super().__init__() which creates the decoder.
+        # TODO: Add a defensive assertion after super().__init__() to verify
+        # the decoder uses OmniEngineCoreOutputs, catching import-order regressions.
+        _vllm_engine_module.EngineCoreOutput = OmniEngineCoreOutput
+        _vllm_engine_module.EngineCoreOutputs = OmniEngineCoreOutputs
+        _vllm_core_client_module.EngineCoreOutputs = OmniEngineCoreOutputs
+        logger.debug(
+            "[StageEngineCoreClient] Patched EngineCoreOutputs -> %s",
+            _vllm_core_client_module.EngineCoreOutputs,
+        )
+
         try:
             super().__init__(
                 vllm_config,

vllm_omni/engine/stage_engine_core_proc.py

@@ -14,6 +14,7 @@ from multiprocessing.process import BaseProcess
 from typing import TYPE_CHECKING, Any
 
 import msgspec
+import vllm.v1.engine.core as _vllm_engine_core_module
 import zmq
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import (
@@ -36,6 +37,7 @@ from vllm.v1.engine.utils import (
 from vllm.v1.utils import shutdown
 
 from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
+from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.stage_init_utils import set_death_signal
 
 if TYPE_CHECKING:
@@ -122,6 +124,18 @@ class StageEngineCoreProc(EngineCoreProc):
             os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
             os.environ["VLLM_OMNI_REPLICA_ID"] = str(max(int(omni_replica_id), 0))
 
+            # Patch the decoder type so process_input_sockets (started
+            # during __init__) decodes OmniEngineCoreRequest (which
+            # carries additional_information) instead of the base
+            # EngineCoreRequest.  Must happen BEFORE __init__ because
+            # the IO thread creates MsgpackDecoder(EngineCoreRequest)
+            # during __init__.
+            _vllm_engine_core_module.EngineCoreRequest = OmniEngineCoreRequest
+            logger.debug(
+                "[StageEngineCoreProc] Patched EngineCoreRequest -> OmniEngineCoreRequest: %s",
+                _vllm_engine_core_module.EngineCoreRequest,
+            )
+
             engine_core = StageEngineCoreProc(
                 *args,
                 engine_index=dp_rank,

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import asyncio
 import base64
 import json
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Mapping
 from typing import cast
 from uuid import uuid4
 
@@ -82,7 +82,10 @@ class RealtimeConnection(VllmRealtimeConnection):
 
     def _extract_audio_chunks(self, output) -> tuple[list[np.ndarray], int]:
         mm = getattr(output, "multimodal_output", None)
-        if not isinstance(mm, dict):
+        if mm is None:
+            return [], 24000
+        # Support both MultimodalPayload and plain dict
+        if not isinstance(mm, Mapping):
             return [], 24000
 
         sr = mm.get("sr") or mm.get("sample_rate") or mm.get("audio_sample_rate") or 24000

vllm_omni/entrypoints/openai/serving_video.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import copy
 import time
+from collections.abc import Mapping
 from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any, cast
@@ -620,7 +621,7 @@ class OmniOpenAIServingVideo:
     def _resolve_fps(result: Any) -> int | None:
         """Extract fps from multimodal_output if the model reported it."""
         multimodal_output = getattr(result, "multimodal_output", None)
-        if isinstance(multimodal_output, dict):
+        if isinstance(multimodal_output, Mapping):
             fps = multimodal_output.get("fps")
             if fps is not None:
                 try:
@@ -633,7 +634,7 @@ class OmniOpenAIServingVideo:
         request_output = getattr(result, "request_output", None)
         if isinstance(request_output, dict):
             mm = request_output.get("multimodal_output") or {}
-            if isinstance(mm, dict):
+            if isinstance(mm, Mapping):
                 fps = mm.get("fps")
                 if fps is not None:
                     try:
@@ -644,7 +645,7 @@ class OmniOpenAIServingVideo:
                         pass
         elif hasattr(request_output, "multimodal_output"):
             mm = getattr(request_output, "multimodal_output", None)
-            if isinstance(mm, dict):
+            if isinstance(mm, Mapping):
                 fps = mm.get("fps")
                 if fps is not None:
                     try:
@@ -659,7 +660,7 @@ class OmniOpenAIServingVideo:
     @classmethod
     def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None:
         multimodal_output = getattr(result, "multimodal_output", None)
-        if isinstance(multimodal_output, dict):
+        if isinstance(multimodal_output, Mapping):
             sample_rate = cls._coerce_audio_sample_rate(
                 multimodal_output.get("audio_sample_rate")
                 or multimodal_output.get("sample_rate")
@@ -672,7 +673,7 @@ class OmniOpenAIServingVideo:
         request_output = getattr(result, "request_output", None)
         if isinstance(request_output, dict):
             multimodal_output = request_output.get("multimodal_output") or {}
-            if isinstance(multimodal_output, dict):
+            if isinstance(multimodal_output, Mapping):
                 return cls._coerce_audio_sample_rate(
                     multimodal_output.get("audio_sample_rate")
                     or multimodal_output.get("sample_rate")
@@ -681,7 +682,7 @@ class OmniOpenAIServingVideo:
                 )
         elif hasattr(request_output, "multimodal_output"):
             multimodal_output = getattr(request_output, "multimodal_output", None)
-            if isinstance(multimodal_output, dict):
+            if isinstance(multimodal_output, Mapping):
                 return cls._coerce_audio_sample_rate(
                     multimodal_output.get("audio_sample_rate")
                     or multimodal_output.get("sample_rate")

vllm_omni/entrypoints/openai/serving_video_stream.py

@@ -32,6 +32,7 @@ import json
 import time as _time
 import uuid
 import wave
+from collections.abc import Mapping
 from typing import Any
 
 import torch
@@ -796,7 +797,7 @@ class OmniStreamingVideoHandler:
         if not isinstance(outputs, list) or not outputs:
             return None
         mm_output = getattr(outputs[0], "multimodal_output", None)
-        if not isinstance(mm_output, dict):
+        if not isinstance(mm_output, Mapping):
             return None
         return mm_output.get("audio")
 

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from dataclasses import replace
 from functools import cached_property
 from pathlib import Path
@@ -220,7 +220,7 @@ class VoxtralTTSForConditionalGeneration(
 
     def tts_postprocess(self, hidden_states: torch.Tensor, multimodal_outputs: object, **info_dict: object | None):
         update_dict = {}
-        if isinstance(multimodal_outputs, dict) and "audio" in multimodal_outputs:
+        if isinstance(multimodal_outputs, Mapping) and "audio" in multimodal_outputs:
             assert self.post_process_idx < len(multimodal_outputs["audio"]), (
                 f"Expect {self.post_process_idx=} < {len(multimodal_outputs['audio'])=}"
             )

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections import defaultdict
+from collections.abc import Mapping
 from contextlib import nullcontext
 from typing import Any
 
@@ -162,7 +163,7 @@ def text2flow(
 
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
-    pooling_output: dict[str, Any] | None,
+    multimodal_output: dict[str, Any] | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -198,12 +199,14 @@ def talker2code2wav_async_chunk(
                     value = _to_cpu_tensor(info_embed.get(key))
                     if value is not None:
                         prompt_payload[key] = value
-                if isinstance(pooling_output, dict):
-                    po_embed = pooling_output.get("embed", {}) if isinstance(pooling_output.get("embed"), dict) else {}
+                if isinstance(multimodal_output, Mapping):
+                    mm_embed = multimodal_output.get("embed", {})
+                    if not isinstance(mm_embed, Mapping):
+                        mm_embed = multimodal_output
                     for key in cond_keys:
                         if key in prompt_payload:
                             continue
-                        value = _to_cpu_tensor(po_embed.get(key))
+                        value = _to_cpu_tensor(mm_embed.get(key))
                         if value is not None:
                             prompt_payload[key] = value
                 # Drop any right-padding carried from batched talker emission so

vllm_omni/model_executor/stage_input_processors/fish_speech.py

@@ -1,5 +1,6 @@
 """Stage input processor for Fish Speech S2 Pro: Slow AR → DAC Decoder."""
 
+from collections.abc import Mapping
 from typing import Any
 
 import torch
@@ -91,14 +92,14 @@ def _select_backlog_chunk_size(
     return backlog_chunk_size
 
 
-def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
-    """Extract the last frame of audio codes from the pooling output."""
-    audio_codes = pooling_output.get("audio_codes")
+def _extract_last_frame(multimodal_output: dict[str, Any]) -> torch.Tensor | None:
+    """Extract the last frame of audio codes from the multimodal output."""
+    audio_codes = multimodal_output.get("audio_codes")
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
         return None
     if audio_codes.ndim == 2:
         frame = audio_codes[-1]
-        valid = pooling_output.get("audio_code_valid")
+        valid = multimodal_output.get("audio_code_valid")
         if isinstance(valid, torch.Tensor) and valid.numel() > 0:
             is_valid = bool(valid.reshape(-1)[-1].item())
         elif valid is not None:
@@ -146,7 +147,7 @@ def slow_ar_to_dac_decoder(
 
 def slow_ar_to_dac_decoder_async_chunk(
     transfer_manager: Any,
-    pooling_output: dict[str, Any] | None,
+    multimodal_output: dict[str, Any] | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -160,8 +161,8 @@ def slow_ar_to_dac_decoder_async_chunk(
     finished = bool(is_finished or request.is_finished())
     cfg = _get_connector_extra(transfer_manager)
 
-    if isinstance(pooling_output, dict):
-        frame = _extract_last_frame(pooling_output)
+    if isinstance(multimodal_output, Mapping):
+        frame = _extract_last_frame(multimodal_output)
         if frame is not None:
             transfer_manager.code_prompt_token_ids[request_id].append(frame.detach())
     elif not finished:

vllm_omni/model_executor/stage_input_processors/glm_image.py

@@ -4,6 +4,7 @@
 
 import math
 import time
+from collections.abc import Mapping
 from typing import Any
 
 import torch
@@ -23,7 +24,7 @@ def _has_source_image(mm_data: Any) -> bool:
     - `img2img`: legacy single-image key
     - `images`: list or single image
     """
-    if not isinstance(mm_data, dict):
+    if not isinstance(mm_data, Mapping):
         return False
     if mm_data.get("image") is not None:
         return True
@@ -35,7 +36,7 @@ def _has_source_image(mm_data: Any) -> bool:
 
 def _first_source_image(mm_data: Any) -> Any:
     """Get first source image from normalized multimodal keys."""
-    if not isinstance(mm_data, dict):
+    if not isinstance(mm_data, Mapping):
         return None
 
     image = mm_data.get("image")
@@ -285,7 +286,7 @@ def ar2diffusion(
 
         if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
             mm_output = ar_output.multimodal_output
-            if isinstance(mm_output, dict) and mm_output.get("ids", {}).get("prior_image") is not None:
+            if isinstance(mm_output, Mapping) and mm_output.get("ids", {}).get("prior_image") is not None:
                 is_i2i = True
         _dt_mode = (time.perf_counter() - _t_mode) * 1000
 
@@ -322,7 +323,7 @@ def ar2diffusion(
         # Check ar_output (RequestOutput) for multimodal_output - this is the correct location
         if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
             mm_output = ar_output.multimodal_output
-            if isinstance(mm_output, dict):
+            if isinstance(mm_output, Mapping):
                 raw_prior_image_ids = mm_output.get("ids", {}).get("prior_image")
                 if raw_prior_image_ids is not None:
                     # Handle different formats:
@@ -348,7 +349,7 @@ def ar2diffusion(
             if hasattr(output, "multimodal_output") and output.multimodal_output:
                 mm_output = output.multimodal_output
                 logger.debug(f"[ar2diffusion] Request {i}: found multimodal_output on CompletionOutput (fallback)")
-                if isinstance(mm_output, dict):
+                if isinstance(mm_output, Mapping):
                     raw_prior_image_ids = mm_output.get("ids", {}).get("prior_image")
                     if raw_prior_image_ids is not None:
                         if isinstance(raw_prior_image_ids, torch.Tensor):

vllm_omni/model_executor/stage_input_processors/glm_tts.py

@@ -188,7 +188,7 @@ def _to_cpu_tensor(value: Any) -> torch.Tensor | None:
 
 def ar_to_dit_async_chunk(
     transfer_manager: Any,
-    pooling_output: dict[str, Any] | None,
+    multimodal_output: dict[str, Any] | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -216,6 +216,7 @@ def ar_to_dit_async_chunk(
     if request_id is None:
         raise ValueError("GLM-TTS async chunk request is missing request id")
     finished = bool(is_finished or request.is_finished())
+    pooling_output = multimodal_output
 
     # Read connector chunk config (supports progressive list or single int)
     connector = getattr(transfer_manager, "connector", None)

vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py

@@ -176,7 +176,7 @@ def _extract_last_frame(pooling_output: OmniPayload) -> torch.Tensor | None:
 
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
-    pooling_output: OmniPayload | None,
+    multimodal_output: OmniPayload | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -188,6 +188,7 @@ def talker2code2wav_async_chunk(
     """
     request_id = request.external_req_id
     finished = bool(is_finished or request.is_finished())
+    pooling_output = multimodal_output
 
     if isinstance(pooling_output, dict):
         frame = _extract_last_frame(pooling_output)

vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py

@@ -214,7 +214,7 @@ def _extract_last_step_row(pooling_output: OmniPayload) -> torch.Tensor | None:
 
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
-    pooling_output: OmniPayload | None,
+    multimodal_output: OmniPayload | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -244,6 +244,7 @@ def talker2code2wav_async_chunk(
     """
     request_id = request.external_req_id
     finished = bool(is_finished or request.is_finished())
+    pooling_output = multimodal_output
 
     # Per-request rolling state. ``code_prompt_token_ids`` is the
     # framework-provided buffer; ``higgs_v3_emitted_frames`` tracks how many

vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py

@@ -12,6 +12,7 @@ signature pattern as glm_image.ar2diffusion.
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from functools import lru_cache
 from typing import Any
 
@@ -208,7 +209,7 @@ def ar2diffusion(
         # Forward multimodal output from AR (if any)
         if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
             mm_output = ar_output.multimodal_output
-            if isinstance(mm_output, dict):
+            if isinstance(mm_output, Mapping):
                 diffusion_input["extra"]["ar_multimodal_output"] = mm_output
 
         # Forward sampling params

vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py

@@ -1,5 +1,6 @@
 """Stage input processor for MammothModa2 (AR -> DiT)."""
 
+from collections.abc import Mapping
 from typing import Any
 
 import torch
@@ -35,7 +36,7 @@ def ar2dit(
         full_token_ids = prompt_token_ids + gen_token_ids
 
         mm_output = getattr(completion_output, "multimodal_output", None)
-        if not isinstance(mm_output, dict) or "latent" not in mm_output:
+        if not isinstance(mm_output, Mapping) or "latent" not in mm_output:
             raise ValueError(
                 "AR stage output missing latent multimodal output. "
                 f"request_id={getattr(ar_output, 'request_id', None)}, "

vllm_omni/model_executor/stage_input_processors/mimo_audio.py

@@ -137,19 +137,19 @@ def _to_code_tensor(codes: Any) -> torch.Tensor | None:
 
 def llm2code2wav_async_chunk(
     transfer_manager: Any,
-    pooling_output: OmniPayload,
+    multimodal_output: OmniPayload | dict[str, Any],
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
     """
-    Async chunk version: convert stage-0 pooling_output to code2wav payload (pooling / connector accumulation).
+    Async chunk version: convert stage-0 multimodal_output to code2wav payload (pooling / connector accumulation).
 
     Accumulates codes in connector per request_id,
     returns payload only when chunk_size is full or request is finished; returns None when waiting.
     """
     # Null guard: chunk_transfer_adapter calls this every emit step
-    # including no-output steps where pooling_output is None.
-    if pooling_output is None or not isinstance(pooling_output, dict):
+    # including no-output steps where multimodal_output is None.
+    if multimodal_output is None or not isinstance(multimodal_output, dict):
         if is_finished:
             connector = getattr(transfer_manager, "connector", None)
             raw_cfg = getattr(connector, "config", {}) or {}
@@ -187,10 +187,8 @@ def llm2code2wav_async_chunk(
 
     # Text-only paths (e.g. modalities=["text"]) yield no codec pooling output;
     # stage-0 still drives the chunk transfer adapter, so treat None as "no codes
-    # this step" rather than letting `.get()` raise AttributeError — an unhandled
-    # error here drops the chunk, starves stage-1 of the finished payload, and
-    # the stage subprocesses die before the final token is emitted.
-    po_codes = pooling_output.get("codes", {}) if pooling_output is not None else {}
+    # this step" rather than letting `.get()` raise AttributeError.
+    po_codes = multimodal_output.get("codes", {}) if multimodal_output is not None else {}
     if "audio" not in po_codes:
         if is_finished:
             return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size)

vllm_omni/model_executor/stage_input_processors/ming_tts.py

@@ -167,10 +167,11 @@ def _build_chunk_observability(
 
 def llm2audio_vae_async_chunk(
     transfer_manager: Any,
-    pooling_output: dict[str, Any] | None,
+    multimodal_output: dict[str, Any] | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
+    pooling_output = multimodal_output
     request_id = request.external_req_id
     chunk_id = int(transfer_manager.put_req_chunk[request_id])
     finished = bool(is_finished or request.is_finished())

vllm_omni/model_executor/stage_input_processors/moss_tts.py

@@ -92,7 +92,7 @@ def talker2codec(
 
 def talker2codec_async_chunk(
     transfer_manager: Any,
-    pooling_output: dict[str, Any] | None,
+    multimodal_output: dict[str, Any] | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -107,6 +107,7 @@ def talker2codec_async_chunk(
     signal "not enough data yet — wait for more frames".
     """
     req_id: str = str(getattr(request, "request_id", id(request)))
+    pooling_output = multimodal_output
 
     # Initialise per-request accumulation state
     if not hasattr(transfer_manager, "_moss_tts_state"):

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -4,6 +4,7 @@
 """Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition."""
 
 import logging
+from collections.abc import Mapping
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -190,7 +191,7 @@ def _get_prefill_multimodal_output(
     if not isinstance(by_req, dict):
         return None
     prefill_mm = by_req.get(request_id)
-    return prefill_mm if isinstance(prefill_mm, dict) else None
+    return prefill_mm if isinstance(prefill_mm, Mapping) else None
 
 
 def _merge_pd_embeddings(
@@ -215,7 +216,7 @@ def _merge_pd_embeddings(
         p_emb = p_layers[int(_EMBED_LAYER_KEY)].detach().to(device=device, dtype=torch.float)
         p_hid = p_layers[int(_HIDDEN_LAYER_KEY)].detach().to(device=device, dtype=torch.float)
     except (KeyError, AttributeError, TypeError) as exc:
-        available_keys = list(prefill_mm.keys()) if isinstance(prefill_mm, dict) else type(prefill_mm).__name__
+        available_keys = list(prefill_mm.keys()) if isinstance(prefill_mm, Mapping) else type(prefill_mm).__name__
         logger.error(
             "_merge_pd_embeddings: failed to extract prefill embeddings (%s). "
             "Expected keys %r and %r, got: %s. "
@@ -433,7 +434,7 @@ def _get_streaming_codec_delta_len(
 
 def thinker2talker_async_chunk(
     transfer_manager: Any,
-    pooling_output: OmniPayload,
+    multimodal_output: OmniPayload | dict[str, Any],
     request: OmniEngineCoreRequest,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -446,13 +447,13 @@ def thinker2talker_async_chunk(
 
     request_id = request.external_req_id
     chunk_id = transfer_manager.put_req_chunk[request_id]
-    if not isinstance(pooling_output, dict):
-        logger.debug("thinker2talker_async_chunk: skip non-dict pooling_output for req=%s", request_id)
+    if not isinstance(multimodal_output, Mapping):
+        logger.debug("thinker2talker_async_chunk: skip non-dict multimodal_output for req=%s", request_id)
         return None
 
-    thinker_hs = pooling_output.get("hidden_states", {})
+    thinker_hs = multimodal_output.get("hidden_states", {})
     thinker_layers = thinker_hs.get("layers", {}) if isinstance(thinker_hs, dict) else {}
-    thinker_embed_raw = pooling_output.get("embed", {})
+    thinker_embed_raw = multimodal_output.get("embed", {})
     thinker_embed = thinker_embed_raw if isinstance(thinker_embed_raw, dict) else {}
     thinker_emb = _layer_tensor(thinker_layers, _EMBED_LAYER_KEY)
     thinker_hid = _layer_tensor(thinker_layers, _HIDDEN_LAYER_KEY)
@@ -534,7 +535,7 @@ def thinker2talker_full_payload(
 ) -> dict[str, Any] | None:
     """Pack complete thinker output for the non-async connector path."""
     rid = getattr(request, "request_id", None)
-    if not isinstance(pooling_output, dict):
+    if not isinstance(pooling_output, Mapping):
         logger.warning(
             "thinker2talker_full_payload: pooling_output not a dict (type=%s) for req=%s; consumer wait gate may hang.",
             type(pooling_output).__name__,
@@ -657,12 +658,12 @@ def thinker2talker(
         thinker_input_ids = prompt_token_ids
         new_seq_length = len(prompt_token_ids + output_ids) - 1
         thinker_mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(thinker_mm_raw, dict):
+        if not isinstance(thinker_mm_raw, Mapping):
             logger.debug("thinker2talker: skip req=%s due to empty multimodal_output", req_id)
             continue
         thinker_mm: OmniPayload = thinker_mm_raw
         mm_hs = thinker_mm.get("hidden_states", {})
-        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, dict) else {}
+        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
         emb_layer = _layer_tensor(mm_layers, _EMBED_LAYER_KEY)
         hid_layer = _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY)
         if emb_layer is None or hid_layer is None:
@@ -748,11 +749,11 @@ def thinker2talker_token_only(
         # hidden-state layers), do the same here so the worker buffer
         # presence agrees with the orchestrator's scheduling decision.
         thinker_mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(thinker_mm_raw, dict):
+        if not isinstance(thinker_mm_raw, Mapping):
             logger.debug("thinker2talker_token_only: skip req=%s due to empty multimodal_output", req_id)
             continue
         mm_hs = thinker_mm_raw.get("hidden_states", {})
-        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, dict) else {}
+        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
         if _layer_tensor(mm_layers, _EMBED_LAYER_KEY) is None or _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) is None:
             logger.debug("thinker2talker_token_only: skip req=%s due to missing hidden-state layers", req_id)
             continue
@@ -801,16 +802,16 @@ def thinker2talker_token_only(
 
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
-    pooling_output: OmniPayload,
+    multimodal_output: OmniPayload | dict[str, Any],
     request: OmniEngineCoreRequest,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
     """
-    Pooling version.
+    Multimodal output version.
     """
-    if not isinstance(pooling_output, dict):
+    if not isinstance(multimodal_output, Mapping):
         return None
-    talker_codes = pooling_output.get("codes", {})
+    talker_codes = multimodal_output.get("codes", {})
     if not isinstance(talker_codes, dict):
         return None
     code_predictor_codes = talker_codes.get("audio")
@@ -884,7 +885,7 @@ def talker2code2wav_full_payload(
 ) -> dict[str, Any] | None:
     """Pack complete talker codec output for the non-async connector path."""
     rid = getattr(request, "request_id", None)
-    if not isinstance(pooling_output, dict):
+    if not isinstance(pooling_output, Mapping):
         logger.warning(
             "talker2code2wav_full_payload: pooling_output not a dict "
             "(type=%s) for req=%s; consumer wait gate may hang.",
@@ -983,7 +984,7 @@ def talker2code2wav(
         if is_streaming_session:
             seq_len = _get_streaming_codec_delta_len(cur_seq_len, req_id, talker_output, streaming_context)
         mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(mm_raw, dict):
+        if not isinstance(mm_raw, Mapping):
             logger.debug("talker2code2wav: skip req=%s due to empty multimodal_output", req_id)
             continue
         mm: OmniPayload = mm_raw

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -1,5 +1,6 @@
 """Stage input processor for Qwen3-TTS: Talker -> Code2Wav."""
 
+from collections.abc import Mapping
 from typing import Any
 
 import torch
@@ -123,8 +124,8 @@ def talker2code2wav(
     return code2wav_inputs
 
 
-def _extract_last_frame(pooling_output: OmniPayload) -> torch.Tensor | None:
-    audio_codes = pooling_output.get("codes", {}).get("audio")
+def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
+    audio_codes = multimodal_output.get("codes", {}).get("audio")
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
         return None
     if audio_codes.ndim == 2:
@@ -139,7 +140,7 @@ def _extract_last_frame(pooling_output: OmniPayload) -> torch.Tensor | None:
 
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
-    pooling_output: OmniPayload | None,
+    multimodal_output: OmniPayload | dict[str, Any] | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -150,12 +151,12 @@ def talker2code2wav_async_chunk(
         request_payload = {}
         transfer_manager.request_payload = request_payload
 
-    if isinstance(pooling_output, dict):
-        frame = _extract_last_frame(pooling_output)
+    if isinstance(multimodal_output, Mapping):
+        frame = _extract_last_frame(multimodal_output)
         if frame is not None:
             codec_codes = frame.cpu().tolist()
             transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
-        ref_code = pooling_output.get("codes", {}).get("ref")
+        ref_code = multimodal_output.get("codes", {}).get("ref")
         if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0 and request_payload.get(request_id) is None:
             request_payload[request_id] = ref_code.to(torch.long).cpu().contiguous()
     elif not finished:

vllm_omni/model_executor/stage_input_processors/voxtral_tts.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from typing import Any
 
 import torch
@@ -35,8 +36,8 @@ def generator2tokenizer(
     return tokenizer_inputs
 
 
-def _extract_last_frame(pooling_output: OmniPayload) -> torch.Tensor | None:
-    audio = pooling_output.get("audio")
+def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
+    audio = multimodal_output.get("audio")
     if not isinstance(audio, torch.Tensor) or audio.numel() == 0:
         return None
     return audio.flatten()
@@ -44,20 +45,20 @@ def _extract_last_frame(pooling_output: OmniPayload) -> torch.Tensor | None:
 
 def generator2tokenizer_async_chunk(
     transfer_manager: Any,
-    pooling_output: OmniPayload,
+    multimodal_output: OmniPayload | dict[str, Any],
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
     request_id = request.external_req_id
     finished = bool(is_finished or request.is_finished())
 
-    if isinstance(pooling_output, dict):
-        frame = _extract_last_frame(pooling_output)
+    if isinstance(multimodal_output, Mapping):
+        frame = _extract_last_frame(multimodal_output)
         if frame is not None:
             codec_codes = frame.cpu().tolist()
             transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
     elif not finished:
-        # Some steps may not produce pooling_output. Only flush on finish.
+        # Some steps may not produce multimodal_output. Only flush on finish.
         return None
 
     connector = getattr(transfer_manager, "connector", None)

vllm_omni/outputs.py

@@ -36,6 +36,7 @@ class OmniConnectorOutput:
     has_pending_kv_work: bool = False
 
 
+@dataclass
 class OmniModelRunnerOutput(ModelRunnerOutput):
     """Model runner output for omni models.
 
@@ -43,11 +44,12 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
     that may be produced by non-autoregressive stages.
 
     Attributes:
-        multimodal_outputs: Optional dictionary mapping modality names to
-            output tensors (e.g., {"image": tensor, "audio": tensor})
+        multimodal_outputs: Optional per-request list of multimodal output
+            dicts, indexed by req_index. Each element is a dict mapping
+            output keys to tensors/values (e.g., [{"audio": tensor}, ...]).
     """
 
-    multimodal_outputs: dict[str, torch.Tensor] | None = None
+    multimodal_outputs: list[dict[str, object]] | None = None
     # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU.
     # The Scheduler can safely free the block tables for these requests.
     kv_extracted_req_ids: list[str] | None = None
@@ -213,11 +215,13 @@ class OmniRequestOutput:
         )
 
     @property
-    def multimodal_output(self) -> dict[str, Any]:
+    def multimodal_output(self) -> Any:
         """Return multimodal output from the underlying request output or local field.
 
         For pipeline outputs, this checks completion outputs first, then request_output.
         For diffusion outputs, this returns 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

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import time
+from collections.abc import Mapping
 from copy import copy, deepcopy
 from typing import Any, NamedTuple
 
@@ -163,7 +164,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
         num_tokens_padded: int,
     ):
         if self.omni_prefix_cache is not None and get_pp_group().is_last_rank:
-            if multimodal_outputs is not None and not isinstance(multimodal_outputs, dict):
+            if multimodal_outputs is not None and not isinstance(multimodal_outputs, Mapping):
                 logger.warning_once(
                     "prefix caching expects mm outputs to be a dict, but got %s",
                     type(multimodal_outputs),
@@ -639,7 +640,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
             if multimodal_outputs is not None:
                 keys_or_type = (
                     list(multimodal_outputs.keys())
-                    if isinstance(multimodal_outputs, dict)
+                    if isinstance(multimodal_outputs, Mapping)
                     else type(multimodal_outputs)
                 )
                 logger.debug(f"[AR] execute_model: multimodal_outputs keys = {keys_or_type}")

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import gc
 import time
+from collections.abc import Mapping
 from copy import copy, deepcopy
 
 import numpy as np
@@ -448,7 +449,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                 pooler_output.append(
                     {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None}
                 )
-        elif isinstance(multimodal_outputs, dict):
+        elif isinstance(multimodal_outputs, Mapping):
             num_reqs = self.input_batch.num_reqs
             for i in range(num_reqs):
                 mm_payload = {}

vllm_omni/utils/mm_outputs.py

@@ -2,6 +2,8 @@
 payloads, most of which are shared by the prefix cache / no prefix cache path.
 """
 
+from collections.abc import Mapping
+
 import torch
 from vllm.logger import init_logger
 
@@ -23,7 +25,7 @@ def build_mm_cpu(multimodal_outputs: dict) -> dict[str, object]:
     mm_cpu: dict[str, object] = {}
     # Currently there are some cases where this is true at the
     # moment, which should be fixed.
-    if not isinstance(multimodal_outputs, dict):
+    if not isinstance(multimodal_outputs, Mapping):
         logger.warning("Multimodal outputs are not a dict and will not be passed")
 
     if multimodal_outputs:

vllm_omni/worker/gpu_ar_model_runner.py

@@ -6,6 +6,7 @@ and also outputs sampled tokens.
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from contextlib import nullcontext
 from copy import copy
 from dataclasses import replace
@@ -46,6 +47,39 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorMode
 logger = init_logger(__name__)
 
 
+def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]:
+    """Convert a flattened payload to strictly ``dict[str, torch.Tensor]``.
+
+    Non-tensor scalars (int, float) are wrapped with ``torch.tensor()``.
+    Values that cannot be safely converted are dropped with a warning.
+    This enforces the tensor-only invariant required by the
+    ``OmniEngineCoreOutput.multimodal_output`` wire field and msgspec
+    serialization.
+    """
+    result: dict[str, torch.Tensor] = {}
+    for key, val in payload.items():
+        if isinstance(val, torch.Tensor):
+            result[key] = val
+        elif isinstance(val, (int, float, bool)):
+            result[key] = torch.tensor(val)
+        elif isinstance(val, (list, tuple)):
+            try:
+                result[key] = torch.tensor(val)
+            except (ValueError, TypeError, RuntimeError):
+                logger.warning(
+                    "Dropping non-tensorizable multimodal output key '%s' (type=%s) from wire payload.",
+                    key,
+                    type(val).__name__,
+                )
+        else:
+            logger.warning(
+                "Dropping non-tensor multimodal output key '%s' (type=%s) from wire payload.",
+                key,
+                type(val).__name__,
+            )
+    return result
+
+
 class ExecuteModelState(NamedTuple):
     scheduler_output: SchedulerOutput
     logits: torch.Tensor | None
@@ -329,7 +363,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         if self.omni_prefix_cache is not None and is_last_pp_rank:
             # If this happens, it generally means the model is not following the correct
             # interface yet and is therefore currently not compatible with prefix cache.
-            if multimodal_outputs is not None and not isinstance(multimodal_outputs, dict):
+            if multimodal_outputs is not None and not isinstance(multimodal_outputs, Mapping):
                 logger.warning_once(
                     "prefix caching expects mm outputs to be a dict, but got %s",
                     type(multimodal_outputs),
@@ -835,21 +869,45 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
 
     @staticmethod
     def _resolve_req_hidden_states(
-        hidden_states_cpu: torch.Tensor,
+        hidden_states_cpu: torch.Tensor | None,
         combined_hidden_states: dict[str, torch.Tensor] | None,
         rid: str,
         start: int,
         end: int,
-    ):
+    ) -> torch.Tensor | None:
         if combined_hidden_states is not None:
             # We always have all request IDs for prefix cache, even for
             # partial cache misses, so this should never happen.
             if rid not in combined_hidden_states:
                 raise RuntimeError("Request IDs in the batch are missing from the merged states!")
             return combined_hidden_states[rid]
-        # Prefix caching is disabled
+        # Prefix caching is disabled. hidden_states_cpu may legitimately be
+        # None (e.g. sparse audio output or no scheduled hidden payload);
+        # callers must omit the "hidden" key in that case.
+        if hidden_states_cpu is None:
+            return None
         return hidden_states_cpu[start:end]
 
+    def _build_multimodal_outputs(
+        self,
+        per_req_payloads: list[dict[str, object]] | None,
+    ) -> list[dict[str, torch.Tensor]] | None:
+        """Build per-request multimodal output payloads (dedicated channel).
+
+        Reuses the per-request payloads assembled by the pooler-payload loop
+        in sample_tokens() (which already handles prefix-cache merging,
+        sparse audio output, and partial downstream batches) so the wire
+        channel stays consistent with the full-payload accumulation path.
+        Enforces the tensor-only invariant required by msgspec: scalars and
+        lists are wrapped into tensors, and anything that cannot be safely
+        converted is dropped.
+        """
+        if self.vllm_config.model_config.engine_output_type == "text":
+            return None
+        if per_req_payloads is None:
+            return None
+        return [_ensure_tensor_values(payload) if payload else {} for payload in per_req_payloads]
+
     @torch.inference_mode()
     def sample_tokens(
         self,
@@ -1057,6 +1115,10 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     deferred_mm_cache_keys=deferred_mm_cache_keys,
                 )
 
+        # NOTE: pooler_output here is used only for the full-payload accumulation
+        # path (accumulate_full_payload_output) and is NOT passed on the wire via
+        # OmniModelRunnerOutput.pooler_output (which is set to None below).
+        # The actual multimodal wire transport uses multimodal_outputs instead.
         pooler_output: list[dict[str, object]] | None = None
         if needs_pooler_payload:
             mm_cpu = None
@@ -1112,7 +1174,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                             start,
                             end,
                         )
-                    payload["hidden"] = req_hidden_states
+                    if req_hidden_states is not None:
+                        payload["hidden"] = req_hidden_states
 
                 mm_payload: dict[str, object] = {}
                 if combined_multimodal_outputs or mm_cpu:
@@ -1127,7 +1190,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
 
                         for mm_key in combined_multimodal_outputs.keys():
                             mm_payload[mm_key] = _unwrap_lists(combined_multimodal_outputs[mm_key][rid])
-
                     else:
                         for mm_key, mm_val in mm_cpu.items():
                             if mm_key in {"meta.req_id", "meta.sparse_audio"}:
@@ -1166,6 +1228,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 if req_state is not None and pooler_output[i]:
                     self.accumulate_full_payload_output(rid, pooler_output[i], req_state)
 
+        multimodal_outputs = self._build_multimodal_outputs(pooler_output)
         with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
             routed_experts_lists = None
             if self.routed_experts_initialized:
@@ -1176,7 +1239,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 sampled_token_ids=valid_sampled_token_ids,
                 logprobs=logprobs_lists,
                 prompt_logprobs_dict=prompt_logprobs_dict,
-                pooler_output=(pooler_output if engine_output_type != "text" and needs_pooler_payload else None),
+                pooler_output=None,
+                multimodal_outputs=multimodal_outputs,
                 kv_connector_output=kv_connector_output,
                 ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
                 num_nans_in_logits=num_nans_in_logits,

vllm_omni/worker/gpu_generation_model_runner.py

@@ -8,6 +8,7 @@ from __future__ import annotations
 
 import gc
 import logging
+from collections.abc import Mapping
 from dataclasses import replace
 
 import numpy as np
@@ -36,7 +37,7 @@ from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState
+from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
@@ -404,7 +405,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             aux_hidden_states,
             ec_connector_output,
             cudagraph_stats,
-            multimodal_outputs,
+            multimodal_outputs_raw,
             slot_mappings,  # OMNI: unpack slot_mappings for upstream v1 API compatibility
         ) = self.execute_model_state
         self.execute_model_state = None
@@ -414,27 +415,29 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         if self.speculative_config is not None:
             self.finalize_kv_connector()
 
-        pooler_output: list[object] = []
-        if isinstance(multimodal_outputs, torch.Tensor):
-            assert multimodal_outputs.shape[0] == 1, (
+        # Build per-request multimodal_outputs list (dedicated channel).
+        # pooler_output is no longer used for multimodal data.
+        multimodal_outputs: list[dict[str, object]] = []
+        if isinstance(multimodal_outputs_raw, torch.Tensor):
+            assert multimodal_outputs_raw.shape[0] == 1, (
                 "model should return a single tensor, to return multiple tensors, use a dict"
             )
-            assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
+            assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs
             for i in range(self.input_batch.num_reqs):
-                pooler_output.append({"model_outputs": multimodal_outputs[i].detach().to("cpu").contiguous()})
-        elif isinstance(multimodal_outputs, list):
-            assert len(multimodal_outputs) == 1, (
+                multimodal_outputs.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()})
+        elif isinstance(multimodal_outputs_raw, list):
+            assert len(multimodal_outputs_raw) == 1, (
                 "model should return a single list, to return multiple lists, use a dict"
             )
-            for out in multimodal_outputs:
-                pooler_output.append(
+            for out in multimodal_outputs_raw:
+                multimodal_outputs.append(
                     {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None}
                 )
-        elif isinstance(multimodal_outputs, dict):
+        elif isinstance(multimodal_outputs_raw, Mapping):
             num_reqs = self.input_batch.num_reqs
             for i in range(num_reqs):
                 mm_payload = {}
-                for key, out in multimodal_outputs.items():
+                for key, out in multimodal_outputs_raw.items():
                     if isinstance(out, list):
                         if len(out) != num_reqs:
                             raise ValueError(
@@ -446,7 +449,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                         mm_payload[key] = out.detach().to("cpu").contiguous()
                     else:
                         logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}")
-                pooler_output.append(mm_payload)
+                multimodal_outputs.append(_ensure_tensor_values(mm_payload))
         else:
             raise RuntimeError("Unsupported diffusion output type")
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
@@ -458,8 +461,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         if self._should_accumulate_full_payload_output():
             for i, rid in enumerate(req_ids_output_copy):
                 req_state = self.requests.get(rid)
-                if req_state is not None and pooler_output[i]:
-                    self.accumulate_full_payload_output(rid, pooler_output[i], req_state)
+                if req_state is not None and multimodal_outputs[i]:
+                    self.accumulate_full_payload_output(rid, multimodal_outputs[i], req_state)
 
         output = OmniModelRunnerOutput(
             req_ids=req_ids_output_copy,
@@ -467,7 +470,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             sampled_token_ids=[],
             logprobs=None,
             prompt_logprobs_dict={},
-            pooler_output=pooler_output,
+            pooler_output=None,
+            multimodal_outputs=multimodal_outputs,
             kv_connector_output=kv_connector_output,
             num_nans_in_logits={},
             cudagraph_stats=cudagraph_stats,