[Enhancement] Custom chunk_size for mimo-audio model (#1964)

Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>

tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID
+from vllm_omni.model_executor.models.mimo_audio.mimo_audio_code2wav import (
+    AudioStreamerConfig,
+    MiMoAudioToken2WavForConditionalGenerationVLLM,
+    flat_codec_group_element_count,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_GROUP = 4
+_AC = 8
+_GROUP_WIDTH = flat_codec_group_element_count(_GROUP, _AC)
+_FTP = 2 * 2 * 240  # frames_per_token from mocked tokenizer.config
+
+
+def _codes_ns(empty: int = 555, eostm: int = 666):
+    return SimpleNamespace(empty=empty, eostm=eostm)
+
+
+def _make_dummy_code_tensor() -> torch.Tensor:
+    """Pad-only talker dummy; matches _check_dummy_code_tensor."""
+    t = torch.zeros(flat_codec_group_element_count(_GROUP, _AC), dtype=torch.long)
+    t = t.view(_GROUP, _AC + 1)
+    t[:, 0] = TALKER_CODEC_PAD_TOKEN_ID
+    return t.view(-1)
+
+
+def _make_valid_flat_codes(num_groups: int = 1, *, empty_id: int = 555) -> torch.Tensor:
+    """Flat layout that extract_audio_code_tensor accepts (text col == empty)."""
+    parts = []
+    for _ in range(num_groups):
+        g = torch.zeros(_GROUP, _AC + 1, dtype=torch.long)
+        g[0, 0] = empty_id
+        g[:, 1:] = torch.randint(1, 50, (_GROUP, _AC))
+        parts.append(g.reshape(-1))
+    return torch.cat(parts)
+
+
+def _make_invalid_flat_immediate_eostm(eostm_id: int = 666) -> torch.Tensor:
+    g = torch.zeros(_GROUP, _AC + 1, dtype=torch.long)
+    g[0, 0] = eostm_id
+    return g.reshape(-1)
+
+
+def _minimal_model():
+    """Avoid __init__ (HF tokenizer paths); only fields used by _batch_decode_waveforms."""
+    model = object.__new__(MiMoAudioToken2WavForConditionalGenerationVLLM)
+    model.device = torch.device("cpu")
+    model.config = SimpleNamespace(group_size=_GROUP, audio_channels=_AC)
+    model.streamer_config = AudioStreamerConfig(group_size=_GROUP, audio_channels=_AC)
+    model.codes = _codes_ns()
+
+    decode_vq = Mock(
+        side_effect=lambda audio_codes: torch.ones(
+            audio_codes.shape[1],
+            7,
+            dtype=torch.float32,
+            device=audio_codes.device,
+        )
+    )
+    decoder = Mock()
+
+    audio_tok = SimpleNamespace(
+        encoder=SimpleNamespace(decode_vq=decode_vq),
+        decoder=decoder,
+        config=SimpleNamespace(avg_pooler=2, stride_size=2, hop_length=240),
+    )
+    model._tokenizer_service = SimpleNamespace(audio_tokenizer=audio_tok)
+    return model, audio_tok
+
+
+def test_batch_decode_waveforms_empty_input_list():
+    """Empty input list returns a single zero-length float32 tensor on model device."""
+    model, _ = _minimal_model()
+    out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, [])
+    assert len(out) == 1
+    assert out[0].dtype == torch.float32
+    assert out[0].numel() == 0
+    assert out[0].device == model.device
+
+
+def test_batch_decode_waveforms_single_vs_multiple_decoder_shapes():
+    """Single and multi-request batches produce correctly shaped packed hidden states and trimmed waveforms."""
+    model, audio_tok = _minimal_model()
+    decoder = audio_tok.decoder
+
+    # Single valid request: decoder output rank-3 for double squeeze path
+    flat = _make_valid_flat_codes(1)
+    decoder.return_value = torch.ones(1, 1, 4 * _FTP + 100, dtype=torch.float32)
+    out1 = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, [flat])
+    decoder.assert_called_once()
+    packed_hs, input_lengths = decoder.call_args[0]
+    assert packed_hs.shape == (4, 7)  # T from decode_vq mock
+    assert input_lengths.tolist() == [4]
+    assert out1[0].shape == (4 * _FTP,)
+    assert out1[0].dtype == torch.float32
+
+    decoder.reset_mock()
+    a = _make_valid_flat_codes(1)
+    b = _make_valid_flat_codes(2)
+    decoder.return_value = torch.ones(2, 1, 8 * _FTP + 100, dtype=torch.float32)
+    out2 = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, [a, b])
+    decoder.assert_called_once()
+    packed_hs2, input_lengths2 = decoder.call_args[0]
+    assert packed_hs2.shape == (4 + 8, 7)  # 1 group -> T=4; 2 groups -> T=8
+    assert input_lengths2.tolist() == [4, 8]
+    assert len(out2) == 2
+    assert out2[0].shape == (4 * _FTP,)
+    assert out2[1].shape == (8 * _FTP,)
+
+
+def test_batch_decode_waveforms_mixed_valid_invalid_requests():
+    """Mixed valid and invalid requests: invalid slots get empty tensors, valid slots get decoded waveforms."""
+    model, audio_tok = _minimal_model()
+    valid_a = _make_valid_flat_codes(1)
+    valid_b = _make_valid_flat_codes(1)
+    dummy = _make_dummy_code_tensor()
+    bad_extract = _make_invalid_flat_immediate_eostm(eostm_id=model.codes.eostm)
+    too_short = torch.tensor([1, 2, 3], dtype=torch.long)
+
+    audio_tok.decoder.return_value = torch.ones(2, 1, 5000, dtype=torch.float32)
+
+    inputs = [
+        None,
+        torch.tensor([], dtype=torch.long),
+        dummy,
+        bad_extract,
+        too_short,
+        valid_a,
+        valid_b,
+    ]
+    out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, inputs)
+
+    assert len(out) == len(inputs)
+    for i in range(5):
+        assert out[i].numel() == 0, f"index {i} should be empty"
+    assert out[5].shape == (4 * _FTP,)
+    assert out[6].shape == (4 * _FTP,)
+    audio_tok.decoder.assert_called_once()
+    packed_hs, input_lengths = audio_tok.decoder.call_args[0]
+    assert packed_hs.shape[0] == 8
+    assert input_lengths.tolist() == [4, 4]
+
+
+def test_batch_decode_waveforms_all_invalid_returns_per_request_empty():
+    """All-invalid batch skips decoder entirely and returns empty tensors for every slot."""
+    model, audio_tok = _minimal_model()
+    out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(
+        model,
+        [None, _make_dummy_code_tensor(), torch.tensor([], dtype=torch.long)],
+    )
+    assert len(out) == 3
+    assert all(t.numel() == 0 for t in out)
+    audio_tok.decoder.assert_not_called()
+
+
+def test_batch_decode_waveforms_output_shape_trim_when_decoder_returns_extra_samples():
+    """Decoder output longer than valid_len is trimmed to the exact expected waveform length."""
+    model, audio_tok = _minimal_model()
+    flat = _make_valid_flat_codes(1)
+    # Longer than valid_len so branch wav = wav[:valid_len] runs
+    audio_tok.decoder.return_value = torch.ones(1, 1, 10_000, dtype=torch.float32)
+    out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, [flat])
+    assert out[0].dim() == 1
+    assert out[0].numel() == 4 * _FTP
+    assert out[0].dtype == torch.float32
+
+
+def test_batch_decode_waveforms_multi_request_trims_each_row_when_decoder_returns_extra():
+    """Else-branch split: per-request wav[:valid_len] when decoder pads each batch row."""
+    model, audio_tok = _minimal_model()
+    a = _make_valid_flat_codes(1)
+    b = _make_valid_flat_codes(2)
+    audio_tok.decoder.return_value = torch.ones(2, 1, 10_000, dtype=torch.float32)
+    out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, [a, b])
+    assert len(out) == 2
+    assert out[0].shape == (4 * _FTP,)
+    assert out[1].shape == (8 * _FTP,)
+    assert out[0].dtype == torch.float32
+    assert out[1].dtype == torch.float32
+
+
+def test_batch_decode_waveforms_valid_only_at_edges_maps_to_correct_indices():
+    """Tensor packing order must match valid_indices when invalid requests are in the middle."""
+    model, audio_tok = _minimal_model()
+    first = _make_valid_flat_codes(1)
+    last = _make_valid_flat_codes(2)
+    inputs = [
+        first,
+        None,
+        _make_dummy_code_tensor(),
+        last,
+    ]
+    audio_tok.decoder.return_value = torch.ones(2, 1, 10_000, dtype=torch.float32)
+    out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, inputs)
+    assert len(out) == 4
+    assert out[0].numel() == 4 * _FTP
+    assert out[1].numel() == 0
+    assert out[2].numel() == 0
+    assert out[3].numel() == 8 * _FTP
+    packed_hs, input_lengths = audio_tok.decoder.call_args[0]
+    assert packed_hs.shape[0] == 12
+    assert input_lengths.tolist() == [4, 8]
+
+
+def test_batch_decode_waveforms_output_shapes_1d_float32_for_all_slots():
+    """Every slot is a 1-D float32 vector (empty or waveform), matching downstream expectations."""
+    model, audio_tok = _minimal_model()
+    inputs = [_make_valid_flat_codes(1), None, _make_valid_flat_codes(1)]
+    audio_tok.decoder.return_value = torch.ones(2, 1, 5000, dtype=torch.float32)
+    out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, inputs)
+    assert len(out) == 3
+    for i, t in enumerate(out):
+        assert t.dim() == 1, f"slot {i}"
+        assert t.dtype == torch.float32, f"slot {i}"
+    assert out[0].numel() == 4 * _FTP
+    assert out[1].numel() == 0
+    assert out[2].numel() == 4 * _FTP

tests/model_executor/stage_input_processors/test_mimo_audio_flush_remaining_codes.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.mimo_audio import _flush_remaining_codes
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _sentinel():
+    return {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
+
+
+def test_flush_remaining_codes_when_no_codes_accumulated_missing_request_id():
+    """No entry for request_id: treat as empty, return finished sentinel with empty codes."""
+    tm = SimpleNamespace(code_prompt_token_ids={})
+    out = _flush_remaining_codes(tm, "missing", chunk_size=3, left_context_size=3)
+    assert out["code_predictor_codes"] == _sentinel()["code_predictor_codes"]
+    assert out["finished"].equal(_sentinel()["finished"])
+
+
+def test_flush_remaining_codes_when_no_codes_accumulated_empty_list():
+    """Explicit empty accumulation list returns the same sentinel."""
+    tm = SimpleNamespace(code_prompt_token_ids={"r": []})
+    out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
+    assert out["code_predictor_codes"] == []
+    assert out["finished"].item() is True
+
+
+def test_flush_remaining_codes_partial_chunk_remaining():
+    """
+    length=7, chunk_size=3 -> remainder 1 frame in last chunk; context window uses
+    left_context + that partial (3+1), capped by length -> last 4 scalar entries.
+    """
+    tm = SimpleNamespace(
+        code_prompt_token_ids={"r": [[1], [2], [3], [4], [5], [6], [7]]},
+    )
+    out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
+    assert out["finished"].item() is True
+    assert out["code_predictor_codes"] == [4, 5, 6, 7]
+
+
+def test_flush_remaining_codes_when_length_is_exact_multiple_of_chunk_size():
+    """length % chunk_size == 0 uses full last chunk as context_length (not remainder)."""
+    tm = SimpleNamespace(
+        code_prompt_token_ids={"r": [[1], [2], [3], [4], [5], [6]]},
+    )
+    out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
+    # context_length = chunk_size = 3, end_index = min(6, 6) -> all 6
+    assert out["code_predictor_codes"] == [1, 2, 3, 4, 5, 6]
+
+
+@pytest.mark.parametrize(
+    "length,chunk_size,left_context,expected_end_index",
+    [
+        (2, 3, 3, 2),  # chunk_length=2 -> context 2, min(2,5)=2
+        (5, 3, 3, 5),  # chunk_length=2 -> min(5,5)=5
+        (7, 3, 3, 4),  # chunk_length=1 -> min(7,4)=4
+        (6, 3, 3, 6),  # chunk_length=0 -> context 3, min(6,6)=6
+        (10, 3, 3, 4),  # chunk_length=1 -> min(10,4)=4
+        (1, 5, 10, 1),  # chunk_length=1 -> min(1,11)=1
+    ],
+)
+def test_flush_remaining_codes_context_window_end_index(
+    length: int, chunk_size: int, left_context: int, expected_end_index: int
+) -> None:
+    """Mirror _flush_remaining_codes context_length and end_index rules."""
+    accumulated = [[i] for i in range(length)]
+    tm = SimpleNamespace(code_prompt_token_ids={"r": accumulated})
+    out = _flush_remaining_codes(tm, "r", chunk_size=chunk_size, left_context_size=left_context)
+    expected_flat = list(range(length - expected_end_index, length))
+    assert out["code_predictor_codes"] == expected_flat
+    assert out["finished"].item() is True

vllm_omni/model_executor/models/mimo_audio/mimo_audio.py

@@ -814,11 +814,13 @@ class MiMoAudioForConditionalGeneration(
                 )
             )
 
-            audio_tensor = self.generate_audio(code)
+            audio_result = self.token2wav(codes=code, **kwargs)
+            if isinstance(audio_result, OmniOutput):
+                return audio_result
             return OmniOutput(
                 text_hidden_states=None,
                 multimodal_outputs={
-                    "model_outputs": audio_tensor.reshape(1, -1) if audio_tensor is not None else audio_tensor
+                    "model_outputs": audio_result.reshape(1, -1) if audio_result is not None else audio_result
                 },
             )
 

vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py

@@ -4,6 +4,7 @@ import os
 import time
 from collections.abc import Iterable
 from dataclasses import dataclass
+from typing import Any
 
 import torch
 import torchaudio
@@ -20,11 +21,21 @@ from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID, MiMoAudioConfig
 from vllm_omni.model_executor.models.mimo_audio.modeling_audio_tokenizer import MiMoAudioTokenizer
+from vllm_omni.model_executor.models.output_templates import OmniOutput
 
 logger = logging.getLogger(__name__)
 
 
-DUMMY_CODE_SHAPE = 36
+def flat_codec_group_element_count(group_size: int, audio_channels: int) -> int:
+    """Flat token count for one MiMo talker codec group on the code2wav wire.
+
+    The talker flattens a ``(group_size, audio_channels + 1)`` layout in column-major
+    order: one leading text/special column per row plus ``audio_channels`` RVQ code
+    columns. Each group therefore occupies ``group_size * (audio_channels + 1)``
+    consecutive ids. This matches ``group_width`` in :func:`extract_audio_code_tensor`
+    and is used to count how many such groups are present in a flat ``codes`` tensor.
+    """
+    return group_size * (audio_channels + 1)
 
 
 class MiMoAudioTokenizerWorker:
@@ -308,7 +319,7 @@ def extract_audio_code_tensor(
     if flat_codes.numel() == 0:
         return None
 
-    group_width = group_size * (audio_channels + 1)
+    group_width = flat_codec_group_element_count(group_size, audio_channels)
     usable = (flat_codes.numel() // group_width) * group_width
     if usable == 0:
         return None
@@ -377,6 +388,9 @@ def get_tokenizer_worker(
 class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
     """Decode MiMo audio codes to waveform for the code2wav stage."""
 
+    have_multimodal_outputs = True
+    enable_update_additional_information = True
+
     def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
@@ -429,6 +443,23 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
             * getattr(audio_tokenizer_config, "hop_length", 240)
         ) * self.config.group_size
 
+        connector_cfg = getattr(vllm_config.model_config, "stage_connector_config", None)
+        extra_cfg = (
+            (
+                connector_cfg.get("extra", connector_cfg)
+                if isinstance(connector_cfg, dict)
+                else getattr(connector_cfg, "extra", None)
+            )
+            if connector_cfg
+            else None
+        )
+        self._codec_chunk_frames = int(extra_cfg.get("codec_chunk_frames", 3)) if isinstance(extra_cfg, dict) else 3
+        if self._codec_chunk_frames <= 0:
+            raise ValueError(f"codec_chunk_frames must be positive, got {self._codec_chunk_frames}")
+        self._codec_left_context_frames = (
+            int(extra_cfg.get("codec_left_context_frames", 3)) if isinstance(extra_cfg, dict) else 3
+        )
+
     def load_weights(
         self,
         weights: Iterable[tuple[str, torch.Tensor]],
@@ -438,6 +469,61 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         # Decoder has no trainable weights to load in this stage.
         return set()
 
+    @staticmethod
+    def _split_flat_codes_for_requests(
+        ids: torch.Tensor,
+        seq_token_counts: list[int] | None,
+        runtime_additional_information: list[dict[str, Any]] | None,
+    ) -> list[torch.Tensor]:
+        """Split flat codec token ids per request.
+
+        Prefer ``code_flat_numel`` from async-chunk / connector payload (same idea as
+        qwen3 omni/tts passing decode metadata in additional_information). Fall back
+        to runner-provided ``seq_token_counts`` (no forward_context / ubatch slicing).
+        """
+        n = ids.numel()
+        if n == 0:
+            return [ids]
+
+        if runtime_additional_information and all(
+            isinstance(info.get("code_flat_numel"), int) and int(info["code_flat_numel"]) > 0
+            for info in runtime_additional_information
+        ):
+            sizes = [int(info["code_flat_numel"]) for info in runtime_additional_information]
+            if sum(sizes) == n:
+                parts: list[torch.Tensor] = []
+                offset = 0
+                for sz in sizes:
+                    parts.append(ids[offset : offset + sz])
+                    offset += sz
+                return parts
+
+        if seq_token_counts is not None and len(seq_token_counts) > 1:
+            boundaries = [0]
+            for count in seq_token_counts:
+                boundaries.append(boundaries[-1] + count)
+            return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))]
+
+        return [ids]
+
+    @staticmethod
+    def _mimo_codec_runtime_lists(
+        num_req: int,
+        runtime_additional_information: list[dict[str, Any]] | None,
+    ) -> tuple[list[int | None], list[int | None]]:
+        """Per-request ``left_context_size`` / ``codec_chunk_frames`` from runtime buffer."""
+        left_frames: list[int | None] = [None] * num_req
+        chunk_frames: list[int | None] = [None] * num_req
+        if not runtime_additional_information:
+            return left_frames, chunk_frames
+        for i in range(min(num_req, len(runtime_additional_information))):
+            info = runtime_additional_information[i]
+            if "left_context_size" in info:
+                left_frames[i] = int(info["left_context_size"])
+            if "codec_chunk_frames" in info:
+                chunk_frames[i] = int(info["codec_chunk_frames"])
+        return left_frames, chunk_frames
+
     def chunked_decode_streaming(
         self,
         codes: torch.Tensor,
@@ -453,9 +539,12 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         wav_chunk = self._decode_waveform_from_codes(codes)
         if wav_chunk.numel() == 0:
             return wav_chunk
-        # Infer number of codec frames: flat format has 36 elements per frame (9*4 with prepend)
         num_flat = codes.numel() if codes.ndim == 1 else codes.shape[-1]
-        num_chunks = num_flat // 36
+        elt_per_group = flat_codec_group_element_count(
+            self.streamer_config.group_size,
+            self.streamer_config.audio_channels,
+        )
+        num_chunks = num_flat // elt_per_group
         if num_chunks <= chunk_size:
             context_size = 0
         else:
@@ -469,50 +558,65 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         self,
         input_ids: torch.Tensor | None = None,
         codes: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
         **kwargs,
-    ) -> torch.Tensor | IntermediateTensors:
-        # Support both input_ids (from vLLM) and codes (from mimo_audio.py)
+    ) -> OmniOutput | torch.Tensor | IntermediateTensors:
+        if runtime_additional_information is None:
+            runtime_additional_information = kwargs.get("model_intermediate_buffer") or kwargs.get(
+                "runtime_additional_information"
+            )
         code_tensor = codes if codes is not None else input_ids
