From 1b31f96240229b9366219b3f97fff7ec00844843 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Sun, 5 Jul 2026 03:57:55 +0800 Subject: [PATCH] [Bugfix] Fix Qwen3-TTS no-async-chunk payload splitting (#4870) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: Yueqian Lin Co-authored-by: Yueqian Lin --- .../openai_api/test_serving_speech.py | 42 ++++++++++++++++ tests/worker/test_gpu_ar_model_runner.py | 50 +++++++++++++++++-- .../entrypoints/openai/serving_speech.py | 14 ++++-- vllm_omni/worker/gpu_ar_model_runner.py | 6 +-- 4 files changed, 99 insertions(+), 13 deletions(-) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 8f2797bc7..efed9980e 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -3450,6 +3450,48 @@ class TestTTSAsyncOffloading: mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True) + def test_prepare_speech_generation_no_async_chunk_stream_uses_final_only( + self, qwen3_tts_server, mocker: MockerFixture + ): + """Full-payload TTS streaming should not request delta multimodal outputs.""" + qwen3_tts_server.engine_client.model_config.async_chunk = False + qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) + qwen3_tts_server._build_tts_params = mocker.MagicMock( + return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]} + ) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + mock_coerce = mocker.patch( + "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types", + return_value=qwen3_tts_server.engine_client.default_sampling_params_list, + ) + request = OpenAICreateSpeechRequest(input="hello", stream_format="audio", response_format="pcm") + + asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, False) + + def test_prepare_speech_generation_no_async_chunk_stream_keeps_delta_for_non_qwen3( + self, voxtral_server, mocker: MockerFixture + ): + """FINAL_ONLY streaming for async_chunk=False is scoped to qwen3_tts only.""" + voxtral_server.engine_client.model_config.async_chunk = False + mocker.patch.object(voxtral_server._get_tts_adapter(), "validate", return_value=None) + voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock( + return_value={ + "prompt_token_ids": [1, 2, 3], + "additional_information": {"voice": ["test"]}, + } + ) + mock_coerce = mocker.patch( + "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types", + return_value=voxtral_server.engine_client.default_sampling_params_list, + ) + request = OpenAICreateSpeechRequest(input="hello", voice="test", stream_format="audio", response_format="pcm") + + asyncio.run(voxtral_server._prepare_speech_generation(request)) + + mock_coerce.assert_called_once_with(voxtral_server.engine_client.default_sampling_params_list, True) + def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false( self, qwen3_tts_server, mocker: MockerFixture ): diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 0e3a39abc..b56917fdb 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -266,7 +266,6 @@ def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monke ec_connector_output=None, cudagraph_stats=None, kv_extracted_req_ids=["r2"], - seq_len=3, num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(), query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long), ) @@ -312,7 +311,6 @@ def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatc ec_connector_output=None, cudagraph_stats=None, kv_extracted_req_ids=None, - seq_len=6, num_scheduled_tokens_np=np.array([1, 2, 3], dtype=np.int32), query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long), ) @@ -366,7 +364,6 @@ def test_process_additional_information_uses_snapshot_request_order(monkeypatch) ec_connector_output=None, cudagraph_stats=None, kv_extracted_req_ids=None, - seq_len=3, num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(), query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long), ) @@ -427,7 +424,6 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch): ec_connector_output=None, cudagraph_stats=None, kv_extracted_req_ids=None, - seq_len=2, num_scheduled_tokens_np=np.array([2], dtype=np.int32), query_start_loc_cpu=torch.tensor([0], dtype=torch.long), ) @@ -439,6 +435,51 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch): assert output.multimodal_outputs is None +def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_only(monkeypatch): + runner = _make_async_output_runner(engine_output_type="latent") + runner.model.omni_pooler_payload_include_hidden = False + runner._async_chunk = False + runner.requests = {"r1": object(), "r2": object(), "r3": object()} + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + lambda self, req_ids: ("latent", req_ids), + ) + monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False) + monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None) + monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None) + + codes = torch.arange(48, dtype=torch.long).reshape(3, 16) + output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot( + runner, + scheduler_output=SimpleNamespace( + total_num_scheduled_tokens=3, + num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1}, + ), + hidden_states=torch.tensor([[1.0]]), + staged_hidden_states_cpu=None, + multimodal_outputs={"codes": {"audio": codes}}, + req_ids_output_copy=["r1", "r2", "r3"], + req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2}, + valid_sampled_token_ids=[[101], [102], [103]], + logprobs_lists=None, + prompt_logprobs_dict={}, + num_nans_in_logits=None, + kv_connector_output=None, + ec_connector_output=None, + cudagraph_stats=None, + kv_extracted_req_ids=None, + num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32), + query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long), + ) + + assert output.inter_stage_outputs is not None + assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1]) + assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2]) + assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], codes[2:3]) + + def test_async_snapshot_payload_omits_hidden_when_model_opts_out(): runner = _make_async_output_runner() runner.model.omni_pooler_payload_include_hidden = False @@ -627,7 +668,6 @@ def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch ec_connector_output=None, cudagraph_stats=None, kv_extracted_req_ids=None, - seq_len=2, num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32), query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long), ) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index ae1a7c1fd..6bfedccba 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -3495,12 +3495,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id = request_id or f"speech-{random_uuid()}" qwen3_ref_audio_warmup_artifact_key: str | None = None - # 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. + # If this is a streaming request with real async chunks, we need to + # coerce cumulative outputs to delta outputs; this ensures we don't + # emit redundant MM data & drain after emitting. Qwen3-TTS full-payload + # (async_chunk=False) has no incremental audio chunks, so keep + # FINAL_ONLY semantics and let the streaming response send the final + # waveform once. Scoped to qwen3_tts: other async_chunk=False models + # keep the DELTA coercion they stream with today. # list() makes a copy to avoid mutating the params. sampling_params_list = list(self.engine_client.default_sampling_params_list) - is_streaming_request = request.is_streaming() + async_chunk = getattr(self.model_config, "async_chunk", True) + qwen3_full_payload = self._tts_model_type == "qwen3_tts" and not bool(async_chunk) + is_streaming_request = request.is_streaming() and not qwen3_full_payload sampling_params_list = coerce_param_message_types(sampling_params_list, is_streaming_request) # Build prompt + tts_params via the per-model adapter (RFC #4327). Every diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 816246410..2b44e7b45 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -1676,7 +1676,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): ec_connector_output: Any, cudagraph_stats: Any, kv_extracted_req_ids: list[str] | None, - seq_len: int, num_scheduled_tokens_np: np.ndarray, query_start_loc_cpu: Any, postprocess_already_applied: bool = False, @@ -1732,6 +1731,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # The actual multimodal wire transport uses multimodal_outputs instead. pooler_output: list[dict[str, object]] | None = None if needs_pooler_payload: + mm_seq_len = int(scheduler_output.total_num_scheduled_tokens) mm_cpu = None if self.omni_prefix_cache is not None: ( @@ -1787,7 +1787,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): mm_cpu=mm_cpu, audio_sparse_output=audio_sparse_output, sparse_mm_index=sparse_mm_index, - seq_len=seq_len, + seq_len=mm_seq_len, ) pooler_output.append(flatten_payload(payload)) @@ -1873,7 +1873,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): slot_mappings, # OMNI: unpack slot_mappings for drafter ) = self.execute_model_state self.execute_model_state = None - seq_len = hidden_states.shape[0] # Apply structured output bitmasks if present. if grammar_output is not None: @@ -2040,7 +2039,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): ec_connector_output=ec_connector_output, cudagraph_stats=cudagraph_stats, kv_extracted_req_ids=kv_extracted_req_ids, - seq_len=seq_len, num_scheduled_tokens_np=num_scheduled_tokens_np, query_start_loc_cpu=query_start_loc_cpu, postprocess_already_applied=omni_postprocess_already_applied,