From c9b31a4c1adc359470a99ecb0ccfb58e70c59687 Mon Sep 17 00:00:00 2001 From: Xu Lu <572605156@qq.com> Date: Thu, 26 Feb 2026 18:48:25 -0800 Subject: [PATCH] feat(qwen3-tts): Add CUDA Graph support for speech tokenizer decoder (#1205) Signed-off-by: xulusjb Co-authored-by: xulusjb --- .../models/qwen3_tts/__init__.py | 0 .../qwen3_tts/test_cuda_graph_decoder.py | 272 ++++++++++++++++++ .../qwen3_tts/cuda_graph_decoder_wrapper.py | 169 +++++++++++ .../models/qwen3_tts/qwen3_tts.py | 27 ++ .../modeling_qwen3_tts_tokenizer_v2.py | 37 +++ 5 files changed, 505 insertions(+) create mode 100644 tests/model_executor/models/qwen3_tts/__init__.py create mode 100644 tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py diff --git a/tests/model_executor/models/qwen3_tts/__init__.py b/tests/model_executor/models/qwen3_tts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py new file mode 100644 index 000000000..373968760 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py new file mode 100644 index 000000000..de278edf7 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py index 418322672..8c17fa50b 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py +++ b/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, diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index d3450a485..bbec6a525 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/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]: