Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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"]
16 changes: 12 additions & 4 deletions vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
Original file line number Diff line number Diff line change
Expand Up @@ -538,10 +538,12 @@ def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: A
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] = []

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could pre-size this to len(info_dicts) and assign by index rather than appending — len(ref_code_list) == batch size then holds by construction, and both append(torch.empty(0)) sites further down drop out:

ref_code_list: list[torch.Tensor] = [
    torch.empty(0, dtype=torch.long) for _ in info_dicts
]
...
for i, info in enumerate(info_dicts):
    if not isinstance(info, dict):
        continue
    ...
    if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
        ref_code_list[i] = ref_code
        has_ref_code = True

Minor — but it means a future continue added to this loop can't silently desync the list from the batch.

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", {})
Expand All @@ -555,7 +557,10 @@ def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: A
)
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))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[minor] The placeholder is 1-D (0,) while real refs are 2-D (N, num_code_groups). Every current consumer only checks numel(), so this works today, but a (0, num_code_groups) placeholder would be shape-consistent if a downstream ever concatenates or reshapes these. Fine to leave as is.

ref_len = meta.get("ref_code_len")
if ref_len is None:
continue
Expand Down Expand Up @@ -585,8 +590,11 @@ def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: A
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)
Expand Down
Loading