From a3834e32cfc6acb231b061e9f2d1c5f31afae351 Mon Sep 17 00:00:00 2001 From: yangyonggit Date: Fri, 12 Jun 2026 00:17:31 +0200 Subject: [PATCH] feat(moss-tts): add CUDA Graph support for codec decoder (#4157) Signed-off-by: leo.yang --- .../models/moss_tts/__init__.py | 0 .../models/moss_tts/test_codec_cudagraph.py | 361 ++++++++++++++++++ vllm_omni/deploy/moss_tts.yaml | 3 +- .../moss_tts/modeling_moss_tts_codec.py | 84 +++- .../models/moss_tts/moss_codec_cudagraph.py | 207 ++++++++++ 5 files changed, 637 insertions(+), 18 deletions(-) create mode 100644 tests/model_executor/models/moss_tts/__init__.py create mode 100644 tests/model_executor/models/moss_tts/test_codec_cudagraph.py create mode 100644 vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py diff --git a/tests/model_executor/models/moss_tts/__init__.py b/tests/model_executor/models/moss_tts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/moss_tts/test_codec_cudagraph.py b/tests/model_executor/models/moss_tts/test_codec_cudagraph.py new file mode 100644 index 000000000..49d6fa411 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_codec_cudagraph.py @@ -0,0 +1,361 @@ +"""Tests for MossTTSCUDAGraphCodecWrapper numerical equivalence. + +Verifies that CUDA Graph-accelerated decoding produces results equivalent +to eager mode, with special attention to the two-argument _decode interface +(codes [NQ, 1, T] + lengths [1]) and the NQ-first input convention. +""" + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import ( + MossAudioTokenizerDecoderOutput, +) +from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import ( + MossTTSCUDAGraphCodecWrapper, +) + +pytestmark = [pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")] + +DEVICE = torch.device("cuda:0") +NUM_QUANTIZERS = 8 +DOWNSAMPLE_RATE = 4 # synthetic; real checkpoint uses 1920 + + +# --------------------------------------------------------------------------- +# Synthetic codec model +# --------------------------------------------------------------------------- + + +class SyntheticCodecModel(nn.Module): + """Minimal stand-in for MossAudioTokenizerModel. + + Exposes the same two-argument _decode(codes, lengths) interface and + returns a MossAudioTokenizerDecoderOutput, mirroring the real model. + + Input: + codes: [NQ, B, T] long — RVQ codes + lengths: [B] long — valid frame counts + Output: + MossAudioTokenizerDecoderOutput(audio=[B, 1, T*upsample], audio_lengths=[B]) + """ + + def __init__( + self, + num_quantizers: int = NUM_QUANTIZERS, + downsample_rate: int = DOWNSAMPLE_RATE, + ): + super().__init__() + self.downsample_rate = downsample_rate + hidden = 32 + # embed: treat NQ quantizers as input channels for Conv1d + self.embed = nn.Conv1d(num_quantizers, hidden, kernel_size=3, padding=1) + self.conv = nn.Conv1d(hidden, hidden, kernel_size=3, padding=1) + self.upsample = nn.ConvTranspose1d(hidden, 1, kernel_size=downsample_rate, stride=downsample_rate) + + def _decode(self, codes: torch.Tensor, lengths: torch.Tensor) -> MossAudioTokenizerDecoderOutput: + """codes: [NQ, B, T], lengths: [B] → audio [B, 1, T*upsample].""" + nq, b, t = codes.shape + # treat NQ as channel dim → [B, NQ, T] for Conv1d (expects [B, C, T]) + x = codes.permute(1, 0, 2).to(dtype=self.embed.weight.dtype) # [B, NQ, T] + x = torch.relu(self.embed(x)) # [B, hidden, T] + x = torch.relu(self.conv(x)) # [B, hidden, T] + audio = self.upsample(x) # [B, 1, T * upsample] + audio_lengths = lengths * self.downsample_rate + return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths) + + def batch_decode( + self, + codes_list: list[torch.Tensor], + num_quantizers: int | None = None, + ) -> MossAudioTokenizerDecoderOutput: + """Mirrors MossAudioTokenizerModel.batch_decode (fallback path).""" + device = codes_list[0].device + nq = num_quantizers or codes_list[0].shape[0] + max_t = max(c.shape[-1] for c in codes_list) + codes = torch.zeros(nq, len(codes_list), max_t, device=device, dtype=torch.long) + lengths = torch.zeros(len(codes_list), device=device, dtype=torch.long) + for i, c in enumerate(codes_list): + codes[:nq, i, : c.shape[-1]] = c[:nq] + lengths[i] = c.shape[-1] + return self._decode(codes, lengths) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def model(): + torch.manual_seed(0) + return SyntheticCodecModel().to(DEVICE).eval() + + +@pytest.fixture(scope="module") +def wrapper(model): + w = MossTTSCUDAGraphCodecWrapper( + model=model, + capture_sizes=[25, 50, 100], + num_quantizers=NUM_QUANTIZERS, + enabled=True, + ) + w.warmup(DEVICE) + return w + + +def _random_codes(t: int, device: torch.device = DEVICE) -> torch.Tensor: + """Return [NQ, T] long tensor (single-request convention).""" + return torch.randint(0, 64, (NUM_QUANTIZERS, t), dtype=torch.long, device=device) + + +def _eager_decode(model: SyntheticCodecModel, codes_nq_t: torch.Tensor) -> MossAudioTokenizerDecoderOutput: + """Run the model directly in eager mode (reference output).""" + t = codes_nq_t.shape[-1] + codes_nq_1_t = codes_nq_t.unsqueeze(1) # [NQ, 1, T] + lengths = torch.tensor([t], dtype=torch.long, device=codes_nq_t.device) + with torch.no_grad(): + return model._decode(codes_nq_1_t, lengths) + + +# --------------------------------------------------------------------------- +# 1. Exact-size inputs — must be bit-identical +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [25, 50, 100]) +def test_exact_size_bit_identical(model, wrapper, t): + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 2. Padded inputs — output trimmed to actual length +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [10, 30, 47, 73, 99]) +def test_padded_output_shape(model, wrapper, t): + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + expected_len = t * DOWNSAMPLE_RATE + assert out.audio.shape[-1] == expected_len, f"got {out.audio.shape[-1]}, want {expected_len}" + assert out.audio.shape == ref.audio.shape + + +@pytest.mark.parametrize("t", [10, 30, 47, 73, 99]) +def test_padded_interior_positions_close(model, wrapper, t): + """Positions well away from the zero-padding boundary must be numerically close.""" + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + boundary = 2 * DOWNSAMPLE_RATE + if ref.audio.shape[-1] > boundary: + torch.testing.assert_close( + out.audio[..., :-boundary], + ref.audio[..., :-boundary], + atol=1e-5, + rtol=1e-5, + ) + + +# --------------------------------------------------------------------------- +# 3. Fallback to eager (T > all capture sizes) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [101, 150, 200]) +def test_fallback_exact_match(model, wrapper, t): + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 4. Static buffer not aliased across calls +# --------------------------------------------------------------------------- + + +def test_output_not_aliased_after_later_replay(model, wrapper): + """clone() in decode() must prevent later replays from overwriting the result.""" + codes_a = _random_codes(50) + codes_b = _random_codes(50) + with torch.no_grad(): + out_a = wrapper.decode(codes_a) + saved = out_a.audio.clone() + _ = wrapper.decode(codes_b) # replay overwrites static buffer + torch.testing.assert_close(out_a.audio, saved, atol=0, rtol=0) + + +def test_deterministic_across_calls(model, wrapper): + codes = _random_codes(30) + with torch.no_grad(): + out1 = wrapper.decode(codes) + out2 = wrapper.decode(codes) + torch.testing.assert_close(out1.audio, out2.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 5. Disabled wrapper falls back to eager +# --------------------------------------------------------------------------- + + +def test_disabled_falls_back_to_eager(model, wrapper): + codes = _random_codes(30) + ref = _eager_decode(model, codes) + wrapper.enabled = False + with torch.no_grad(): + out = wrapper.decode(codes) + wrapper.enabled = True + torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 6. audio_lengths is consistent with audio shape +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [25, 30, 50, 99, 100]) +def test_audio_lengths_consistent(model, wrapper, t): + codes = _random_codes(t) + with torch.no_grad(): + out = wrapper.decode(codes) + assert out.audio_lengths is not None + assert int(out.audio_lengths[0].item()) == out.audio.shape[-1] + + +# --------------------------------------------------------------------------- +# 7. NQ-first input layout: codes_nq_t is [NQ, T], not [T, NQ] +# --------------------------------------------------------------------------- + + +def test_nq_first_layout_matches_eager(model, wrapper): + """Verify that wrapper correctly interprets [NQ, T] (NQ-first) layout.""" + t = 25 + codes = _random_codes(t) # [NQ, T] + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 8. Non-contiguous input — copy_() must handle strided tensors correctly +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [20, 47]) +def test_noncontiguous_input_matches_contiguous(model, wrapper, t): + """Padded path uses copy_() on a slice; must produce the same result + whether codes_nq_t is contiguous or a non-contiguous view (e.g. a column + slice of a larger buffer, which is the realistic production layout when the + caller slices out one request from a batched tensor).""" + codes_contiguous = _random_codes(t) + # Build a non-contiguous view: embed in a larger [NQ, 2, T] buffer and + # take column 0. The resulting tensor has stride (2*T, 2, 1) → non-contiguous. + buf = torch.zeros(NUM_QUANTIZERS, 2, t, dtype=torch.long, device=DEVICE) + buf[:, 0, :] = codes_contiguous + codes_noncontig = buf[:, 0, :] # [NQ, T], non-contiguous + assert not codes_noncontig.is_contiguous() + + with torch.no_grad(): + out_c = wrapper.decode(codes_contiguous) + out_nc = wrapper.decode(codes_noncontig) + torch.testing.assert_close(out_c.audio, out_nc.audio, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 9. Zero-padding uses token 0 — must be a valid codebook index (no OOB) +# --------------------------------------------------------------------------- + + +def test_zero_padding_no_cuda_error(model, wrapper): + """Padded inputs fill unused frames with code 0. Verify that the zero- + padded region does not corrupt the output shape. + + Note: SyntheticCodecModel uses Conv1d (not nn.Embedding), so it cannot + trigger a real CUDA index-out-of-bounds on large token ids. This test + validates shape correctness only; OOB safety on the real codec requires + an integration test with MossAudioTokenizerModel.""" + t = 7 # well below the 25-frame bucket → large zero-padded region + codes = _random_codes(t) + with torch.no_grad(): + out = wrapper.decode(codes) + torch.accelerator.synchronize() # surface any deferred CUDA errors + assert out.audio.shape[-1] == t * DOWNSAMPLE_RATE + + +# --------------------------------------------------------------------------- +# 10. Shape alignment: graph path and eager fallback return identical shapes +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("t", [25, 30, 101]) +def test_graph_and_eager_shape_identical(model, wrapper, t): + """Graph replay (t<=100) and eager fallback (t>100) must return tensors + with the same number of dimensions and the same batch/channel layout so + downstream consumers can treat both paths uniformly.""" + codes = _random_codes(t) + ref = _eager_decode(model, codes) + with torch.no_grad(): + out = wrapper.decode(codes) + assert out.audio.ndim == ref.audio.ndim, f"ndim mismatch: graph={out.audio.ndim} eager={ref.audio.ndim}" + assert out.audio.shape[:-1] == ref.audio.shape[:-1], ( + f"batch/channel shape mismatch: graph={out.audio.shape} eager={ref.audio.shape}" + ) + + +# --------------------------------------------------------------------------- +# 11. Half-precision weights: graph and eager must agree under fp16/bf16 +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_half_precision_graph_eager_agree(dtype): + """CUDA Graph must produce outputs consistent with eager when model weights + are in half precision. The wrapper itself stores codes as long and lengths + as long; only the model's conv weights change dtype.""" + torch.manual_seed(0) + model_hp = SyntheticCodecModel().to(device=DEVICE, dtype=dtype).eval() + wrapper_hp = MossTTSCUDAGraphCodecWrapper( + model=model_hp, + capture_sizes=[25, 50], + num_quantizers=NUM_QUANTIZERS, + enabled=True, + ) + wrapper_hp.warmup(DEVICE) + + codes = _random_codes(25) + ref = _eager_decode(model_hp, codes) + with torch.no_grad(): + out = wrapper_hp.decode(codes) + # Half-precision accumulation introduces small numerical error; use + # tolerances that match typical fp16/bf16 rounding. + atol = 1e-2 if dtype == torch.float16 else 5e-2 + torch.testing.assert_close(out.audio.float(), ref.audio.float(), atol=atol, rtol=1e-3) + + +# --------------------------------------------------------------------------- +# 12. Empty input (T=0): wrapper must not crash and return zero-length audio +# --------------------------------------------------------------------------- + + +def test_empty_input_t0(model, wrapper): + """T=0 input must be handled gracefully: no exception, audio length == 0. + + The outer MossTTSCodecDecoder guards against empty segments, but the + wrapper should also handle T=0 without crashing so it is safe to call + directly.""" + codes = _random_codes(0) # [NQ, 0] + with torch.no_grad(): + out = wrapper.decode(codes) + assert out.audio.shape[-1] == 0, f"expected empty audio, got shape {out.audio.shape}" diff --git a/vllm_omni/deploy/moss_tts.yaml b/vllm_omni/deploy/moss_tts.yaml index acdb87721..dbab816eb 100644 --- a/vllm_omni/deploy/moss_tts.yaml +++ b/vllm_omni/deploy/moss_tts.yaml @@ -14,6 +14,7 @@ connectors: connector_get_max_wait: 300 codec_chunk_frames: 25 codec_left_context_frames: 0 + decode_cudagraph_capture_sizes: [4, 8, 16, 25, 32, 50, 64, 100, 128, 200, 256] stages: - stage_id: 0 @@ -38,7 +39,7 @@ stages: - stage_id: 1 max_num_seqs: 1 gpu_memory_utilization: 0.12 - enforce_eager: true + enforce_eager: false enable_prefix_caching: false async_scheduling: true # Talker can emit up to max_tokens=4096 frames × n_vq=32 = 131072 codec diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index fa731fa6f..b1b6063d4 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -20,6 +20,9 @@ from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import ( MossAudioTokenizerConfig, MossAudioTokenizerModel, ) +from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import ( + MossTTSCUDAGraphCodecWrapper, +) from vllm_omni.model_executor.models.output_templates import OmniOutput logger = init_logger(__name__) @@ -28,15 +31,16 @@ logger = init_logger(__name__) class MossTTSCodecDecoder(nn.Module): """Stage-1 decoder for all MOSS-TTS variants. - Consumes ``(T, NQ)`` audio VQ codes emitted by Stage 0 and decodes them + Consumes ``(NQ, T)`` audio VQ codes emitted by Stage 0 and decodes them to a 24 kHz mono waveform using the vendored ``MossAudioTokenizerModel``. All five variants share the same codec checkpoint ``OpenMOSS-Team/MOSS-Audio-Tokenizer``. The number of quantizers - (``n_vq``) is passed as a ``num_quantizers`` argument to ``batch_decode`` - so one codec instance handles both ``n_vq=32`` (MOSS-TTS) and ``n_vq=16`` - (all other variants) without swapping weights. + (``n_vq``) is read from ``hf_config`` at construction time and fixed for + the lifetime of the instance; the same checkpoint can be configured as + ``n_vq=32`` (MOSS-TTS) or ``n_vq=16`` (all other variants) without + swapping weights. The codec checkpoint path defaults to the value stored in ``vllm_config.model_config.hf_config.codec_model_name_or_path`` but can @@ -64,6 +68,7 @@ class MossTTSCodecDecoder(nn.Module): # Resolved at load_weights() time self._codec: MossAudioTokenizerModel | None = None + self._cuda_graph_wrapper: MossTTSCUDAGraphCodecWrapper | None = None self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32) # ------------------------------------------------------------------ @@ -100,10 +105,11 @@ class MossTTSCodecDecoder(nn.Module): Stage 0 emits flat codebook-major ``[NQ * T_chunk]`` audio codes. The chunk transfer adapter assigns those to ``request.prompt_token_ids``, - which arrives here as ``input_ids``. Per-request offsets are derived - from ``runtime_additional_information``-attached request metadata; meta - like ``left_context_size`` (overlap context for the causal decoder) - also lives there. + which arrives here as ``input_ids`` concatenated across all requests. + Per-request slice boundaries are computed from + ``kwargs["num_scheduled_tokens"]`` (a list of token counts, one per + request). ``runtime_additional_information`` carries per-request + metadata such as ``left_context_size``. Returns ------- @@ -163,11 +169,10 @@ class MossTTSCodecDecoder(nn.Module): continue t_chunk = int(seg.numel() // self._n_vq) codes_nq_t = seg.reshape(self._n_vq, t_chunk).to(device=device) - # Clamp out-of-range codes (the talker uses ``audio_pad_code`` - # =``codebook_size`` for delay-pattern padding; passing it to the - # codec embedding lookup would be an OOB index. The downstream - # de-delay step is intended to drop those rows, but until that's - # wired we clamp so the kernel does not assert.) + # Clamp out-of-range codes: the talker uses ``audio_pad_code`` + # (= ``codebook_size``) for delay-pattern padding. The stage input + # processor de-delays and drops pad rows before forwarding here, but + # clamp as a defensive guard against any edge-case leakage. codebook_size = self._codec.config.codebook_size codes_nq_t = codes_nq_t.clamp_(0, int(codebook_size) - 1) @@ -183,7 +188,10 @@ class MossTTSCodecDecoder(nn.Module): left_ctx = int(left_ctx.reshape(-1)[0].item()) if left_ctx.numel() else 0 left_ctx = int(left_ctx) - out = self._codec.batch_decode(codes_list=[codes_nq_t], num_quantizers=self._n_vq) + if self._cuda_graph_wrapper is not None: + out = self._cuda_graph_wrapper.decode(codes_nq_t) + else: + out = self._codec.batch_decode(codes_list=[codes_nq_t], num_quantizers=self._n_vq) if out.audio is None: continue @@ -194,7 +202,13 @@ class MossTTSCodecDecoder(nn.Module): # Trim left-context samples. if left_ctx > 0: - trim = left_ctx * self._codec.downsample_rate + trim = min(left_ctx * self._codec.downsample_rate, wav.shape[0]) + if trim < left_ctx * self._codec.downsample_rate: + logger.warning( + "left_ctx trim (%d samples) exceeds wav length (%d); returning empty audio.", + left_ctx * self._codec.downsample_rate, + wav.shape[0], + ) wav = wav[trim:] audios[i] = wav @@ -288,8 +302,7 @@ class MossTTSCodecDecoder(nn.Module): codec_cfg.num_quantizers, ) - # Move sr_tensor to the same device - self._sr_tensor = self._sr_tensor.to(device=device) + self._maybe_enable_decoder_cudagraph(device) # vLLM's track_weights_loading() compares the returned set against # ``self.named_parameters()``. After ``self._codec = codec`` above, @@ -297,5 +310,42 @@ class MossTTSCodecDecoder(nn.Module): # mirror that here. return {f"_codec.{name}" for name, _ in codec.named_parameters()} + def _maybe_enable_decoder_cudagraph(self, device: torch.device) -> None: + """Capture CUDA Graphs for the codec decoder if enforce_eager is False.""" + if getattr(self.vllm_config.model_config, "enforce_eager", True): + return + if self._codec is None: + return + + # Read capture sizes from the connector's extra config (same convention + # as Qwen3-TTS), falling back to a sensible default covering common + # codec_chunk_frames values used in moss_tts.yaml. + capture_sizes: list[int] = [4, 8, 16, 25, 32, 50, 64, 100, 128, 200, 256] + model_cfg = getattr(self.vllm_config, "model_config", None) + connector_cfg = getattr(model_cfg, "stage_connector_config", None) + if isinstance(connector_cfg, dict): + extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg) + else: + extra_cfg = getattr(connector_cfg, "extra", None) + if isinstance(extra_cfg, dict): + raw = extra_cfg.get("decode_cudagraph_capture_sizes") + if raw is not None: + if isinstance(raw, (list, tuple)): + parsed = sorted({int(v) for v in raw if int(v) > 0}) + elif isinstance(raw, str): + parsed = sorted({int(v.strip()) for v in raw.split(",") if v.strip()}) + else: + parsed = [int(raw)] + if parsed: + capture_sizes = parsed + + self._cuda_graph_wrapper = MossTTSCUDAGraphCodecWrapper( + model=self._codec, + capture_sizes=capture_sizes, + num_quantizers=self._n_vq, + enabled=True, + ) + self._cuda_graph_wrapper.warmup(device) + __all__ = ["MossTTSCodecDecoder"] diff --git a/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py b/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py new file mode 100644 index 000000000..ed03c018e --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py @@ -0,0 +1,207 @@ +"""CUDA Graph acceleration for the MOSS Audio Tokenizer codec decoder. + +Captures MossAudioTokenizerModel._decode for a set of fixed frame-count +bucket sizes, then replays the captured graph at inference time to eliminate +kernel-launch overhead. Inputs that exceed all captured sizes fall back to +eager execution transparently. +""" + +from __future__ import annotations + +import time + +import torch +from torch.cuda import CUDAGraph +from vllm.logger import init_logger +from vllm.platforms import current_platform + +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import ( + MossAudioTokenizerDecoderOutput, + MossAudioTokenizerModel, +) + +logger = init_logger(__name__) + + +class MossTTSCUDAGraphCodecWrapper: + """CUDA Graph wrapper for MossAudioTokenizerModel._decode. + + Graphs are keyed by padded_T (int). On each call the actual T is + bucket-matched to the smallest pre-captured size >= T. The static code + buffer [NQ, 1, padded_T] is filled left-aligned (right-zero-padded) and + the graph is replayed. The output audio is sliced to the correct length + by scaling from the captured audio shape (actual_T / padded_T * captured_len), + avoiding any assumption about downsample_rate vs effective decoder upsample. + The slice is cloned before returning so the static buffer can be reused. + + Usage:: + + wrapper = MossTTSCUDAGraphCodecWrapper(codec_model, capture_sizes, nq) + wrapper.warmup(device) + + # per-request decode: + out = wrapper.decode(codes_nq_t) # codes_nq_t: [NQ, T] + """ + + def __init__( + self, + model: MossAudioTokenizerModel, + capture_sizes: list[int], + num_quantizers: int, + enabled: bool = True, + ) -> None: + self.model = model + self.capture_sizes: list[int] = sorted(capture_sizes) + self.num_quantizers = num_quantizers + self.enabled = enabled + + # All dicts keyed by padded_T. + self.graphs: dict[int, CUDAGraph] = {} + self.static_codes: dict[int, torch.Tensor] = {} # [NQ, 1, padded_T] + # static_lengths is kept alive here — the captured graph holds a + # reference to the underlying storage and must not be GC'd. + self.static_lengths: dict[int, torch.Tensor] = {} # [1] + self.static_audio: dict[int, torch.Tensor] = {} # [1, 1, padded_T * effective_upsample] + + self._warmed_up = False + + # ------------------------------------------------------------------ + # Size helpers + # ------------------------------------------------------------------ + + def _get_padded_size(self, actual_t: int) -> int | None: + """Return the smallest capture size >= actual_t, or None if too large.""" + for s in self.capture_sizes: + if actual_t <= s: + return s + return None + + # ------------------------------------------------------------------ + # Warmup / capture + # ------------------------------------------------------------------ + + def warmup(self, device: torch.device) -> None: + """Allocate static buffers and capture CUDA Graphs for all sizes.""" + if device.type != "cuda" or not self.enabled or self._warmed_up: + return + + nq = self.num_quantizers + logger.info( + "MOSS-TTS codec CUDA Graph warmup: nq=%d capture_sizes=%s", + nq, + self.capture_sizes, + ) + t0 = time.perf_counter() + + # One eager run per size to let cuDNN / CUDA allocate memory before + # the capture window (graph capture forbids new CUDA allocs during it). + for size in self.capture_sizes: + dummy_codes = torch.zeros(nq, 1, size, dtype=torch.long, device=device) + dummy_lengths = torch.tensor([size], dtype=torch.long, device=device) + with torch.no_grad(): + _ = self.model._decode(dummy_codes, dummy_lengths) + + torch.accelerator.synchronize(device) + + for size in self.capture_sizes: + try: + self._capture(size, device) + logger.info(" Captured CUDA Graph for size=%d", size) + except Exception: + logger.warning( + " Failed to capture CUDA Graph for size=%d; this size will fall back to eager decode", + size, + exc_info=True, + ) + + self._warmed_up = True + elapsed_ms = (time.perf_counter() - t0) * 1000.0 + logger.info( + "MOSS-TTS codec CUDA Graph warmup complete: %d/%d captured in %.1f ms", + len(self.graphs), + len(self.capture_sizes), + elapsed_ms, + ) + + def _capture(self, size: int, device: torch.device) -> None: + nq = self.num_quantizers + static_codes = torch.zeros(nq, 1, size, dtype=torch.long, device=device) + # lengths holds the number of valid code frames; set to full size at + # capture time so the decoder emits a full-size audio buffer. + static_lengths = torch.tensor([size], dtype=torch.long, device=device) + + # Extra eager warmup inside capture to ensure all kernels are compiled. + with torch.no_grad(): + _ = self.model._decode(static_codes, static_lengths) + torch.accelerator.synchronize(device) + + graph = CUDAGraph() + with torch.no_grad(): + with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + static_out = self.model._decode(static_codes, static_lengths) + + self.graphs[size] = graph + self.static_codes[size] = static_codes + self.static_lengths[size] = static_lengths + # static_out.audio is a static buffer reused every replay; hold a + # reference so it is not garbage-collected. + self.static_audio[size] = static_out.audio # [1, 1, size * effective_upsample] + + # ------------------------------------------------------------------ + # Inference + # ------------------------------------------------------------------ + + @torch.no_grad() + def decode(self, codes_nq_t: torch.Tensor) -> MossAudioTokenizerDecoderOutput: + """Decode [NQ, T] codes to waveform using a CUDA Graph when possible. + + Falls back to eager batch_decode when: + - CUDA Graph is disabled or not yet warmed up + - an outer CUDA stream capture is active (e.g. vLLM FULL graph mode) + - actual T exceeds all pre-captured sizes + """ + if not self.enabled or not self._warmed_up: + return self.model.batch_decode(codes_list=[codes_nq_t], num_quantizers=self.num_quantizers) + + # Replaying a graph inside an active stream capture would corrupt the + # outer graph. Fall back to eager so the caller can complete its capture. + if torch.cuda.is_current_stream_capturing(): + return self.model.batch_decode(codes_list=[codes_nq_t], num_quantizers=self.num_quantizers) + + actual_t = int(codes_nq_t.shape[-1]) + padded_size = self._get_padded_size(actual_t) + + if padded_size is None or padded_size not in self.graphs: + return self.model.batch_decode(codes_list=[codes_nq_t], num_quantizers=self.num_quantizers) + + # --- Fill static buffers then replay --- + + static_codes = self.static_codes[padded_size] # [NQ, 1, padded_size] + + if actual_t == padded_size: + # Exact fit: copy the whole buffer at once. + # codes_nq_t is [NQ, T]; unsqueeze(1) → [NQ, 1, T] matching static_codes. + static_codes.copy_(codes_nq_t.unsqueeze(1)) + else: + # Smaller input: zero the buffer first (right-zero-pad), then fill + # left-aligned. static_codes[:, 0, :actual_t] is [NQ, actual_t] + # and codes_nq_t is [NQ, actual_t]. + static_codes.zero_() + static_codes[:, 0, :actual_t].copy_(codes_nq_t) + + self.graphs[padded_size].replay() + + # static_audio[padded_size] is the live output buffer of the graph. + # Slice to the real audio length and clone before returning; without + # the clone the next replay would overwrite the caller's tensor. + # Derive the slice length from the captured audio shape rather than + # downsample_rate: the decoder's effective upsample (product of all + # PatchedPretransform patch sizes) can differ from the config attribute. + captured_len = self.static_audio[padded_size].shape[-1] + actual_wav_len = captured_len * actual_t // padded_size + audio = self.static_audio[padded_size][..., :actual_wav_len].clone() + audio_lengths = torch.tensor([actual_wav_len], dtype=torch.long, device=audio.device) + return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths) + + +__all__ = ["MossTTSCUDAGraphCodecWrapper"]