From 8e12007ac1680d169e68e2601d855f2c7eabb363 Mon Sep 17 00:00:00 2001 From: Juan Pablo Zuluaga <46724788+JuanPZuluaga@users.noreply.github.com> Date: Sat, 21 Mar 2026 12:11:34 +0100 Subject: [PATCH] [BugFix][Qwen3TTS] CodePredictor CudaGraph Pool (#2059) Signed-off-by: JuanPZuluaga Co-authored-by: JuanPZuluaga --- benchmarks/qwen3-tts/README.md | 22 ++--- .../vllm_omni/configs/qwen3_tts_bs1.yaml | 7 +- .../vllm_omni/configs/qwen3_tts_bs16.yaml | 94 +++++++++++++++++++ .../vllm_omni/configs/qwen3_tts_bs4.yaml | 7 +- .../qwen3_tts_code_predictor_vllm.py | 58 ++++++++---- 5 files changed, 152 insertions(+), 36 deletions(-) create mode 100644 benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml diff --git a/benchmarks/qwen3-tts/README.md b/benchmarks/qwen3-tts/README.md index b038669e8..73bc420f9 100644 --- a/benchmarks/qwen3-tts/README.md +++ b/benchmarks/qwen3-tts/README.md @@ -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 diff --git a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml index 5d730f605..ca441d286 100644 --- a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml +++ b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml @@ -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 diff --git a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml new file mode 100644 index 000000000..2cc5cf535 --- /dev/null +++ b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml @@ -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 diff --git a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml index 10dc89564..5de107d49 100644 --- a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml +++ b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml @@ -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 diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 4a90b83af..00114f0cc 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -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: