From 5dc787d5e8d10301a47fde006200f641741f11b1 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Fri, 10 Jul 2026 11:11:29 +0800 Subject: [PATCH] [Perf][Qwen3-TTS] Align talker MTP CUDA graph capture with Qwen3-Omni & Adapt to async output (#4923) Signed-off-by: amy-why-3459 --- vllm_omni/deploy/qwen3_tts.yaml | 13 +++++--- .../qwen3_tts_code_predictor_vllm.py | 5 ++-- .../models/qwen3_tts/qwen3_tts_talker.py | 30 +++++++++++++++---- 3 files changed, 37 insertions(+), 11 deletions(-) diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 4f95c1f48..f27da1b5b 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -35,7 +35,7 @@ connectors: stages: - stage_id: 0 - max_num_seqs: 10 + max_num_seqs: 64 gpu_memory_utilization: 0.3 trust_remote_code: true # Prefix caching on by default. The HS opt-out (talker only consumes the @@ -43,9 +43,9 @@ stages: # PR make per-step bookkeeping negligible, and any text repetition across # requests (same voice or shared ref_audio) is a net throughput win. # Disable explicitly if memory headroom is tight on a single-GPU deploy. - enable_prefix_caching: true + enable_prefix_caching: false async_scheduling: true - max_num_batched_tokens: 512 + max_num_batched_tokens: 32768 max_model_len: 4096 devices: "0" output_connectors: @@ -67,7 +67,7 @@ stages: - stage_id: 1 # Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1. - max_num_seqs: 10 + max_num_seqs: 64 gpu_memory_utilization: 0.3 enforce_eager: false trust_remote_code: true @@ -95,3 +95,8 @@ platforms: stages: - stage_id: 0 enforce_eager: false + compilation_config: + cudagraph_mode: PIECEWISE + - stage_id: 1 + compilation_config: + cudagraph_mode: PIECEWISE 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 8d2f0686a..9d00cc2ba 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 @@ -13,6 +13,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import ( CodePredictorWrapper, CodePredictorWrapperConfig, ) +from vllm_omni.platforms import current_omni_platform from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig @@ -21,7 +22,7 @@ Qwen3TTSTalkerCodePredictorModelVLLM = CodePredictorBaseModel class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapper): - """Qwen3-TTS code predictor (CUDA graphs, per-call sampling, projection).""" + """Qwen3-TTS code predictor (per-call sampling, projection).""" def __init__( self, @@ -35,7 +36,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp vllm_config=vllm_config, cp_config=config, wrapper_config=CodePredictorWrapperConfig( - use_cuda_graphs=True, + use_cuda_graphs=current_omni_platform.is_npu(), use_parallel_embedding=False, use_projection=(config.hidden_size != talker_config.hidden_size), return_proj_buf=False, diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index 2d28a0730..3daeb4cf2 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -324,13 +324,33 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # OmniGPUModelRunner will store talker_mtp output under this key in # per-request additional_information. self.talker_mtp_output_key = ("codes", "audio") + self.talker_mtp_graph_safe = True # talker_mtp samples with per-row generators, so explicitly-seeded # requests stay batched instead of one scalar forward per row (#4883). - # Only valid while talker_mtp receives the unpadded active batch (this - # talker is not graph-wrapped); a padded batch would need the runner to - # pad the generators list as well. - self.talker_mtp_accepts_per_row_generators = True - + # This is only valid while talker_mtp receives the *unpadded* active + # batch, i.e. when it is NOT graph-wrapped. The runner wraps talker_mtp + # (padded batches + a single captured RNG stream) whenever full + # cudagraphs are enabled, so disable per-row generators in that case. + # Consequence: under full cudagraphs, per-request ``tts_local_seed`` is + # not reproducible (matching Qwen3-Omni, which has no per-row seeding); + # eager mode keeps the per-row generator fast-path. + # TODO(#4923): the model should not read ``cudagraph_mode`` — it is an + # upstream (runner) concern, and the runner already re-derives the same + # wrap decision in ``_init_talker_mtp``. Once all TTS models declare + # ``talker_mtp_graph_safe`` we should move this gating into the model + # runner (models only expose the static capability flag, and users can + # then configure ``cudagraph_mode`` freely). This also means mtp seeding + # is silently dropped when the decode batch > 1 under full cudagraphs; + # the follow-up should refactor the per-row generator in the runner to + # be cudagraph-safe so seeded batches stay reproducible. + cudagraph_mode = getattr(vllm_config.compilation_config, "cudagraph_mode", None) + talker_mtp_graph_wrapped = ( + self.talker_mtp_graph_safe and cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs() + ) + self.talker_mtp_accepts_per_row_generators = not talker_mtp_graph_wrapped + self.use_async_omni_output = True + self.eager_omni_postprocess_before_async_output = True + self.omni_pooler_payload_include_hidden = False self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) if get_pp_group().is_last_rank: