[Bugfix] Fix cross-request codes.ref leak in Qwen3-TTS make_omni_output (#4373)

Signed-off-by: jjw <m@jjw.in>
Co-authored-by: jjw <m@jjw.in>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for cross-request ``codes.ref`` leakage (#4370).
+
+``make_omni_output`` used to collapse every request's reference codec
+frames into a single last-writer-wins slot and emit it as a length-1
+list. ``to_payload_element`` indexes list payloads with
+``element[idx] if idx < len(element) else element[0]``, so that length-1
+list was broadcast to every request in the batch: each utterance's first
+vocoder chunks then decoded with another request's reference voice as
+context (audible onset timbre deformation at any concurrency > 1).
+"""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+    Qwen3TTSTalkerForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_NUM_CODE_GROUPS = 16
+
+
+def _make_talker() -> Qwen3TTSTalkerForConditionalGeneration:
+    # make_omni_output reads no instance state, so skip __init__ (no real
+    # config / checkpoint needed), same as the preprocess tests.
+    return Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+
+
+def _info(span_frames: int, ref_code: torch.Tensor | None) -> dict:
+    codes: dict = {"audio": torch.zeros((span_frames, _NUM_CODE_GROUPS), dtype=torch.long)}
+    meta: dict = {}
+    if ref_code is not None:
+        codes["ref"] = ref_code
+        meta["ref_code_len"] = int(ref_code.shape[0])
+    return {"codes": codes, "meta": meta}
+
+
+def test_make_omni_output_keeps_ref_codes_per_request() -> None:
+    ref_a = torch.ones((3, _NUM_CODE_GROUPS), dtype=torch.long)
+    ref_b = torch.full((5, _NUM_CODE_GROUPS), 2, dtype=torch.long)
+
+    out = _make_talker().make_omni_output(
+        torch.zeros((4, 8)),
+        model_intermediate_buffer=[_info(2, ref_a), _info(2, ref_b)],
+    )
+
+    ref_list = out.multimodal_outputs["codes"]["ref"]
+    assert len(ref_list) == 2, "codes.ref must stay batch-aligned (one entry per request)"
+    assert torch.equal(ref_list[0], ref_a)
+    assert torch.equal(ref_list[1], ref_b)
+
+
+def test_make_omni_output_pads_requests_without_ref_code() -> None:
+    ref_b = torch.full((5, _NUM_CODE_GROUPS), 2, dtype=torch.long)
+
+    out = _make_talker().make_omni_output(
+        torch.zeros((4, 8)),
+        model_intermediate_buffer=[_info(2, None), _info(2, ref_b)],
+    )
+
+    ref_list = out.multimodal_outputs["codes"]["ref"]
+    assert len(ref_list) == 2
+    assert ref_list[0].numel() == 0, "no-ref request must get an empty placeholder, not a neighbor's ref"
+    assert torch.equal(ref_list[1], ref_b)
+
+
+def test_make_omni_output_omits_ref_key_when_no_request_has_one() -> None:
+    out = _make_talker().make_omni_output(
+        torch.zeros((4, 8)),
+        model_intermediate_buffer=[_info(2, None), _info(2, None)],
+    )
+
+    assert "ref" not in out.multimodal_outputs["codes"]

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -538,10 +538,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
         audio_codes_list: list[torch.Tensor] = []
         ref_code_len_list: list[torch.Tensor] = []
-        ref_code_tensor: torch.Tensor | None = None
+        ref_code_list: list[torch.Tensor] = []
+        has_ref_code = False
         codec_streaming_list: list[torch.Tensor] = []
         for info in info_dicts:
             if not isinstance(info, dict):
+                ref_code_list.append(torch.empty(0, dtype=torch.long))
                 continue
             codes = info.get("codes", {})
             meta = info.get("meta", {})
@@ -555,7 +557,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     )
             ref_code = codes.get("ref")
             if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
-                ref_code_tensor = ref_code
+                ref_code_list.append(ref_code)
+                has_ref_code = True
+            else:
+                ref_code_list.append(torch.empty(0, dtype=torch.long))
             ref_len = meta.get("ref_code_len")
             if ref_len is None:
                 continue
@@ -585,8 +590,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         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]
-        if ref_code_tensor is not None:
-            mm.setdefault("codes", {})["ref"] = [ref_code_tensor]
+        if has_ref_code:
+            # Batch-aligned, one entry per request: ``to_payload_element``
+            # indexes ``element[idx]`` per request, and a shorter list would
+            # silently broadcast one request's reference codes to the others.
+            mm.setdefault("codes", {})["ref"] = ref_code_list
         if codec_streaming_list:
             mm.setdefault("meta", {})["codec_streaming"] = torch.cat(codec_streaming_list, dim=0)[:span_len]
         return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)