From e1088026faa5e4ee7b27ae3cd835fdc74f6431c0 Mon Sep 17 00:00:00 2001 From: Baoyuan Qi Date: Mon, 11 May 2026 10:05:46 +0800 Subject: [PATCH] [Performance] Improve MiMo-Audio tokenizer decoding performance (#2183) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 齐保元 Co-authored-by: Jialong Liu <88185941+Galleons2029@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- .../test_mimo_audio_code2wav_batch_decode.py | 2 +- .../mimo_audio/cuda_graph_decoder_wrapper.py | 396 ++++++++++++++++++ .../models/mimo_audio/mimo_audio.py | 8 +- .../models/mimo_audio/mimo_audio_code2wav.py | 116 ++++- .../mimo_audio/modeling_audio_tokenizer.py | 279 +++++++++++- .../models/mimo_audio/quantization.py | 7 +- 6 files changed, 775 insertions(+), 33 deletions(-) create mode 100644 vllm_omni/model_executor/models/mimo_audio/cuda_graph_decoder_wrapper.py diff --git a/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py b/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py index 0262eaa02..194a024db 100644 --- a/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py +++ b/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 diff --git a/vllm_omni/model_executor/models/mimo_audio/cuda_graph_decoder_wrapper.py b/vllm_omni/model_executor/models/mimo_audio/cuda_graph_decoder_wrapper.py new file mode 100644 index 000000000..a88e3548a --- /dev/null +++ b/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 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 9a09ab030..ecf094c87 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/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 diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py index e6d70915c..449425a2f 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py +++ b/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) diff --git a/vllm_omni/model_executor/models/mimo_audio/modeling_audio_tokenizer.py b/vllm_omni/model_executor/models/mimo_audio/modeling_audio_tokenizer.py index ddd683e11..0093da541 100644 --- a/vllm_omni/model_executor/models/mimo_audio/modeling_audio_tokenizer.py +++ b/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, diff --git a/vllm_omni/model_executor/models/mimo_audio/quantization.py b/vllm_omni/model_executor/models/mimo_audio/quantization.py index 2ce7beb04..574224aa5 100644 --- a/vllm_omni/model_executor/models/mimo_audio/quantization.py +++ b/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