From dacb8e6f10f06ffee05f6705f61f084e186470a4 Mon Sep 17 00:00:00 2001 From: Juan Pablo Zuluaga <46724788+JuanPZuluaga@users.noreply.github.com> Date: Mon, 20 Jul 2026 19:22:45 +0200 Subject: [PATCH] [Bugfix][Perf] mimo_audio: fix per-request speech-code routing under continuous batching (#5070) Signed-off-by: JuanPZuluaga Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../test_mimo_audio_per_request_codes.py | 120 ++++++++++++++++++ vllm_omni/entrypoints/openai/serving_chat.py | 5 +- .../models/mimo_audio/mimo_audio.py | 14 +- .../models/mimo_audio/mimo_audio_llm.py | 53 +++++--- 4 files changed, 172 insertions(+), 20 deletions(-) create mode 100644 tests/model_executor/models/mimo_audio/test_mimo_audio_per_request_codes.py diff --git a/tests/model_executor/models/mimo_audio/test_mimo_audio_per_request_codes.py b/tests/model_executor/models/mimo_audio/test_mimo_audio_per_request_codes.py new file mode 100644 index 000000000..4dc365ef0 --- /dev/null +++ b/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) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index eec6a03b9..3f6250be8 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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: diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index ecf094c87..2af6abd42 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/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": diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index 48715fb98..19a339056 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/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: