[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>

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)