[Bugfix][Perf] mimo_audio: fix per-request speech-code routing under continuous batching (#5070)

Signed-off-by: JuanPZuluaga <juanz9312@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/model_executor/models/mimo_audio/test_mimo_audio_per_request_codes.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiMo-Audio per-request speech-code routing.
+
+Guards the cross-request audio-corruption class (same class as RFC #4316
+Issue 1 for MOSS-TTS): the fused thinker+talker emits one speech-code row per
+request, so ``codes["audio"]`` must be a batch-ordered **list** for the
+generic splitter (``to_payload_element``) to route ``element[idx]`` to
+request ``idx``. A batched ``(num_reqs, 1, 8, 4)`` tensor is only split
+correctly by the token-offset path when every request has exactly one
+scheduled token; in a mixed prefill+decode step the tensor's first dim
+matches neither token length, the splitter falls through to ``clone()``, and
+every request receives the whole batch — leaking codes across requests and
+crashing the async-chunk ragged-list conversion
+(``expected sequence of length 36 at dim 1 (got 144)``).
+
+CPU-only and weight-free: ``forward`` is driven on a bare model with
+``generate_codes`` stubbed to return row-indexed codes, so assertions on the
+payload prove which request got which row.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from vllm_omni.utils.mm_outputs import to_payload_element  # noqa: E402
+
+NUM_CODEBOOKS = 8
+GROUP_SIZE = 4
+
+
+@functools.lru_cache(maxsize=1)
+def _model_cls():
+    """Defer the import (pulls vLLM model_executor) until first use."""
+    from vllm_omni.model_executor.models.mimo_audio.mimo_audio import (
+        MiMoAudioForConditionalGeneration,
+    )
+
+    return MiMoAudioForConditionalGeneration
+
+
+def _row_indexed_codes(num_reqs: int) -> torch.Tensor:
+    """(num_reqs, 1, 8, 4) tensor where request r's codes all hold r."""
+    codes = torch.empty((num_reqs, 1, NUM_CODEBOOKS, GROUP_SIZE), dtype=torch.long)
+    for r in range(num_reqs):
+        codes[r] = r
+    return codes
+
+
+def _make_bare_model(next_speech_tokens, num_tokens: int):
+    """A model whose only behaviour is the emit logic under test."""
+    cls = _model_cls()
+    model = cls.__new__(cls)
+    model.model_stage = "fused_thinker_talker"
+    hidden = torch.zeros((num_tokens, 16), dtype=torch.float32)
+    model.generate_codes = lambda **kwargs: (next_speech_tokens, hidden)
+    return model
+
+
+def _emit(next_speech_tokens, num_tokens: int):
+    model = _make_bare_model(next_speech_tokens, num_tokens)
+    result = model.forward(
+        input_ids=torch.zeros((num_tokens,), dtype=torch.long),
+        positions=torch.zeros((num_tokens,), dtype=torch.long),
+    )
+    return result.multimodal_outputs["codes"]["audio"]
+
+
+class TestPerRequestCodeRouting:
+    def test_emit_is_batch_aligned_list(self):
+        audio = _emit(_row_indexed_codes(3), num_tokens=3)
+
+        assert isinstance(audio, list)
+        assert len(audio) == 3  # one entry per request, batch-aligned
+        for r, codes in enumerate(audio):
+            assert codes.shape == (1, 1, NUM_CODEBOOKS, GROUP_SIZE)
+            assert torch.all(codes == r)
+
+    def test_no_codes_step_emits_none(self):
+        audio = _emit(None, num_tokens=2)
+        assert audio is None
+
+    def test_mixed_prefill_decode_routes_by_request_index(self):
+        """req 1 prefills 7 tokens while reqs 0 and 2 decode: the splitter
+        must hand each request its own row, not the whole batch."""
+        audio = _emit(_row_indexed_codes(3), num_tokens=9)
+
+        seq_len = 9  # 1 (decode) + 7 (prefill) + 1 (decode) scheduled tokens
+        spans = [(0, 1), (1, 8), (8, 9)]
+        for idx, (start, end) in enumerate(spans):
+            routed = to_payload_element(
+                audio,
+                idx,
+                start,
+                end,
+                seq_len=seq_len,
+                scheduled_seq_len=seq_len,
+            )
+            assert routed.shape == (1, 1, NUM_CODEBOOKS, GROUP_SIZE)
+            assert torch.all(routed == idx), f"req {idx} received another request's codes"
+
+    def test_pure_decode_routes_by_request_index(self):
+        """One token per request (the only case the old tensor emit handled)."""
+        audio = _emit(_row_indexed_codes(4), num_tokens=4)
+
+        for idx in range(4):
+            routed = to_payload_element(audio, idx, idx, idx + 1, seq_len=4, scheduled_seq_len=4)
+            assert torch.all(routed == idx)
+
+    def test_routed_codes_do_not_alias_the_emit(self):
+        """The splitter clones list elements; mutating the routed payload must
+        not corrupt another request's view of the emit."""
+        audio = _emit(_row_indexed_codes(2), num_tokens=5)
+        routed = to_payload_element(audio, 0, 0, 4, seq_len=5, scheduled_seq_len=5)
+        routed.fill_(99)
+        assert torch.all(audio[0] == 0)

vllm_omni/entrypoints/openai/serving_chat.py

@@ -2539,7 +2539,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         final_res = omni_outputs.request_output
         # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput
         # Reference: examples/offline_inference/qwen3_omni/end2end.py line 421
-        mm_output = final_res.outputs[0].multimodal_output
+        # The attribute is attached dynamically when stage audio arrives; fall
+        # back to the no-audio error response instead of an AttributeError 500
+        # when the pipeline produced no audio for this request.
+        mm_output = getattr(final_res.outputs[0], "multimodal_output", None) or {}
         audio_data = mm_output.get("audio")
         if isinstance(audio_data, list):
             if not audio_data:

vllm_omni/model_executor/models/mimo_audio/mimo_audio.py

@@ -793,9 +793,21 @@ class MiMoAudioForConditionalGeneration(
                 **kwargs,
             )
 
+            # next_speech_tokens is request-indexed (num_reqs, 1, 8, 4), not
+            # token-indexed, so ship it as a per-request list. The payload
+            # builder slices batched tensors by token offsets (start:end),
+            # which only coincides with the request index in pure-decode
+            # steps; in a mixed prefill+decode batch every request would
+            # receive the whole batch tensor, leaking codes across requests
+            # and crashing the downstream ragged-list conversion.
+            if next_speech_tokens is not None:
+                speech_token_payload = list(next_speech_tokens.split(1, dim=0))
+            else:
+                speech_token_payload = None
+
             return OmniOutput(
                 text_hidden_states=text_hidden_states.reshape(-1, text_hidden_states.shape[-1]),
-                multimodal_outputs={"codes": {"audio": next_speech_tokens}},
+                multimodal_outputs={"codes": {"audio": speech_token_payload}},
             )
 
         if self.model_stage == "code2wav":

vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py

@@ -897,20 +897,28 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         input_ids: torch.Tensor,
         *,
         request_ids: list[str],
-        query_start_loc: torch.Tensor,
+        query_start_loc_list: list[int],
         kwargs: dict,
     ) -> tuple[dict[str, any], dict]:
         has_merge_mm_embedding = False
         merge_mm_embedding_info: dict[str, any] = {}
         seq_len = input_ids.shape[1] if input_ids.ndim == 2 else input_ids.shape[0]
 
