From 9fdb3bcb89ba82617422d1779b69232bef90daf3 Mon Sep 17 00:00:00 2001 From: nperraud <6399466+nperraud@users.noreply.github.com> Date: Fri, 29 May 2026 17:18:25 +0200 Subject: [PATCH] [Bugfix][Model] Qwen3-TTS: don't collapse 2D ref_code list when estimating prompt length (#3940) Signed-off-by: nperraud --- .../test_qwen3_tts_talker_prompt_len.py | 129 ++++++++++++++++++ .../models/qwen3_tts/qwen3_tts_talker.py | 14 +- 2 files changed, 141 insertions(+), 2 deletions(-) create mode 100644 tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py new file mode 100644 index 000000000..c63622eee --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py @@ -0,0 +1,129 @@ +"""Regression tests for `estimate_prompt_len_from_additional_information`. + +Pins the 2D `voice_clone_prompt.ref_code` shape behaviour. Applying the +singleton-batch unwrapper `_first(...)` to that value strips its outer +dimension and reports `len(ref_code) == num_codebooks` instead of +`num_frames`, which silently truncates `inputs_embeds` downstream in +`_build_prompt_embeds`. +""" + +import pytest + +from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( + Qwen3TTSTalkerForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _fake_tokenize(text: str, **_kwargs): + return [0] * (8 + max(1, len(text.split()))) + + +def test_estimate_prompt_len_uses_full_ref_code_length() -> None: + num_frames = 318 + num_codebooks = 8 + info = { + "task_type": ["Base"], + "text": ["hello"], + "ref_text": ["world"], + "voice_clone_prompt": [ + { + "ref_spk_embedding": [0.0] * 512, + "ref_code": [[0] * num_codebooks for _ in range(num_frames)], + "icl_mode": True, + } + ], + "non_streaming_mode": [True], + "language": ["English"], + } + + est = Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + additional_information=info, + task_type="Base", + tokenize_prompt=_fake_tokenize, + codec_language_id={"english": 0}, + spk_is_dialect=None, + ) + + # codec_lens = 1 + num_frames = 319; plus text-side and codec-prefix + # terms ~20. Would be ~30 if `_first` collapses ref_code to its first row. + assert est > 100, f"got {est}; expected ~339. Did `_first(ref_code)` collapse the 2D list again?" + + +def test_estimate_prompt_len_handles_1d_ref_code() -> None: + num_frames = 50 + info = { + "task_type": ["Base"], + "text": ["hello"], + "ref_text": ["world"], + "voice_clone_prompt": [ + { + "ref_spk_embedding": [0.0] * 512, + "ref_code": list(range(num_frames)), + "icl_mode": True, + } + ], + "non_streaming_mode": [True], + "language": ["English"], + } + + est = Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + additional_information=info, + task_type="Base", + tokenize_prompt=_fake_tokenize, + codec_language_id={"english": 0}, + spk_is_dialect=None, + ) + + assert est > 50, f"got {est}; 1D ref_code must contribute its own length" + + +def test_estimate_prompt_len_uses_list_of_int_ref_ids_from_voice_clone_prompt() -> None: + """`voice_clone_prompt.ref_ids` as a list-of-int (the natural pre-tokenized shape) + must be read as a sequence, not unwrapped to its first element. Without the fix + `_first` collapsed it to an int, the `isinstance(list)` / `hasattr("shape")` + branches fell through, and `ref_ids_len` defaulted to 0 — silently dropping the + reference-token contribution to the prefix length budget. + """ + + def make_info(num_ref_ids: int) -> dict: + return { + "task_type": ["Base"], + "text": ["hello"], + "ref_text": ["world"], + "voice_clone_prompt": [ + { + "ref_spk_embedding": [0.0] * 512, + "ref_code": [[0] * 8 for _ in range(200)], + "ref_ids": list(range(num_ref_ids)), + "icl_mode": True, + } + ], + "non_streaming_mode": [True], + "language": ["English"], + } + + est_30 = Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + additional_information=make_info(30), + task_type="Base", + tokenize_prompt=_fake_tokenize, + codec_language_id={"english": 0}, + spk_is_dialect=None, + ) + est_50 = Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + additional_information=make_info(50), + task_type="Base", + tokenize_prompt=_fake_tokenize, + codec_language_id={"english": 0}, + spk_is_dialect=None, + ) + + # Each additional ref_id contributes exactly one token to the prefix budget + # (the estimator strips a fixed 5-token tail from ref_ids regardless of length). + # Under the previous bug `_first` collapsed both lists to a single int, so both + # estimates would have been identical (delta = 0). + assert est_50 - est_30 == 20, ( + f"got delta={est_50 - est_30}; expected 20 (= 50 - 30 extra ref_ids). " + f"Did `_first(ref_ids)` collapse the list to its first int?" + ) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index dbede17a3..bd5ec6108 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -1041,7 +1041,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): if in_context_mode: ref_code = None if isinstance(voice_clone_prompt, dict): - ref_code = _first(voice_clone_prompt.get("ref_code"), None) + # Do NOT apply `_first(...)` here. `_first` unwraps the singleton-list + # batching convention used at the top level of `additional_information`; + # values inside `voice_clone_prompt` are per-item payloads. `ref_code` + # is a `[num_frames, num_codebooks]` 2D list, so `_first` would strip + # its outer dim and report `num_codebooks` (~8) instead of `num_frames` + # (~hundreds), silently truncating the prefill embeddings downstream. + ref_code = voice_clone_prompt.get("ref_code") ref_code_len: int | None = None if isinstance(ref_code, list): @@ -1071,7 +1077,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # text_embed = ref_ids + text_ids + eos. ref_ids = _first(info.get("ref_ids"), None) if isinstance(voice_clone_prompt, dict) and ref_ids is None: - ref_ids = _first(voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id"), None) + # Same reasoning as ref_code above: values inside `voice_clone_prompt` + # are per-item payloads, not singleton-wrapped batches. Applying + # `_first` to a list-of-int `ref_ids` collapses it to an int and the + # branches below fall through to `ref_ids_len = 0`. + ref_ids = voice_clone_prompt.get("ref_ids") or voice_clone_prompt.get("ref_id") if ref_ids is None: ref_text = _first(info.get("ref_text"), "")