[BugFix] code2wav supports disabling CUDA graph. (#3732)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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"

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

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

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.",

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

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",