Feat: non_streaming_mode for Qwen3-TTS Base Models During Online Inference (#4198)

Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Sy03 <1370724210@qq.com>

docs/serving/speech_api.md

@@ -114,6 +114,7 @@ Content-Type: application/json
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
 | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
+| `non_streaming_mode` | bool | null | Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. |
 | `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
 
 **Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
@@ -332,7 +333,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \
         "input": "Hello, this is a cloned voice",
         "task_type": "Base",
         "ref_audio": "https://example.com/reference.wav",
-        "ref_text": "Original transcript of the reference audio"
+        "ref_text": "Original transcript of the reference audio",
+        "non_streaming_mode": true
     }' --output cloned.wav
 ```
 
@@ -454,6 +456,7 @@ Content-Type: application/json
 | `ref_audio` | string | null | Default reference audio (Base task) |
 | `ref_text` | string | null | Default reference transcript (Base task) |
 | `max_new_tokens` | integer | null | Default max tokens |
+| `non_streaming_mode` | bool | null | Default Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. |
 
 Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:
 
@@ -469,6 +472,7 @@ Each item in the `items` array requires only `input` (the text). All other field
 | `ref_audio` | string | Override reference audio |
 | `ref_text` | string | Override reference transcript |
 | `max_new_tokens` | integer | Override max tokens |
+| `non_streaming_mode` | bool | Override Qwen3-TTS prompt construction mode. Does not affect HTTP response streaming or async-chunk pipelining. When null, inherits the batch-level value (then the model default). |
 
 ### Response Format
 

tests/entrypoints/openai_api/test_serving_speech.py

@@ -31,6 +31,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
     CreateAudio,
     OpenAICreateSpeechRequest,
     SpeechBatchItem,
+    StreamingSpeechSessionConfig,
 )
 from vllm_omni.entrypoints.openai.serving_speech import (
     _TTS_LANGUAGES,
@@ -1235,6 +1236,49 @@ class TestTTSMethods:
         assert params["language"] == ["English"]
         assert params["task_type"] == ["CustomVoice"]
 
+    def test_build_tts_params_base_non_streaming_mode_true(self, speech_server):
+        """Base task should pass through an explicit non_streaming_mode override."""
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            task_type="Base",
+            ref_audio="data:audio/wav;base64,abc",
+            ref_text="reference",
+            non_streaming_mode=True,
+        )
+
+        params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["Base"]
+        assert params["non_streaming_mode"] == [True]
+
+    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."""
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            task_type="Base",
+            ref_audio="data:audio/wav;base64,abc",
+            ref_text="reference",
+        )
+
+        params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["Base"]
+        assert "non_streaming_mode" 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."""
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            task_type="VoiceDesign",
+            instructions="warm and calm",
+            non_streaming_mode=False,
+        )
+
+        params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["VoiceDesign"]
+        assert params["non_streaming_mode"] == [False]
+
     def test_load_supported_speakers(self, mocker: MockerFixture):
         """Test _load_supported_speakers."""
         mock_engine_client = mocker.MagicMock()
@@ -2023,6 +2067,34 @@ class TestMergeBatchItem:
         assert merged.task_type == "CustomVoice"
         assert merged.max_new_tokens == 512
 
+    def test_non_streaming_mode_batch_default_used(self):
+        """Batch-level non_streaming_mode should be used when item doesn't specify one."""
+        batch = BatchSpeechRequest(
+            items=[SpeechBatchItem(input="hi")],
+            non_streaming_mode=True,
+        )
+
+        merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
+
+        assert merged.non_streaming_mode is True
+
+    def test_non_streaming_mode_item_override_wins(self):
+        """Per-item false should override a true batch-level default."""
+        batch = BatchSpeechRequest(
+            items=[SpeechBatchItem(input="hi", non_streaming_mode=False)],
+            non_streaming_mode=True,
+        )
+
+        merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
+
+        assert merged.non_streaming_mode is False
+
+
+def test_streaming_speech_session_config_accepts_non_streaming_mode():
+    config = StreamingSpeechSessionConfig(non_streaming_mode=True)
+
+    assert config.non_streaming_mode is True
+
 
 class TestAsyncOmniSupportedTasks:
     """Test that AsyncOmni reports correct supported tasks based on output modalities."""
@@ -2811,6 +2883,52 @@ class TestTTSAsyncOffloading:
         qwen3_tts_server._build_tts_params.assert_called_once()
         qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
 
