[Bugfix] Qwen3-TTS talker handles multi-token replay on resume (#4559)

Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -34,6 +34,7 @@ def _make_minimal_talker(
     """
     model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
     model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
+    model._embedding_dtype = torch.bfloat16
     if tts_pad_embed is None:
         tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16)
     model._tts_pad_embed = tts_pad_embed
@@ -82,6 +83,7 @@ def _make_minimal_builder(
     builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
     builder._residual_code_embeddings = lambda: []
     builder._speaker_encoder = None
+    builder._embedding_dtype = torch.bfloat16
     builder._tts_pad_embed_buffer = (
         tts_pad_embed if tts_pad_embed is not None else torch.zeros((1, 4), dtype=torch.bfloat16)
     )
@@ -254,6 +256,36 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
     assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text[65:])
 
 
+def test_decode_replay_span_embeds_all_tokens_without_mutating_decode_state():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_text = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+
+    out_ids, out_embeds, update = model.preprocess(
+        input_ids=torch.tensor([101, 202, 303], dtype=torch.long),
+        input_embeds=None,
+        text=["hello"],
+        task_type=["Base"],
+        hidden_states={"trailing_text": trailing_text, "last": last_hidden},
+        meta={"talker_text_offset": 1, "codec_streaming": True},
+        _omni_is_prefill=False,
+        _omni_num_computed_tokens=2400,
+        _omni_prompt_len=213,
+    )
+
+    assert out_ids.tolist() == [101, 202, 303]
+    assert torch.equal(
+        out_embeds.cpu(),
+        torch.tensor([[101.0] * 4, [202.0] * 4, [303.0] * 4], dtype=torch.bfloat16),
+    )
+    assert update == {"meta": {"codec_streaming": True}}
+
+
 def test_decode_batch_preprocess_matches_decode_state_updates():
     tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
     model = _make_minimal_talker(tts_pad_embed=tts_pad)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py

@@ -74,3 +74,18 @@ def test_make_omni_output_omits_ref_key_when_no_request_has_one() -> None:
     )
 
     assert "ref" not in out.multimodal_outputs["codes"]
+
+
+def test_make_omni_output_keeps_hidden_states_for_replay_spans_without_audio_codes() -> None:
+    hidden = torch.arange(6 * 8, dtype=torch.float32).reshape(6, 8)
+
+    out = _make_talker().make_omni_output(
+        hidden,
+        model_intermediate_buffer=[
+            {"meta": {"codec_streaming": True}},
+            _info(1, None),
+        ],
+    )
+
+    assert torch.equal(out.text_hidden_states, hidden)
+    assert out.multimodal_outputs["codes"]["audio"].shape == (1, _NUM_CODE_GROUPS)

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -592,7 +592,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         audio_codes = torch.cat(audio_codes_list, dim=0)
         span_len = int(audio_codes.shape[0])
-        hidden = hidden[:span_len]
         mm: OmniPayload = {"codes": {"audio": audio_codes}}
         if ref_code_len_list:
             mm.setdefault("meta", {})["ref_code_len"] = torch.cat(ref_code_len_list, dim=0)[:span_len]
@@ -718,6 +717,14 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             info_update.setdefault("codes", {})["audio"] = zeros
             return input_ids_out, prompt_embeds, info_update
 
+        if span_len > 1:
+            inputs_embeds_out = (
+                self.embed_input_ids(input_ids.reshape(-1, 1).to(torch.long))
+                .to(device=input_ids.device, dtype=dtype)
+                .reshape(span_len, -1)
+            )
+            return input_ids, inputs_embeds_out, {"meta": {"codec_streaming": codec_streaming}}
+
         # Decode: span_len == 1
         # Pop one text-step vector from tailing_text_hidden queue.
         # ``tts_pad_embed`` was materialized above from :attr:`_tts_pad_embed`