anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[BugFix][Qwen3TTS] CodePredictor CudaGraph Pool (#2059)

Browse Source 
Signed-off-by: JuanPZuluaga <juanz9312@gmal.com>
Co-authored-by: JuanPZuluaga <juanz9312@gmal.com>
This commit is contained in:
Juan Pablo ZuluagaandJuanPZuluaga authored and GitHub committed 2026-03-21 19:11:34 +08:00
1 parent da21e991cc
commit 8e12007ac1
5 files changed
+152 -36
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

benchmarks/qwen3-tts/README.md
+11 -11
View File
@@ -32,8 +32,8 @@ bash run_benchmark.sh --hf-only
# Use a different model (e.g. 1.7B)
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
# Use batch_size=4 config for higher throughput
STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
# Use bs16 config for higher throughput
STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only
# Custom GPU, prompt count, concurrency levels
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
@@ -54,7 +54,7 @@ CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
### 2) Run online serving benchmark
```bash
python vllm_omni/bench_tts_serve.py \
python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
--port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
@@ -65,7 +65,7 @@ python vllm_omni/bench_tts_serve.py \
### 3) Run HuggingFace baseline
```bash
python transformers/bench_tts_hf.py \
python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--num-prompts 50 \
--gpu-device 0 \
@@ -75,7 +75,7 @@ python transformers/bench_tts_hf.py \
### 4) Generate comparison plots
```bash
python plot_results.py \
python benchmarks/qwen3-tts/plot_results.py \
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
--labels "vllm-omni" "hf_transformers" \
--output results/comparison.png
@@ -83,14 +83,14 @@ python plot_results.py \
## Stage Configs
| Config | Batch Size | Description |
|--------|:----------:|-------------|
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
| Config | max_num_seqs | Description |
|--------|:------------:|-------------|
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
## Metrics
benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml
+4 -3
View File
@@ -4,6 +4,7 @@ async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
is_comprehension: true
runtime:
devices: "0"
engine_args:
@@ -14,7 +15,7 @@ stage_args:
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: false
async_scheduling: true
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
@@ -45,10 +46,10 @@ stage_args:
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: false
async_scheduling: true
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 8192
max_model_len: 32768
benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml
+94
View File
@@ -0,0 +1,94 @@
# Qwen3-TTS max_num_seqs=16 config (streaming with async_chunk)
# High-throughput concurrent request processing
# 2-stage pipeline: Talker -> Code2Wav
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
is_comprehension: true
runtime:
devices: "0"
engine_args:
max_num_seqs: 16
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: true
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 4096
max_model_len: 4096
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 4096
seed: 42
detokenize: false
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
engine_args:
max_num_seqs: 16
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: true
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
distributed_executor_backend: "mp"
max_num_batched_tokens: 16384
max_model_len: 32768
engine_input_source: [0]
final_output: true
final_output_type: audio
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 65536
seed: 42
detokenize: true
repetition_penalty: 1.0
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 16
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 25
edges:
- from: 0
to: 1
window_size: -1
benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml
+4 -3
View File
@@ -1,10 +1,11 @@
# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
# Enables concurrent request processing with max_inflight=4
# Enables concurrent request processing
# 2-stage pipeline: Talker -> Code2Wav
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
is_comprehension: true
runtime:
devices: "0"
engine_args:
@@ -15,7 +16,7 @@ stage_args:
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: false
async_scheduling: true
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
@@ -46,7 +47,7 @@ stage_args:
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: false
async_scheduling: true
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py
+39 -19
View File
@@ -353,6 +353,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
self._bucket_pos_ids: dict[int, torch.Tensor] = {}
self._lm_heads_list: list[nn.Module] | None = None
self._codec_embeds_list: list[nn.Module] | None = None
self._cuda_graphs: dict[int, tuple[torch.cuda.CUDAGraph, torch.Tensor]] = {}
def get_input_embeddings(self) -> nn.ModuleList:
return self.model.get_input_embeddings()
@@ -398,12 +399,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
)
def _setup_compile(self) -> None:
"""Lazily set up torch.compiled model forward for kernel fusion.
Uses ``mode="reduce-overhead"`` with ``dynamic=False`` so Inductor
captures internal CUDA graphs for fixed shapes, eliminating kernel
launch overhead entirely.
"""
"""Lazily set up torch.compile with manual CUDA graph capture."""
if self._compiled_model_fwd is not None:
return
self._lm_heads_list = list(self.lm_head)
@@ -412,13 +408,11 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
logger.warning_once("code_predictor: torch.compile disabled")
self._compiled_model_fwd = self.model.forward
return
self._compiled_model_fwd = torch.compile(
self.model.forward,
mode="reduce-overhead",
dynamic=False,
)
logger.info("code_predictor: torch.compile enabled (mode=reduce-overhead, dynamic=False)")
self._warmup_compile()
self._compiled_model_fwd = torch.compile(self.model.forward, mode="default", dynamic=False)
self._warmup_buckets()
self._capture_cuda_graphs()
logger.info("code_predictor: torch.compile (mode=default) + CUDA graphs")
def _padded_bsz(self, bsz: int) -> int:
for bucket in self._bucket_sizes:
@@ -426,8 +420,8 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
return bucket
return bsz
def _warmup_compile(self) -> None:
"""Warmup power-of-2 batch-size buckets to front-load compilation."""
def _warmup_buckets(self) -> None:
"""Warmup power-of-2 batch-size buckets to front-load Inductor compilation."""
max_bsz = self._vllm_config.scheduler_config.max_num_seqs
bucket_sizes = [1 << i for i in range(max_bsz.bit_length()) if (1 << i) <= max_bsz]
if max_bsz not in bucket_sizes:
@@ -444,7 +438,28 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
self._bucket_pos_ids[bsz] = pos_ids
for _ in range(3):
self._compiled_model_fwd(proj_buf[:bsz, :max_seq, :], pos_ids)
logger.info("code_predictor: warmup done for bucket sizes %s", self._bucket_sizes)
logger.info("code_predictor: warmup done for buckets %s", self._bucket_sizes)
def _capture_cuda_graphs(self) -> None:
"""Capture a CUDA graph per bucket using vLLM's global graph pool."""
from vllm.platforms import current_platform
pool = current_platform.get_global_graph_pool()
max_seq = self._num_groups + 1
proj_buf = self._proj_buf
for bsz in self._bucket_sizes:
static_input = proj_buf[:bsz, :max_seq, :]
pos_ids = self._bucket_pos_ids[bsz]
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g, pool=pool):
static_output = self._compiled_model_fwd(static_input, pos_ids)
self._cuda_graphs[bsz] = (g, static_output)
logger.info("code_predictor: captured CUDA graphs for buckets %s", self._bucket_sizes)
# ------------------------------------------------------------------
# Optimized forward: re-prefill + torch.compile + projection cache
@@ -505,10 +520,15 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
base_pos = torch.arange(max_seq, device=device, dtype=torch.long)
full_pos_ids = base_pos if padded_bsz == 1 else base_pos.repeat(padded_bsz)
for step in range(1, num_groups):
projected = proj_buf[:padded_bsz, :max_seq, :]
# Use captured CUDA graph if available, otherwise call compiled fn.
cuda_graph_entry = self._cuda_graphs.get(padded_bsz)
hidden_out = model_fwd(projected, full_pos_ids)
for step in range(1, num_groups):
if cuda_graph_entry is not None:
cuda_graph_entry[0].replay()
hidden_out = cuda_graph_entry[1]
else:
hidden_out = model_fwd(proj_buf[:padded_bsz, :max_seq, :], full_pos_ids)
logits = lm_heads[step - 1](hidden_out[:bsz, step, :])
if use_sampling: