[Bugfix] Truncate mimo-audio code2wav prompt to MAX_CODE2WAV_TOKENS (#2693)

Signed-off-by: lishunyang <lishunyang12@163.com>

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

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,