feat(moss-tts): add CUDA Graph support for codec decoder (#4157)
Signed-off-by: leo.yang <leo.yang.engineer@gmail.com>
This commit is contained in:
yangyonggit authored and GitHub committed
1 parent
152343d5fb
commit
a3834e32cf
5 files changed
+637
-18
No files matched your search
Whitespace-only changes.
@@ -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}"
|
||||
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
Reference in new issue
Block a user