-        for req_idx, req_id in enumerate(request_ids):
-            query_start_loc_by_req = int(query_start_loc[req_idx].item())
-            query_end_loc_by_req = int(query_start_loc[req_idx + 1].item())
-            input_ids_by_req = input_ids[query_start_loc_by_req:query_end_loc_by_req]
-            seq_len_by_req = input_ids_by_req.shape[1] if input_ids_by_req.ndim == 2 else input_ids_by_req.shape[0]
+        # Only each request's start token is ever read below, and only for
+        # single-token requests. Gather just those start offsets (num_reqs elements)
+        # to host instead of copying the whole flattened sequence every forward.
+        # query_start_loc entries are flat token offsets, so we index the flattened
+        # ids by the same offsets and derive per-request length from the boundaries.
+        flat_input_ids = input_ids.reshape(-1)
+        start_positions = query_start_loc_list[:-1]
+        start_is_empty_cpu = (flat_input_ids[start_positions] == self.empty_token_id).cpu().tolist()
 
-            if seq_len_by_req == 1 and bool(input_ids_by_req == self.empty_token_id):
+        for req_idx, req_id in enumerate(request_ids):
+            query_start_loc_by_req = query_start_loc_list[req_idx]
+            query_end_loc_by_req = query_start_loc_list[req_idx + 1]
+            seq_len_by_req = query_end_loc_by_req - query_start_loc_by_req
+
+            if seq_len_by_req == 1 and start_is_empty_cpu[req_idx]:
                 merge_mm_embedding_info[req_id] = {
                     "query_start_loc": query_start_loc_by_req,
                     "query_end_loc": query_end_loc_by_req,
@@ -1035,13 +1043,19 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         inputs_embeds: torch.Tensor | None = None,
         **kwargs: object,
     ) -> torch.Tensor | IntermediateTensors:
+        # Materialize the request boundaries on the host once instead of one
+        # `.item()` sync per boundary in the loops below.
         _forward_context = get_forward_context()
-        _default_query_start_loc = torch.tensor([0, input_ids.shape[-1]], device=input_ids.device)
-        query_start_loc = (
-            next(iter(_forward_context.attn_metadata.values())).query_start_loc
-            if _forward_context.attn_metadata is not None
-            else _default_query_start_loc
-        )
+        if _forward_context.attn_metadata is not None:
+            query_start_loc = next(iter(_forward_context.attn_metadata.values())).query_start_loc
+            query_start_loc_list = query_start_loc.tolist()
+        else:
+            # Single implicit request over all tokens. Build the host boundaries
+            # directly (no .tolist() D2H); the small device tensor is still needed
+            # for logits_indices below.
+            n_tokens = int(input_ids.shape[-1])
+            query_start_loc_list = [0, n_tokens]
+            query_start_loc = torch.tensor([0, n_tokens], device=input_ids.device)
 
         runtime_additional_information = kwargs.get("runtime_additional_information", [])
         if runtime_additional_information:
@@ -1054,7 +1068,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         merge_mm_embedding_info, has_merge_mm_embedding, kwargs = self._collect_merge_mm_embedding_info(
             input_ids,
             request_ids=request_ids,
-            query_start_loc=query_start_loc,
+            query_start_loc_list=query_start_loc_list,
             kwargs=kwargs,
         )
 
@@ -1077,15 +1091,18 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         batch_next_speech_tokens: torch.Tensor | None = None
 
         if not is_capturing and next_ids is not None and num_reqs > 0:
-            if (next_ids == self.empty_token_id).any():
+            # Copy the empty-token mask to host once; the `.any()` gate and the
+            # per-request checks below then avoid a GPU sync each.
+            next_is_empty_cpu = (next_ids == self.empty_token_id).cpu()
+            if bool(next_is_empty_cpu.any()):
                 batch_hs_list = []
                 valid_mask = []
 
                 for req_idx in range(num_reqs):
-                    start = int(query_start_loc[req_idx].item())
-                    end = int(query_start_loc[req_idx + 1].item())
+                    start = query_start_loc_list[req_idx]
+                    end = query_start_loc_list[req_idx + 1]
                     hs_req = hidden_states[start:end][-1:, :]
-                    is_empty = bool(next_ids[req_idx] == self.empty_token_id)
+                    is_empty = bool(next_is_empty_cpu[req_idx])
                     valid_mask.append(is_empty)
 
                     if not is_empty: