feat(qwen3-tts): Add CUDA Graph support for speech tokenizer decoder (#1205)

Signed-off-by: xulusjb <fdukeshik@gmail.com>
Co-authored-by: xulusjb <fdukeshik@gmail.com>

tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py

@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for CUDA Graph decoder wrapper numerical equivalence.
+
+Verifies that CUDA Graph-accelerated decoding produces results equivalent
+to eager mode, with special attention to padding cases where zero-padding
+may introduce small numerical differences due to attention and convolution.
+
+Architecture note: the real Qwen3TTSTokenizerV2Decoder uses causal
+convolutions, so zero-padding on the right has minimal impact (~2e-3).
+The synthetic decoder here uses standard (non-causal) Conv1d for a
+worst-case test of the wrapper mechanism.
+"""
+
+import importlib.util
+import os
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")]
+
+DEVICE = torch.device("cuda:0")
+NUM_QUANTIZERS = 8
+TOTAL_UPSAMPLE = 4
+
+# Load CUDAGraphDecoderWrapper: try package import first, fall back to direct file load
+try:
+    from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+except Exception:
+    _WRAPPER_PATH = os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "model_executor",
+        "models",
+        "qwen3_tts",
+        "cuda_graph_decoder_wrapper.py",
+    )
+    _spec = importlib.util.spec_from_file_location("cuda_graph_decoder_wrapper", os.path.abspath(_WRAPPER_PATH))
+    _mod = importlib.util.module_from_spec(_spec)
+    _spec.loader.exec_module(_mod)
+    CUDAGraphDecoderWrapper = _mod.CUDAGraphDecoderWrapper
+
+
+class SyntheticDecoder(nn.Module):
+    """A small decoder mimicking Qwen3TTSTokenizerV2Decoder's interface.
+
+    Uses Conv1d layers so that zero-padding can affect neighboring positions
+    via the receptive field, providing a worst-case test for padding effects.
+    """
+
+    def __init__(self, num_quantizers=NUM_QUANTIZERS, total_upsample=TOTAL_UPSAMPLE):
+        super().__init__()
+        hidden = 32
+        self.total_upsample = total_upsample
+        self.embed = nn.Conv1d(num_quantizers, hidden, kernel_size=3, padding=1)
+        self.conv1 = nn.Conv1d(hidden, hidden, kernel_size=5, padding=2)
+        self.conv2 = nn.Conv1d(hidden, hidden, kernel_size=3, padding=1)
+        self.upsample = nn.ConvTranspose1d(hidden, hidden, kernel_size=total_upsample, stride=total_upsample)
+        self.out = nn.Conv1d(hidden, 1, kernel_size=1)
+
+    def forward(self, codes):
+        x = codes.float()
+        x = torch.relu(self.embed(x))
+        x = torch.relu(self.conv1(x))
+        x = torch.relu(self.conv2(x))
+        x = self.upsample(x)
+        return self.out(x).clamp(min=-1, max=1)
+
+
+@pytest.fixture(scope="module")
+def decoder():
+    """Create a synthetic decoder on CUDA with fixed weights."""
+    torch.manual_seed(42)
+    return SyntheticDecoder().to(DEVICE).eval()
+
+
+@pytest.fixture(scope="module")
+def wrapper(decoder):
+    """Create a warmed-up CUDAGraphDecoderWrapper."""
+    w = CUDAGraphDecoderWrapper(
+        decoder=decoder,
+        capture_sizes=[25, 50, 100],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    w.warmup(DEVICE)
+    return w
+
+
+def _random_codes(seq_len, device=DEVICE):
+    return torch.randint(0, 100, (1, NUM_QUANTIZERS, seq_len), dtype=torch.long, device=device)
+
+
+# ──────────────────────────────────────────────────────────────────
+# 1. Exact-size inputs (no padding needed) → bit-identical
+# ──────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("seq_len", [25, 50, 100])
+def test_exact_size_numerical_equivalence(decoder, wrapper, seq_len):
+    """When input exactly matches a capture size, output must be bit-identical."""
+    codes = _random_codes(seq_len)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ──────────────────────────────────────────────────────────────────
+# 2. Padded inputs (zero-padding to nearest capture size)
+# ──────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
+def test_padded_output_shape_and_length(decoder, wrapper, seq_len):
+    """Padded decode must return output trimmed to actual input length."""
+    codes = _random_codes(seq_len)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+    expected_len = seq_len * TOTAL_UPSAMPLE
+    assert graph_out.shape == eager_out.shape
+    assert graph_out.shape[-1] == expected_len
+
+
+@pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
+def test_padded_interior_positions_close(decoder, wrapper, seq_len):
+    """Interior positions (away from padding boundary) should be very close.
+
+    The conv receptive field is at most 5 (kernel_size=5), so positions
+    more than 2 timesteps from the end (times the upsample factor) should
+    be nearly identical between eager and graph modes.
+    """
+    codes = _random_codes(seq_len)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+
+    # Exclude the last (receptive_field * upsample) positions from strict check
+    boundary = 3 * TOTAL_UPSAMPLE  # conservative: 3 positions * 4x upsample
+    if eager_out.shape[-1] > boundary:
+        interior_eager = eager_out[..., :(-boundary)]
+        interior_graph = graph_out[..., :(-boundary)]
+        torch.testing.assert_close(interior_graph, interior_eager, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
+def test_padded_output_bounded(decoder, wrapper, seq_len):
+    """Padded output values must remain in [-1, 1] and max diff should be bounded."""
+    codes = _random_codes(seq_len)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+
+    assert graph_out.min() >= -1.0 and graph_out.max() <= 1.0
+    max_diff = (graph_out - eager_out).abs().max().item()
+    # With non-causal conv, boundary diffs can be large (~0.5).
+    # The real causal decoder shows ~2e-3.
+    assert max_diff < 1.0, f"Max diff {max_diff} exceeds bound"
+
+
+# ──────────────────────────────────────────────────────────────────
+# 3. Fallback to eager (size exceeds all capture sizes) → bit-identical
+# ──────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("seq_len", [101, 150, 200])
+def test_fallback_eager_exact_match(decoder, wrapper, seq_len):
+    """Input larger than all capture sizes falls back to eager -> bit-identical."""
+    codes = _random_codes(seq_len)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ──────────────────────────────────────────────────────────────────
+# 4. Chunked decode equivalence
+# ──────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("total_len", [60, 100, 150, 250])
+def test_chunked_decode_shape_match(decoder, wrapper, total_len):
+    """Chunked decode output shape must match between eager and graph modes."""
+    codes = _random_codes(total_len)
+    chunk_size, ctx = 50, 10
+
+    with torch.no_grad():
+        eager_out = _eager_chunked(decoder, codes, chunk_size, ctx)
+        graph_out = wrapper.chunked_decode_with_cudagraph(codes, chunk_size=chunk_size, left_context_size=ctx)
+
+    assert eager_out.shape == graph_out.shape
+
+
+@pytest.mark.parametrize("total_len", [50, 100])
+def test_chunked_decode_exact_size_equivalence(decoder, wrapper, total_len):
+    """Chunked decode with chunks matching capture sizes should be bit-identical."""
+    codes = _random_codes(total_len)
+    # chunk_size=50 matches a capture size exactly, no context overlap
+    chunk_size, ctx = 50, 0
+
+    with torch.no_grad():
+        eager_out = _eager_chunked(decoder, codes, chunk_size, ctx)
+        graph_out = wrapper.chunked_decode_with_cudagraph(codes, chunk_size=chunk_size, left_context_size=ctx)
+
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+def _eager_chunked(decoder, codes, chunk_size, left_context_size):
+    """Eager chunked decode matching the real decoder's chunked_decode logic."""
+    wavs = []
+    start = 0
+    total_len = codes.shape[-1]
+    while start < total_len:
+        end = min(start + chunk_size, total_len)
+        ctx = left_context_size if start - left_context_size > 0 else start
+        chunk = codes[..., start - ctx : end]
+        wav = decoder(chunk)
+        wavs.append(wav[..., ctx * decoder.total_upsample :])
+        start = end
+    return torch.cat(wavs, dim=-1)
+
+
+# ──────────────────────────────────────────────────────────────────
+# 5. Edge cases and control tests
+# ──────────────────────────────────────────────────────────────────
+
+
+def test_single_frame(decoder, wrapper):
+    """Single-frame input (seq_len=1) should work with padding."""
+    codes = _random_codes(1)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+    assert graph_out.shape == eager_out.shape
+    assert graph_out.shape[-1] == TOTAL_UPSAMPLE
+
+
+def test_disabled_wrapper_matches_eager(decoder, wrapper):
+    """Disabled wrapper should produce bit-identical output to eager."""
+    codes = _random_codes(30)
+    wrapper.enabled = False
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+    wrapper.enabled = True
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+def test_batch_size_gt1_falls_back(decoder, wrapper):
+    """Batch size > 1 should fall back to eager (bit-identical)."""
+    codes = torch.randint(0, 100, (2, NUM_QUANTIZERS, 25), dtype=torch.long, device=DEVICE)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+def test_deterministic_across_calls(decoder, wrapper):
+    """Same input should produce identical CUDA graph output across calls."""
+    codes = _random_codes(30)
+    with torch.no_grad():
+        out1 = wrapper.decode(codes)
+        out2 = wrapper.decode(codes)
+    torch.testing.assert_close(out1, out2, atol=0, rtol=0)

vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py

@@ -0,0 +1,169 @@
+# Copyright 2026 The Alibaba Qwen team.
+# SPDX-License-Identifier: Apache-2.0
+"""
+CUDA Graph wrapper for Qwen3TTSTokenizerV2Decoder.
+
+This module provides CUDA Graph acceleration for the speech tokenizer decoder,
+reducing kernel launch overhead during inference.
+"""
+
+import torch
+from torch.cuda import CUDAGraph
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class CUDAGraphDecoderWrapper:
+    """
+    CUDA Graph wrapper for Qwen3TTSTokenizerV2Decoder.
+
+    This wrapper captures the decoder forward pass for fixed input sizes
+    and replays them during inference to reduce kernel launch overhead.
+
+    Usage:
+        wrapper = CUDAGraphDecoderWrapper(decoder, capture_sizes=[25, 50, 100, 200, 300])
+        wrapper.warmup(device)
+
+        # During inference:
+        output = wrapper.decode(codes)  # Automatically uses CUDA graph if possible
+    """
+
+    DEFAULT_CAPTURE_SIZES = [25, 50, 100, 150, 200, 250, 300]
+
+    def __init__(
+        self,
+        decoder: torch.nn.Module,
+        capture_sizes: list[int] | None = None,
+        num_quantizers: int = 8,
+        enabled: bool = True,
+    ):
+        self.decoder = decoder
+        self.capture_sizes = capture_sizes or self.DEFAULT_CAPTURE_SIZES
+        self.num_quantizers = num_quantizers
+        self.enabled = enabled
+
+        self.graphs: dict[int, CUDAGraph] = {}
+        self.static_inputs: dict[int, torch.Tensor] = {}
+        self.static_outputs: dict[int, torch.Tensor] = {}
+
+        self._warmed_up = False
+        self._device = None
+
+    def _get_padded_size(self, actual_size: int) -> int | None:
+        for size in self.capture_sizes:
+            if actual_size <= size:
+                return size
+        return None
+
+    def warmup(self, device: torch.device, dtype: torch.dtype = torch.long):
+        if device.type != "cuda":
+            logger.info("CUDA Graph warmup skipped: device %s is not CUDA", device)
+            return
+
+        if not self.enabled:
+            logger.info("CUDA Graph is disabled, skipping warmup")
+            return
+
+        if self._warmed_up:
+            logger.warning("CUDA Graph already warmed up, skipping")
+            return
+
+        self._device = device
+        self.decoder.eval()
+
+        logger.info("Starting CUDA Graph warmup for %d sizes: %s", len(self.capture_sizes), self.capture_sizes)
+
+        # Warmup runs to ensure CUDA memory is allocated
+        for size in self.capture_sizes:
+            dummy_codes = torch.zeros(
+                1,
+                self.num_quantizers,
+                size,
+                dtype=dtype,
+                device=device,
+            )
+            with torch.no_grad():
+                _ = self.decoder(dummy_codes)
+
+        torch.cuda.synchronize(device)
+
+        for size in self.capture_sizes:
+            try:
+                self._capture_graph_for_size(size, device, dtype)
+                logger.info("  Captured CUDA Graph for size=%d", size)
+            except Exception:
+                logger.warning("  Failed to capture CUDA Graph for size=%d", size, exc_info=True)
+
+        self._warmed_up = True
+        logger.info("CUDA Graph warmup complete. Captured %d graphs.", len(self.graphs))
+
+    def _capture_graph_for_size(self, size: int, device: torch.device, dtype: torch.dtype):
+        static_input = torch.zeros(
+            1,
+            self.num_quantizers,
+            size,
+            dtype=dtype,
+            device=device,
+        )
+
+        with torch.no_grad():
+            _ = self.decoder(static_input)
+
+        torch.cuda.synchronize(device)
+
+        graph = CUDAGraph()
+        with torch.no_grad():
+            with torch.cuda.graph(graph):
+                static_output = self.decoder(static_input)
+
+        self.graphs[size] = graph
+        self.static_inputs[size] = static_input
+        self.static_outputs[size] = static_output
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        if not self.enabled or not self._warmed_up:
+            return self.decoder(codes)
+
+        if codes.shape[0] != 1:
+            return self.decoder(codes)
+
+        actual_size = codes.shape[-1]
+        padded_size = self._get_padded_size(actual_size)
+
+        if padded_size is None or padded_size not in self.graphs:
+            return self.decoder(codes)
+
+        self.static_inputs[padded_size].zero_()
+        self.static_inputs[padded_size][:, :, :actual_size] = codes
+
+        self.graphs[padded_size].replay()
+
+        output = self.static_outputs[padded_size]
+        total_upsample = self.decoder.total_upsample
+        actual_output_len = actual_size * total_upsample
+
+        return output[..., :actual_output_len].clone()
+
+    def chunked_decode_with_cudagraph(
+        self,
+        codes: torch.Tensor,
+        chunk_size: int = 300,
+        left_context_size: int = 25,
+    ) -> torch.Tensor:
+        wavs = []
+        start_index = 0
+        total_len = codes.shape[-1]
+        total_upsample = self.decoder.total_upsample
+
+        while start_index < total_len:
+            end_index = min(start_index + chunk_size, total_len)
+            context_size = left_context_size if start_index - left_context_size > 0 else start_index
+
+            codes_chunk = codes[..., start_index - context_size : end_index]
+            wav_chunk = self.decode(codes_chunk)
+
+            wavs.append(wav_chunk[..., context_size * total_upsample :])
+            start_index = end_index
+
+        return torch.cat(wavs, dim=-1)

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py

@@ -100,6 +100,33 @@ class Qwen3TTSModelForGeneration(nn.Module):
         # Store vllm_config for potential future use
         self.vllm_config = vllm_config
 
+        # Enable CUDA Graph for decoder
+        self._enable_decoder_cudagraph()
+
+    def _enable_decoder_cudagraph(self):
+        # Respect --enforce-eager flag
+        model_cfg = getattr(self.vllm_config, "model_config", None)
+        if model_cfg and getattr(model_cfg, "enforce_eager", False):
+            logger.info("CUDA Graph not enabled: --enforce-eager is set")
+            return
+        try:
+            inner_model = getattr(self.model, "model", None)
+            if inner_model is None or not hasattr(inner_model, "speech_tokenizer"):
+                return
+            tokenizer = inner_model.speech_tokenizer
+            if not (hasattr(tokenizer, "model") and hasattr(tokenizer.model, "decoder")):
+                return
+            decoder = tokenizer.model.decoder
+            device = next(decoder.parameters()).device
+            if device.type != "cuda":
+                logger.info("CUDA Graph not enabled: decoder is on %s", device)
+                return
+            if hasattr(decoder, "enable_cudagraph"):
+                decoder.enable_cudagraph()
+                logger.info("CUDA Graph enabled for speech tokenizer decoder")
+        except Exception:
+            logger.warning("Failed to enable CUDA Graph for decoder", exc_info=True)
+
     def forward(
         self,
         input_ids: torch.Tensor | None = None,

vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py

@@ -848,6 +848,38 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
 
         self.post_init()
 
+        # CUDA Graph support
+        self._cudagraph_enabled = False
+        self._cudagraph_wrapper = None
+
+    def enable_cudagraph(
+        self,
+        capture_sizes: list[int] | None = None,
+        device: torch.device | None = None,
+    ):
+        from ..cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+
+        if device is None:
+            device = next(self.parameters()).device
+        if device.type != "cuda":
+            logger.warning("Cannot enable CUDA Graph: decoder is not on a CUDA device (got %s)", device)
+            return
+        self._cudagraph_wrapper = CUDAGraphDecoderWrapper(
+            decoder=self,
+            capture_sizes=capture_sizes,
+            num_quantizers=self.config.num_quantizers,
+            enabled=True,
+        )
+        self._cudagraph_wrapper.warmup(device, dtype=torch.long)
+        self._cudagraph_enabled = True
+        sizes = self._cudagraph_wrapper.capture_sizes
+        logger.info("CUDA Graph enabled for decoder with sizes: %s", sizes)
+
+    def disable_cudagraph(self):
+        self._cudagraph_enabled = False
+        self._cudagraph_wrapper = None
+        logger.info("CUDA Graph disabled for decoder")
+
     def forward(self, codes):
         if codes.shape[1] != self.config.num_quantizers:
             raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}")
@@ -866,6 +898,11 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
         return wav.clamp(min=-1, max=1)
 
     def chunked_decode(self, codes, chunk_size=300, left_context_size=25):
+        # Use CUDA graph if enabled
+        if self._cudagraph_enabled and self._cudagraph_wrapper is not None:
+            return self._cudagraph_wrapper.chunked_decode_with_cudagraph(codes, chunk_size, left_context_size)
+
+        # Original implementation (eager mode)
         wavs = []
         start_index = 0
         while start_index < codes.shape[-1]: