[Bugfix][Model] Qwen3-TTS: don't collapse 2D ref_code list when estimating prompt length (#3940)

Signed-off-by: nperraud <nathanael.perraudin@gmail.com>

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?"
+    )

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"), "")