From e03ab82dd237c9cbab71b7cce5dd9ad295927dcd Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Sat, 25 Apr 2026 21:11:53 -0600 Subject: [PATCH] [Refactor/Bugfix] Use Delta Messages for Streaming (#2911) Signed-off-by: Alex Brooks Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: Yueqian Lin Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Co-authored-by: Sy03 <1370724210@qq.com> Co-authored-by: Yueqian Lin Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../qwen3_omni/end2end_async_chunk.py | 9 +- tests/e2e/online_serving/test_voxtral_tts.py | 24 ++ tests/engine/test_output_processor.py | 211 ++++++++++++++++++ tests/entrypoints/test_async_omni.py | 134 +++++++++++ tests/entrypoints/test_async_omni_abort.py | 85 ------- tests/entrypoints/test_utils.py | 43 ++++ .../test_cosyvoice3_stage_input_processors.py | 2 +- .../stage_input_processors/test_glm_image.py | 2 +- .../test_qwen3_tts_async_chunk.py | 2 + vllm_omni/engine/output_modality.py | 23 +- vllm_omni/engine/output_processor.py | 105 +++++---- vllm_omni/engine/stage_engine_core_client.py | 2 +- vllm_omni/entrypoints/async_omni.py | 17 +- vllm_omni/entrypoints/omni_base.py | 7 +- .../entrypoints/openai/realtime_connection.py | 50 +++-- vllm_omni/entrypoints/openai/serving_chat.py | 13 +- .../entrypoints/openai/serving_speech.py | 10 +- vllm_omni/entrypoints/utils.py | 39 ++++ .../stage_input_processors/cosyvoice3.py | 2 +- .../stage_input_processors/dynin_omni.py | 2 +- .../stage_input_processors/glm_image.py | 2 +- .../stage_input_processors/hunyuan_image3.py | 2 +- .../stage_input_processors/mammoth_moda2.py | 2 +- .../stage_input_processors/omnivoice.py | 2 +- .../stage_input_processors/qwen2_5_omni.py | 2 +- .../stage_input_processors/qwen3_omni.py | 4 +- .../stage_input_processors/qwen3_tts.py | 4 +- 27 files changed, 623 insertions(+), 177 deletions(-) create mode 100644 tests/engine/test_output_processor.py create mode 100644 tests/entrypoints/test_async_omni.py delete mode 100644 tests/entrypoints/test_async_omni_abort.py diff --git a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py index f38922e94..85c2da20b 100644 --- a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py +++ b/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 diff --git a/tests/e2e/online_serving/test_voxtral_tts.py b/tests/e2e/online_serving/test_voxtral_tts.py index 649f822ba..91f62bb71 100644 --- a/tests/e2e/online_serving/test_voxtral_tts.py +++ b/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) diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py new file mode 100644 index 000000000..4576da5a1 --- /dev/null +++ b/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 diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py new file mode 100644 index 000000000..c8ef5eddf --- /dev/null +++ b/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 diff --git a/tests/entrypoints/test_async_omni_abort.py b/tests/entrypoints/test_async_omni_abort.py deleted file mode 100644 index b34652162..000000000 --- a/tests/entrypoints/test_async_omni_abort.py +++ /dev/null @@ -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()) diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 248629d51..6e52e4c6c 100644 --- a/tests/entrypoints/test_utils.py +++ b/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 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 5a58ca266..9c2366646 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 @@ -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)], ) diff --git a/tests/model_executor/stage_input_processors/test_glm_image.py b/tests/model_executor/stage_input_processors/test_glm_image.py index 94b52d420..542cdd5fb 100644 --- a/tests/model_executor/stage_input_processors/test_glm_image.py +++ b/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, ) 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 ea00917ea..07e343bf0 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 @@ -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)], diff --git a/vllm_omni/engine/output_modality.py b/vllm_omni/engine/output_modality.py index c4a1288a8..c2615ca8c 100644 --- a/vllm_omni/engine/output_modality.py +++ b/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. diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index ac18597ee..84016fdb4 100644 --- a/vllm_omni/engine/output_processor.py +++ b/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 diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 844be7194..37b0e538a 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/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 diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 119232ce7..056f56c00 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 75d3ee290..4147c8027 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/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: diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index 1d5470f56..9fc2a1ee3 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/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: diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index fd6484e6d..42d421db8 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 4946368e9..eb8a12c66 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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. diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 5541fde30..d728e7641 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index 736d3a4bd..dc1e12dfe 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/dynin_omni.py b/vllm_omni/model_executor/stage_input_processors/dynin_omni.py index 9ec849799..6ce8881a9 100644 --- a/vllm_omni/model_executor/stage_input_processors/dynin_omni.py +++ b/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" 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 b21cb29c2..828f7f218 100644 --- a/vllm_omni/model_executor/stage_input_processors/glm_image.py +++ b/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 {} 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 89a7a28f6..0c0e6d7b3 100644 --- a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py +++ b/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 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 1239aafaf..5eec3d045 100644 --- a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py +++ b/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) diff --git a/vllm_omni/model_executor/stage_input_processors/omnivoice.py b/vllm_omni/model_executor/stage_input_processors/omnivoice.py index b7f5c102e..b866e9779 100644 --- a/vllm_omni/model_executor/stage_input_processors/omnivoice.py +++ b/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] diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index c5ae43c82..072869eaf 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/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"] 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 db87465d2..8018df343 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/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: 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 a0c62cba3..95a771534 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/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.