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