[Performance] Improve MiMo-Audio tokenizer decoding performance (#2183)

Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>
Co-authored-by: Jialong Liu <88185941+Galleons2029@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py

@@ -91,7 +91,7 @@ def _minimal_model(mocker: MockerFixture):
         decoder=decoder,
         config=SimpleNamespace(avg_pooler=2, stride_size=2, hop_length=240),
     )
-    model._tokenizer_service = SimpleNamespace(audio_tokenizer=audio_tok)
+    model._tokenizer_service = SimpleNamespace(audio_tokenizer=audio_tok, cuda_graph_wrapper=None)
     return model, audio_tok
 
 

vllm_omni/model_executor/models/mimo_audio/cuda_graph_decoder_wrapper.py

@@ -0,0 +1,396 @@
+# Copyright 2025 Xiaomi Corporation.
+"""
+CUDA Graph wrapper for MiMo Audio Tokenizer decoder.
+
+Captures the full decode path (VQ dequantize -> AudioDecoder -> Vocoder -> Waveform)
+into CUDA Graphs for fixed bucket sizes, eliminating kernel-launch overhead that
+dominates small-chunk streaming decoding.
+
+The vocoder uses non-causal sliding-window attention, so a pre-built attention mask
+(static tensor) is updated in-place before each graph replay to properly handle
+padding positions.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch.cuda import CUDAGraph
+from vllm.logger import init_logger
+
+from .modeling_audio_tokenizer import StreamingCache, StreamingConfig, TransformerVocos
+
+logger = init_logger(__name__)
+
+
+class CUDAGraphMiMoDecoderWrapper:
+    """CUDA Graph wrapper for MiMoAudioTokenizer.decode_fixed().
+
+    Usage::
+
+        wrapper = CUDAGraphMiMoDecoderWrapper(tokenizer)
+        wrapper.warmup(device)
+        wav = wrapper.decode(codes)  # auto graph-replay or eager fallback
+
+    When ``code_rows`` lists multiple RVQ depths (e.g. talker emits 8 rows but
+    config.num_quantizers is 20), one CUDA Graph is captured per (depth, bucket).
+    """
+
+    DEFAULT_CAPTURE_SIZES = [4, 6, 8, 12, 16, 24, 32, 64, 128, 256]
+
+    def __init__(
+        self,
+        tokenizer,
+        capture_sizes: list[int] | None = None,
+        enabled: bool = True,
+        code_rows: list[int] | None = None,
+    ):
+        self.tokenizer = tokenizer
+
+        # num_quantizers = RVQ layers in checkpoint; talker may emit fewer rows
+        self.audio_channels = tokenizer.config.num_quantizers
+        nq = self.audio_channels
+        if code_rows is None:
+            rows = {nq}
+        else:
+            rows = {r for r in code_rows if 1 <= r <= nq}
+            rows.add(nq)
+        self._capture_code_rows: list[int] = sorted(rows)
+        if len(self._capture_code_rows) > 1:
+            logger.info(
+                "CUDAGraphMiMoDecoderWrapper: capturing graphs for code depths %s (num_quantizers=%d)",
+                self._capture_code_rows,
+                nq,
+            )
+
+        self._explicit_sizes = capture_sizes is not None
+        self.capture_sizes: list[int] = sorted(capture_sizes) if capture_sizes else []
+        self.enabled = enabled
+
+        self.graphs: dict[tuple[int, int], CUDAGraph] = {}
+        self.static_code_inputs: dict[tuple[int, int], torch.Tensor] = {}
+        self.static_wav_outputs: dict[tuple[int, int], torch.Tensor] = {}
+        self.static_vocoder_masks: dict[int, torch.Tensor] = {}
+
+        self._warmed_up = False
+        self._device: torch.device | None = None
+
+        cfg = tokenizer.config
+        self.frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length
+
+    @staticmethod
+    def compute_capture_sizes(
+        codec_chunk_frames: int = 0,
+        codec_left_context_frames: int = 0,
+    ) -> list[int]:
+        """Derive bucket sizes from streaming chunk config for high hit rate."""
+        sizes: set[int] = set()
+        if codec_chunk_frames > 0:
+            sizes.add(codec_chunk_frames)
+            if codec_left_context_frames > 0:
+                sizes.add(codec_chunk_frames + codec_left_context_frames)
+
+        for p2 in CUDAGraphMiMoDecoderWrapper.DEFAULT_CAPTURE_SIZES:
+            sizes.add(p2)
+        return sorted(sizes)
+
+    def _get_padded_size(self, actual_size: int) -> int | None:
+        """Round up to the nearest captured bucket size."""
+        for size in self.capture_sizes:
+            if actual_size <= size:
+                return size
+        return None
+
+    def _compute_output_wav_len(self, code_len: int) -> int:
+        """Estimate the output waveform length for a given code length."""
+        return self.tokenizer._compute_vocoder_seq_len(code_len) * self.tokenizer.config.hop_length
+
+    def _compute_vocoder_seq_len(self, code_len: int) -> int:
+        return self.tokenizer._compute_vocoder_seq_len(code_len)
+
+    def warmup(
+        self,
+        device: torch.device,
+        codec_chunk_frames: int = 0,
+        codec_left_context_frames: int = 0,
+    ):
+        """Eager warmup then CUDA Graph capture for all bucket sizes."""
+        if device.type != "cuda" or not self.enabled or self._warmed_up:
+            return
+
+        self._device = device
+        self.tokenizer.eval()
+
+        if not self._explicit_sizes:
+            self.capture_sizes = self.compute_capture_sizes(
+                codec_chunk_frames=codec_chunk_frames,
+                codec_left_context_frames=codec_left_context_frames,
+            )
+
+        logger.info(
+            "CUDAGraphMiMoDecoderWrapper: starting warmup for %d sizes: %s",
+            len(self.capture_sizes),
+            self.capture_sizes,
+        )
+
+        model_dtype = next(self.tokenizer.parameters()).dtype
+
+        for size in self.capture_sizes:
+            voc_len = self._compute_vocoder_seq_len(size)
+            mask = TransformerVocos.build_vocoder_attn_mask(
+                voc_len,
+                voc_len,
+                tuple(self.tokenizer.config.vocoder_attn_window_size),
+                device,
+                model_dtype,
+            )
+            self.static_vocoder_masks[size] = mask
+
+            for n_rows in self._capture_code_rows:
+                dummy = torch.zeros(n_rows, size, dtype=torch.long, device=device)
+                with torch.no_grad():
+                    _ = self.tokenizer.decode_fixed(dummy, size, vocoder_attn_mask=mask)
+
+        torch.accelerator.synchronize(device)
+
+        expected = len(self.capture_sizes) * len(self._capture_code_rows)
+        for n_rows in self._capture_code_rows:
+            for size in self.capture_sizes:
+                try:
+                    self._capture(size, device, n_code_rows=n_rows)
+                    logger.info(
+                        "  Captured CUDA Graph for code_rows=%d code_len=%d",
+                        n_rows,
+                        size,
+                    )
+                except Exception:
+                    logger.warning(
+                        "  Failed to capture graph for code_rows=%d code_len=%d",
+                        n_rows,
+                        size,
+                        exc_info=True,
+                    )
+
+        self._warmed_up = True
+        logger.info(
+            "CUDAGraphMiMoDecoderWrapper warmup complete: %d/%d captured",
+            len(self.graphs),
+            expected,
+        )
+
+    def _capture(self, size: int, device: torch.device, n_code_rows: int):
+        key = (n_code_rows, size)
+        static_codes = torch.zeros(n_code_rows, size, dtype=torch.long, device=device)
+        mask = self.static_vocoder_masks[size]
+
+        with torch.no_grad():
+            _ = self.tokenizer.decode_fixed(static_codes, size, vocoder_attn_mask=mask)
+        torch.accelerator.synchronize(device)
+
+        graph = CUDAGraph()
+        with torch.no_grad():
+            with torch.cuda.graph(graph):
+                static_output = self.tokenizer.decode_fixed(
+                    static_codes,
+                    size,
+                    vocoder_attn_mask=mask,
+                )
+
+        self.graphs[key] = graph
+        self.static_code_inputs[key] = static_codes
+        self.static_wav_outputs[key] = static_output
+
+    def _update_vocoder_mask(self, padded_size: int, actual_code_shape: int):
+        """Update the static vocoder mask in-place for the given actual code length."""
+        mask = self.static_vocoder_masks[padded_size]
+        voc_padded = self._compute_vocoder_seq_len(padded_size)
+        voc_actual = self._compute_vocoder_seq_len(actual_code_shape)
+
+        new_mask = TransformerVocos.build_vocoder_attn_mask(
+            voc_padded,
+            voc_actual,
+            tuple(self.tokenizer.config.vocoder_attn_window_size),
+            mask.device,
+            mask.dtype,
+        )
+        mask.copy_(new_mask)
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode codes with CUDA Graph replay when possible.
+
+        Args:
+            codes: [n_rvq_rows, T] on CUDA device (n_rvq_rows may be < num_quantizers).
+        Returns:
+            Waveform tensor (same shape as MiMoAudioTokenizer.decode output).
+        """
+        if not self.enabled or not self._warmed_up:
+            return self.tokenizer.decode(codes)
+
+        actual_code_shape = codes.shape[-1]
+        n_rows = codes.shape[0]
+        padded_size = self._get_padded_size(actual_code_shape)
+        if padded_size is None:
+            return self.tokenizer.decode(codes)
+
+        key = (n_rows, padded_size)
+        if key not in self.graphs:
+            return self.tokenizer.decode(codes)
+
+        self.static_code_inputs[key].zero_()
+        self.static_code_inputs[key][:, :actual_code_shape] = codes
+
+        self._update_vocoder_mask(padded_size, actual_code_shape)
+
+        self.graphs[key].replay()
+
+        actual_wav_len = self._compute_output_wav_len(actual_code_shape)
+        output = self.static_wav_outputs[key]
+        return output[..., :actual_wav_len].clone()
+
+    @torch.no_grad()
+    def streaming_decode(
+        self,
+        codes_chunks: torch.Tensor,
+        chunk_input_lengths: list[int],
+        history_cache: StreamingCache | None = None,
+        streaming_config: StreamingConfig | None = None,
+        last_chunk: bool = False,
+    ) -> tuple[list[torch.Tensor | None], StreamingCache]:
+        """Streaming decode with CUDA Graph acceleration.
+
+        Mirrors MiMoAudioTokenizer.streaming_decode(): identical cache /
+        overlap logic, but replaces ``self.decoder(hs, lengths)`` with
+        ``decoder.forward_fixed`` + vocoder attention mask so the
+        fixed-shape path (and optionally CUDA Graph) can be used.
+
+        Args:
+            codes_chunks: [num_quantizers, sum(chunk_input_lengths)] — codes
+                for all samples in the batch, concatenated along dim-1.
+            chunk_input_lengths: per-sample code lengths in this chunk.
+            history_cache: StreamingCache from the previous call.
+            streaming_config: overlap / segment config.
+            last_chunk: whether this is the final chunk.
+
+        Returns:
+            (return_wavs, history_cache) — same contract as the original.
+        """
+        if history_cache is None:
+            history_cache = StreamingCache()
+        if streaming_config is None:
+            streaming_config = StreamingConfig()
+        tokenizer = self.tokenizer
+        hidden_states = tokenizer.encoder.decode_vq(codes_chunks)
+
+        input_lengths: list[int] = []
+        input_hidden_states: list[torch.Tensor] = []
+        cache_hidden_states: list[torch.Tensor] = []
+        start_idx = 0
+
+        for i, length in enumerate(chunk_input_lengths):
+            sample_hs = hidden_states[start_idx : start_idx + length]
+            start_idx += length
+            if history_cache.hidden_states is not None:
+                sample_hs = torch.cat(
+                    [history_cache.hidden_states[i], sample_hs],
+                    dim=0,
+                )
+                length += history_cache.hidden_states[i].size(0)
+            input_hidden_states.append(sample_hs)
+            cache_hidden_states.append(sample_hs.clone())
+            input_lengths.append(length)
+
+        output = self._batch_decode_hidden_states(input_hidden_states, input_lengths)
+
+        return_wavs: list[torch.Tensor | None] = []
+        frames_per_token = self.frames_per_token
+        processed_lengths: list[int] = []
+
+        for i, wav in enumerate(output):
+            wav = wav.float().detach().cpu()
+            prev_processed = history_cache.processed_lengths[i] if history_cache.processed_lengths is not None else 0
+
+            if last_chunk:
+                return_wavs.append(wav[:, prev_processed * frames_per_token :])
+                new_processed_length = input_lengths[i]
+            elif input_lengths[i] <= streaming_config.right_overlap:
+                return_wavs.append(None)
+                new_processed_length = 0
+            else:
+                end_idx = input_lengths[i] - streaming_config.right_overlap
+                wav = wav[:, prev_processed * frames_per_token : end_idx * frames_per_token]
+                return_wavs.append(wav)
+                new_processed_length = end_idx
+                if input_lengths[i] > streaming_config.left_overlap:
+                    cache_hidden_states[i] = cache_hidden_states[i][-streaming_config.left_overlap :]
+                    new_processed_length -= input_lengths[i] - streaming_config.left_overlap
+
+            processed_lengths.append(new_processed_length)
+
+        history_cache.hidden_states = cache_hidden_states
+        history_cache.processed_lengths = processed_lengths
+
+        return return_wavs, history_cache
+
+    def _batch_decode_hidden_states(
+        self,
+        hidden_states_list: list[torch.Tensor],
+        input_lengths: list[int],
+    ) -> list[torch.Tensor]:
+        """Decode a list of per-sample hidden_states using the fixed-shape path.
+
+        Each sample is decoded independently with padding + vocoder mask.
+        Returns a list of wav tensors, each shaped [1, wav_len] (matching the
+        per-sample shape from the original batched decoder output).
+        """
+        results: list[torch.Tensor] = []
+        decoder = self.tokenizer.decoder
+        cfg = self.tokenizer.config
+        window_size = tuple(cfg.vocoder_attn_window_size)
+
+        mask_dtype = decoder.vocoder.embeddings.weight.dtype
+
+        for sample_hs, actual_code_shape in zip(hidden_states_list, input_lengths):
+            padded_size = self._get_padded_size(actual_code_shape)
+
+            if padded_size is None:
+                wav = decoder(
+                    sample_hs,
+                    torch.tensor([actual_code_shape], device=sample_hs.device),
+                )
+                results.append(wav.squeeze(0))
+                continue
+
+            if actual_code_shape < padded_size:
+                pad = torch.zeros(
+                    padded_size - actual_code_shape,
+                    sample_hs.size(-1),
+                    device=sample_hs.device,
+                    dtype=sample_hs.dtype,
+                )
+                padded_hs = torch.cat([sample_hs, pad], dim=0)
+            else:
+                padded_hs = sample_hs
+
+            voc_padded = self._compute_vocoder_seq_len(padded_size)
+            voc_actual = self._compute_vocoder_seq_len(actual_code_shape)
+            vocoder_mask = TransformerVocos.build_vocoder_attn_mask(
+                voc_padded,
+                voc_actual,
+                window_size,
+                padded_hs.device,
+                mask_dtype,
+            )
+
+            wav = decoder.forward_fixed(
+                padded_hs.unsqueeze(0),
+                vocoder_attn_mask=vocoder_mask,
+            )
+
+            actual_wav_len = self._compute_output_wav_len(actual_code_shape)
+            wav = wav[..., :actual_wav_len].squeeze(0)
+            results.append(wav)
+
+        return results
+
+    @property
+    def is_ready(self) -> bool:
+        return self.enabled and self._warmed_up and len(self.graphs) > 0

vllm_omni/model_executor/models/mimo_audio/mimo_audio.py

@@ -459,13 +459,7 @@ class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProces
 
             num_features = audio_output_lengths[item_idx] // 4
             if num_features == 0:
-                try:
-                    audios = mm_items.get_items("audio", AudioProcessorItems)
-                    audio_len = audios.get_audio_length(item_idx)
-                    raise ValueError(f"The audio (len={audio_len}) is too short to be represented inside the model")
-                except (AttributeError, KeyError):
-                    # If AudioProcessorItems is not available, use default
-                    num_features = 1
+                num_features = 1
 
             audio_tokens = [audio_token_id] * num_features
 

vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py

@@ -20,6 +20,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
 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.cuda_graph_decoder_wrapper import CUDAGraphMiMoDecoderWrapper
 from vllm_omni.model_executor.models.mimo_audio.modeling_audio_tokenizer import MiMoAudioTokenizer
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 
@@ -114,6 +115,8 @@ class MiMoAudioTokenizerWorker:
         self.audio_channels = self.config.audio_channels
         self.sample_rate = self.audio_tokenizer.config.sampling_rate
 
+        self.cuda_graph_wrapper: CUDAGraphMiMoDecoderWrapper | None = None
+
         # Warmup (skip for CPU due to potential shape mismatch issues)
         if device_str != "cpu":
             logger.info("[tokenizer worker] Running warmup encode/decode...")
@@ -128,6 +131,30 @@ class MiMoAudioTokenizerWorker:
                 )
             except Exception as e:
                 logger.warning("[tokenizer worker] Warmup failed (non-critical): %s", str(e))
+
+            cuda_graph_enabled = os.environ.get("MIMO_AUDIO_TOKENIZER_CUDA_GRAPH", "1") == "1"
+            if cuda_graph_enabled:
+                try:
+                    logger.info("[tokenizer worker] Initializing CUDA Graph decoder wrapper...")
+                    cg_start = time.monotonic()
+
+                    n_q = self.audio_tokenizer.config.num_quantizers
+                    cg_code_rows = sorted({self.audio_channels, n_q})
+                    self.cuda_graph_wrapper = CUDAGraphMiMoDecoderWrapper(
+                        self.audio_tokenizer,
+                        enabled=True,
+                        code_rows=cg_code_rows,
+                    )
+
+                    self.cuda_graph_wrapper.warmup(torch.device(device_str))
+                    logger.info(
+                        "[tokenizer worker] CUDA Graph decoder ready in %.2fs",
+                        time.monotonic() - cg_start,
+                    )
+                except Exception as e:
+                    logger.warning("[tokenizer worker] CUDA Graph warmup failed (non-critical): %s", str(e))
+                    self.cuda_graph_wrapper = None
+
         else:
             logger.info("[tokenizer worker] Skipping warmup for CPU device")
 
@@ -266,8 +293,13 @@ class MiMoAudioTokenizerWorker:
     ) -> torch.Tensor:
         """Decode audio tokens to waveform using the tokenizer's decoder"""
         tokens = tokens.to(self.device)
-        with torch.no_grad():
-            decoded_audio: torch.Tensor = self.audio_tokenizer.decode(tokens)
+
+        if self.cuda_graph_wrapper is not None and self.cuda_graph_wrapper.is_ready:
+            decoded_audio = self.cuda_graph_wrapper.decode(tokens)
+        else:
+            with torch.no_grad():
+                decoded_audio = self.audio_tokenizer.decode(tokens)
+
         decoded_audio = decoded_audio.float().reshape(-1).detach().cpu()
         return decoded_audio  # [samples] cpu
 
@@ -673,9 +705,10 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         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] = []
+        cg_wrapper = self._tokenizer_service.cuda_graph_wrapper
+        cg_ready = cg_wrapper is not None and cg_wrapper.is_ready
+
+        extracted: list[tuple[int, torch.Tensor]] = []
 
         for i, req_codes in enumerate(request_codes_list):
             if req_codes is None or req_codes.numel() == 0:
@@ -693,13 +726,33 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
             if audio_codes is None or audio_codes.numel() == 0:
                 continue
 
