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