Repository navigation
[Bugfix] Fix cross-request codes.ref leak in Qwen3-TTS make_omni_output #4373
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| 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"] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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] = [] | ||
| 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 @@ 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)) | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [minor] The placeholder is 1-D |
||
| ref_len = meta.get("ref_code_len") | ||
| if ref_len is None: | ||
| continue | ||
|
|
@@ -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) | ||
|
|
||
There was a problem hiding this comment.
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 sizethen holds by construction, and bothappend(torch.empty(0))sites further down drop out:Minor — but it means a future
continueadded to this loop can't silently desync the list from the batch.