+            extracted.append((i, audio_codes))
+
+        if not extracted:
+            return [empty] * num_req
+
+        valid_indices = [t[0] for t in extracted]
+        # CUDA graph decode is single-batch only; avoid duplicate decode_vq for len>1.
+        use_cuda_graph = cg_ready and len(extracted) == 1
+
+        if use_cuda_graph:
+            wav_out = cg_wrapper.decode(extracted[0][1].to(self.device))
+            wav = wav_out.squeeze(0).squeeze(0)
+            cfg = tokenizer.config
+            frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length
+            valid_len = extracted[0][1].shape[-1] * frames_per_token
+            if wav.numel() > valid_len:
+                wav = wav[:valid_len]
+            result: list[torch.Tensor] = [empty] * num_req
+            result[valid_indices[0]] = wav.to(dtype=torch.float32).reshape(-1)
+            return result
+
+        hidden_list: list[torch.Tensor] = []
+        lengths: list[int] = []
+        for _, audio_codes in extracted:
             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]
@@ -749,9 +802,10 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         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] = []
+        cg_wrapper = self._tokenizer_service.cuda_graph_wrapper
+        cg_ready = cg_wrapper is not None and cg_wrapper.is_ready
+
+        extracted: list[tuple[int, torch.Tensor]] = []
         context_sizes: list[int] = []
 
         for i, req_codes in enumerate(request_codes_list):
@@ -770,10 +824,7 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
             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)
+            extracted.append((i, audio_codes))
 
             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)
@@ -793,9 +844,37 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
                 ctx = 0 if num_chunks <= chunk_sz else default_left_context_size
             context_sizes.append(ctx)
 
-        if not hidden_list:
+        if not extracted:
             return [empty] * num_req
 
+        valid_indices = [t[0] for t in extracted]
+        use_cuda_graph = cg_ready and len(extracted) == 1
+
+        cfg = tokenizer.config
+        frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length
+
+        if use_cuda_graph:
+            wav_out = cg_wrapper.decode(extracted[0][1].to(self.device))
+            wav = wav_out.squeeze(0).squeeze(0)
+            valid_len = extracted[0][1].shape[-1] * 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: list[torch.Tensor] = [empty] * num_req
+            result[valid_indices[0]] = wav.to(dtype=torch.float32).reshape(-1)
+            return result
+
+        hidden_list: list[torch.Tensor] = []
+        lengths: list[int] = []
+        for _, audio_codes in extracted:
+            hs = tokenizer.encoder.decode_vq(audio_codes.to(self.device))
+            hidden_list.append(hs)
+            lengths.append(hs.size(0))
+
         if len(hidden_list) == 1:
             packed_hs = hidden_list[0]
         else:
