From 748470c1fb2db8fba95eb7d4b2498656097865b2 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Sun, 24 May 2026 00:05:55 +0800 Subject: [PATCH] [BugFix] code2wav supports disabling CUDA graph. (#3732) Signed-off-by: amy-why-3459 Co-authored-by: Hongsheng Liu --- tests/test_config_factory.py | 2 +- vllm_omni/deploy/qwen3_omni_moe.yaml | 2 +- vllm_omni/deploy/qwen3_tts.yaml | 6 +- .../models/qwen3_omni/qwen3_omni.py | 43 ++++++---- .../models/qwen3_omni/qwen3_omni_code2wav.py | 4 +- .../models/qwen3_tts/qwen3_tts_code2wav.py | 80 +++++++++++++------ 6 files changed, 91 insertions(+), 46 deletions(-) diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index be3a012cc..9ac3d859c 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -1453,7 +1453,7 @@ class TestPlatformOverrides: assert deploy.stages[0].tensor_parallel_size == 2 assert deploy.stages[0].devices == "0,1" # Stage 2 unaffected fields stay at base - assert deploy.stages[2].enforce_eager is True + assert deploy.stages[2].enforce_eager is False def test_xpu_overrides(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index 445437c0f..b77f4caf2 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe.yaml @@ -56,7 +56,7 @@ stages: max_num_batched_tokens: 51200 max_num_seqs: 64 gpu_memory_utilization: 0.1 - enforce_eager: true + enforce_eager: false trust_remote_code: true enable_prefix_caching: false async_scheduling: false diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 80c316849..0c976d090 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -9,8 +9,8 @@ # * Stage 1 (code2wav) sets max_num_batched_tokens=65536 for correctness: # codec prefill length (Q * num_frames) exceeds the 32k default. # * Stage 0 does not set enforce_eager — talker runs cudagraph by default. -# Stage 1 sets true because its codec generation loop is not -# cudagraph-compatible. NPU platform flips stage 0 to true where +# Stage 1 defaults to false so Code2Wav can use its inner cudagraph; +# set true to disable that path in eager mode. NPU platform flips stage 0 to true where # cudagraph is not yet verified. async_chunk: true @@ -56,7 +56,7 @@ stages: - stage_id: 1 max_num_seqs: 10 gpu_memory_utilization: 0.3 - enforce_eager: true + enforce_eager: false trust_remote_code: true enable_prefix_caching: false async_scheduling: true diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 3c658fc3f..ba924a1ea 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -420,13 +420,16 @@ class Qwen3OmniMoeForConditionalGeneration( else: codes = input_ids.reshape(1, 16, -1) else: - logger.warning( - ( - "Input_ids length: %s is not divisible by 16, padding " - "with zeros. This should only happen in warm up." - ), - input_ids.shape[0], - ) + if seq_token_counts is None: + logger.debug( + "Code2Wav warmup input length %s is not divisible by 16; padding with zeros.", + input_ids.shape[0], + ) + else: + logger.warning_once( + "Code2Wav input length is not divisible by 16; padding with zeros. " + "This is expected only during cudagraph warmup." + ) input_ids_flatten = input_ids.reshape(-1) input_ids_flatten = torch.cat( [ @@ -1283,6 +1286,22 @@ class Qwen3OmniMoeForConditionalGeneration( left_frames = int(extra.get("codec_left_context_frames", 0) or 0) return chunk_frames, left_frames + def _maybe_enable_code2wav_cudagraph(self) -> None: + """Enable the inner Code2Wav CUDA graph unless this stage runs in eager mode.""" + if not self.code2wav or not hasattr(self.code2wav, "enable_cudagraph"): + return + + model_cfg = getattr(self.vllm_config, "model_config", None) + if getattr(model_cfg, "enforce_eager", False): + logger.info("Code2Wav CUDA Graph disabled because enforce_eager is set") + return + + chunk_frames, left_frames = self._get_codec_frame_config() + self.code2wav.enable_cudagraph( + codec_chunk_frames=chunk_frames, + codec_left_context_frames=left_frames, + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights for all components of the omni model.""" loaded_weights = set() @@ -1319,15 +1338,11 @@ class Qwen3OmniMoeForConditionalGeneration( code2wav_loaded = add_prefix_to_loaded_weights(code2wav_loaded, "code2wav") loaded_weights.update(code2wav_loaded) - # Precompute SnakeBeta caches and enable CUDA graph for Code2Wav decoder + # Precompute SnakeBeta caches; Code2Wav CUDA graph follows the stage's + # enforce_eager setting, the same switch vLLM uses for outer graphs. try: self.code2wav.precompute_snake_caches() - if hasattr(self.code2wav, "enable_cudagraph"): - chunk_frames, left_frames = self._get_codec_frame_config() - self.code2wav.enable_cudagraph( - codec_chunk_frames=chunk_frames, - codec_left_context_frames=left_frames, - ) + self._maybe_enable_code2wav_cudagraph() except Exception: logger.warning( "Failed to enable CUDA Graph for Code2Wav; falling back to eager.", diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 8bcba4340..2b1a7ccf2 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -297,9 +297,9 @@ class Qwen3OmniMoeCode2Wav(nn.Module): single tensor with shape ``[1, waveform_len]``. """ if not (left_context_size and seq_token_counts and len(left_context_size) == len(seq_token_counts)): - logger.warning( + logger.warning_once( "chunked_decode_streaming: missing/invalid left_context_size or seq_token_counts; " - "defaulting to left_context_size=zeros(len=codes.shape[0])." + "defaulting to left_context_size=zeros(len(codes)). This is expected during cudagraph warmup." ) left_context_size = [0] * codes.shape[0] # Decode chunk diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index 8c200141a..b46018e36 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -111,6 +111,56 @@ class Qwen3TTSCode2Wav(nn.Module): return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)] return [ids] + def _maybe_enable_decoder_cudagraph( + self, + *, + device: torch.device, + codec_chunk_frames: int, + codec_left_context_frames: int, + decode_cudagraph_capture_sizes: list[int] | None, + decode_cudagraph_batch_sizes: list[int] | None, + decode_cudagraph_extra_capture_shapes: list[tuple[int, int]] | None, + decode_compile_shapes: list[tuple[int, int]] | None, + ) -> None: + """Enable inner Code2Wav CUDA graph unless stage is enforce_eager.""" + if not hasattr(self.decoder, "enable_cudagraph") or device.type != "cuda": + return + + model_cfg = getattr(self.vllm_config, "model_config", None) + if getattr(model_cfg, "enforce_eager", False): + logger.info("Qwen3-TTS Code2Wav CUDA Graph disabled because enforce_eager is set") + return + + if ( + codec_chunk_frames > 0 + and codec_left_context_frames > 0 + and self._decoder_sliding_window + and codec_left_context_frames < self._decoder_sliding_window + ): + logger.warning( + "Qwen3-TTS streaming codec_left_context_frames=%d " + "is smaller than decoder sliding_window=%d; " + "chunk-boundary distortion may occur. " + "Increase codec_left_context_frames to at least " + "%d for streaming.", + codec_left_context_frames, + self._decoder_sliding_window, + self._decoder_sliding_window, + ) + + self.decoder.enable_cudagraph( + capture_sizes=decode_cudagraph_capture_sizes, + capture_batch_sizes=decode_cudagraph_batch_sizes, + extra_capture_shapes=decode_cudagraph_extra_capture_shapes, + compile_shapes=decode_compile_shapes, + device=device, + codec_chunk_frames=codec_chunk_frames, + codec_left_context_frames=codec_left_context_frames, + decode_chunk_size=self._decode_chunk_frames, + decode_left_context=self._decode_left_context_frames, + ) + logger.info("Code2Wav decoder CUDA Graph enabled") + def _get_decode_batch_bucket_frames(self, actual_frames: int) -> int: for bucket_frames in self._decode_batch_bucket_frames: if actual_frames <= bucket_frames: @@ -581,35 +631,15 @@ class Qwen3TTSCode2Wav(nn.Module): if hasattr(self.decoder, "enable_cudagraph") and device.type == "cuda": try: - if ( - codec_chunk_frames > 0 - and codec_left_context_frames > 0 - and self._decoder_sliding_window - and codec_left_context_frames < self._decoder_sliding_window - ): - logger.warning( - "Qwen3-TTS streaming codec_left_context_frames=%d " - "is smaller than decoder sliding_window=%d; " - "chunk-boundary distortion may occur. " - "Increase codec_left_context_frames to at least " - "%d for streaming.", - codec_left_context_frames, - self._decoder_sliding_window, - self._decoder_sliding_window, - ) - - self.decoder.enable_cudagraph( - capture_sizes=decode_cudagraph_capture_sizes, - capture_batch_sizes=decode_cudagraph_batch_sizes, - extra_capture_shapes=decode_cudagraph_extra_capture_shapes, - compile_shapes=decode_compile_shapes, + self._maybe_enable_decoder_cudagraph( device=device, codec_chunk_frames=codec_chunk_frames, codec_left_context_frames=codec_left_context_frames, - decode_chunk_size=self._decode_chunk_frames, - decode_left_context=self._decode_left_context_frames, + decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes, + decode_cudagraph_batch_sizes=decode_cudagraph_batch_sizes, + decode_cudagraph_extra_capture_shapes=decode_cudagraph_extra_capture_shapes, + decode_compile_shapes=decode_compile_shapes, ) - logger.info("Code2Wav decoder CUDA Graph enabled") except Exception: logger.warning( "Failed to enable CUDA Graph for Code2Wav decoder",