From a693ae67f6eec6df1d443581857a9c016850a0a2 Mon Sep 17 00:00:00 2001
From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Date: Mon, 15 Jun 2026 03:00:03 -0700
Subject: [PATCH] [Bugfix] Qwen3-TTS: trim reference audio in no_async_chunk
 voice clone (#4429)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
---
 .../test_qwen3_tts_async_chunk.py             | 47 +++++++++++++++++++
 .../stage_input_processors/qwen3_tts.py       | 14 +++++-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
index bbb83b484..da25ee7bd 100644
--- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
+++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
@@ -14,6 +14,7 @@ from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
 from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
     talker2code2wav,
     talker2code2wav_async_chunk,
+    talker2code2wav_full_payload,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -439,3 +440,49 @@ def test_non_async_processor_filters_out_of_range_codec_values():
     # Only ref_code (1 frame) + 2 valid frames = 3 frames * 4 quantizers = 12 codes
     assert len(prompt["prompt_token_ids"]) == 4 * 3
     assert prompt["additional_information"] == {"meta": {"left_context_size": 1}}
+
+
+def test_full_payload_emits_left_context_size_for_ref_clone():
+    """Regression for #4421.
+
+    The worker-side ``talker2code2wav_full_payload`` producer is the
+    authoritative channel that prepends ``ref_code`` to the codec stream.
+    The orchestrator-side ``talker2code2wav_token_only`` can no longer derive
+    ``left_context_size`` (the stage-0 RequestOutput multimodal_output no longer
+    carries the talker codec since the separated mm-output channel landed), so
+    full_payload MUST emit the matching ``left_context_size`` in its connector
+    meta or Code2Wav trims nothing and the reference audio leaks into the
+    output.
+    """
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)  # 2 ref frames
+    audio_codes = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [1, 1, 1, 1]], dtype=torch.long)  # 3 generated frames
+    pooling_output = {
+        "codes.audio": audio_codes,
+        "codes.ref": ref_code,
+        "meta.ref_code_len": 2,
+    }
+    request = SimpleNamespace(request_id="r", output_token_ids=list(range(4)))  # seq_len=3
+
+    payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request)
+
+    assert payload is not None
+    assert payload["meta"]["finished"].item() is True
+    # The fix: trim length is co-located with the ref prepend it describes.
+    assert payload["meta"]["left_context_size"] == 2
+    # ref(2) + generated(3) frames, codebook-major flat = Q * 5.
+    assert len(payload["codes"]["audio"]) == _Q * 5
+
+
+def test_full_payload_omits_left_context_size_without_ref():
+    """Without a reference (non-clone tasks) nothing is prepended, so no
+    ``left_context_size`` is emitted and Code2Wav trims nothing."""
+    audio_codes = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=torch.long)
+    pooling_output = {"codes.audio": audio_codes}
+    request = SimpleNamespace(request_id="r", output_token_ids=list(range(3)))  # seq_len=2
+
+    payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request)
+
+    assert payload is not None
+    assert payload["meta"]["finished"].item() is True
+    assert "left_context_size" not in payload["meta"]
+    assert len(payload["codes"]["audio"]) == _Q * 2
diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
index 25afa9b9e..c67e4a1af 100644
--- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
@@ -526,7 +526,19 @@ def talker2code2wav_full_payload(
         audio = torch.cat([ref_code.to(audio.device), audio], dim=0)
 
     codec_codes = audio.transpose(0, 1).to(device="cpu", dtype=torch.long).reshape(-1).contiguous()
+    meta: dict[str, Any] = {"finished": torch.tensor(True, dtype=torch.bool)}
+    # Co-locate the Code2Wav trim length with the ref prepend it describes.
+    # The orchestrator-side ``talker2code2wav_token_only`` channel derives
+    # left_context_size from the stage-0 RequestOutput ``multimodal_output``,
+    # which no longer carries the talker codec (ref reads as absent, so it
+    # emits left_context_size=0).  This producer is the authoritative side
+    # that actually prepends ``ref_code``; if it does not also emit the
+    # matching ``left_context_size`` the consumer trims nothing and the
+    # reference audio leaks into the output (issue #4421).  Mirrors the
+    # async-chunk path, which already ships left_context_size in-band.
+    if ref_code is not None and ref_frames > 0:
+        meta["left_context_size"] = ref_frames
     return {
         "codes": {"audio": codec_codes},
-        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+        "meta": meta,
     }