@@ -804,9 +883,6 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
 
         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)

vllm_omni/model_executor/models/mimo_audio/modeling_audio_tokenizer.py

@@ -243,6 +243,44 @@ class ISTFT(nn.Module):
 
         return y
 
+    def forward_fixed(self, spec: torch.Tensor) -> torch.Tensor:
+        """CUDA-Graph-safe ISTFT: identical to forward() but without assert."""
+        if self.padding == "center":
+            return torch.istft(
+                spec,
+                self.n_fft,
+                self.hop_length,
+                self.win_length,
+                self.window,
+                center=True,
+            )
+
+        pad = (self.win_length - self.hop_length) // 2
+        B, N, T = spec.shape
+
+        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+        ifft = ifft * self.window[None, :, None]
+
+        output_size = (T - 1) * self.hop_length + self.win_length
+        y = torch.nn.functional.fold(
+            ifft,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        )[:, 0, 0, pad:-pad]
+
+        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_envelope = torch.nn.functional.fold(
+            window_sq,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        ).squeeze()[pad:-pad]
+
+        y = y / window_envelope
+
+        return y
+
 
 class ISTFTHead(nn.Module):
     """
@@ -291,6 +329,20 @@ class ISTFTHead(nn.Module):
         audio = audio.to(original_dtype)
         return audio
 
+    def forward_fixed(self, x: torch.Tensor) -> torch.Tensor:
+        """CUDA-Graph-safe ISTFTHead: uses istft.forward_fixed (no assert)."""
+        x = self.out(x).transpose(1, 2)
+        mag, p = x.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(mag, max=1e2)
+        x = torch.cos(p)
+        y = torch.sin(p)
+        original_dtype = x.dtype
+        S = mag.float() * (x.float() + 1j * y.float())
+        audio = self.istft.forward_fixed(S)
+        audio = audio.to(original_dtype)
+        return audio
+
 
 class RotaryEmbedding(nn.Module):
     inv_freq: torch.Tensor
@@ -446,6 +498,44 @@ class Attention(nn.Module):
         attn_output = self.out_proj(attn_output)
         return attn_output
 
+    def forward_fixed(
+        self,
+        hidden_states: torch.Tensor,
+        rope_position_embeddings=None,
+        attn_mask: torch.Tensor | None = None,
+    ):
+        """Fixed-shape forward for CUDA Graph (batch_size=1, 3D tensor [B, L, D]).
+
+        Uses F.scaled_dot_product_attention instead of flash_attn_varlen_func.
+
+        Args:
+            attn_mask: optional [B, 1, L, L] bool/float mask for SDPA.
+                       Needed when padding is present and attention is non-causal,
+                       or when sliding-window behavior must be replicated.
+        """
+        B, L, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if rope_position_embeddings is not None:
+            cos, sin = rope_position_embeddings
+            query_states = apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=1)
+            key_states = apply_rotary_pos_emb(key_states, cos, sin, unsqueeze_dim=1)
+
+        attn_output = F.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attn_mask,
+            is_causal=(self.causal and attn_mask is None),
+        )
+
+        attn_output = attn_output.transpose(1, 2).reshape(B, L, self.embed_dim)
+        attn_output = self.out_proj(attn_output)
+        return attn_output
+
 
 class TransformerLayer(nn.Module):
     def __init__(
@@ -493,6 +583,33 @@ class TransformerLayer(nn.Module):
             hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
         return hidden_states
 
+    def forward_fixed(
+        self,
+        hidden_states: torch.Tensor,
+        rope_position_embeddings: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Fixed-shape forward for CUDA Graph (3D input [B, L, D])."""
+        residual = hidden_states
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+        hidden_states = self.self_attn.forward_fixed(
+            hidden_states,
+            rope_position_embeddings=rope_position_embeddings,
+            attn_mask=attn_mask,
+        )
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = residual + hidden_states
+
+        if hidden_states.dtype == torch.float16 or hidden_states.dtype == torch.bfloat16:
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+        return hidden_states
+
 
 class TransformerVocos(nn.Module):
     def __init__(self, config: MiMoAudioTokenizerConfig):
@@ -548,6 +665,65 @@ class TransformerVocos(nn.Module):
         output_length = input_length * self.hop_size
         return x[:, None, :], output_length
 