+        empty = torch.zeros((0,), dtype=torch.float32, device=self.device)
 
-        if code_tensor is None:
-            raise ValueError("Either input_ids or codes must be provided")
-
-        if code_tensor.numel() == 0:
-            raise ValueError("code_tensor is empty.")
-
-        if getattr(self.vllm_config.model_config, "async_chunk", False):
-            waveform = self.chunked_decode_streaming(code_tensor, chunk_size=10, left_context_size=10)
-        else:
-            waveform = self._decode_waveform_from_codes(code_tensor)
-
-        # For warmup/dummy run: if waveform is empty, return dummy hidden_states
-        # _dummy_run expects a tensor that can be processed by extract_multimodal_outputs
-        if waveform.numel() == 0:
-            # Return dummy hidden_states with shape [seq_len, hidden_size]
-            # This is needed for warmup to avoid NoneType errors
-            hidden_size = getattr(self.config, "hidden_size", None)
-            if hidden_size is None:
-                # Fallback: try to get from vllm_config.model_config
-                hidden_size = getattr(
-                    self.vllm_config.model_config,
-                    "hidden_size",
-                    4096,  # Default fallback
-                )
-
-            # Return [seq_len=1, hidden_size] tensor for warmup
-            dummy_hidden = torch.zeros(
-                (1, hidden_size),
-                dtype=torch.bfloat16,
-                device=self.device,
+        if code_tensor is None or code_tensor.numel() == 0:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [empty]},
             )
 
-            return dummy_hidden
+        is_async_chunk = getattr(self.vllm_config.model_config, "async_chunk", False)
+        ids = code_tensor.reshape(-1).to(dtype=torch.long)
+        request_ids_list = self._split_flat_codes_for_requests(
+            ids, kwargs.get("seq_token_counts"), runtime_additional_information
+        )
+        per_left, per_chunk = self._mimo_codec_runtime_lists(len(request_ids_list), runtime_additional_information)
 
-        if waveform.numel() == 1 and waveform.item() == 0:
-            # Designed for Text output only
-            return None
+        is_capturing = torch.cuda.is_current_stream_capturing()
+        if is_capturing:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [empty] * len(request_ids_list)},
+            )
 
