From f11b32a42edfcf753be682b93f8bfa4e5317873c Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Thu, 2 Jul 2026 12:46:05 +0800 Subject: [PATCH] [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) Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../test_qwen3_tts_talker_preprocess.py | 32 +++++++++++++++++++ .../test_qwen3_tts_talker_ref_codes.py | 15 +++++++++ .../models/qwen3_tts/qwen3_tts_talker.py | 9 +++++- 3 files changed, 55 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py index 9d6041f41..7358c7893 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py +++ b/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) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py index 2adfa17a4..3c145e4ca 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py +++ b/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) 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 105672399..c1c7f4264 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 @@ -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`