[Fix][Qwen3-TTS] Preserve ref_code decoder context for Base ICL (#1731)

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -7,7 +7,10 @@ from types import SimpleNamespace
 import pytest
 import torch
 
-from vllm_omni.model_executor.stage_input_processors.qwen3_tts import talker2code2wav_async_chunk
+from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+    talker2code2wav,
+    talker2code2wav_async_chunk,
+)
 
 _FRAME = [1, 2, 3, 4]  # 4-codebook frame
 _Q = len(_FRAME)  # num quantizers
@@ -29,6 +32,7 @@ def _tm(*, chunk_frames=25, left_context=25, initial_chunk=0):
     return SimpleNamespace(
         code_prompt_token_ids=defaultdict(list),
         put_req_chunk=defaultdict(int),
+        request_payload={},
         connector=SimpleNamespace(
             config={
                 "extra": {
@@ -152,3 +156,113 @@ def test_per_request_override_wins_over_stage_config():
     tm = _tm(initial_chunk=5)
     payload = _call(tm, "r-override2", n_frames=10, put_req=0, req_ic=15)
     assert payload is None
+
+
+def test_first_streaming_chunk_prepends_ref_code_context():
+    tm = _tm(initial_chunk=10)
+    rid = "r-ref"
+    tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(10)]
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"audio_codes": torch.zeros((0,)), "ref_code": ref_code},
+        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+        is_finished=False,
+    )
+
+    assert payload is not None
+    assert payload["left_context_size"] == 2
+    assert len(payload["code_predictor_codes"]) == _Q * 12
+
+
+def test_ref_code_context_only_applies_to_first_streaming_chunk():
+    tm = _tm(initial_chunk=10)
+    rid = "r-ref2"
+    tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(20)]
+    tm.put_req_chunk[rid] = 1
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"audio_codes": torch.zeros((0,)), "ref_code": ref_code},
+        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+        is_finished=False,
+    )
+
+    assert payload is not None
+    assert payload["left_context_size"] == 10
+    assert len(payload["code_predictor_codes"]) == _Q * 20
+
+
+def test_ref_code_context_can_be_buffered_before_first_emit():
+    tm = _tm(initial_chunk=10)
+    rid = "r-ref-buffered"
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+
+    first_payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]]), "ref_code": ref_code},
+        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+        is_finished=False,
+    )
+    assert first_payload is None
+    assert rid in tm.request_payload
+
+    for _ in range(8):
+        talker2code2wav_async_chunk(
+            transfer_manager=tm,
+            pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]])},
+            request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+            is_finished=False,
+        )
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]])},
+        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
+        is_finished=False,
+    )
+
+    assert payload is not None
+    assert payload["left_context_size"] == 2
+    assert len(payload["code_predictor_codes"]) == _Q * 12
+    assert rid not in tm.request_payload
+
+
+def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+    audio_codes = torch.tensor(
+        [
+            [0, 0, 0, 0],
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+        ],
+        dtype=torch.long,
+    )
+    output = SimpleNamespace(multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code})
+    stage = SimpleNamespace(engine_outputs=[SimpleNamespace(outputs=[output])])
+
+    prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
+
+    assert len(prompts) == 1
+    prompt = prompts[0]
+    assert prompt["additional_information"] == {"left_context_size": [2]}
+    assert prompt["prompt_token_ids"] == [
+        9,
+        8,
+        1,
+        5,
+        9,
+        8,
+        2,
+        6,
+        9,
+        8,
+        3,
+        7,
+        9,
+        8,
+        4,
+        8,
+    ]