-        return waveform
+        if is_async_chunk:
+            audios = self._batch_chunked_decode_streaming(
+                request_ids_list,
+                default_chunk_size=self._codec_chunk_frames,
+                default_left_context_size=self._codec_left_context_frames,
+                per_req_left_context_frames=per_left,
+                per_req_codec_chunk_frames=per_chunk,
+            )
+        else:
+            audios = self._batch_decode_waveforms(request_ids_list)
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": audios},
+        )
+
+    def make_omni_output(self, model_output: torch.Tensor | OmniOutput, **kwargs) -> OmniOutput:
+        """Convert raw model output to OmniOutput if needed."""
+        if isinstance(model_output, OmniOutput):
+            return model_output
+        empty = torch.zeros((0,), dtype=torch.float32, device=self.device)
+        if model_output is None:
+            return OmniOutput(text_hidden_states=None, multimodal_outputs={"model_outputs": [empty]})
+        if isinstance(model_output, torch.Tensor):
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [model_output.to(dtype=torch.float32).reshape(-1)]},
+            )
+        raise TypeError(f"Unexpected model output type: {type(model_output)}")
 
     def _get_full_code_sequence(
         self,
@@ -545,8 +649,194 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
 
         return code_tensor
 
+    @torch.inference_mode()
+    def _batch_decode_waveforms(
+        self,
+        request_codes_list: list[torch.Tensor],
+    ) -> list[torch.Tensor]:
+        """Batch-decode multiple requests' codes in a single decoder forward pass.
+
+        Instead of calling _decode_waveform_from_codes per request (which incurs
+        4 GPU↔CPU round-trips each), this method:
+          1. Extracts audio_codes for all valid requests on CPU (cheap).
+          2. Runs quantizer.decode_vq (embedding lookup) for each on GPU.
+          3. Packs all hidden-states into one tensor and calls
+             decoder(packed_hs, input_lengths) once.
+          4. Splits the output waveforms back to per-request tensors.
+        """
+        empty = torch.zeros((0,), dtype=torch.float32, device=self.device)
+        num_req = len(request_codes_list)
+        if num_req == 0:
+            return [empty]
+
+        tokenizer = self._tokenizer_service.audio_tokenizer
+        group_size = self.streamer_config.group_size
+        audio_channels = self.streamer_config.audio_channels
+
+        hidden_list: list[torch.Tensor] = []
+        lengths: list[int] = []
+        valid_indices: list[int] = []
+
+        for i, req_codes in enumerate(request_codes_list):
+            if req_codes is None or req_codes.numel() == 0:
+                continue
+            if self._check_dummy_code_tensor(req_codes):
+                continue
+
+            flat_codes = req_codes.detach().to(torch.long).cpu()
+            audio_codes = extract_audio_code_tensor(
+                flat_codes,
+                group_size,
+                audio_channels,
+                self.codes,
+            )
+            if audio_codes is None or audio_codes.numel() == 0:
+                continue
+
+            hs = tokenizer.encoder.decode_vq(audio_codes.to(self.device))
+            hidden_list.append(hs)
+            lengths.append(hs.size(0))
+            valid_indices.append(i)
+
+        if not hidden_list:
+            return [empty] * num_req
+
+        if len(hidden_list) == 1:
+            packed_hs = hidden_list[0]
+        else:
+            packed_hs = torch.cat(hidden_list, dim=0)
+        input_lengths = torch.tensor(lengths, device=self.device)
+
+        recon_wav = tokenizer.decoder(packed_hs, input_lengths)
+
+        cfg = tokenizer.config
+        frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length
+
+        result: list[torch.Tensor] = [empty] * num_req
+        if len(valid_indices) == 1:
+            wav = recon_wav.squeeze(0).squeeze(0)
+            valid_len = lengths[0] * frames_per_token
+            if wav.numel() > valid_len:
+                wav = wav[:valid_len]
+            result[valid_indices[0]] = wav.to(dtype=torch.float32).reshape(-1)
+        else:
+            for j, idx in enumerate(valid_indices):
+                wav = recon_wav[j].squeeze(0)
+                valid_len = lengths[j] * frames_per_token
+                if wav.numel() > valid_len:
+                    wav = wav[:valid_len]
+                result[idx] = wav.to(dtype=torch.float32).reshape(-1)
+
+        return result
+
+    @torch.inference_mode()
+    def _batch_chunked_decode_streaming(
+        self,
+        request_codes_list: list[torch.Tensor],
+        default_chunk_size: int,
+        default_left_context_size: int,
+        *,
+        per_req_left_context_frames: list[int | None] | None = None,
+        per_req_codec_chunk_frames: list[int | None] | None = None,
+    ) -> list[torch.Tensor]:
+        """Batch version of chunked_decode_streaming for async_chunk mode."""
+        empty = torch.zeros((0,), dtype=torch.float32, device=self.device)
+        num_req = len(request_codes_list)
+        if num_req == 0:
+            return [empty]
+
+        tokenizer = self._tokenizer_service.audio_tokenizer
+        group_size = self.streamer_config.group_size
+        audio_channels = self.streamer_config.audio_channels
+
+        hidden_list: list[torch.Tensor] = []
+        lengths: list[int] = []
+        valid_indices: list[int] = []
+        context_sizes: list[int] = []
+
+        for i, req_codes in enumerate(request_codes_list):
+            if req_codes is None or req_codes.numel() == 0:
+                continue
+            if self._check_dummy_code_tensor(req_codes):
+                continue
+
+            flat_codes = req_codes.detach().to(torch.long).cpu()
+            audio_codes = extract_audio_code_tensor(
+                flat_codes,
+                group_size,
+                audio_channels,
+                self.codes,
+            )
+            if audio_codes is None or audio_codes.numel() == 0:
+                continue
+
+            hs = tokenizer.encoder.decode_vq(audio_codes.to(self.device))
+            hidden_list.append(hs)
+            lengths.append(hs.size(0))
+            valid_indices.append(i)
+
+            num_flat = req_codes.numel() if req_codes.ndim == 1 else req_codes.shape[-1]
+            elt_per_group = flat_codec_group_element_count(group_size, audio_channels)
+            num_chunks = num_flat // elt_per_group
+            chunk_sz = default_chunk_size
+            if per_req_codec_chunk_frames is not None and i < len(per_req_codec_chunk_frames):
+                c = per_req_codec_chunk_frames[i]
+                if c is not None:
+                    chunk_sz = c
+            if (
+                per_req_left_context_frames is not None
+                and i < len(per_req_left_context_frames)
+                and per_req_left_context_frames[i] is not None
+            ):
+                ctx = int(per_req_left_context_frames[i])
+            else:
+                ctx = 0 if num_chunks <= chunk_sz else default_left_context_size
+            context_sizes.append(ctx)
+
+        if not hidden_list:
+            return [empty] * num_req
+
+        if len(hidden_list) == 1:
+            packed_hs = hidden_list[0]
+        else:
+            packed_hs = torch.cat(hidden_list, dim=0)
+        input_lengths = torch.tensor(lengths, device=self.device)
+
+        recon_wav = tokenizer.decoder(packed_hs, input_lengths)
+
+        cfg = tokenizer.config
+        frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length
+
+        result: list[torch.Tensor] = [empty] * num_req
+        if len(valid_indices) == 1:
+            wav = recon_wav.squeeze(0).squeeze(0)
+            valid_len = lengths[0] * frames_per_token
+            if wav.numel() > valid_len:
+                wav = wav[:valid_len]
+            drop = context_sizes[0] * self.total_upsample
+            if drop > 0 and drop < wav.numel():
+                wav = wav[drop:]
+            elif drop >= wav.numel():
+                wav = wav[:0]
+            result[valid_indices[0]] = wav.to(dtype=torch.float32).reshape(-1)
+        else:
+            for j, idx in enumerate(valid_indices):
+                wav = recon_wav[j].squeeze(0)
+                valid_len = lengths[j] * frames_per_token
+                if wav.numel() > valid_len:
+                    wav = wav[:valid_len]
+                drop = context_sizes[j] * self.total_upsample
+                if drop > 0 and drop < wav.numel():
+                    wav = wav[drop:]
+                elif drop >= wav.numel():
+                    wav = wav[:0]
+                result[idx] = wav.to(dtype=torch.float32).reshape(-1)
+
+        return result
+
     def _check_dummy_code_tensor(self, code_tensor: torch.Tensor) -> bool:
-        if code_tensor is not None and code_tensor.numel() == DUMMY_CODE_SHAPE:
+        expected = flat_codec_group_element_count(self.config.group_size, self.config.audio_channels)
+        if code_tensor is not None and code_tensor.numel() == expected:
             code_groups = code_tensor.view(self.config.group_size, self.config.audio_channels + 1)
             return (
                 (code_groups[:, 0] == TALKER_CODEC_PAD_TOKEN_ID).all() and (code_groups[:, 1:].sum() == 0).all()

vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml

@@ -24,6 +24,8 @@ stage_args:
       max_model_len: 8192
       max_num_batched_tokens: 8192
       custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
     is_comprehension: true
     final_output: true
     final_output_type: text
@@ -57,6 +59,8 @@ stage_args:
       max_model_len: 8192
       max_num_batched_tokens: 8192
     engine_input_source: [ 0 ]
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
     is_comprehension: false
     final_output: true
     final_output_type: audio
@@ -67,3 +71,26 @@ stage_args:
       max_tokens: 18192
       seed: 42
       detokenize: false
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 365536
+        codec_streaming: true
+        connector_get_sleep_s: 0.001
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 3
+        codec_left_context_frames: 3
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

vllm_omni/model_executor/stage_input_processors/mimo_audio.py

@@ -50,6 +50,58 @@ def _make_finished_sentinel() -> dict[str, Any]:
     return {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
 
 
+def _flush_remaining_codes(
+    transfer_manager: Any,
+    request_id: str,
+    chunk_size: int,
+    left_context_size: int,
+) -> dict[str, Any]:
+    """Flush any accumulated but unsent codes when the request finishes."""
+    accumulated = transfer_manager.code_prompt_token_ids.get(request_id, [])
+    if not accumulated:
+        return _make_finished_sentinel()
+
+    length = len(accumulated)
+    chunk_length = length % chunk_size
+    context_length = chunk_length if chunk_length != 0 else chunk_size
+    end_index = min(length, left_context_size + context_length)
+
+    # Align with qwen3_omni talker2code2wav_async_chunk: decoder strip uses explicit frame count.
+    left_ctx_frames = max(0, min(length - context_length, left_context_size))
+    flat_codes = torch.tensor(accumulated[-end_index:]).reshape(-1).tolist()
+
+    return {
+        "code_predictor_codes": flat_codes,
+        "left_context_size": left_ctx_frames,
+        "codec_chunk_frames": chunk_size,
+        "codec_left_context_frames": left_context_size,
+        "code_flat_numel": len(flat_codes),
+        "finished": torch.tensor(True, dtype=torch.bool),
+    }
+
+
+def _is_codes_empty(codes: Any) -> bool:
+    """Check whether code_predictor_codes should be treated as empty / invalid."""
+    if codes is None:
+        return True
+    if isinstance(codes, torch.Tensor):
+        return codes.numel() == 0 or not codes.any()
+    if hasattr(codes, "__len__") and len(codes) == 0:
+        return True
+    t = torch.tensor(codes, dtype=torch.long) if not isinstance(codes, torch.Tensor) else codes
+    return not t.any()
+
+
+def _to_code_tensor(codes: Any) -> torch.Tensor | None:
+    """Convert codes to a (B, 1, 8, 4) long tensor, or return None if shape is invalid."""
+    code_tensor = codes.to(torch.long) if isinstance(codes, torch.Tensor) else torch.tensor(codes, dtype=torch.long)
+    if code_tensor.ndim == 3:
+        code_tensor = code_tensor.unsqueeze(0)
+    if code_tensor.ndim != 4 or code_tensor.shape[-2:] != (8, 4):
+        return None
+    return code_tensor
+
+
 def llm2code2wav_async_chunk(
     transfer_manager: Any,
     pooling_output: dict[str, Any],
@@ -62,57 +114,39 @@ def llm2code2wav_async_chunk(
     Accumulates codes in connector per request_id,
     returns payload only when chunk_size is full or request is finished; returns None when waiting.
     """
-    if "code_predictor_codes" not in pooling_output:
+
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    chunk_size = int(cfg.get("codec_chunk_frames", 3))
+    left_context_size = int(cfg.get("codec_left_context_frames", 3))
+
+    request_id = getattr(request, "external_req_id", None)
+
+    codes = pooling_output.get("code_predictor_codes")
+
+    if _is_codes_empty(codes):
         if is_finished:
-            return _make_finished_sentinel()
+            return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size)
         return None
 
-    code_predictor_codes = pooling_output["code_predictor_codes"]
-
-    if code_predictor_codes is None:
+    code_tensor = _to_code_tensor(codes)
+    if code_tensor is None:
         if is_finished:
-            return _make_finished_sentinel()
-        return None
-    if isinstance(code_predictor_codes, torch.Tensor):
-        if code_predictor_codes.numel() == 0:
-            if is_finished:
-                return _make_finished_sentinel()
-            return None
-    elif hasattr(code_predictor_codes, "__len__"):
-        if len(code_predictor_codes) == 0:
-            return None
-
-    if isinstance(code_predictor_codes, torch.Tensor):
-        if not code_predictor_codes.any():
-            if is_finished:
-                return _make_finished_sentinel()
-            return None
-        code_tensor = code_predictor_codes.to(torch.long)
-    else:
-        code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long)
-        if not code_tensor.any():
-            if is_finished:
-                return _make_finished_sentinel()
-            return None
-
-    if code_tensor.ndim == 3:
-        code_tensor = code_tensor.unsqueeze(0)
-    if code_tensor.ndim != 4 or code_tensor.shape[-2:] != (8, 4):
+            return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size)
         return None
 
     pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4, device=code_tensor.device, dtype=code_tensor.dtype)
-    code_final = prepend_and_flatten_colmajor(code_tensor, pad_vec)
-    code_list = code_final.tolist()
+    code_list = prepend_and_flatten_colmajor(code_tensor, pad_vec).tolist()
+
     if sum(code_list) == 0:
         if is_finished:
-            return _make_finished_sentinel()
+            return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size)
         return None
 
-    request_id = getattr(request, "external_req_id", None)
     if request_id is None:
         return None
 
-    chunk_size = left_context_size = 10
     transfer_manager.code_prompt_token_ids[request_id].append(code_list)
     length = len(transfer_manager.code_prompt_token_ids[request_id])
     chunk_length = length % chunk_size
@@ -121,14 +155,17 @@ def llm2code2wav_async_chunk(
 
     context_length = chunk_length if chunk_length != 0 else chunk_size
     end_index = min(length, left_context_size + context_length)
+    left_ctx_frames = max(0, min(length - context_length, left_context_size))
+    flat_codes = torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]).reshape(-1).tolist()
 
-    info = {
-        "code_predictor_codes": (
-            torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]).reshape(-1).tolist()
-        ),
+    return {
+        "code_predictor_codes": flat_codes,
+        "left_context_size": left_ctx_frames,
+        "codec_chunk_frames": chunk_size,
+        "codec_left_context_frames": left_context_size,
+        "code_flat_numel": len(flat_codes),
         "finished": torch.tensor(is_finished, dtype=torch.bool),
     }
-    return info
 
 
 def llm2code2wav(