From 8f869c8ea5de06d1ab0d1914b0e1c5b62833c90d Mon Sep 17 00:00:00 2001 From: Megha Agarwal Date: Fri, 12 Jun 2026 12:49:14 +0530 Subject: [PATCH] [Refactor] Output Processor Phase 2: separate multimodal output channel (#1601) (#2744) Signed-off-by: Megha Agarwal Signed-off-by: meghaagr13 Signed-off-by: Taichang Zhou Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Hongsheng Liu Co-authored-by: meghaagr13 Co-authored-by: Boao Shi Co-authored-by: Taichang Zhou Co-authored-by: Cursor --- pyproject.toml | 1 + .../test_chunk_transfer_adapter.py | 18 +- .../e2e/offline_inference/test_voxcpm2_tts.py | 3 +- tests/engine/test_ensure_tensor_values.py | 59 +++ tests/engine/test_output_modality.py | 4 +- tests/engine/test_output_processor.py | 61 +++ tests/engine/test_wire_multimodal_output.py | 88 ++++ tests/helpers/assertions.py | 32 +- .../test_cosyvoice3_stage_input_processors.py | 24 +- .../test_qwen3_tts_async_chunk.py | 20 +- .../test_voxtral_tts_async_chunk.py | 28 +- .../higgs_audio_v3/test_higgs_audio_v3.py | 28 +- tests/worker/test_gpu_ar_model_runner.py | 7 +- .../test_gpu_generation_model_runner.py | 20 +- vllm_omni/core/sched/omni_ar_scheduler.py | 7 +- .../core/sched/omni_generation_scheduler.py | 6 +- .../chunk_transfer_adapter.py | 14 +- .../omni_connectors/utils/serialization.py | 14 +- vllm_omni/engine/__init__.py | 5 +- vllm_omni/engine/async_omni_engine.py | 2 +- vllm_omni/engine/mm_outputs.py | 55 +- vllm_omni/engine/output_modality.py | 2 + vllm_omni/engine/output_processor.py | 480 +++++++++++++----- vllm_omni/engine/stage_engine_core_client.py | 18 + vllm_omni/engine/stage_engine_core_proc.py | 14 + .../entrypoints/openai/realtime_connection.py | 7 +- vllm_omni/entrypoints/openai/serving_video.py | 13 +- .../openai/serving_video_stream.py | 3 +- .../models/voxtral_tts/voxtral_tts.py | 4 +- .../stage_input_processors/cosyvoice3.py | 11 +- .../stage_input_processors/fish_speech.py | 15 +- .../stage_input_processors/glm_image.py | 11 +- .../stage_input_processors/glm_tts.py | 3 +- .../stage_input_processors/higgs_audio_v2.py | 3 +- .../stage_input_processors/higgs_audio_v3.py | 3 +- .../stage_input_processors/hunyuan_image3.py | 3 +- .../stage_input_processors/mammoth_moda2.py | 3 +- .../stage_input_processors/mimo_audio.py | 14 +- .../stage_input_processors/ming_tts.py | 3 +- .../stage_input_processors/moss_tts.py | 3 +- .../stage_input_processors/qwen3_omni.py | 37 +- .../stage_input_processors/qwen3_tts.py | 13 +- .../stage_input_processors/voxtral_tts.py | 13 +- vllm_omni/outputs.py | 12 +- .../npu/worker/npu_ar_model_runner.py | 5 +- .../npu/worker/npu_generation_model_runner.py | 3 +- vllm_omni/utils/mm_outputs.py | 4 +- vllm_omni/worker/gpu_ar_model_runner.py | 78 ++- .../worker/gpu_generation_model_runner.py | 38 +- 49 files changed, 985 insertions(+), 327 deletions(-) create mode 100644 tests/engine/test_ensure_tensor_values.py create mode 100644 tests/engine/test_wire_multimodal_output.py diff --git a/pyproject.toml b/pyproject.toml index 1bb8fb2f1..21b0ff42e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -281,6 +281,7 @@ extend-ignore-identifiers-re = [ ".*[Oo]no_[Aa]nna.*", ".*cann.*", ".*[Nn][Dd]3.*", + ".*updator.*", "byt5.*", ] extend-ignore-re = [ diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index ab671ec66..9c6a86ed3 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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 diff --git a/tests/e2e/offline_inference/test_voxcpm2_tts.py b/tests/e2e/offline_inference/test_voxcpm2_tts.py index a53c80e4d..817be4064 100644 --- a/tests/e2e/offline_inference/test_voxcpm2_tts.py +++ b/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") diff --git a/tests/engine/test_ensure_tensor_values.py b/tests/engine/test_ensure_tensor_values.py new file mode 100644 index 000000000..f8b71f507 --- /dev/null +++ b/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) diff --git a/tests/engine/test_output_modality.py b/tests/engine/test_output_modality.py index 7a9c76502..0abfe6ea0 100644 --- a/tests/engine/test_output_modality.py +++ b/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({})}") diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 4576da5a1..56447264a 100644 --- a/tests/engine/test_output_processor.py +++ b/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 diff --git a/tests/engine/test_wire_multimodal_output.py b/tests/engine/test_wire_multimodal_output.py new file mode 100644 index 000000000..47da26938 --- /dev/null +++ b/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) diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index f0d515d01..f79c50173 100644 --- a/tests/helpers/assertions.py +++ b/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(), diff --git a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py index 6debabb0e..c46c49d97 100644 --- a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py +++ b/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, ) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 9c4270c1e..bbb83b484 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/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, ) diff --git a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py index 7e63ede9e..5235d7c68 100644 --- a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py +++ b/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, ) diff --git a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py index 066d0c6a0..33be067db 100644 --- a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py +++ b/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 ---- diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 8c6accad5..3dbbc4ac4 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/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]]), ) diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index 7bc0f0ba5..be7fadcb6 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/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) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 4c4ae361a..736ba549c 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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 diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 9f5d3e104..2e99c8e1a 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/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(), diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index eb1c3eec6..04368842c 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/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, diff --git a/vllm_omni/distributed/omni_connectors/utils/serialization.py b/vllm_omni/distributed/omni_connectors/utils/serialization.py index c2145dbff..1166ede19 100644 --- a/vllm_omni/distributed/omni_connectors/utils/serialization.py +++ b/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 diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index aa08fcde6..89bb41908 100644 --- a/vllm_omni/engine/__init__.py +++ b/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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 05da3d3c3..f53cb26ac 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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 diff --git a/vllm_omni/engine/mm_outputs.py b/vllm_omni/engine/mm_outputs.py index f2a8927f7..3c9da2496 100644 --- a/vllm_omni/engine/mm_outputs.py +++ b/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: diff --git a/vllm_omni/engine/output_modality.py b/vllm_omni/engine/output_modality.py index 5887bb910..ac72910dd 100644 --- a/vllm_omni/engine/output_modality.py +++ b/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", } diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 5c0b752b3..5780c574c 100644 --- a/vllm_omni/engine/output_processor.py +++ b/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, diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index b5aecf124..2c6c00665 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/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, diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index eb989db0e..62c5b2a8f 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/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, diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index 5cf936e48..0f0efbef8 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index f05ce204b..ac6038f92 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/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") diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index a76b241c5..0c1fb52e9 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/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") diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py index eb8b5c14a..d12034c24 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/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'])=}" ) diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index 2a0fdd869..072fc9cf6 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py index ac12e6035..668ff2ed4 100644 --- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py +++ b/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: diff --git a/vllm_omni/model_executor/stage_input_processors/glm_image.py b/vllm_omni/model_executor/stage_input_processors/glm_image.py index 39533913e..bf644abd0 100644 --- a/vllm_omni/model_executor/stage_input_processors/glm_image.py +++ b/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): diff --git a/vllm_omni/model_executor/stage_input_processors/glm_tts.py b/vllm_omni/model_executor/stage_input_processors/glm_tts.py index caafbf86c..af037aee7 100644 --- a/vllm_omni/model_executor/stage_input_processors/glm_tts.py +++ b/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) diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py index 0d6781586..c7582f9ff 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py +++ b/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) diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py index 6d71649e3..353be03e8 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py b/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py index d0e9b17fa..9a81a8cc3 100644 --- a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py index 6b699008e..7c821c243 100644 --- a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py +++ b/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)}, " diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py index 2443ad2c4..7c0db4131 100644 --- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py +++ b/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) diff --git a/vllm_omni/model_executor/stage_input_processors/ming_tts.py b/vllm_omni/model_executor/stage_input_processors/ming_tts.py index 46aa5308c..92f25eea0 100644 --- a/vllm_omni/model_executor/stage_input_processors/ming_tts.py +++ b/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()) diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py index 562318d73..48c0f98cb 100644 --- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py +++ b/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"): diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index c2b33368d..d624b4a49 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 158963c6e..25afa9b9e 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/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: diff --git a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py index be251f9c9..2928a7cf3 100644 --- a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py +++ b/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) diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 7b3aab2cb..2dfa1699f 100644 --- a/vllm_omni/outputs.py +++ b/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 diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 9ea236bb7..38808648e 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/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}") diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 8e27026d4..341b79d36 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/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 = {} diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py index 43e8fa110..3009e21d4 100644 --- a/vllm_omni/utils/mm_outputs.py +++ b/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: diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index b68cf3800..c7cca20ea 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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, diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index e4b7944d1..793228158 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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,