[Bugfix][Qwen3-TTS] Handle codec generations that exhaust their token budget

Signed-off-by: Sy03 <1370724210@qq.com>

recipes/Qwen/Qwen3-TTS.md

@@ -18,11 +18,11 @@ examples in this repository.
 
 Qwen3-TTS supports three task types, each backed by a dedicated model checkpoint:
 
-| Task Type     | Model                                    | Description                                                  |
-| ------------- | ---------------------------------------- | ------------------------------------------------------------ |
-| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice`  | Predefined speaker voices with optional style/emotion control |
-| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign`  | Generate speech from a natural language voice description     |
-| `Base`        | `Qwen/Qwen3-TTS-12Hz-1.7B-Base`         | Voice cloning from reference audio + transcript              |
+| Task Type     | Model                                    | Description                                                   |
+| ------------- | ---------------------------------------- | ------------------------------------------------------------- |
+| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice`   | Predefined speaker voices with optional style/emotion control |
+| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign`   | Generate speech from a natural language voice description     |
+| `Base`        | `Qwen/Qwen3-TTS-12Hz-1.7B-Base`          | Voice cloning from reference audio + transcript               |
 
 Smaller 0.6B variants are also available for `CustomVoice` and `Base`.
 
@@ -166,6 +166,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
 - Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config.
 - Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`.
 - Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail.
+- Base codec termination: Base requests without an explicit `max_new_tokens`
+  use a text-scaled safety ceiling (at least 192 codec frames and no more than
+  the configured model limit). If the Talker reaches that ceiling without
+  codec EOS, non-streaming serving discards the incomplete audio and retries
+  once with a fresh seed; an explicit `seed` or `max_new_tokens` disables the
+  retry. SSE and WebSocket clients receive structured errors and must discard
+  previously emitted audio when the error contains `"action":"discard"`;
+  raw PCM streams terminate with a connection error after any partial bytes
+  already sent.
 - Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model.
 
 ## Hardware Support

tests/engine/test_orchestrator.py

@@ -2216,7 +2216,12 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None:
     )
     output = SimpleNamespace(
         request_id="req-stream",
-        outputs=[SimpleNamespace(cumulative_token_ids=list(range(11)))],
+        outputs=[
+            SimpleNamespace(
+                cumulative_token_ids=list(range(11)),
+                finish_reason=FinishReason.LENGTH,
+            )
+        ],
     )
 
     metrics = pool.build_stage_metrics(
@@ -2228,6 +2233,7 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None:
 
     assert metrics.num_tokens_out == 3
     assert metrics.output_unit_count == 3
+    assert metrics.finish_reason == "length"
 
 
 def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None:

tests/entrypoints/openai_api/test_serving_speech.py

@@ -45,9 +45,15 @@ from vllm_omni.entrypoints.openai.serving_speech import (
     OmniOpenAIServingSpeech,
     _create_wav_header,
 )
-from vllm_omni.entrypoints.openai.tts_adapters.base import DEFAULT_TTS_LANGUAGES, PreparedRequest, SpeechServingContext
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    DEFAULT_TTS_LANGUAGES,
+    PreparedRequest,
+    SpeechServingContext,
+    TTSGenerationError,
+)
 from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers
 from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
+from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSCodecLimitError
 from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     FISH_TEXT_ONLY_SYSTEM_PROMPT,
@@ -2928,6 +2934,40 @@ class TestStreamingResponse:
         assert "text/event-stream" not in response.headers["content-type"]
         assert len(response.content) > 0
 
+    def test_raw_stream_passes_tts_params_to_common_guard(self, streaming_app, monkeypatch):
+        """The raw route must not bypass the shared Qwen3-TTS EOS guard."""
+        speech_server = streaming_app.state.speech_server
+        finalized_tts_params = {"_qwen3_tts_effective_max_tokens": [192]}
+        captured: dict = {}
+
+        async def prepare(_request, request_id=None):
+            return request_id, object(), finalized_tts_params
+
+        async def generate_chunks(
+            _generator,
+            _request_id,
+            _response_format="pcm",
+            raw_request=None,
+            request_start_s=None,
+            include_sample_rate=False,
+            usage_acc=None,
+            tts_params=None,
+            collect=None,
+        ):
+            captured["tts_params"] = tts_params
+            yield b"\x00\x00"
+
+        monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare)
+        monkeypatch.setattr(speech_server, "_generate_audio_chunks", generate_chunks)
+
+        response = TestClient(streaming_app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
+        )
+
+        assert response.status_code == 200
+        assert captured["tts_params"] is finalized_tts_params
+
     def test_sse_streaming(self, streaming_app):
         """stream_format=sse without stream=True returns audio deltas as SSE."""
         client = TestClient(streaming_app)
@@ -4848,6 +4888,57 @@ class TestTTSAsyncOffloading:
         assert ("artifact-a", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
         assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()}
 
+    @pytest.mark.asyncio
+    async def test_qwen3_tts_nonstream_retries_codec_limit_once(self, qwen3_tts_server, mocker: MockerFixture) -> None:
+        qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+        qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(
+            side_effect=[
+                Qwen3TTSCodecLimitError("codec limit"),
+                (b"RIFF" + b"\x00" * 32, "audio/wav"),
+            ]
+        )
+        request = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
+
+        response = await qwen3_tts_server.create_speech(request)
+
+        assert response.status_code == 200
+        assert qwen3_tts_server._generate_audio_bytes.await_count == 2
+        first_call, retry_call = qwen3_tts_server._generate_audio_bytes.await_args_list
+        assert first_call.args[0] is request
+        assert first_call.kwargs["request_id"] != retry_call.kwargs["request_id"]
+        assert retry_call.args[0].seed is not None
+        assert request.seed is None
+
+    @pytest.mark.asyncio
+    async def test_nonstream_does_not_retry_nonretryable_generation_error(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ) -> None:
+        qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+        qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=TTSGenerationError("invalid generation"))
+
+        response = await qwen3_tts_server.create_speech(OpenAICreateSpeechRequest(input="Hello", task_type="Base"))
+
+        assert response.status_code == 500
+        qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("request_kwargs", [{"seed": 100423}, {"max_new_tokens": 192}])
+    async def test_qwen3_tts_nonstream_does_not_retry_explicit_sampling_controls(
+        self,
+        qwen3_tts_server,
+        mocker: MockerFixture,
+        request_kwargs: dict[str, int],
+    ) -> None:
+        qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+        qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=Qwen3TTSCodecLimitError("codec limit"))
+
+        response = await qwen3_tts_server.create_speech(
+            OpenAICreateSpeechRequest(input="Hello", task_type="Base", **request_kwargs)
+        )
+
+        assert response.status_code == 500
+        qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
+
     @pytest.mark.asyncio
     async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server):
         async def failing_generator():
@@ -4862,6 +4953,80 @@ class TestTTSAsyncOffloading:
         assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys
         assert ("artifact-fail", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
 
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_rejects_qwen3_codec_limit_before_success(self, qwen3_tts_server):
+        async def length_limited_generator():
+            yield SimpleNamespace(
+                multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
+                metrics={
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 191,
+                            "finish_reason": "length",
+                        }
+                    }
+                },
+            )
+
+        chunks = qwen3_tts_server._generate_audio_chunks(
+            length_limited_generator(),
+            "req-codec-limit",
+            tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
+        )
+        with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"):
+            async for _ in chunks:
+                pass
+
+    @pytest.mark.asyncio
+    async def test_sse_codec_limit_marks_partial_audio_for_discard(self, qwen3_tts_server):
+        async def length_limited_generator():
+            yield SimpleNamespace(
+                multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
+                metrics={
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 191,
+                            "finish_reason": "length",
+                        }
+                    }
+                },
+            )
+
+        events = [
+            event
+            async for event in qwen3_tts_server._generate_audio_sse_events(
+                length_limited_generator(),
+                "req-sse-codec-limit",
+                tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
+            )
+        ]
+
+        assert any("event: speech.audio.delta" in event for event in events)
+        assert not any("event: speech.audio.done" in event for event in events)
+        error_event = next(event for event in events if "event: speech.audio.error" in event)
+        payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
+        assert payload["error"]["partial_audio"] is True
+        assert payload["error"]["action"] == "discard"
+
+    @pytest.mark.asyncio
+    async def test_sse_error_before_audio_is_not_marked_partial(self, qwen3_tts_server):
+        async def failing_generator():
+            raise RuntimeError("boom")
+            yield  # pragma: no cover
+
+        events = [
+            event
+            async for event in qwen3_tts_server._generate_audio_sse_events(
+                failing_generator(),
+                "req-sse-no-audio",
+            )
+        ]
+
+        error_event = next(event for event in events if "event: speech.audio.error" in event)
+        payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
+        assert "partial_audio" not in payload["error"]
+        assert "action" not in payload["error"]
+
     @pytest.mark.asyncio
     async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
         async def pcm_generator():

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 import base64
 from types import SimpleNamespace
@@ -44,7 +47,9 @@ def _build_test_app(
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
         speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+        ):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
                 yield (chunk, 24000) if include_sample_rate else chunk
 
@@ -244,6 +249,7 @@ class TestStreamingSpeechWebSocket:
 
     def test_streaming_multiple_binary_frames(self, mocker: MockerFixture):
         captured_requests = []
+        captured_tts_params = []
 
         speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
@@ -253,11 +259,12 @@ class TestStreamingSpeechWebSocket:
 
         async def mock_prepare_speech_generation(request):
             captured_requests.append(request)
-            return "req-stream", object(), {}
+            return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]}
 
         speech_service._prepare_speech_generation = mock_prepare_speech_generation
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
+            captured_tts_params.append(tts_params)
             for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
                 yield (chunk, 24000) if include_sample_rate else chunk
 
@@ -302,6 +309,7 @@ class TestStreamingSpeechWebSocket:
         assert captured_requests[0].stream is True
         assert captured_requests[0].response_format == "pcm"
         assert captured_requests[0].initial_codec_chunk_frames == 12
+        assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}]
         assert speech_service._generate_audio_bytes.await_count == 0
 
     def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture):
@@ -344,7 +352,9 @@ class TestStreamingSpeechWebSocket:
 
         # The forced-aligner stage rides the same generator: its pooling output
         # is surfaced via the ``collect`` channel once the audio has streamed.
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+        ):
             for chunk in (first_chunk, second_chunk):
                 yield (chunk, 1000) if include_sample_rate else chunk
             if collect is not None:
@@ -427,7 +437,9 @@ class TestStreamingSpeechWebSocket:
         speech_service.forced_aligner_enabled = True
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+        ):
             chunk = b"\x01" * 1000
             yield (chunk, 1000) if include_sample_rate else chunk
             if collect is not None:
@@ -608,7 +620,7 @@ class TestStreamingSpeechWebSocket:
         speech_service.engine_client.abort = mocker.AsyncMock()
         speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
             yield b"\x01\x02"
             raise RuntimeError("stream boom")
 
@@ -633,6 +645,8 @@ class TestStreamingSpeechWebSocket:
                 assert ws.receive_json() == {
                     "type": "error",
                     "message": "Generation failed for utterance 0, sentence 0: stream boom",
+                    "partial_audio": True,
+                    "action": "discard",
                 }
                 assert ws.receive_json() == {
                     "type": "audio.done",
@@ -702,7 +716,7 @@ class TestStreamingSpeechWebSocket:
         speech_service.engine_client.abort = mocker.AsyncMock()
         speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
             yield b"\x01\x02"
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks

tests/entrypoints/openai_api/test_speech_usage.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for speech token-usage accounting (issue #4646).
 
 These exercise the pure usage logic without a model/server: input-token
@@ -205,6 +206,14 @@ def test_output_tokens_from_stage_metrics():
     assert acc.total() == 77
 
 
+def test_output_tokens_preserves_stage0_finish_reason():
+    acc = SpeechOutputTokenCounter()
+    res = _final_res(stage0_out=128)
+    res.metrics["stage_metrics"]["0"]["finish_reason"] = "length"
+    acc.observe(res)
+    assert acc.stage0_finish_reason == "length"
+
+
 def test_output_tokens_takes_max_across_stages():
     acc = SpeechOutputTokenCounter()
     acc.observe(_final_res(stage0_out=42, stage1_out=0))

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -32,7 +32,11 @@ from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import (
     MossTTSAdapter,
     MossTTSNanoAdapter,
 )
-from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
+from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import (
+    QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY,
+    Qwen3TTSAdapter,
+    Qwen3TTSCodecLimitError,
+)
 from vllm_omni.model_executor.models.indextts2 import prompt_utils
 from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
     INDEXTTS25_TOKENIZER_FILE,
@@ -162,6 +166,56 @@ def test_qwen3_tts_metadata():
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
 
 
+@pytest.mark.parametrize(
+    ("task_type", "text_tokens", "request_cap", "expected_cap"),
+    [
+        ("Base", 0, None, 4096),
+        ("Base", 10, None, 192),
+        ("Base", 23, None, 276),
+        ("Base", 23, 128, 128),
+        ("Base", 23, 512, 512),
+        ("Base", 400, None, 4096),
+        ("CustomVoice", 10, None, 4096),
+        ("CustomVoice", 10, 256, 256),
+    ],
+)
+def test_qwen3_tts_applies_text_scaled_codec_safety_limit(task_type, text_tokens, request_cap, expected_cap, mocker):
+    server = mocker.Mock()
+    server._count_usage_text_tokens.return_value = text_tokens
+    adapter = Qwen3TTSAdapter(SpeechServingContext(server=server))
+    stage_defaults = [SamplingParams(max_tokens=4096, min_tokens=2)]
+    request = OpenAICreateSpeechRequest(
+        input="test text",
+        task_type=task_type,
+        max_new_tokens=request_cap,
+    )
+    prompt: dict[str, Any] = {"additional_information": {}}
+
+    overridden = adapter.apply_sampling_overrides(stage_defaults, request, prompt)
+
+    assert overridden[0].max_tokens == expected_cap
+    assert prompt["additional_information"][QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] == [expected_cap]
+    assert stage_defaults[0].max_tokens == 4096
+
+
+def test_qwen3_tts_rejects_only_length_finished_base_audio(mocker):
+    adapter = Qwen3TTSAdapter(SpeechServingContext(server=mocker.Mock()))
+    params = {
+        "task_type": ["Base"],
+        QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY: [192],
+    }
+
+    # EOS at the exact budget is valid; the token count alone is not a
+    # sufficient failure signal.
+    adapter.validate_generation(params, stage0_finish_reason="stop", output_tokens=192)
+    adapter.validate_generation(params, stage0_finish_reason=None, output_tokens=192)
+
+    # Some frontends expose 191 decoded frames for max_new_tokens=192. The
+    # engine terminal reason still makes this an unambiguous limit failure.
+    with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"):
+        adapter.validate_generation(params, stage0_finish_reason="length", output_tokens=191)
+
+
 def test_indextts_adapters_are_versioned():
     assert resolve_adapter("indextts2") is IndexTTS2Adapter
     assert resolve_adapter("indextts2_5") is IndexTTS25Adapter

vllm_omni/engine/stage_pool.py

@@ -633,6 +633,15 @@ class StagePool:
             if callable(pop_native_text_metrics):
                 native_text_metrics = pop_native_text_metrics(request_id)
         native_generation_tokens = native_text_metrics.get("num_generation_tokens")
+        finish_reason = next(
+            (
+                str(reason)
+                for request_output in reversed(request_outputs)
+                for completion in reversed(getattr(request_output, "outputs", []) or [])
+                if (reason := getattr(completion, "finish_reason", None)) is not None
+            ),
+            None,
+        )
         num_tokens_out = (
             max(int(native_generation_tokens), 0)
             if isinstance(native_generation_tokens, int) and not isinstance(native_generation_tokens, bool)
@@ -701,6 +710,7 @@ class StagePool:
             # happens inside the model runner and is not observable here.
             batch_size=1,
             replica_id=replica_id,
+            finish_reason=finish_reason,
             rx_decode_time_ms=0.0,
             rx_transfer_bytes=0,
             rx_in_flight_time_ms=0.0,

vllm_omni/entrypoints/openai/serving_speech.py

@@ -56,6 +56,7 @@ from vllm_omni.entrypoints.openai.speech_usage import (
 )
 from vllm_omni.entrypoints.openai.tts_adapters import (
     SpeechServingContext,
+    TTSGenerationError,
     all_tts_stage_keys,
     detect_tts_model_type,
     resolve_adapter,
@@ -78,6 +79,7 @@ from vllm_omni.utils.speaker_cache import get_speaker_cache
 
 logger = init_logger(__name__)
 
+
 _SPEECH_USAGE_INPUT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TOKENS"
 _SPEECH_USAGE_OUTPUT_TOKENS_HEADER = "X-VLLM-OMNI-OUTPUT-TOKENS"
 _SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS"
@@ -750,6 +752,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         details = self._compute_speech_input_details(request, tts_params)
         return build_speech_usage(details, output_tokens)
 
+    def _validate_tts_generation(
+        self,
+        tts_params: dict[str, Any],
+        usage_acc: SpeechOutputTokenCounter,
+    ) -> None:
+        adapter = self._get_tts_adapter()
+        if adapter is None:
+            return
+        adapter.validate_generation(
+            tts_params,
+            stage0_finish_reason=usage_acc.stage0_finish_reason,
+            output_tokens=usage_acc.total(),
+        )
+
     @staticmethod
     def _build_speech_usage_headers(usage: SpeechTokenUsage | None) -> dict[str, str]:
         """Map speech usage into non-streaming response headers.
@@ -2173,6 +2189,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_start_s: float | None = None,
         include_sample_rate: bool = False,
         usage_acc: SpeechOutputTokenCounter | None = None,
+        tts_params: dict[str, Any] | None = None,
         collect: dict | None = None,
     ):
         """Generate audio chunks for streaming response.
@@ -2198,6 +2215,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         stream_start_s = request_start_s if request_start_s is not None else time.perf_counter()
         artifact_ready = False
 
+        # SSE supplies an accumulator for usage output. Raw-audio and WebSocket
+        # streams retain terminal metrics only when their model adapter needs
+        # generation validation.
+        adapter = self._get_tts_adapter()
+        if tts_params is not None and usage_acc is None and adapter is not None and adapter.validates_generation:
+            usage_acc = SpeechOutputTokenCounter()
+
         try:
             async for res in generator:
                 # Tally generated codec tokens for usage (reads per-stage metrics
@@ -2275,6 +2299,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 # Audex contract: zero codec tokens must abort the stream, not
                 # complete it cleanly with zero audio bytes.
                 raise ValueError("Audex produced no audio (the thinker emitted zero or invalid codec tokens)")
+            # Check before committing the reference-audio artifact or logging
+            # success. Streaming protocols may already have emitted partial
+            # bytes, but they must terminate as an error rather than cleanly.
+            if tts_params is not None and usage_acc is not None:
+                self._validate_tts_generation(tts_params, usage_acc)
             self._mark_ref_audio_artifact_ready_for_request(request_id)
             artifact_ready = True
             total_ms = (time.perf_counter() - stream_start_s) * 1000.0
@@ -2359,6 +2388,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         ``input_tokens`` is computed from the request text + reference audio.
         """
         usage_acc = SpeechOutputTokenCounter()
+        emitted_audio = False
         try:
             async for chunk in self._generate_audio_chunks(
                 generator,
@@ -2367,6 +2397,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 raw_request=raw_request,
                 request_start_s=request_start_s,
                 usage_acc=usage_acc,
+                tts_params=tts_params,
             ):
                 payload = {
                     "type": "speech.audio.delta",
@@ -2374,6 +2405,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     "response_format": response_format,
                 }
                 data = json.dumps(payload, separators=(",", ":"))
+                emitted_audio = True
                 yield f"event: speech.audio.delta\ndata: {data}\n\n"
             done_payload: dict[str, Any] = {"type": "speech.audio.done"}
             if request is not None:
@@ -2385,16 +2417,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         except asyncio.CancelledError:
             raise
         except Exception as e:
-            payload = {
-                "type": "speech.audio.error",
-                "error": {
-                    "message": str(e),
-                    "type": "server_error",
-                    "param": None,
-                    "code": HTTPStatus.INTERNAL_SERVER_ERROR.value,
-                },
+            error: dict[str, Any] = {
+                "message": str(e),
+                "type": "server_error",
+                "param": None,
+                "code": HTTPStatus.INTERNAL_SERVER_ERROR.value,
             }
-            data = json.dumps(payload, separators=(",", ":"))
+            if emitted_audio:
+                error.update(partial_audio=True, action="discard")
+            error_payload: dict[str, Any] = {
+                "type": "speech.audio.error",
+                "error": error,
+            }
+            data = json.dumps(error_payload, separators=(",", ":"))
             yield f"event: speech.audio.error\ndata: {data}\n\n"
 
     @staticmethod
@@ -3109,7 +3144,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         return request_id, generator, tts_params
 
     async def _generate_pcm_chunks(
-        self, generator, request_id: str, *, include_sample_rate: bool = False, collect: dict | None = None
+        self,
+        generator,
+        request_id: str,
+        *,
+        include_sample_rate: bool = False,
+        tts_params: dict[str, Any] | None = None,
+        collect: dict | None = None,
     ):
         """Yield raw PCM byte chunks from the engine generator.
 
@@ -3123,15 +3164,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             request_id,
             response_format="pcm",
             include_sample_rate=include_sample_rate,
+            tts_params=tts_params,
             collect=collect,
         ):
             yield chunk
 
     async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest):
         """Yield raw PCM bytes for a speech request as soon as chunks are decoded."""
-        request_id, generator, _ = await self._prepare_speech_generation(request)
+        request_id, generator, tts_params = await self._prepare_speech_generation(request)
         try:
-            async for chunk in self._generate_pcm_chunks(generator, request_id):
+            async for chunk in self._generate_pcm_chunks(generator, request_id, tts_params=tts_params):
                 yield chunk
         finally:
             self._discard_ref_audio_artifact_warmup(request_id)
@@ -3203,6 +3245,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if final_output is None:
                 raise ValueError("No output generated from the model.")
 
+            self._validate_tts_generation(bytes_tts_params or {}, usage_acc)
+
             # Extract audio from the audio-bearing res (not necessarily the last
             # yielded one, which may be the aligner's timestamps output).
             audio_source = audio_res if audio_res is not None else final_output
@@ -3527,7 +3571,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     return error
 
                 media_type = "audio/wav" if response_format == "wav" else "audio/pcm"
-                _, generator, _ = await self._prepare_speech_generation(request, request_id=request_id)
+                _, generator, raw_tts_params = await self._prepare_speech_generation(request, request_id=request_id)
                 return StreamingResponse(
                     self._generate_audio_chunks(
                         generator,
@@ -3535,6 +3579,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         response_format,
                         raw_request=raw_request,
                         request_start_s=request_start_s,
+                        tts_params=raw_tts_params,
                     ),
                     media_type=media_type,
                 )
@@ -3563,9 +3608,35 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
             collect: dict = {}
             usage_box: list[SpeechTokenUsage] = []
-            audio_bytes, media_type = await self._generate_audio_bytes(
-                request, request_id=request_id, usage_out=usage_box, collect=collect
-            )
+            try:
+                audio_bytes, media_type = await self._generate_audio_bytes(
+                    request, request_id=request_id, usage_out=usage_box, collect=collect
+                )
+            except TTSGenerationError as error:
+                # An adapter can reject otherwise completed audio. Retry only
+                # retryable, stochastic, server-budgeted non-streaming
+                # requests: changing an explicit seed would violate
+                # reproducibility, and changing a caller-provided token limit
+                # would violate the requested budget.
+                if not error.retryable or request.seed is not None or request.max_new_tokens is not None:
+                    raise
+                retry_seed = int(random_uuid()[:8], 16)
+                retry_request = request.model_copy(update={"seed": retry_seed})
+                retry_request_id = f"speech-{random_uuid()}"
+                collect.clear()
+                usage_box.clear()
+                logger.warning(
+                    "TTS request %s failed generation validation; retrying once as %s with seed=%d",
+                    request_id,
+                    retry_request_id,
+                    retry_seed,
+                )
+                audio_bytes, media_type = await self._generate_audio_bytes(
+                    retry_request,
+                    request_id=retry_request_id,
+                    usage_out=usage_box,
+                    collect=collect,
+                )
             total_ms = (time.perf_counter() - request_start_s) * 1000.0
             logger.info(
                 "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d",
@@ -3618,6 +3689,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 e,
             )
             return self.create_error_response(e)
+        except TTSGenerationError as e:
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.error(
+                "[SpeechE2E] request_id=%s stream=%s status=invalid_generation total_ms=%.2f error=%s",
+                request_id,
+                bool(request.stream),
+                total_ms,
+                e,
+            )
+            return self._diffusion_error_response(str(e), status_code=500)
         except Exception as e:
             total_ms = (time.perf_counter() - request_start_s) * 1000.0
             logger.exception(

vllm_omni/entrypoints/openai/serving_speech_stream.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """WebSocket handler for streaming text input TTS.
 
 Accepts text incrementally via WebSocket, buffers it until input.done, and
@@ -329,7 +332,7 @@ class OmniStreamingSpeechHandler:
         request_id = None
         try:
             if config.stream_audio:
-                request_id, generator, _ = await self._speech_service._prepare_speech_generation(request)
+                request_id, generator, tts_params = await self._speech_service._prepare_speech_generation(request)
                 if config.word_timestamps:
                     total_bytes = await self._stream_audio_with_alignments(
                         websocket=websocket,
@@ -339,9 +342,16 @@ class OmniStreamingSpeechHandler:
                         utterance_index=utterance_index,
                         sentence_index=sentence_index,
                         language=config.language,
+                        tts_params=tts_params,
                     )
                 else:
-                    async with aclosing(self._speech_service._generate_pcm_chunks(generator, request_id)) as stream:
+                    async with aclosing(
+                        self._speech_service._generate_pcm_chunks(
+                            generator,
+                            request_id,
+                            tts_params=tts_params,
+                        )
+                    ) as stream:
                         async for chunk in stream:
                             total_bytes += len(chunk)
                             await websocket.send_bytes(chunk)
@@ -367,6 +377,7 @@ class OmniStreamingSpeechHandler:
             await self._send_error(
                 websocket,
                 f"Generation failed for utterance {utterance_index}, sentence {sentence_index}: {e}",
+                partial_audio=total_bytes > 0,
             )
         finally:
             try:
@@ -392,6 +403,7 @@ class OmniStreamingSpeechHandler:
         utterance_index: int,
         sentence_index: int,
         language: str | None = None,
+        tts_params: dict | None = None,
     ) -> int:
         """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence.
 
@@ -438,6 +450,7 @@ class OmniStreamingSpeechHandler:
                 generator,
                 request_id,
                 include_sample_rate=True,
+                tts_params=tts_params,
                 collect=collect,
             )
         ) as stream:
@@ -463,14 +476,15 @@ class OmniStreamingSpeechHandler:
         return total_bytes
 
     @staticmethod
-    async def _send_error(websocket: WebSocket, message: str) -> None:
+    async def _send_error(websocket: WebSocket, message: str, *, partial_audio: bool = False) -> None:
         """Send an error message to the client."""
         try:
-            await websocket.send_json(
-                {
-                    "type": "error",
-                    "message": message,
-                }
-            )
+            payload: dict[str, object] = {
+                "type": "error",
+                "message": message,
+            }
+            if partial_audio:
+                payload.update(partial_audio=True, action="discard")
+            await websocket.send_json(payload)
         except Exception:
             pass  # Connection may already be closed; safe to ignore

vllm_omni/entrypoints/openai/speech_usage.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Token-usage accounting for the Speech (``/v1/audio/speech``) API.
 
 Why this module exists (issue #4646)
@@ -160,6 +161,7 @@ class SpeechOutputTokenCounter:
     """
 
     output_tokens: int = 0
+    stage0_finish_reason: str | None = None
 
     def observe(self, res: Any) -> None:
         metrics = getattr(res, "metrics", None)
@@ -168,6 +170,11 @@ class SpeechOutputTokenCounter:
         stage_metrics = metrics.get("stage_metrics")
         if not isinstance(stage_metrics, dict):
             return
+        stage0 = stage_metrics.get("0", stage_metrics.get(0))
+        if isinstance(stage0, dict):
+            reason = stage0.get("finish_reason")
+            if reason is not None:
+                self.stage0_finish_reason = str(reason)
         for stage in stage_metrics.values():
             if not isinstance(stage, dict):
                 continue

vllm_omni/entrypoints/openai/tts_adapters/__init__.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Registry of TTS serving adapters, and the model detection built on it.
 
 Adapters register themselves by their ``name`` (the model-type discriminator)
@@ -19,6 +20,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import (
     OutputPolicy,
     PreparedRequest,
     SpeechServingContext,
+    TTSGenerationError,
     TTSModelAdapter,
 )
 
@@ -126,6 +128,7 @@ __all__ = [
     "OutputPolicy",
     "PreparedRequest",
     "SpeechServingContext",
+    "TTSGenerationError",
     "TTSModelAdapter",
     "TTS_ADAPTER_REGISTRY",
     "all_tts_model_types",

vllm_omni/entrypoints/openai/tts_adapters/base.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Base contract for per-model TTS serving adapters.
 
 This package factors the per-model ``if self._tts_model_type == ...`` dispatch
@@ -11,7 +12,7 @@ See the RFC for the full design (issue #4327).
 """
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -69,6 +70,14 @@ def apply_max_new_tokens(
     return sampling_params_list
 
 
+class TTSGenerationError(RuntimeError):
+    """A completed TTS generation that must not be returned as valid audio."""
+
+    def __init__(self, message: str, *, retryable: bool = False) -> None:
+        super().__init__(message)
+        self.retryable = retryable
+
+
 @dataclass
 class OutputPolicy:
     """How the orchestrator aggregates engine output for a model.
@@ -153,6 +162,8 @@ class TTSModelAdapter(ABC):
     detect_priority: ClassVar[int] = 100
     #: Serving backend: ``"ar"`` (engine_client) or ``"diffusion"``.
     backend: ClassVar[str] = "ar"
+    #: Whether streaming entrypoints must retain terminal metrics for validation.
+    validates_generation: ClassVar[bool] = False
     #: Whether the model consumes ``request.speed`` in its native parameters.
     native_speed_control: ClassVar[bool] = False
 
@@ -233,6 +244,21 @@ class TTSModelAdapter(ABC):
         """
         return sampling_params_list
 
+    def validate_generation(
+        self,
+        tts_params: Mapping[str, object],
+        *,
+        stage0_finish_reason: str | None,
+        output_tokens: int,
+    ) -> None:
+        """Reject a completed generation that is invalid for this model.
+
+        Adapters that override this hook must also set
+        :attr:`validates_generation` so streaming entrypoints retain the
+        terminal metrics needed by the validation without charging that cost
+        to unrelated TTS models.
+        """
+
     async def warmup(self) -> None:
         return
 
@@ -304,5 +330,6 @@ __all__ = [
     "OutputPolicy",
     "PreparedRequest",
     "SpeechServingContext",
+    "TTSGenerationError",
     "TTSModelAdapter",
 ]

vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Qwen3-TTS serving adapter."""
 
 from collections.abc import Mapping
@@ -12,7 +13,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import (
     DEFAULT_TTS_LANGUAGES,
     ARTTSAdapter,
     PreparedRequest,
-    apply_max_new_tokens,
+    TTSGenerationError,
 )
 from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_precomputed_speakers
 from vllm_omni.utils.speaker_cache import validate_qwen3_tts_profile
@@ -22,11 +23,23 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY = "_qwen3_tts_effective_max_tokens"
+_MIN_CODEC_FRAMES = 192
+_MAX_CODEC_FRAMES_PER_TEXT_TOKEN = 12
+
+
+class Qwen3TTSCodecLimitError(TTSGenerationError):
+    """Qwen3-TTS Base exhausted its codec budget without emitting EOS."""
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message, retryable=True)
+
 
 @register_tts_adapter
 class Qwen3TTSAdapter(ARTTSAdapter):
     """Adapter for Qwen3-TTS (AR ``engine_client`` backend)."""
 
+    validates_generation = True
     stage_keys = frozenset({"qwen3_tts"})
     name = "qwen3_tts"
 
@@ -226,4 +239,95 @@ class Qwen3TTSAdapter(ARTTSAdapter):
         prompt: dict[str, Any] | None = None,
         request_id: str | None = None,
     ) -> list:
-        return apply_max_new_tokens(sampling_params_list, request)
+        """Apply a text-scaled safety ceiling to Base codec generation.
+
+        Qwen3-TTS can rarely enter a repetitive state in which codec EOS is no
+        longer reachable through top-k sampling. A fixed 4096-frame ceiling
+        turns that into several minutes of unusable audio. Bound the default
+        Base-task budget by text length, while preserving an explicit caller
+        ``max_new_tokens`` override and the configured budget for other tasks.
+        """
+        import copy
+
+        del request_id
+        server = self.ctx.server
+        # Only scalar fields on stage 0 are changed below. Shallow-copy each
+        # stage so the shared defaults stay immutable without deep-copying the
+        # complete sampling configuration on every request.
+        sampling_params_list = [copy.copy(params) for params in sampling_params_list]
+        configured_cap = getattr(sampling_params_list[0], "max_tokens", None)
+        task_type = request.task_type or "CustomVoice"
+        text_tokens = None
+        dynamic_cap = None
+        effective_cap: int | None
+
+        if request.max_new_tokens is not None:
+            # An explicit request budget is an opt-out from the automatic
+            # ceiling. It remains an upper bound, and a length finish is still
+            # surfaced as an incomplete generation rather than valid audio.
+            effective_cap = int(request.max_new_tokens)
+        elif task_type == "Base":
+            counted_text_tokens = server._count_usage_text_tokens(request.input)
+            if counted_text_tokens > 0:
+                text_tokens = counted_text_tokens
+                dynamic_cap = max(_MIN_CODEC_FRAMES, text_tokens * _MAX_CODEC_FRAMES_PER_TEXT_TOKEN)
+                effective_cap = min(dynamic_cap, int(configured_cap)) if configured_cap is not None else dynamic_cap
+            else:
+                # Token counting is best-effort. If the tokenizer is missing or
+                # rejects the input, preserve the configured budget instead of
+                # truncating an otherwise valid request at the minimum ceiling.
+                effective_cap = int(configured_cap) if configured_cap is not None else None
+        else:
+            effective_cap = int(configured_cap) if configured_cap is not None else None
+
+        if effective_cap is not None:
+            effective_cap = max(1, effective_cap)
+            sampling_params_list[0].max_tokens = effective_cap
+            sampling_params_list[0].min_tokens = min(
+                int(getattr(sampling_params_list[0], "min_tokens", 0) or 0),
+                effective_cap,
+            )
+
+        if isinstance(prompt, dict):
+            additional_information = prompt.get("additional_information")
+            if isinstance(additional_information, dict) and effective_cap is not None:
+                additional_information[QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] = [effective_cap]
+
+        logger.debug(
+            "Qwen3-TTS codec budget: task_type=%s text_tokens=%s dynamic_cap=%s "
+            "configured_cap=%s request_cap=%s effective_cap=%s",
+            task_type,
+            text_tokens,
+            dynamic_cap,
+            configured_cap,
+            request.max_new_tokens,
+            effective_cap,
+        )
+        return sampling_params_list
+
+    def validate_generation(
+        self,
+        tts_params: Mapping[str, object],
+        *,
+        stage0_finish_reason: str | None,
+        output_tokens: int,
+    ) -> None:
+        if QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY not in tts_params:
+            return
+        task_type = tts_params.get("task_type")
+        if isinstance(task_type, (list, tuple)):
+            task_type = task_type[0] if task_type else None
+        if task_type != "Base" or stage0_finish_reason != "length":
+            return
+
+        raw_limit = tts_params.get(QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY)
+        if isinstance(raw_limit, (list, tuple)):
+            raw_limit = raw_limit[0] if raw_limit else None
+        try:
+            limit = int(raw_limit) if isinstance(raw_limit, (str, bytes, bytearray, int, float)) else 0
+        except (TypeError, ValueError):
+            limit = 0
+        raise Qwen3TTSCodecLimitError(
+            "Qwen3-TTS Base did not emit codec EOS before its token budget "
+            f"({output_tokens}/{limit} codec tokens); the generated audio is incomplete."
+        )

vllm_omni/metrics/stats.py

@@ -45,6 +45,7 @@ class StageRequestStats:
     stage_id: int | None = None
     replica_id: int | None = None
     final_output_type: str | None = None
+    finish_reason: str | None = None
     request_id: str | None = None
     postprocess_time_ms: float = 0.0
     diffusion_metrics: dict[str, float] = None
@@ -126,6 +127,7 @@ STAGE_EXCLUDE = {
     "rx_decode_time_ms",
     "rx_in_flight_time_ms",
     "final_output_type",
+    "finish_reason",
     "pipeline_timings",
 }
 TRANSFER_EXCLUDE = {"from_stage", "to_stage", "request_id", "used_shm"}
@@ -473,6 +475,7 @@ class OrchestratorAggregator:
             current = {
                 "stage_id": sid,
                 "final_output_type": evt.final_output_type,
+                "finish_reason": evt.finish_reason,
                 defs.NUM_TOKENS_IN: int(evt.num_tokens_in),
                 defs.NUM_TOKENS_OUT: int(evt.num_tokens_out),
                 defs.STAGE_GEN_TIME_MS: float(evt.stage_gen_time_ms),
@@ -498,6 +501,8 @@ class OrchestratorAggregator:
 
         current[defs.NUM_TOKENS_IN] = int(current.get(defs.NUM_TOKENS_IN, 0)) + int(evt.num_tokens_in)
         current[defs.NUM_TOKENS_OUT] = int(current.get(defs.NUM_TOKENS_OUT, 0)) + int(evt.num_tokens_out)
+        if evt.finish_reason is not None:
+            current["finish_reason"] = evt.finish_reason
         current[defs.STAGE_GEN_TIME_MS] = float(current.get(defs.STAGE_GEN_TIME_MS, 0.0)) + float(evt.stage_gen_time_ms)
         current[defs.POSTPROCESS_TIME_MS] = float(current.get(defs.POSTPROCESS_TIME_MS, 0.0)) + float(
             evt.postprocess_time_ms