From 618268d6b9e802545829e92bf5d779367c625f61 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Tue, 21 Apr 2026 00:40:29 +0800 Subject: [PATCH] [Bugfix] Truncate mimo-audio code2wav prompt to MAX_CODE2WAV_TOKENS (#2693) Signed-off-by: lishunyang --- .../test_mimo_audio_llm2code2wav.py | 72 +++++++++++++++++++ .../stage_input_processors/mimo_audio.py | 22 ++++++ 2 files changed, 94 insertions(+) create mode 100644 tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py diff --git a/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py b/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py new file mode 100644 index 000000000..1c14ab0cf --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import logging +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors import mimo_audio as sip +from vllm_omni.model_executor.stage_input_processors.mimo_audio import ( + MAX_CODE2WAV_TOKENS, + llm2code2wav, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_stage_list(codec_codes: torch.Tensor, request_id: str = "req-0"): + """Build a minimal stage_list[0] with a single talker output carrying codec_codes.""" + output = SimpleNamespace(multimodal_output={"code_predictor_codes": codec_codes}) + talker_output = SimpleNamespace(outputs=[output], request_id=request_id) + stage0 = SimpleNamespace(engine_outputs=[talker_output]) + return [stage0] + + +def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog): + """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated, not passed through.""" + # prepend_and_flatten_colmajor produces 36 ids per (8, 4) codec frame: + # pad adds one row -> (9, 4) per frame, permuted and flattened. + # Pick enough frames to comfortably exceed the cap. + frames = (MAX_CODE2WAV_TOKENS // 36) + 100 + codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long) + + stage_list = _make_stage_list(codec_codes, request_id="req-long") + + # Attach caplog's handler directly to the module logger so the warning is + # captured regardless of propagation (vllm's logger configuration can + # interact badly with caplog.at_level's default root-handler path). + target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.mimo_audio") + target_logger.addHandler(caplog.handler) + prev_level = target_logger.level + target_logger.setLevel(logging.WARNING) + try: + prompts = llm2code2wav(stage_list, engine_input_source=[0]) + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(prev_level) + + assert len(prompts) == 1 + assert len(prompts[0]["prompt_token_ids"]) == MAX_CODE2WAV_TOKENS + assert any("truncating" in rec.getMessage() for rec in caplog.records), ( + f"Expected a 'truncating' warning; captured records: {[r.getMessage() for r in caplog.records]}" + ) + + +def test_llm2code2wav_short_sequence_unchanged(): + """Short codec sequences are returned without truncation.""" + codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long) + stage_list = _make_stage_list(codec_codes, request_id="req-short") + + prompts = llm2code2wav(stage_list, engine_input_source=[0]) + + assert len(prompts) == 1 + # 4 frames + 1 pad row, flattened col-major → well below the cap + assert 0 < len(prompts[0]["prompt_token_ids"]) <= MAX_CODE2WAV_TOKENS + + +def test_llm2code2wav_truncation_boundary_constant_matches_yaml(): + """MAX_CODE2WAV_TOKENS must match the stage-1 max_model_len in mimo_audio.yaml and end2end.py.""" + assert sip.MAX_CODE2WAV_TOKENS == 18192 diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py index 71573c0e2..39e17b2fe 100644 --- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py +++ b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py @@ -9,6 +9,14 @@ from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_ logger = init_logger(__name__) +# Maximum tokens supported by the code2wav stage. The flattened talker codec +# sequence fed to stage-1 must not exceed this, otherwise gpu_input_batch +# add_request will fail with a broadcast error when copying prompt_token_ids +# into token_ids_cpu. Keep in sync with the stage-1 ``max_model_len`` in +# ``vllm_omni/model_executor/stage_configs/mimo_audio.yaml`` and the offline +# example ``examples/offline_inference/mimo_audio/end2end.py``. +MAX_CODE2WAV_TOKENS = 18192 + def prepend_and_flatten_colmajor(x: torch.Tensor, pad_vec: torch.Tensor) -> torch.Tensor: """ @@ -235,6 +243,20 @@ def llm2code2wav( code_final = prepend_and_flatten_colmajor(codec_codes, pad_vec) code_final = code_final.tolist() + # Guard against flattened sequences longer than code2wav's max_model_len. + # Without this, add_request raises ``could not broadcast input array + # from shape (N,) into shape (max_model_len,)`` and kills the engine + # core (see issue #2683). Mirrors the offline end2end.py safeguard. + if len(code_final) > MAX_CODE2WAV_TOKENS: + request_id = getattr(talker_output, "request_id", f"unknown_{i}") + logger.warning( + "Request %s: code_final len=%d > MAX_CODE2WAV_TOKENS=%d, truncating.", + request_id, + len(code_final), + MAX_CODE2WAV_TOKENS, + ) + code_final = code_final[:MAX_CODE2WAV_TOKENS] + code2wav_inputs.append( OmniTokensPrompt( prompt_token_ids=code_final,