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

[xulusjb]

Purpose

Add CUDA Graph acceleration for the Qwen3-TTS speech tokenizer decoder to reduce kernel launch overhead during the audio decoding phase.

Changes

• Add class that captures and replays CUDA graphs for fixed input sizes
• Add / methods to
• Modify to route through CUDA graph when enabled, with automatic fallback to eager execution for unsupported sizes
• Auto-enable CUDA graphs during model initialization in
• Respect flag to allow users to disable CUDA graphs
• Use the model's existing device placement (no hardcoded )

Captured graph sizes: [25, 50, 100, 150, 200, 250, 300]

Test Plan

1. Start Qwen3-TTS service with specified configuration
2. Run benchmark with 40 sentences at concurrency 1

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4886635c1a

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Avoid forcing decoder onto cuda:0 during init

This unconditionally selects cuda:0 when CUDA is available and moves the decoder there. In vLLM workers that bind each model replica to a specific device (e.g., GPU1 in multi-GPU or tensor-parallel setups), inputs will be on the worker’s assigned device, while the decoder has been moved to cuda:0, causing cross-device tensor errors at decode time. The CUDA-graph warmup should use the model’s current device (or the worker’s configured device) and avoid moving the module in __init__.

Useful? React with 👍 / 👎.

[Gaohan123]

@linyueqian PTAL

[Copilot]

Pull request overview

Adds CUDA Graph acceleration for the Qwen3-TTS (12Hz) speech tokenizer decoder to reduce kernel-launch overhead during audio decoding, with eager fallback when graphs can’t be used.

Changes:

• Add CUDAGraphDecoderWrapper to capture/replay decoder forward graphs for fixed code lengths.
• Add enable_cudagraph() / disable_cudagraph() to Qwen3TTSTokenizerV2Decoder and route chunked_decode() through the wrapper when enabled.
• Auto-enable decoder CUDA Graphs during Qwen3TTSModelForGeneration initialization.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.

File	Description
vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py	Adds decoder-level CUDA Graph toggle and integrates wrapper into chunked_decode().
vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py	Auto-enables CUDA Graphs for the speech tokenizer decoder during model init.
vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py	New wrapper that captures and replays CUDA Graphs for decoder inference.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Defaulting to torch.device("cuda:0") and then potentially calling decoder.to(device) here can break multi-GPU / non-zero local-rank setups and can conflict with vLLM’s device placement (only the decoder submodule is moved). Prefer using the decoder/model’s existing device (e.g., next(decoder.parameters()).device / inner_model.device) and enable CUDA Graph only after the model is already on the intended CUDA device, without forcing cuda:0.

[Copilot]

The broad except Exception here logs only the exception message, losing the traceback needed to debug CUDA Graph capture/enable failures. Use logger.exception(...) (or logger.warning(..., exc_info=True)) so failures include the full stack trace, especially since this runs during model initialization.

Suggested change

	logger.warning(f"Failed to enable CUDA Graph: {e}")
	logger.warning("Failed to enable CUDA Graph", exc_info=True)

[Copilot]

enable_cudagraph() warms up/captures graphs unconditionally. If the decoder is on CPU (or CUDA is unavailable), warmup() will hit torch.cuda.* APIs and raise. Add an explicit guard (e.g., if device.type != "cuda" or not torch.cuda.is_available(): log and return without enabling) before calling warmup.

[Copilot]

torch.cuda.synchronize(device) (and subsequent capture) assumes device is CUDA, but warmup() doesn’t validate that. If warmup is called with a CPU device (or CUDA is unavailable), this will raise. Add a guard at the start of warmup() to no-op or raise a clear error when device.type != "cuda" / torch.cuda.is_available() is false.

[Copilot]

Import of 'Tuple' is not used.

Suggested change

	from typing import Dict, List, Optional, Tuple
	from typing import Dict, List, Optional

[linyueqian]

I tested the CUDA graph decoder locally with actual Qwen3-TTS-12Hz-1.7B-CustomVoice weights on H200.

Correctness looks good for exact sizes. Zero diff when input matches a captured graph size. For chunked_decode, diffs are ~2.5e-5 since chunk_size=300 hits a captured size exactly. However, when input needs padding (e.g. size 75 padded to 100), the zero-padding bleeds through the conv layers and introduces diffs up to ~2e-3. Probably inaudible but worth documenting.

Speedup was only ~1.02x on H200 (36.3ms eager vs 35.4ms CUDA graph for size=300). The 26% e2e improvement in the PR description likely compounds across many calls in the full pipeline. What GPU was the benchmark run on? H200 has very low kernel launch overhead so older GPUs would benefit more.

Other things to address:

1. Codex already flagged the hardcoded cuda:0, will break in multi-GPU/TP. enable_cudagraph() on the decoder already does the right thing with next(self.parameters()).device, so the outer method should just delegate to it
2. Graph capture at line 158 is missing torch.no_grad() (the warmup runs have it but the capture itself doesn't)
3. No way to disable without modifying code, consider respecting --enforce-eager or an env var
4. patch_decoder_with_cudagraph is dead code, never called
5. Needs rebase

[xulusjb]

Thanks for the thorough review @linyueqian! All great points. Pushed a fix commit and rebased onto latest main. Here is what was addressed:

1. **Hardcoded ** — Removed. Now uses to respect vLLM's device placement. Removed the call entirely.
2. **Missing ** — Added around both the warmup run and graph capture in .
3. ** support** — now checks and skips if set.
4. Dead code — Removed .
5. Unused import — Removed .
6. Logging — Switched to for full tracebacks instead of bare with f-string.
7. Device guards — Added checks in and to gracefully handle non-CUDA scenarios.

Rebase onto latest main is done. PTAL when you get a chance!

[linyueqian]

LGTM, all code-level concerns addressed. @Gaohan123 @hsliuustc0106

[linyueqian]

Please fix pre-commit and DCO checks.

[JuanPZuluaga]

@linyueqian should the cuda graph capture, be also added in qwen3_tts_talker.py? Introduced by: #1161?

[hsliuustc0106]

🤖 Code Review: PR #1205 — CUDA Graph Support for Speech Tokenizer Decoder 🔴

Verdict: Request Changes

Summary: Adds CUDAGraphDecoderWrapper that captures CUDA graphs at fixed input sizes for the 12Hz decoder, replaying them during inference. Claims ~26% latency reduction and ~36% throughput improvement on H200, directly targeting RTF < 1.0. However, the PR has inconsistent capture size defaults, duplicated chunking logic, no thread safety guards, and zero test coverage. These must be addressed before merge.

Key Concerns:

• ⚠️ DEFAULT_CAPTURE_SIZES in the wrapper class ([25,50,100,150,200,250,300,400,500]) is inconsistent with the hardcoded list passed by enable_cudagraph() ([25,50,100,150,200,250,300]). The wrapper default is never used because enable_cudagraph() always passes an explicit list. Reconcile these.
• ⚠️ chunked_decode_with_cudagraph() duplicates the chunking logic from Qwen3TTSTokenizerV2Decoder.chunked_decode() (line 868). If the original changes, the copy silently diverges. Extract shared logic.
• ⚠️ No unit tests at all. Need at minimum: _get_padded_size() edge cases, torch.allclose equivalence test (eager vs graph output), and an integration test verifying audio quality preservation.
• ⚠️ Static input/output tensors are shared across calls with no concurrency guard. Stage config has max_batch_size: 1 but there is no assertion. A concurrent call would corrupt the static buffers silently.

Detailed Feedback

📄 vllm_omni/model_executor/models/qwen3_tts.py

🟠 HIGH (architecture)

DEFAULT_CAPTURE_SIZES in CUDAGraphDecoderWrapper is [25,50,100,150,200,250,300,400,500] but enable_cudagraph() passes [25,50,100,150,200,250,300]. The wrapper default is dead code since enable_cudagraph() always provides an explicit list. This is confusing and error-prone.

# Reconcile by using a single source of truth:
CUDA_GRAPH_CAPTURE_SIZES = [25, 50, 100, 150, 200, 250, 300, 400, 500]

# In enable_cudagraph():
def enable_cudagraph(self, capture_sizes=None):
    sizes = capture_sizes or CUDA_GRAPH_CAPTURE_SIZES
    self.decoder = CUDAGraphDecoderWrapper(self.decoder, capture_sizes=sizes)

🟠 HIGH (architecture) — Line 868

chunked_decode_with_cudagraph() duplicates the chunking logic from Qwen3TTSTokenizerV2Decoder.chunked_decode(). If the original logic changes, the wrapper's copy will silently diverge. This violates DRY and creates a maintenance hazard.

# Have the wrapper delegate to the original chunking logic:
def chunked_decode_with_cudagraph(self, codes, chunk_size=None):
    # Reuse the base chunking strategy, only override the per-chunk decode call
    chunks = self._split_into_chunks(codes, chunk_size)
    return torch.cat([self.decode(chunk) for chunk in chunks], dim=-1)

🟠 HIGH (security)

Static input/output tensors are shared across calls with no thread safety guard. Stage config sets max_batch_size: 1 but there is no runtime assertion. A concurrent call would corrupt the static buffers silently, producing incorrect audio output.

# Add a guard at the top of decode():
assert not self._in_use, 'CUDAGraphDecoderWrapper does not support concurrent calls'
self._in_use = True
try:
    # ... existing decode logic ...
finally:
    self._in_use = False

🟡 MEDIUM (performance)

9 capture sizes x graph capture = significant GPU memory for intermediate activations. Stage config sets gpu_memory_utilization: 0.1 which is very low. No logging of memory consumed by captured graphs during warmup.

# Log memory usage during warmup:
import torch
for size in capture_sizes:
    mem_before = torch.cuda.memory_allocated()
    self._capture_graph(size)
    mem_after = torch.cuda.memory_allocated()
    logger.info(f'CUDA graph for size {size}: {(mem_after - mem_before) / 1024**2:.1f} MB')

🟡 MEDIUM (docs)

No documentation on how to enable CUDA graphs. Is it a stage config flag? An API parameter? The relationship to enforce_eager: true in qwen3_tts.yaml is unclear.

📄 tests/

🟠 HIGH (tests)

No unit tests at all. Need at minimum: (1) _get_padded_size() edge cases including size=0, size=max, size>max, (2) torch.allclose numerical equivalence test comparing eager vs graph output, (3) integration test verifying audio quality is preserved after CUDA graph replay.

# Example test structure:
def test_get_padded_size_exact_match():
    wrapper = CUDAGraphDecoderWrapper(mock_decoder, capture_sizes=[25, 50, 100])
    assert wrapper._get_padded_size(25) == 25

def test_get_padded_size_rounds_up():
    assert wrapper._get_padded_size(30) == 50

def test_get_padded_size_exceeds_max():
    assert wrapper._get_padded_size(101) is None

def test_numerical_equivalence():
    eager_out = decoder(codes)
    graph_out = wrapper.decode(codes)
    assert torch.allclose(eager_out, graph_out, atol=1e-5)

📄 configs/qwen3_tts.yaml

🟢 LOW (docs)

enforce_eager: true is set in the config but the relationship to the new CUDA graph feature is not documented. Should this be set to false when CUDA graphs are enabled? Or are they independent?

Suggestions

1. Reconcile DEFAULT_CAPTURE_SIZES vs the hardcoded list in enable_cudagraph() -- use a single constant.
2. Extract shared chunking logic to avoid duplication with chunked_decode().
3. Make capture sizes configurable via stage config YAML.
4. Add torch.allclose equivalence test.
5. Log GPU memory consumed during warmup.

Part of Issue #938 — Qwen3-TTS Production Ready

---

Review generated by vLLM-Omni PR Review Agent

[linyueqian]

@linyueqian should the cuda graph capture, be also added in qwen3_tts_talker.py? Introduced by: #1161?

The talker only uses the speech tokenizer for encoding (tok.encode()), not decoding, so no CUDA graph needed there.

But qwen3_tts_code2wav.py from #1161 does call chunked_decode() through its own separate tokenizer instance — the CUDA graph enablement here won't carry over to it. Worth adding similar support there too (or at least a TODO).

Also minor: DEFAULT_CAPTURE_SIZES in the wrapper vs the hardcoded list in enable_cudagraph() are inconsistent.

[xulusjb]

Hi team,
For the new commit 98fa76a, I:

1. Rebased onto latest main, squashed into a single commit,
2. Fixed all pre-commit/DCO issues.
3. Unified the capture sizes to a single source of truth in CUDAGraphDecoderWrapper.DEFAULT_CAPTURE_SIZES.

Ready for re-review @linyueqian.

[zsl24]

@xulusjb Hi team,

I noticed that the andimarafioti/faster-qwen3-tts implementation achieves significant speedups for Qwen3-TTS by applying CUDA Graph optimization to the Talker and MTP modules, reducing CUDA kernel launch overhead.

This PR currently focuses on optimizing the Codec module via CUDA Graph capture. To align our approach with the most impactful strategy:

Could we evaluate whether prioritizing Talker/MTP optimization (as in andimarafioti/faster-qwen3-tts) might yield greater performance gains?
Are there architectural constraints that make Codec optimization preferable in our case?

[linyueqian]

@xulusjb Hi team,

I noticed that the andimarafioti/faster-qwen3-tts implementation achieves significant speedups for Qwen3-TTS by applying CUDA Graph optimization to the Talker and MTP modules, reducing CUDA kernel launch overhead.

This PR currently focuses on optimizing the Codec module via CUDA Graph capture. To align our approach with the most impactful strategy:

Could we evaluate whether prioritizing Talker/MTP optimization (as in andimarafioti/faster-qwen3-tts) might yield greater performance gains? Are there architectural constraints that make Codec optimization preferable in our case?

I am working on a pr right now for this.

[linyueqian]

LGTM. For follow-up, please add tests for eager vs CUDA graph numerical equivalence, especially padding cases.

[xulusjb]

Thanks for the review! I've added the requested tests in the follow-up commit:

31 tests covering eager vs CUDA graph numerical equivalence:

• Exact-size inputs (no padding): verified bit-identical output
• Padded inputs: verified interior positions match within 1e-5 atol, output shape/length correct, values bounded
• Fallback to eager (oversized input): verified bit-identical
• Chunked decode: shape match and bit-identical for aligned chunks
• Edge cases: single frame, disabled wrapper, batch > 1, determinism

Tests use a synthetic Conv1d decoder as a worst-case proxy (non-causal convs amplify boundary effects). The real causal
decoder shows ~2e-3 max diff vs ~0.5 in the synthetic case, confirming padding impact is negligible for actual
inference.

Branch has also been rebased onto latest main.