[Refactor/Bugfix] Use Delta Messages for Streaming (#2911)

Signed-off-by: Alex Brooks <albrooks@redhat.com>
Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Sy03 <1370724210@qq.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -242,8 +242,7 @@ async def run_single_request(
                 if stage_0_first_output_ts is None:
                     stage_0_first_output_ts = time.perf_counter()
                 text_output = output.outputs[0].text
-                if output.finished:
-                    text_parts.append(text_output)
+                text_parts.append(text_output)
             elif omni_output.final_output_type == "audio":
                 mm_out = output.outputs[0].multimodal_output
                 if mm_out and "audio" in mm_out:
@@ -293,7 +292,7 @@ async def run_single_request(
     if text_parts:
         text_file = os.path.join(output_dir, f"{request_id}.txt")
         with open(text_file, "w", encoding="utf-8") as f:
-            f.write("\n".join(text_parts))
+            f.write("".join(text_parts))
         result["saved_files"].append(text_file)
         print(
             f"[Request {request_id}] Text saved to {text_file} "
@@ -393,6 +392,10 @@ async def run_all(args):
 
         # Use default sampling params from stage config (they are pre-configured
         # in the YAML for each stage).
+        #
+        # NOTE: Since we do not set the sampling params directly, .generate in
+        # will automatically set the output kind to delta, since this is what
+        # makes sense for most multimodal use-cases.
         sampling_params_list = None
 
         output_dir = args.output_dir

tests/e2e/online_serving/test_voxtral_tts.py

@@ -82,6 +82,30 @@ class TestVoxtralTTSFixedVoice:
             f"Audio content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes"
         )
 
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_english_streaming(self, omni_server) -> None:
+        """Test basic streaming English TTS generation."""
+        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
+        payload = {
+            "input": "Hello, how are you?",
+            "voice": "casual_female",
+            "language": "English",
+            "stream": True,
+            "response_format": "pcm",
+        }
+
+        with httpx.Client(timeout=120.0) as client:
+            with client.stream("POST", url, json=payload) as response:
+                assert response.status_code == 200
+                assert response.headers.get("content-type") == "audio/pcm"
+                total = sum(len(c) for c in response.iter_bytes())
+
+        assert total > MIN_AUDIO_BYTES, (
+            f"Streamed audio too small ({total} bytes), expected at least {MIN_AUDIO_BYTES} bytes"
+        )
+
     @pytest.mark.advanced_model
     @pytest.mark.omni
     @hardware_test(res={"cuda": "H100"}, num_cards=1)

tests/engine/test_output_processor.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from vllm.outputs import PoolingRequestOutput
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine import FinishReason
+
+from vllm_omni.engine.output_modality import OutputModalityNames
+from vllm_omni.engine.output_processor import OmniRequestState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Audio is explicitly listed as a drainable modality
+AUDIO = OutputModalityNames.AUDIO
+
+# Latent is explicitly not drainable, but the choice doesn't matter here as
+# long as isn't listed as drainable. I.e., could also be arbitrary keys for
+# the purposes of these tests
+LATENT = OutputModalityNames.LATENT
+
+# NOTE: detokenizer and logprobs aren't really used here, but we mock them since
+# some of the utils called in vLLM superclass assert require them to be None.
+_DETOK = MagicMock(
+    output_token_ids=[0],
+    get_next_output_text=MagicMock(return_value=""),
+    num_output_tokens=MagicMock(return_value=1),
+)
+_LOGPROBS = MagicMock(logprobs=None, cumulative_logprob=None, prompt_logprobs=None)
+
+_DEFAULT_STATE_KWARGS = dict(
+    request_id="r",
+    external_req_id="r",
+    parent_req=None,
+    request_index=0,
+    lora_request=None,
+    prompt=None,
+    prompt_token_ids=[0],
+    prompt_embeds=None,
+    logprobs_processor=_LOGPROBS,
+    detokenizer=_DETOK,
+    max_tokens_param=None,
+    arrival_time=0.0,
+    queue=None,
+    log_stats=False,
+    stream_interval=1,
+)
+
+
+def _make_state(output_kind: RequestOutputKind):
+    return OmniRequestState(**_DEFAULT_STATE_KWARGS, output_kind=output_kind)
+
+
+def test_init_empty_dict():
+    """Ensure mm_accumulated is initially empty."""
+    assert _make_state(RequestOutputKind.CUMULATIVE).mm_accumulated == {}
+    assert _make_state(RequestOutputKind.DELTA).mm_accumulated == {}
+
+
+def test_delta_drains_output_modality_per_step():
+    """DELTA drains the mm_type key (output modality) but preserves hidden-state keys."""
+    s = _make_state(RequestOutputKind.DELTA)
+    audio1, audio2, hs1, hs2 = [torch.ones(num_elem) for num_elem in range(1, 5)]
+
+    # Add audio and hidden state tensors
+    s.add_multimodal_tensor(audio1, mm_type=AUDIO)  # should be drained
+    s.add_multimodal_tensor(hs1, mm_type=LATENT)  # shouldn't be drained
+
+    out1 = s._new_completion_output([1], None, None)
+    out1_audio = out1.multimodal_output[AUDIO]
+    out1_hidden = out1.multimodal_output[LATENT]
+    assert isinstance(out1_audio, torch.Tensor)
+    assert torch.equal(out1.multimodal_output[AUDIO], audio1)
+    assert isinstance(out1_hidden, torch.Tensor)
+    assert torch.equal(out1.multimodal_output[LATENT], hs1)
+
+    # After emission, hidden states should remain, but audio is drained
+    assert set(s.mm_accumulated.keys()) == {LATENT}
+
+    s.add_multimodal_tensor(audio2, AUDIO)
+    s.add_multimodal_tensor(hs2, mm_type=LATENT)
+    out2 = s._new_completion_output([2], None, None)
+    out2_audio = out2.multimodal_output[AUDIO]
+    out2_hidden = out2.multimodal_output[LATENT]
+    assert isinstance(out2_audio, torch.Tensor)
+    assert torch.equal(out2_audio, audio2)
+    # Since hidden isn't drained, it's grown to a list
+    assert isinstance(out2_hidden, list) and len(out2_hidden) == 2
+    assert torch.equal(hs1, out2_hidden[0])
+    assert torch.equal(hs2, out2_hidden[1])
+
+
+def test_cumulative_emits_consolidated_audio_each_step():
+    """Ensure cumulative accumulates and consolidates modality keys every step."""
+    s = _make_state(RequestOutputKind.CUMULATIVE)
+    # NOTE: audio is usually emitted as (1, size) chunks; we need to be sure
+    # to not change the tensor dimension when we consolidate
+    audio1 = torch.ones(1, 500)
+    s.add_multimodal_tensor(audio1, mm_type=AUDIO)
+    req_out = s.make_request_output([1], None, None, None)
+    assert req_out is not None
+    cons_audio = req_out.outputs[0].multimodal_output[AUDIO]
+    # Single chunk keeps original shape [1, 500]
+    assert isinstance(cons_audio, torch.Tensor) and cons_audio.shape == audio1.shape
+
+    audio2 = torch.ones(1, 300)
+    s.add_multimodal_tensor(audio2, mm_type=AUDIO)
+    req_out = s.make_request_output([2], None, None, None)
+    assert req_out is not None
+    cons_audio = req_out.outputs[0].multimodal_output[AUDIO]
+    # After consolidation, audio chunks are concatenated on last axis,
+    # preserving the [1, N] channel dimension
+    total_audio_len = audio1.shape[-1] + audio2.shape[-1]
+    assert isinstance(audio2, torch.Tensor) and cons_audio.shape == (1, total_audio_len)
+
+    assert "audio" in s.mm_accumulated
+
+
+def test_finish_consolidates_hidden_states():
+    """Ensure consolidation merges hidden-state tensor lists on finish."""
+    s = _make_state(RequestOutputKind.CUMULATIVE)
+    s.add_multimodal_tensor(torch.ones(5, 4), mm_type=LATENT)
+    s.add_multimodal_tensor(torch.ones(3, 4), mm_type=LATENT)
+
+    result = s.make_request_output([1], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+
+    hs = result.outputs[0].multimodal_output[LATENT]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 8
+
+
+def test_finish_consolidation_for_hs_delta():
+    """Ensure finish doesn't drop the accumulated hidden states."""
+    s = _make_state(RequestOutputKind.DELTA)
+    # hidden state accumulation (nothing drained)
+    s.add_multimodal_tensor({"foo": torch.ones(5, 4)}, mm_type=LATENT)
+    result = s.make_request_output([0], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+    hs = result.outputs[0].multimodal_output["foo"]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 5
+
+    # Since we don't drain the hidden states, if we add 3 elements, we should get 8
+    s.add_multimodal_tensor({"foo": torch.ones(3, 4)}, mm_type=LATENT)
+    result = s.make_request_output([0], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+    hs = result.outputs[0].multimodal_output["foo"]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 8
+    assert "foo" in s.mm_accumulated
+
+
+def test_finish_consolidation_drains_mm_delta():
+    """Ensure making the request output drains modality deltas (e.g., audio)."""
+    s = _make_state(RequestOutputKind.DELTA)
+    # multimodal data accumulation (drained)
+    s.add_multimodal_tensor({AUDIO: torch.ones(5, 4)}, mm_type=AUDIO)
+    result = s.make_request_output([0], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+    hs = result.outputs[0].multimodal_output[AUDIO]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 5
+
+    # Since we did drain the hidden states, we no longer get the 5 back
+    s.add_multimodal_tensor({AUDIO: torch.ones(3, 4)}, mm_type=AUDIO)
+    result = s.make_request_output([0], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+    hs = result.outputs[0].multimodal_output[AUDIO]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 3
+    assert AUDIO not in s.mm_accumulated  # drained
+
+
+@pytest.mark.parametrize("mm_type", [AUDIO, "hidden"])
+def test_final_only_consolidates_drainable_keys(mm_type):
+    """FINAL_ONLY never drains per-step, so modality keys and hidden state
+    keys both accumulate and are consolidated on finish."""
+    s = _make_state(RequestOutputKind.FINAL_ONLY)
+
+    # NOTE: Currently there is brittlness in the tensor stacking, so we just
+    # test a 1D tensor here. The intention is just to ensure audio /hidden
+    # behave the same.
+    s.add_multimodal_tensor(torch.ones(500), mm_type=mm_type)
+    # Non-finish step returns None without calling _new_completion_output
+    assert s.make_request_output([1], None, None, None) is None
+    assert mm_type in s.mm_accumulated
+
+    s.add_multimodal_tensor(torch.ones(300), mm_type=mm_type)
+    result = s.make_request_output([2], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+
+    audio = result.outputs[0].multimodal_output[mm_type]
+    assert isinstance(audio, torch.Tensor)
+    assert audio.shape == (800,)
+
+
+def test_cumulative_token_ids_always_set():
+    """cumulative_token_ids is set for all output kinds."""
+    for kind in (RequestOutputKind.DELTA, RequestOutputKind.CUMULATIVE, RequestOutputKind.FINAL_ONLY):
+        s = _make_state(kind)
+        out = s._new_completion_output([42], None, None)
+        assert hasattr(out, "cumulative_token_ids")
+        # The mock detokenizer has output_token_ids=[0]
+        assert list(out.cumulative_token_ids) == [0]
+
+
+def test_cumulative_token_ids_is_a_copy():
+    """cumulative_token_ids must be a snapshot, not a live reference."""
+    s = _make_state(RequestOutputKind.DELTA)
+    out = s._new_completion_output([42], None, None)
+    assert out.cumulative_token_ids is not _DETOK.output_token_ids

tests/entrypoints/test_async_omni.py

@@ -0,0 +1,134 @@
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+from vllm.sampling_params import RequestOutputKind, SamplingParams
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+async def _noop(**kw):
+    pass
+
+
+def get_fake_add_request(submitted_request_ids):
+    async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
+        del prompt, sampling_params_list, final_stage_id, kwargs
+        submitted_request_ids.append(request_id)
+
+    return fake_add_request_async
+
+
+def get_fake_abort(aborted_request_batches):
+    async def fake_abort_async(request_ids):
+        aborted_request_batches.append(list(request_ids))
+
+    return fake_abort_async
+
+
+async def fake_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts):
+    del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts
+    if request_id.startswith("cancel-"):
+        await asyncio.Future()
+        return
+    yield SimpleNamespace(
+        stage_id=0,
+        request_output=SimpleNamespace(outputs=[]),
+        finished=True,
+    )
+
+
+def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) -> AsyncOmni:
+    omni = object.__new__(AsyncOmni)
+    omni._pause_cond = asyncio.Condition()
+    omni._paused = False
+    omni.engine = SimpleNamespace(
+        num_stages=1,
+        add_request_async=fake_add_request,
+        abort_async=fake_abort_request,
+    )
+    omni.log_stats = False
+    omni.request_states = {}
+    omni._final_output_handler = lambda: None
+    omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params
+    omni._compute_final_stage_id = lambda output_modalities: 0
+    omni._process_orchestrator_results = fake_process_results
+    omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
+    return omni
+
+
+def test_generate_accepts_request_after_repeated_cancellations():
+    async def run_test():
+        submitted_request_ids = []
+        aborted_request_batches = []
+
+        async def collect_outputs(request_id):
+            outputs = []
+            async for output in AsyncOmni.generate(
+                omni,
+                prompt={"prompt": "prompt"},
+                request_id=request_id,
+                sampling_params_list=[SimpleNamespace()],
+                output_modalities=["image"],
+            ):
+                outputs.append(output)
+            return outputs
+
+        omni = get_async_omni_instance(
+            fake_add_request=get_fake_add_request(submitted_request_ids),
+            fake_abort_request=get_fake_abort(aborted_request_batches),
+        )
+
+        assert len(await collect_outputs("baseline")) == 1
+
+        for idx in range(3):
+            task = asyncio.create_task(collect_outputs(f"cancel-{idx}"))
+            await asyncio.sleep(0)
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+
+        assert len(await collect_outputs("after-cancel")) == 1
+        assert submitted_request_ids == [
+            "baseline",
+            "cancel-0",
+            "cancel-1",
+            "cancel-2",
+            "after-cancel",
+        ]
+        assert aborted_request_batches == [
+            ["cancel-0"],
+            ["cancel-1"],
+            ["cancel-2"],
+        ]
+
+    asyncio.run(run_test())
+
+
+@pytest.mark.parametrize(
+    "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE]
+)
+def test_output_kind_is_preserved_with_explicit_sampling_params(output_kind):
+    """Ensure we don't change the output kind in async generate if params are provided directly."""
+
+    captured_params = []
+
+    async def capturing_add_request(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
+        del prompt, final_stage_id, kwargs
+        captured_params.extend(sampling_params_list)
+
+    async def run():
+        omni = get_async_omni_instance(fake_add_request=capturing_add_request)
+        sp = SamplingParams(output_kind=output_kind)
+        async for _ in omni.generate(
+            prompt={"prompt": "test"},
+            request_id="test-req",
+            sampling_params_list=[sp],
+            output_modalities=["text"],
+        ):
+            pass
+
+    asyncio.run(run())
+    assert captured_params[0].output_kind == output_kind

tests/entrypoints/test_async_omni_abort.py

@@ -1,85 +0,0 @@
-import asyncio
-from types import SimpleNamespace
-
-import pytest
-
-from vllm_omni.entrypoints.async_omni import AsyncOmni
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_generate_accepts_request_after_repeated_cancellations():
-    async def run_test():
-        submitted_request_ids = []
-        aborted_request_batches = []
-
-        async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
-            del prompt, sampling_params_list, final_stage_id, kwargs
-            submitted_request_ids.append(request_id)
-
-        async def fake_abort_async(request_ids):
-            aborted_request_batches.append(list(request_ids))
-
-        async def fake_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts):
-            del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts
-            if request_id.startswith("cancel-"):
-                await asyncio.Future()
-                return
-            yield SimpleNamespace(
-                stage_id=0,
-                request_output=SimpleNamespace(outputs=[]),
-                finished=True,
-            )
-
-        async def collect_outputs(request_id):
-            outputs = []
-            async for output in AsyncOmni.generate(
-                omni,
-                prompt={"prompt": "prompt"},
-                request_id=request_id,
-                sampling_params_list=[SimpleNamespace()],
-                output_modalities=["image"],
-            ):
-                outputs.append(output)
-            return outputs
-
-        omni = object.__new__(AsyncOmni)
-        omni._pause_cond = asyncio.Condition()
-        omni._paused = False
-        omni.engine = SimpleNamespace(
-            num_stages=1,
-            add_request_async=fake_add_request_async,
-            abort_async=fake_abort_async,
-        )
-        omni.log_stats = False
-        omni.request_states = {}
-        omni._final_output_handler = lambda: None
-        omni.resolve_sampling_params_list = lambda params: params
-        omni._compute_final_stage_id = lambda output_modalities: 0
-        omni._process_orchestrator_results = fake_process_results
-        omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
-
-        assert len(await collect_outputs("baseline")) == 1
-
-        for idx in range(3):
-            task = asyncio.create_task(collect_outputs(f"cancel-{idx}"))
-            await asyncio.sleep(0)
-            task.cancel()
-            with pytest.raises(asyncio.CancelledError):
-                await task
-
-        assert len(await collect_outputs("after-cancel")) == 1
-        assert submitted_request_ids == [
-            "baseline",
-            "cancel-0",
-            "cancel-1",
-            "cancel-2",
-            "after-cancel",
-        ]
-        assert aborted_request_batches == [
-            ["cancel-0"],
-            ["cancel-1"],
-            ["cancel-2"],
-        ]
-
-    asyncio.run(run_test())

tests/entrypoints/test_utils.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass
 import pytest
 import torch
 from pytest_mock import MockerFixture
+from vllm.sampling_params import RequestOutputKind, SamplingParams
 
 from vllm_omni.config.yaml_util import create_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -15,6 +16,7 @@ from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.utils import (
     _convert_dataclasses_to_dict,
     _filter_dict_like_object,
+    coerce_param_message_types,
     filter_dataclass_kwargs,
     load_and_resolve_stage_configs,
     load_stage_configs_from_yaml,
@@ -424,3 +426,44 @@ class TestLoadStageConfigsFromYaml:
 
         assert stages[0]["engine_args"]["nested"]["base"] == 1
         assert stages[0]["engine_args"]["nested"]["override"] == 2
+
+
+class TestCumulativeStreamingCoercion:
+    @pytest.mark.parametrize("skip_clone", [True, False])
+    def test_cumulative_default_becomes_delta_if_stream(self, skip_clone):
+        """Ensure cumulative messages are coercible to delta if streaming."""
+        sp = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
+        sp.skip_clone = skip_clone
+        result = coerce_param_message_types([sp], is_streaming=True)[0]
+        assert isinstance(result, SamplingParams)
+        assert result.output_kind == RequestOutputKind.DELTA
+        assert (skip_clone and sp is result) or (not skip_clone and sp is not result)
+
+    @pytest.mark.parametrize("skip_clone", [True, False])
+    def test_cumulative_default_becomes_final_only_if_not_stream(self, skip_clone):
+        """Ensure cumulative messages are coercible to final only if not streaming."""
+        sp = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
+        sp.skip_clone = skip_clone
+        result = coerce_param_message_types([sp], is_streaming=False)[0]
+        assert isinstance(result, SamplingParams)
+        assert result.output_kind == RequestOutputKind.FINAL_ONLY
+        assert (skip_clone and sp is result) or (not skip_clone and sp is not result)
+
+    @pytest.mark.parametrize("is_streaming", [True, False])
+    @pytest.mark.parametrize("output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])
+    def test_non_cumulative_are_coerced(self, output_kind, is_streaming):
+        """Ensure non-cumulative params are coerced to the target type."""
+        sp = SamplingParams(output_kind=output_kind)
+        expected = RequestOutputKind.DELTA if is_streaming else RequestOutputKind.FINAL_ONLY
+        result = coerce_param_message_types([sp], is_streaming=is_streaming)[0]
+        assert isinstance(result, SamplingParams)
+        assert result.output_kind == expected
+
+    def test_coercion_applies_to_all_stages(self):
+        """Ensure all stages are coerced to DELTA for streaming."""
+        sp0 = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
+        sp1 = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
+        result = coerce_param_message_types([sp0, sp1], is_streaming=True)
+        assert all([isinstance(r, SamplingParams) for r in result])
+        assert result[0].output_kind == RequestOutputKind.DELTA
+        assert result[1].output_kind == RequestOutputKind.DELTA

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -13,7 +13,7 @@ def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], m
     return SimpleNamespace(
         request_id=request_id,
         prompt_token_ids=prompt_ids,
-        outputs=[SimpleNamespace(token_ids=out_ids, multimodal_output=mm)],
+        outputs=[SimpleNamespace(token_ids=out_ids, cumulative_token_ids=out_ids, multimodal_output=mm)],
     )
 
 

tests/model_executor/stage_input_processors/test_glm_image.py

@@ -27,7 +27,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def _source_output(token_ids: list[int], mm_output: dict | None = None):
     """Create a minimal AR output mock."""
     return SimpleNamespace(
-        outputs=[SimpleNamespace(token_ids=token_ids)],
+        outputs=[SimpleNamespace(token_ids=token_ids, cumulative_token_ids=token_ids)],
         multimodal_output=mm_output,
     )
 

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -336,6 +336,7 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
     output = SimpleNamespace(
         multimodal_output={"codes": {"audio": audio_codes, "ref": ref_code}},
         token_ids=list(range(3)),
+        cumulative_token_ids=list(range(3)),
     )
     stage = SimpleNamespace(
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
@@ -381,6 +382,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
     output = SimpleNamespace(
         multimodal_output={"codes": {"audio": audio_codes, "ref": ref_code}},
         token_ids=list(range(4)),
+        cumulative_token_ids=list(range(4)),
     )
     stage = SimpleNamespace(
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],

vllm_omni/engine/output_modality.py

@@ -8,7 +8,7 @@ for type-safe multimodal output routing and tensor merging.
 from __future__ import annotations
 
 import re
-from enum import Enum, Flag, auto
+from enum import Enum, Flag, StrEnum, auto
 
 _MODALITY_ALIASES: dict[str, str] = {
     "speech": "audio",
@@ -21,6 +21,27 @@ _MODALITY_ALIASES: dict[str, str] = {
 }
 
 
+class OutputModalityNames(StrEnum):
+    """Keys for output modalities.
+
+    TODO: (Alex) Integrate this with the big-flag enum below + throughout the code
+    for better type safety (currently only used for output processor).
+    """
+
+    TEXT = "text"
+    IMAGE = "image"
+    AUDIO = "audio"
+    LATENT = "latent"
+
+
+# Specify which output modalities may be drained when handling delta messages.
+# For some types, e.g., latents, we need to be careful to ensure the full context
+# is passed as the stream yields due to assumptions in the I/O processing and model
+# when async chunk isn't enabled.
+NON_DRAINABLE_MODALITIES = {OutputModalityNames.TEXT, OutputModalityNames.LATENT}
+DRAINABLE_MODALITIES = {mod for mod in OutputModalityNames if mod not in NON_DRAINABLE_MODALITIES}
+
+
 class OutputModality(Flag):
     """Bit-flag enum for output modalities.
 

vllm_omni/engine/output_processor.py

@@ -17,6 +17,7 @@ from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.data_entry_keys import unflatten_payload
+from vllm_omni.engine.output_modality import DRAINABLE_MODALITIES
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -37,25 +38,29 @@ class OmniRequestState(RequestState):
     ):
         super().__init__(*args, **kwargs)
         # Omni-specific: multimodal output accumulation
-        self.mm_type: str | None = None
-        self.mm_accumulated: dict[str, Any] | None = None
+        # 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] = {}
+
+    @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()
+            except Exception:
+                return x
+        return x
 
     def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None:
         if payload is None:
             return
+
+        mm_type = (mm_type or "").lower()
         try:
-            if mm_type:
-                self.mm_type = (mm_type or "").lower()
-
-            # Normalize incoming payload to dict on CPU
-            def _to_cpu(x):
-                if isinstance(x, torch.Tensor):
-                    try:
-                        return x.detach().to("cpu", non_blocking=True).contiguous()
-                    except Exception:
-                        return x
-                return x
-
             if isinstance(payload, dict):
                 # Keep payload flat (dotted keys like "hidden_states.layer_0")
                 # during accumulation so that all values are tensors/scalars and
@@ -63,7 +68,8 @@ class OmniRequestState(RequestState):
                 # later in _consolidate_multimodal_tensors after concatenation.
 
                 incoming: dict[str, Any] = {}
-                target_key = self.mm_type or "hidden"
+                # TODO (Alex): Clean up and simplify key management
+                target_key = mm_type or "hidden"
 
                 for k, v in payload.items():
                     if k == "model_outputs":
@@ -71,12 +77,12 @@ class OmniRequestState(RequestState):
                     elif k == "hidden" and target_key != "hidden":
                         k = target_key
 
-                    incoming[k] = _to_cpu(v)
+                    incoming[k] = self._to_cpu(v)
             else:
-                key = self.mm_type or "hidden"
-                incoming = {key: _to_cpu(payload)}
+                key = mm_type or "hidden"
+                incoming = {key: self._to_cpu(payload)}
 
-            if self.mm_accumulated is None:
+            if not self.mm_accumulated:
                 self.mm_accumulated = incoming
             else:
                 for k, v in incoming.items():
@@ -105,17 +111,26 @@ class OmniRequestState(RequestState):
             logger.exception("Error accumulating multimodal tensor")
 
     def _consolidate_multimodal_tensors(self) -> None:
-        """Consolidate accumulated tensor lists into single tensors via concatenation."""
-        if self.mm_accumulated is None:
+        """Consolidate accumulated tensor lists into single tensors via concatenation.
+
+        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.
+        """
+        if not self.mm_accumulated:
             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
                 if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
                     try:
                         if k == "audio":
-                            # When the audio tensor shape is inconsistent, torch.cat will fail.
-                            # We need to use torch.cat in -1 dimension.
-                            continue
+                            # Audio chunks are usually 2D (i.e., 1, N); concatenate
+                            # on the last axis to preserve the channel dimension.
+                            self.mm_accumulated[k] = torch.cat(v, dim=-1)
                         elif k == "sr":
                             # Sample rate is a constant scalar, keep last value.
                             self.mm_accumulated[k] = v[-1]
@@ -182,13 +197,13 @@ class OmniRequestState(RequestState):
             )
 
         finished = finish_reason is not None
-        final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
+        is_final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
+        is_delta = self.output_kind == RequestOutputKind.DELTA
 
-        if not finished and final_only:
+        if not finished and is_final_only:
             return None
 
-        # Consolidate accumulated tensors when finishing.
-        if finished:
+        if finished or not is_delta:
             self._consolidate_multimodal_tensors()
 
         if self.stream_interval > 1:
@@ -205,7 +220,7 @@ class OmniRequestState(RequestState):
             ):
                 return None
 
-            if self.output_kind == RequestOutputKind.DELTA:
+            if is_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 :]
@@ -233,18 +248,26 @@ class OmniRequestState(RequestState):
     ) -> Any:
         # Reuse base text/logprobs logic, then annotate with pooling_result.
         base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason, routed_experts)
-        try:
-            if not hasattr(base_output, "multimodal_output"):
-                setattr(base_output, "multimodal_output", {})
-            if self.mm_accumulated is not None:
-                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)
-        except Exception:
-            logger.exception("Error in _new_completion_output")
+
+        # 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)
+
+        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)
+
+        if self.output_kind == RequestOutputKind.DELTA:
+            for modality_key in DRAINABLE_MODALITIES:
+                self.mm_accumulated.pop(modality_key, None)
+
         return base_output
 
 

vllm_omni/engine/stage_engine_core_client.py

@@ -395,7 +395,7 @@ class StageEngineCoreClientBase:
 
         return [
             OmniTokensPrompt(
-                prompt_token_ids=so.outputs[0].token_ids,
+                prompt_token_ids=so.outputs[0].cumulative_token_ids,
                 multi_modal_data=(mm_data[so.request_id] if self.requires_multimodal_data else None),
             )
             for so in source_outputs

vllm_omni/entrypoints/async_omni.py

@@ -31,6 +31,7 @@ from vllm_omni.entrypoints.omni_base import (
     OmniBase,
     OmniEngineDeadError,
 )
+from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
@@ -40,7 +41,7 @@ if TYPE_CHECKING:
     from vllm.tokenizers import TokenizerLike
     from vllm.v1.engine import PauseMode
 
-    from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams
+    from vllm_omni.inputs.data import OmniPromptType
 
 logger = init_logger(__name__)
 _FINAL_OUTPUT_IDLE_SLEEP_S = 0.001
@@ -278,7 +279,13 @@ class AsyncOmni(EngineClient, OmniBase):
                 and not isinstance(sampling_params_list, (str, bytes))
             ):
                 sampling_params_list = self._maybe_expand_sampling_params(list(sampling_params_list))
-            sampling_params_list = self.resolve_sampling_params_list(sampling_params_list)
+
+            # Set the output kind to delta output if sampling params were omitted,
+            # since AsyncOmni is typically used for streaming.
+            sampling_params_list = self.resolve_sampling_params_list(
+                sampling_params_list,
+                allow_delta_coercion=True,
+            )
 
             # Track per-request metrics
             wall_start_ts = time.time()
@@ -366,15 +373,15 @@ class AsyncOmni(EngineClient, OmniBase):
         # only check thinker's sampling params now
         stage0_params = sampling_params_list[0]
         self._validate_streaming_input_sampling_params(stage0_params)
-
         req_state = self.request_states[request_id]
+        has_submitted_first_chunk = False
 
+        # NOTE: InputProcessor in vLLM should generally do this too, but for
+        # now we do it defensively. TODO (Alex) ensure clones/copying are optimized
         if not stage0_params.skip_clone:
             stage0_params = stage0_params.clone()
             stage0_params.skip_clone = True
 
-        has_submitted_first_chunk = False
-
         async def handle_inputs() -> None:
             nonlocal has_submitted_first_chunk
             cancelled = False

vllm_omni/entrypoints/omni_base.py

@@ -16,7 +16,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
-from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e
+from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.outputs import OmniRequestOutput
@@ -232,9 +232,14 @@ class OmniBase(PDDisaggregationMixin):
     def resolve_sampling_params_list(
         self,
         sampling_params_list: Sequence[Any] | Any | None,
+        allow_delta_coercion: bool = False,
     ) -> Sequence[Any]:
         if sampling_params_list is None:
             normalized = self.default_sampling_params_list
+            # Set the output kind to delta since no params were specified
+            if allow_delta_coercion:
+                normalized = coerce_param_message_types(list(normalized), is_streaming=True)
+
         elif isinstance(sampling_params_list, Sequence) and not isinstance(sampling_params_list, (str, bytes)):
             normalized = sampling_params_list
         elif self.num_stages == 1:

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -4,6 +4,7 @@ import asyncio
 import base64
 import json
 from collections.abc import AsyncGenerator
+from typing import cast
 from uuid import uuid4
 
 import numpy as np
@@ -12,6 +13,9 @@ from vllm.entrypoints.openai.realtime.connection import RealtimeConnection as Vl
 from vllm.entrypoints.openai.realtime.protocol import TranscriptionDelta, TranscriptionDone
 from vllm.logger import init_logger
 
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.utils import coerce_param_message_types
+
 logger = init_logger(__name__)
 
 
@@ -24,8 +28,7 @@ class RealtimeConnection(VllmRealtimeConnection):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        # Last audio buffer seen for this realtime generation (cumulative or concatenation
-        # of increments); used to turn server cumulative PCM into true deltas.
+        self.engine = cast(AsyncOmni, self.serving.engine_client)
         self._realtime_audio_ref: np.ndarray | None = None
 
     async def start_generation(self):
@@ -115,41 +118,48 @@ class RealtimeConnection(VllmRealtimeConnection):
         sent_audio = False
         audio_done_sent = False
         full_text = ""
-        sent_text_len = 0
         prompt_token_ids_len = 0
         completion_tokens_len = 0
         self._realtime_audio_ref = None
 
+        # Coerce cumulative outputs to delta outputs; this ensures
+        # we don't emit redundant MM data & drain after emitting.
+        sampling_params_list = list(self.engine.default_sampling_params_list)
+        sampling_params_list = coerce_param_message_types(
+            sampling_params_list,
+            is_streaming=True,
+        )
+
         try:
-            result_gen = self.serving.engine_client.generate(
+            result_gen = self.engine.generate(
                 prompt=streaming_input_gen,
                 request_id=request_id,
+                sampling_params_list=sampling_params_list,
             )
 
             async for output in result_gen:
+                # Handle delta texts; this is very similar to the client from vLLM
                 if output.outputs and len(output.outputs) > 0:
-                    output0 = output.outputs[0]
-                    token_ids = list(output0.token_ids)
-                    if token_ids:
-                        input_stream.put_nowait(token_ids)
-                        # token_ids are cumulative per request
-                        completion_tokens_len = len(token_ids)
+                    first_output = output.outputs[0]
+                    new_token_ids = list(first_output.token_ids)
+                    new_tokens_len = len(new_token_ids)
+
                     if not prompt_token_ids_len and output.prompt_token_ids:
                         prompt_token_ids_len = len(output.prompt_token_ids)
-                    cumulative_text = output0.text or ""
-                    if cumulative_text:
-                        if len(cumulative_text) >= sent_text_len:
-                            delta_text = cumulative_text[sent_text_len:]
-                        else:
-                            delta_text = cumulative_text
-                        sent_text_len = len(cumulative_text)
-                        full_text = cumulative_text
-                    else:
-                        delta_text = ""
 
+                    if new_tokens_len:
+                        input_stream.put_nowait(new_token_ids)
+
+                    delta_text = first_output.text or ""
+                    full_text += delta_text
+
+                    # append output to input if there was any delta text
                     if delta_text:
                         await self.send(TranscriptionDelta(delta=delta_text))
 
+                    completion_tokens_len += new_tokens_len
+
+                # Handle audio chunking; this is Omni specific
                 audio_chunks, sample_rate = self._extract_audio_chunks(output)
 
                 for chunk in audio_chunks:

vllm_omni/entrypoints/openai/serving_chat.py

@@ -16,6 +16,7 @@ from pydantic import TypeAdapter
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse
+from vllm_omni.entrypoints.utils import coerce_param_message_types
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 
 try:
@@ -391,6 +392,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     # Use standard OpenAI API parameters for comprehension stage
                     sampling_params_list = self._build_sampling_params_list_from_request(request)
 
+                # If this is a streaming (output) request, coerce cumulative outputs
+                # to delta to ensure emitted outputs are correctly drained. Otherwise
+                # convert cumulative to Final Only to ensure the output is correct.
+                sampling_params_list = coerce_param_message_types(sampling_params_list, request.stream)
+
                 # Apply user-specified overrides to diffusion stage(s) for image generation
                 for idx, sp in enumerate(sampling_params_list):
                     if hasattr(sp, "height") and _image_gen_height is not None:
@@ -1007,12 +1013,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                             cur_channel = harmony_parser.current_channel
                             cur_recipient = harmony_parser.current_recipient
                         else:
-                            # output.text is cumulative, extract only the delta portion
-                            previous_text = previous_texts[i] if previous_texts else ""
-                            if output.text is not None:
-                                delta_text = output.text[len(previous_text) :]
-                            else:
-                                delta_text = ""
+                            delta_text = output.text or ""
 
                         if not delta_text and not output.token_ids and not previous_num_tokens[i]:
                             # Chunked prefill case, don't return empty chunks

vllm_omni/entrypoints/openai/serving_speech.py

@@ -40,6 +40,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
     SpeechBatchItem,
     SpeechBatchItemResult,
 )
+from vllm_omni.entrypoints.utils import coerce_param_message_types
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     build_fish_text_only_prompt_ids,
     estimate_fish_voice_clone_prompt_len_from_normalized,
@@ -1628,6 +1629,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if self.engine_client.errored:
             raise self.engine_client.dead_error
 
+        # If this is a streaming request, we need to coerce
+        # cumulative outputs to delta outputs; this ensures
+        # we don't emit redundant MM data & drain after emitting.
+        # list() makes a copy to avoid mutating the params.
+        sampling_params_list = list(self.engine_client.default_sampling_params_list)
+        sampling_params_list = coerce_param_message_types(sampling_params_list, request.stream)
+
         if self._is_fish_speech:
             validation_error = self._validate_fish_tts_request(request)
             if validation_error:
@@ -1743,8 +1751,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             model_type,
         )
 
-        sampling_params_list = self.engine_client.default_sampling_params_list
-
         # CosyVoice3: set dynamic min/max tokens based on text length.
         # The official model requires min_token_text_ratio to prevent early
         # EOS and max_token_text_ratio to cap generation length.

vllm_omni/entrypoints/utils.py

@@ -8,12 +8,14 @@ from typing import Any, get_args, get_origin
 
 import yaml
 from vllm.logger import init_logger
+from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
 from vllm.transformers_utils.repo_utils import file_or_path_exists
 
 from vllm_omni.config.stage_config import StageConfigFactory
 from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs
 from vllm_omni.entrypoints.stage_utils import _to_dict
+from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
 # Get the project root directory (2 levels up from this file)
@@ -850,3 +852,40 @@ def detect_pid_host() -> bool:
         return True
 
     return has_pid_host()
+
+
+### Helpers for handling delta messages
+def coerce_param_message_types(params: list[OmniSamplingParams], is_streaming: bool):
+    """Iterate over the sampling params and convert to the message types
+    to DELTA messages, if streaming is enabled, or FINAL_ONLY if
+    it's disabled, while respecting `.skip_clone` on the params.
+
+    This is needed to avoid emitting redundant multimodal data.
+    """
+    # Coerce vLLM's default output kinds as needed to handle streaming
+    # (i.e., DELTA output kind). Note that this is only applied to non
+    # Diffusion sampling params.
+    #
+    # NOTE: Hidden states will still be passed between stages.
+    for idx, sp in enumerate(params):
+        # For OmniDiffusionParams don't set output kind
+        if isinstance(sp, SamplingParams):
+            params[idx] = maybe_coerce_to_message_type(sp, is_streaming)
+    return params
+
+
+def maybe_coerce_to_message_type(params: SamplingParams, is_streaming: bool):
+    """If this is a CUMULATIVE message, coerce it to DELTA if streaming, otherwise FINAL_ONLY."""
+    target_type = RequestOutputKind.DELTA if is_streaming else RequestOutputKind.FINAL_ONLY
+    if params.output_kind == target_type:
+        return params
+    elif is_streaming and params.output_kind == RequestOutputKind.FINAL_ONLY:
+        logger.warning("Request appears to be streaming, but got request type final only!")
+    elif not is_streaming and params.output_kind == RequestOutputKind.DELTA:
+        logger.warning("Request appears to not be streaming, but got request type delta!")
+
+    if not params.skip_clone:
+        params = params.clone()
+        params.skip_clone = True
+    params.output_kind = target_type
+    return params

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -79,7 +79,7 @@ def text2flow(
         if multi_modal_data is None:
             raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
 
-        output_ids = _ensure_list(output.token_ids)
+        output_ids = _ensure_list(output.cumulative_token_ids)
         prefix_ids = _ensure_list(source_output.prompt_token_ids)
         additional_info = dict(multi_modal_data)
         additional_info.setdefault("ids", {})["prompt"] = prefix_ids

vllm_omni/model_executor/stage_input_processors/dynin_omni.py

@@ -117,7 +117,7 @@ def _bridge_tokens(
         if not token_ids:
             token_ids = _to_token_id_list(mm_out.get("text_tokens"))
         if not token_ids:
-            token_ids = list(getattr(output, "token_ids", []) or [])
+            token_ids = list(output.cumulative_token_ids or [])
         if not token_ids:
             raise RuntimeError(
                 f"Stage {source_stage_id} output for request {source_output.request_id} has no token_ids"

vllm_omni/model_executor/stage_input_processors/glm_image.py

@@ -240,7 +240,7 @@ def ar2diffusion(
     for i, ar_output in enumerate(ar_outputs):
         _t_req = time.perf_counter()
         output = ar_output.outputs[0]
-        generated_token_ids = output.token_ids
+        generated_token_ids = output.cumulative_token_ids
 
         # Get original prompt info
         original_prompt = prompt[i] if i < len(prompt) else {}

vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py

@@ -57,7 +57,7 @@ def ar2diffusion(
 
     for i, ar_output in enumerate(ar_outputs):
         output = ar_output.outputs[0]
-        generated_token_ids = output.token_ids
+        generated_token_ids = output.cumulative_token_ids
         generated_text = getattr(output, "text", "") or ""
 
         # Get original prompt info

vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py

@@ -34,7 +34,7 @@ def ar2dit(
         prompt_token_ids = ar_output.prompt_token_ids
         # exclude the last token because it has no corresponding hidden state
         completion_output = ar_output.outputs[0]
-        gen_token_ids = completion_output.token_ids[:-1]
+        gen_token_ids = completion_output.cumulative_token_ids[:-1]
         full_token_ids = prompt_token_ids + gen_token_ids
 
         mm_output = getattr(completion_output, "multimodal_output", None)

vllm_omni/model_executor/stage_input_processors/omnivoice.py

@@ -35,7 +35,7 @@ def tokens2audio(
 
     # Pass audio_tokens from generator to decoder
     engine_input = OmniTokensPrompt(
-        prompt_token_ids=output.token_ids,
+        prompt_token_ids=output.cumulative_token_ids,
         additional_information=multi_modal_data,
     )
     return [engine_input]

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -33,7 +33,7 @@ def thinker2talker(
     for i, thinker_output in enumerate(thinker_outputs):
         output = thinker_output.outputs[0]
         prompt_token_ids = thinker_output.prompt_token_ids
-        thinker_output_ids = output.token_ids
+        thinker_output_ids = output.cumulative_token_ids
         prompt_token_ids_len = len(prompt_token_ids)
         mm: OmniPayload = output.multimodal_output
         latent = mm["latent"]

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -416,7 +416,7 @@ def thinker2talker(
         output = thinker_output.outputs[0]
         req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
         prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
-        output_ids = _ensure_list(output.token_ids)
+        output_ids = _ensure_list(output.cumulative_token_ids)
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
         if is_streaming_session:
             prompt_token_ids, output_ids, thinker_sequences, thinker_input_ids = _get_streaming_talker_tokens(
@@ -595,7 +595,7 @@ def talker2code2wav(
     for i, talker_output in enumerate(talker_outputs):
         output = talker_output.outputs[0]
         req_id = str(getattr(talker_output, "request_id", f"idx-{i}"))
-        cur_seq_len = len(output.token_ids) - 1
+        cur_seq_len = len(output.cumulative_token_ids) - 1
         seq_len = cur_seq_len
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
         if is_streaming_session:

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -39,9 +39,11 @@ def talker2code2wav(
         output = talker_output.outputs[0]
         mm = output.multimodal_output
         mm_codes = mm.get("codes", {})
+
         # audio_codes shape: [num_frames, Q] where Q=num_quantizers (16)
         audio_codes = mm_codes["audio"].to(torch.long)
-        token_ids = output.token_ids
+        token_ids = output.cumulative_token_ids
+
         # token_ids provides an upper bound on the newly generated codec span.
         # audio_codes may still contain zero-padded / invalid rows, so trim only
         # after filtering valid frames instead of trying to align EOS indices.