[Bugfix][Qwen3-TTS] Gate async-chunk Code2Wav on non_streaming_mode (#4371) (#6898)

Signed-off-by: Zhichao Zhang <60429419+smartDream-chao@users.noreply.github.com>

examples/offline_inference/text_to_speech/qwen3_tts/end2end.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Offline inference demo for Qwen3 TTS via vLLM Omni.
 
 Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base
@@ -166,6 +169,9 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
                 "language": [language],
                 "speaker": [speaker],
                 "max_new_tokens": [2048],
+                # Optional Code2Wav emit opt-in under deploy.async_chunk.
+                # Independent of prompt-mode non_streaming_mode (#4198 / #6898).
+                "full_utterance_decode": [True],
             }
             inputs.append(
                 {
@@ -185,6 +191,7 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
             "speaker": [speaker],
             "instruct": [instruct],
             "max_new_tokens": [2048],
+            "full_utterance_decode": [True],
         }
         inputs = {
             "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
@@ -226,6 +233,7 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
                 "instruct": [instruct],
                 "max_new_tokens": [2048],
                 "non_streaming_mode": [True],
+                "full_utterance_decode": [True],
             }
             inputs.append(
                 {
@@ -244,6 +252,7 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
             "instruct": [instruct],
             "max_new_tokens": [2048],
             "non_streaming_mode": [True],
+            "full_utterance_decode": [True],
         }
         inputs = {
             "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
@@ -291,6 +300,7 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que
                 "language": [language],
                 "x_vector_only_mode": [x_vector_only_mode],
                 "max_new_tokens": [2048],
+                "full_utterance_decode": [True],
             }
             inputs.append(
                 {
@@ -307,6 +317,7 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que
             "language": [syn_lang_single],
             "x_vector_only_mode": [x_vector_only_mode],
             "max_new_tokens": [2048],
+            "full_utterance_decode": [True],
         }
         inputs = {
             "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),

tests/entrypoints/openai_api/test_serving_speech.py

@@ -2241,6 +2241,7 @@ class TestTTSMethods:
 
         assert params["task_type"] == ["Base"]
         assert params["non_streaming_mode"] == [True]
+        assert "full_utterance_decode" not in params
 
     def test_build_tts_params_base_omits_non_streaming_mode_by_default(self, speech_server):
         """Base task should keep using the model default when no override is sent."""
@@ -2255,6 +2256,7 @@ class TestTTSMethods:
 
         assert params["task_type"] == ["Base"]
         assert "non_streaming_mode" not in params
+        assert "full_utterance_decode" not in params
 
     def test_build_tts_params_explicit_non_streaming_mode_overrides_voicedesign_default(self, speech_server):
         """Explicit false should not be replaced by the VoiceDesign fallback."""
@@ -2269,6 +2271,38 @@ class TestTTSMethods:
 
         assert params["task_type"] == ["VoiceDesign"]
         assert params["non_streaming_mode"] == [False]
+        assert "full_utterance_decode" not in params
+
+    def test_build_tts_params_streaming_voicedesign_keeps_prompt_mode_not_full_decode(self, speech_server):
+        """Streaming VoiceDesign defaults: prompt-mode True, no full_utterance_decode."""
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            task_type="VoiceDesign",
+            instructions="warm and calm",
+            stream=True,
+            response_format="pcm",
+        )
+
+        params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["VoiceDesign"]
+        assert params["non_streaming_mode"] == [True]
+        assert "full_utterance_decode" not in params
+
+    def test_build_tts_params_streaming_customvoice_explicit_non_streaming_mode_still_windowed(self, speech_server):
+        """Explicit non_streaming_mode=True must not inject full_utterance_decode."""
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            voice="Vivian",
+            non_streaming_mode=True,
+            stream=True,
+            response_format="pcm",
+        )
+
+        params = speech_server._build_tts_params(req)
+
+        assert params["non_streaming_mode"] == [True]
+        assert "full_utterance_decode" not in params
 
     def test_load_supported_speakers(self, mocker: MockerFixture):
         """Test _load_supported_speakers."""
@@ -5287,9 +5321,34 @@ class TestTTSAsyncOffloading:
 
         assert tts_params["task_type"] == ["VoiceDesign"]
         assert tts_params["non_streaming_mode"] == [False]
+        assert "full_utterance_decode" not in tts_params
         prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
         assert prompt["additional_information"] is tts_params
         assert prompt["additional_information"]["non_streaming_mode"] == [False]
+        assert "full_utterance_decode" not in prompt["additional_information"]
+
+    def test_prepare_speech_generation_qwen3_streaming_voicedesign_default_params(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """Streaming VoiceDesign defaults keep prompt-mode True without full_utterance_decode."""
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+
+        request = OpenAICreateSpeechRequest(
+            input="hello",
+            task_type="VoiceDesign",
+            instructions="warm and calm",
+            stream=True,
+            response_format="pcm",
+        )
+        _request_id, _generator, tts_params = asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        assert tts_params["task_type"] == ["VoiceDesign"]
+        assert tts_params["non_streaming_mode"] == [True]
+        assert "full_utterance_decode" not in tts_params
+        prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
+        assert prompt["additional_information"]["non_streaming_mode"] == [True]
+        assert "full_utterance_decode" not in prompt["additional_information"]
 
     def test_prepare_speech_generation_qwen3_base_non_streaming_mode_true(
         self, qwen3_tts_server, mocker: MockerFixture
@@ -5312,9 +5371,11 @@ class TestTTSAsyncOffloading:
         assert tts_params["task_type"] == ["Base"]
         assert tts_params["ref_text"] == ["reference transcript"]
         assert tts_params["non_streaming_mode"] == [True]
+        assert "full_utterance_decode" not in tts_params
         prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
         assert prompt["additional_information"] is tts_params
         assert prompt["additional_information"]["non_streaming_mode"] == [True]
+        assert "full_utterance_decode" not in prompt["additional_information"]
 
     def test_qwen3_repeated_ref_audio_hot_path_sends_cache_key_without_waveform(self, qwen3_tts_server):
         """After a ref artifact is marked ready, repeated requests avoid ref_audio payload IPC."""

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from collections import defaultdict
 from types import SimpleNamespace
@@ -32,11 +32,15 @@ _FRAME = [1, 2, 3, 4]
 _Q = len(_FRAME)
 
 
-def _req(rid, *, finished, initial_codec_chunk_frames=None):
-    ai = None
+def _req(rid, *, finished, initial_codec_chunk_frames=None, non_streaming_mode=None, full_utterance_decode=None):
+    entries = {}
     if initial_codec_chunk_frames is not None:
-        entry = SimpleNamespace(list_data=[initial_codec_chunk_frames])
-        ai = SimpleNamespace(entries={"initial_codec_chunk_frames": entry})
+        entries["initial_codec_chunk_frames"] = SimpleNamespace(list_data=[initial_codec_chunk_frames])
+    if non_streaming_mode is not None:
+        entries["non_streaming_mode"] = SimpleNamespace(list_data=[non_streaming_mode])
+    if full_utterance_decode is not None:
+        entries["full_utterance_decode"] = SimpleNamespace(list_data=[full_utterance_decode])
+    ai = SimpleNamespace(entries=entries) if entries else None
     return SimpleNamespace(
         external_req_id=rid,
         is_finished=lambda: finished,
@@ -123,6 +127,85 @@ def test_flush_on_finish():
     assert len(p.codes.audio) == _Q * 24
 
 
+def test_full_utterance_decode_defers_until_finished():
+    """full_utterance_decode=True should not emit windowed chunks mid-utterance (#4371)."""
+    tm = _tm()
+    tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(50)]
+    p = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        multimodal_output=None,
+        request=_req("r", finished=False, full_utterance_decode=True),
+        is_finished=False,
+    )
+    assert p is None
+
+
+def test_full_utterance_decode_emits_full_utterance_on_finish():
+    tm = _tm()
+    tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(50)]
+    p = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        multimodal_output=None,
+        request=_req("r", finished=True, full_utterance_decode=True),
+        is_finished=True,
+    )
+    assert p is not None
+    assert p.meta.finished.item() is True
+    assert len(p.codes.audio) == _Q * 50
+
+
+def test_full_utterance_decode_false_keeps_windowed_emit():
+    # Dynamic IC defers emit at 25 frames (see _CASES); 41 hits the first
+    # steady-state window boundary and proves streaming path is unchanged.
+    tm = _tm()
+    tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(41)]
+    p = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        multimodal_output={"codes": {"audio": torch.zeros((0,))}},
+        request=_req("r", finished=False, full_utterance_decode=False),
+        is_finished=False,
+    )
+    assert p is not None
+    assert len(p.codes.audio) == _Q * 25
+
+
+def test_non_streaming_mode_alone_does_not_defer_async_chunk():
+    """Prompt-mode non_streaming_mode must not gate Code2Wav emits (#6898).
+
+    VoiceDesign serving defaults non_streaming_mode=True even for streaming
+    responses; those requests still need windowed async-chunk audio for TTFA.
+    """
+    tm = _tm()
+    tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(41)]
+    p = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        multimodal_output={"codes": {"audio": torch.zeros((0,))}},
+        request=_req("r", finished=False, non_streaming_mode=True),
+        is_finished=False,
+    )
+    assert p is not None
+    assert len(p.codes.audio) == _Q * 25
+
+
+def test_voicedesign_prompt_mode_with_streaming_keeps_windowed_emit():
+    """Streaming VoiceDesign-style defaults: prompt-mode True, decode False."""
+    tm = _tm()
+    tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(41)]
+    p = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        multimodal_output={"codes": {"audio": torch.zeros((0,))}},
+        request=_req(
+            "r",
+            finished=False,
+            non_streaming_mode=True,
+            full_utterance_decode=False,
+        ),
+        is_finished=False,
+    )
+    assert p is not None
+    assert len(p.codes.audio) == _Q * 25
+
+
 _CASES = [
     # ── IC boundary rule ──────────────────────────────────────────────
     # initial_codec_chunk_frames only controls the first emitted chunk.

vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py

@@ -374,6 +374,12 @@ class Qwen3TTSAdapter(ARTTSAdapter):
         elif params["task_type"][0] == "VoiceDesign":
             params["non_streaming_mode"] = [True]
 
+        # Do not set ``full_utterance_decode`` here. That flag is an explicit
+        # Code2Wav emit opt-in (offline / advanced callers via
+        # additional_information) and must stay independent of prompt-mode
+        # ``non_streaming_mode`` (#4198 / #6898). Online serving keeps the
+        # default windowed async-chunk path (incl. #5202 incremental decode).
+
         return params
 
     def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int:

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Stage input processor for Qwen3-TTS: Talker -> Code2Wav."""
 
 import time
@@ -24,6 +27,7 @@ from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
     parse_chunk_ramp,
 )
 from vllm_omni.model_executor.stage_input_processors.tts_utils import (
+    extract_full_utterance_decode_from_request,
     extract_language_from_prompt,
     extract_language_from_request,
     extract_speaker_from_prompt,
@@ -180,6 +184,15 @@ def talker2code2wav_async_chunk(
         initial_chunk_size = chunk_size
     length = len(transfer_manager.code_prompt_token_ids[request_id])
 
+    # Full-utterance Code2Wav is gated by ``full_utterance_decode`` (explicit
+    # opt-in via additional_information). Do not reuse ``non_streaming_mode``:
+    # that flag is prompt construction only and VoiceDesign defaults it to True
+    # even for streaming responses (#4198, #6898 review). When True, defer
+    # connector chunks until the talker finishes so Code2Wav decodes once.
+    full_utterance_decode = extract_full_utterance_decode_from_request(request)
+    if full_utterance_decode is True and not finished:
+        return None
+
     if length <= 0:
         if finished:
             return OmniPayloadStruct(
@@ -192,7 +205,10 @@ def talker2code2wav_async_chunk(
             )
         return None
 
-    if adaptive_enabled:
+    if full_utterance_decode is True:
+        first_chunk = int(transfer_manager.put_req_chunk.get(request_id, 0)) <= 0
+        context_length = length
+    elif adaptive_enabled:
         _adaptive_states = getattr(transfer_manager, "_adaptive_states", None)
         if _adaptive_states is None:
             _adaptive_states = {}

vllm_omni/model_executor/stage_input_processors/tts_utils.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The Qwen team.
 """Shared TTS utility functions for speaker and language extraction.
 
@@ -103,6 +103,37 @@ def extract_speaker_from_prompt(
     return None
 
 
+# =============================================================================
+# Code2Wav emit-mode helpers (independent of prompt construction)
+# =============================================================================
+
+
+def extract_full_utterance_decode_from_request(request: Any) -> bool | None:
+    """Extract ``full_utterance_decode`` from a request's additional_information.
+
+    This is a *response / Code2Wav emit* intent flag. It must not be confused
+    with ``non_streaming_mode``, which only controls Qwen3-TTS prompt
+    construction and is independent of HTTP/WebSocket streaming (#4198).
+
+    Returns:
+        True/False when explicitly set, otherwise None.
+    """
+    additional_information = getattr(request, "additional_information", None)
+    if additional_information is None:
+        return None
+    entries = getattr(additional_information, "entries", None)
+    if not isinstance(entries, dict):
+        return None
+    entry = entries.get("full_utterance_decode")
+    if entry is None:
+        return None
+    list_data = getattr(entry, "list_data", None)
+    if isinstance(list_data, list) and list_data:
+        raw = list_data[0]
+        return raw if isinstance(raw, bool) else None
+    return None
+
+
 # =============================================================================
 # Language helpers
 # =============================================================================