+    @staticmethod
+    def build_vocoder_attn_mask(
+        seq_len: int,
+        actual_len: int,
+        window_size: tuple[int, int],
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Build a [1, 1, seq_len, seq_len] attention mask combining padding + sliding window.
+
+        - Padding positions (>= actual_len) are masked out.
+        - Sliding window (left, right) is applied.
+        - Returns a float mask: 0.0 for attend, -inf for mask out.
+        """
+        row = torch.arange(seq_len, device=device)
+        col = torch.arange(seq_len, device=device)
+
+        valid_mask = (col.unsqueeze(0) < actual_len) & (row.unsqueeze(1) < actual_len)
+
+        left_win, right_win = window_size
+        if left_win > 0 or right_win > 0:
+            dist = row.unsqueeze(1) - col.unsqueeze(0)
+            if left_win > 0 and right_win > 0:
+                window_mask = (dist >= -right_win) & (dist <= left_win)
+            elif left_win > 0:
+                window_mask = dist <= left_win
+            else:
+                window_mask = dist >= -right_win
+            valid_mask = valid_mask & window_mask
+
+        mask = torch.zeros(seq_len, seq_len, device=device, dtype=dtype)
+        mask.masked_fill_(~valid_mask, float("-inf"))
+        return mask.unsqueeze(0).unsqueeze(0)
+
+    def forward_fixed(self, x: torch.Tensor, vocoder_attn_mask: torch.Tensor | None = None):
+        """Fixed-shape forward for CUDA Graph. x: [1, n_mels, L].
+
+        Args:
+            vocoder_attn_mask: optional pre-built [1, 1, L, L] attention mask.
+                Caller is responsible for constructing this (e.g. via
+                build_vocoder_attn_mask) and, when used with CUDA Graph,
+                updating the static tensor in-place before replay.
+        """
+
+        x = x.transpose(1, 2)  # [1, L, n_mels]
+        B, L, _ = x.size()
+        x = self.embeddings(x)  # [1, L, vocoder_dim]
+        position_ids = torch.arange(0, L, device=x.device, dtype=torch.long)
+        rope_position_embeddings = self.position_embedding(x.view(L, -1), position_ids)
+        cos, sin = rope_position_embeddings
+        cos = cos.unsqueeze(0)  # [1, L, head_dim]
+        sin = sin.unsqueeze(0)
+        for layer in self.layers:
+            x = layer.forward_fixed(x, rope_position_embeddings=(cos, sin), attn_mask=vocoder_attn_mask)
+        x = self.layer_norm(x)
+        x = self.head.forward_fixed(x)
+        output_length = L * self.hop_size
+        return x[:, None, :], output_length
+
 
 class AudioEncoder(nn.Module):
     def __init__(self, config: MiMoAudioTokenizerConfig):
@@ -755,9 +931,9 @@ class CausalConvTranspose1d(nn.Module):
         hidden_states = self.norm(hidden_states)
         hidden_states = hidden_states.transpose(2, 1)  # (N, C, L) -> (N, L, C)
 
-        casual_padding_right = max(0, kernel_size - stride)
-        hidden_states = hidden_states[:, : hidden_states.shape[1] - casual_padding_right, :]
-        output_length = (input_length - 1) * stride + kernel_size - casual_padding_right
+        causal_padding_right = max(0, kernel_size - stride)
+        hidden_states = hidden_states[:, : hidden_states.shape[1] - causal_padding_right, :]
+        output_length = (input_length - 1) * stride + kernel_size - causal_padding_right
         sequence_mask, _ = get_sequence_mask(hidden_states, output_length)
         if output_dim <= 2:
             hidden_states = torch.masked_select(hidden_states, sequence_mask).view(-1, self.out_channels)
@@ -766,6 +942,19 @@ class CausalConvTranspose1d(nn.Module):
             hidden_states = hidden_states[:, : torch.max(output_length), :]
         return hidden_states, output_length
 
+    def forward_fixed(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Fixed-shape forward for CUDA Graph. Input: [1, L, C], Output: [1, L', C']."""
+        kernel_size = self.conv.kernel_size[0]
+        stride = self.conv.stride[0]
+        hidden_states = hidden_states.transpose(2, 1)  # (1, L, C) -> (1, C, L)
+        hidden_states = self.conv(hidden_states)
+        hidden_states = self.norm(hidden_states)
+        hidden_states = hidden_states.transpose(2, 1)  # (1, C, L') -> (1, L', C)
+        causal_padding_right = max(0, kernel_size - stride)
+        if causal_padding_right > 0:
+            hidden_states = hidden_states[:, : hidden_states.shape[1] - causal_padding_right, :]
+        return hidden_states
+
 
 class AudioDecoder(nn.Module):
     def __init__(self, config: MiMoAudioTokenizerConfig):
@@ -855,6 +1044,49 @@ class AudioDecoder(nn.Module):
 
         return recon_wav
 
+    def forward_fixed(
+        self,
+        audio_embed: torch.Tensor,
+        vocoder_attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Fixed-shape forward for CUDA Graph.
+
+        Input: audio_embed [1, T, d_model] — padded, batch_size=1.
+        Output: recon_wav [1, 1, wav_len].
+
+        Args:
+            vocoder_attn_mask: pre-built attention mask for the vocoder's
+                non-causal layers. Shape [1, 1, L_voc, L_voc].
+        """
+        audio_embed = audio_embed.to(self.layer_norm.weight)
+        B, L, D = audio_embed.size()
+
+        if self.dconv1 is not None:
+            audio_embed = self.dconv1.forward_fixed(audio_embed)  # [1, L', d_model]
+
+        hidden_states = audio_embed
+        _, cur_L, _ = hidden_states.size()
+
+        position_ids = torch.arange(0, cur_L, device=hidden_states.device, dtype=torch.long)
+        rope_position_embeddings = self.position_embedding(hidden_states.view(cur_L, -1), position_ids)
+        cos, sin = rope_position_embeddings
+        cos = cos.unsqueeze(0)  # [1, L', head_dim]
+        sin = sin.unsqueeze(0)
+
+        for layer in self.layers:
+            hidden_states = layer.forward_fixed(hidden_states, rope_position_embeddings=(cos, sin))
+
+        hidden_states = self.layer_norm(hidden_states)
+
+        coarse_mel = self.dconv2.forward_fixed(hidden_states)  # [1, L'', n_mels]
+
+        recon_wav, wav_length = self.vocoder.forward_fixed(
+            x=coarse_mel.transpose(1, 2),
+            vocoder_attn_mask=vocoder_attn_mask,
+        )
+
+        return recon_wav
+
 
 class MiMoAudioTokenizer(PreTrainedModel):
     config_class = MiMoAudioTokenizerConfig
@@ -902,6 +1134,47 @@ class MiMoAudioTokenizer(PreTrainedModel):
         )
         return output
 
+    def _compute_vocoder_seq_len(self, code_len: int) -> int:
+        """Compute the vocoder input sequence length from code_len."""
+        cfg = self.config
+        L = code_len
+        if cfg.avg_pooler != 1:
+            k, s = cfg.avg_pooler, cfg.avg_pooler
+            L = (L - 1) * s + k - max(0, k - s)
+        k2, s2 = cfg.decoder_kernel_size, cfg.decoder_stride_size
+        L = (L - 1) * s2 + k2 - max(0, k2 - s2)
+        return L
+
+    @torch.no_grad()
+    def decode_fixed(
+        self,
+        codes: torch.Tensor,
+        padded_len: int,
+        vocoder_attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Fixed-shape decode for CUDA Graph capture.
+
+        Args:
+            codes: [audio_channels, T] — actual codes (T <= padded_len).
+            padded_len: the fixed sequence length to pad to.
+            vocoder_attn_mask: optional pre-built vocoder attention mask.
+        Returns:
+            recon_wav: [1, 1, wav_len] — full padded output (caller truncates).
+        """
+        hidden_states = self.encoder.decode_vq(codes)  # [T, d_model]
+        actual_len = hidden_states.size(0)
+        if actual_len < padded_len:
+            pad = torch.zeros(
+                padded_len - actual_len,
+                hidden_states.size(1),
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )
+            hidden_states = torch.cat([hidden_states, pad], dim=0)
+        hidden_states = hidden_states.unsqueeze(0)  # [1, padded_len, d_model]
+        output = self.decoder.forward_fixed(hidden_states, vocoder_attn_mask=vocoder_attn_mask)
+        return output
+
     @torch.no_grad()
     def streaming_decode(
         self,

vllm_omni/model_executor/models/mimo_audio/quantization.py

@@ -356,11 +356,14 @@ class ResidualVectorQuantization(nn.Module):
         return out_indices
 
     def decode(self, q_indices: torch.Tensor, st: int = 0) -> torch.Tensor:
-        quantized_out = torch.tensor(0.0, device=q_indices.device)
+        quantized_out = None
         for i, indices in enumerate(q_indices):
             layer = self.layers[st + i]
             quantized = layer.decode(indices)
-            quantized_out = quantized_out + quantized
+            if quantized_out is None:
+                quantized_out = quantized
+            else:
+                quantized_out = quantized_out + quantized
         return quantized_out