+    def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """VoiceDesign explicit false should reach the model prompt additional_information."""
+        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",
+            non_streaming_mode=False,
+        )
+        _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"] == [False]
+        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]
+
+    def test_prepare_speech_generation_qwen3_base_non_streaming_mode_true(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """Base explicit true should reach the model prompt additional_information."""
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000))
+        qwen3_tts_server._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+
+        request = OpenAICreateSpeechRequest(
+            input="hello",
+            task_type="Base",
+            ref_audio="data:audio/wav;base64,abc",
+            ref_text="reference transcript",
+            non_streaming_mode=True,
+        )
+        _request_id, _generator, tts_params = asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        assert tts_params["task_type"] == ["Base"]
+        assert tts_params["ref_text"] == ["reference transcript"]
+        assert tts_params["non_streaming_mode"] == [True]
+        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]
+
     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."""
         wav_list = [0.0] * 48000

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -133,6 +133,14 @@ class OpenAICreateSpeechRequest(BaseModel):
         ge=0,
         description="Per-request initial chunk size override. If null, computed dynamically based on server load.",
     )
+    non_streaming_mode: bool | None = Field(
+        default=None,
+        description=(
+            "Qwen3-TTS prompt construction mode override. This does not "
+            "control HTTP response streaming or async-chunk pipelining. "
+            "When null, use model defaults: Base=False, CustomVoice/VoiceDesign=True."
+        ),
+    )
     extra_params: dict[str, Any] | None = Field(
         default=None,
         description=("Optional model-specific parameters passed directly to the model's extra_args."),
@@ -370,6 +378,7 @@ class SpeechBatchItem(BaseModel):
     x_vector_only_mode: bool | None = None
     max_new_tokens: int | None = None
     initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
+    non_streaming_mode: bool | None = None
 
 
 class BatchSpeechRequest(BaseModel):
@@ -389,6 +398,7 @@ class BatchSpeechRequest(BaseModel):
     x_vector_only_mode: bool | None = None
     max_new_tokens: int | None = None
     initial_codec_chunk_frames: int | None = Field(default=None, ge=0)
+    non_streaming_mode: bool | None = None
 
 
 class SpeechBatchItemResult(BaseModel):
@@ -423,6 +433,14 @@ class StreamingSpeechSessionConfig(BaseModel):
         ge=0,
         description="Initial chunk size for reduced TTFA. Overrides stage config for this session.",
     )
+    non_streaming_mode: bool | None = Field(
+        default=None,
+        description=(
+            "Qwen3-TTS prompt construction mode override. This does not "
+            "control WebSocket audio streaming or async-chunk pipelining. "
+            "When null, use model defaults: Base=False, CustomVoice/VoiceDesign=True."
+        ),
+    )
     ref_audio: str | None = None
     ref_text: str | None = None
     x_vector_only_mode: bool | None = None

vllm_omni/entrypoints/openai/serving_speech.py

@@ -2815,9 +2815,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if request.initial_codec_chunk_frames is not None:
             params["initial_codec_chunk_frames"] = [request.initial_codec_chunk_frames]
 
-        # VoiceDesign requires non_streaming_mode (match offline script behaviour).
-        # CustomVoice and Base rely on the model default (True and False respectively).
-        if params["task_type"][0] == "VoiceDesign":
+        if request.non_streaming_mode is not None:
+            params["non_streaming_mode"] = [request.non_streaming_mode]
+        # Preserve the legacy VoiceDesign fallback when the request omits an
+        # explicit override. CustomVoice and Base rely on model defaults
+        # (True and False respectively).
+        elif params["task_type"][0] == "VoiceDesign":
             params["non_streaming_mode"] = [True]
 
         return params
@@ -3974,6 +3977,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             x_vector_only_mode=_pick("x_vector_only_mode"),
             max_new_tokens=_pick("max_new_tokens"),
             initial_codec_chunk_frames=_pick("initial_codec_chunk_frames"),
+            non_streaming_mode=_pick("non_streaming_mode"),
         )
 
     async def create_speech_batch(

vllm_omni/entrypoints/openai/serving_speech_stream.py

@@ -221,6 +221,7 @@ class OmniStreamingSpeechHandler:
             speed=config.speed,
             max_new_tokens=config.max_new_tokens,
             initial_codec_chunk_frames=config.initial_codec_chunk_frames,
+            non_streaming_mode=config.non_streaming_mode,
             ref_audio=config.ref_audio,
             ref_text=config.ref_text,
             x_vector_only_mode=config.x_vector_only_mode,