Stack torch.compile with CUDA graph decode instead of choosing one

Regionally compile the DiT blocks even when enable_cuda_graph_decode is set
and let graph capture record the compiled (fused) kernels; inductor's own
cudagraphs stay off so the two graph layers never nest. Measured on GB200
(1024x1024, 50 steps, seed 42, BF16): combo 3.04s vs compile-only 3.2s,
graph-only 4.0s, eager 4.3s per image; 44.5-45.6 dB PSNR vs eager.

recipes/Qwen/Qwen-Image-2.1.md

@@ -129,17 +129,23 @@ vllm serve Qwen/Qwen-Image-2.1 --omni --no-enable-cuda-graph-decode
 The same control is available as `Omni(..., enable_cuda_graph_decode=False)`
 and as `enable_cuda_graph_decode: false` on a stage in the deployment YAML.
 `--enforce-eager` / `enforce_eager: true` remains the broader switch: it
-disables both graph capture and automatic `torch.compile`. Otherwise, this
-transformer uses its model-specific decode graphs instead of automatic
-`torch.compile`.
+disables both graph capture and automatic `torch.compile`. The two
+optimizations stack: the DiT blocks are regionally `torch.compile`'d and
+graph capture records the compiled (fused) kernels, while inductor's own
+cudagraphs stay disabled so the two graph layers never nest. Measured on a
+single GB200 (1024×1024, 50 steps, seed 42, BF16, `true_cfg_scale=1.0`,
+warmup 1 + median of 3): compile+graph 3.04 s, compile-only 3.2 s,
+graph-only 4.0 s, eager 4.3 s end-to-end per image; combo output vs eager is
+44.5–45.6 dB PSNR (the same magnitude as pure `torch.compile` fusion
+divergence), and graph-only vs eager is bit-identical.
 
 The autoregressive engine's `compilation_config.cudagraph_mode` does not
 control this diffusion path; use `enable_cuda_graph_decode` (or
 `enforce_eager` for everything) to force eager decode execution.
 
-TP/SP/ring parallelism, HSDP, offload/cache hooks, compiled blocks, quantized
-KV caches, dynamic LoRA, padded text masks, and a second in-flight request
-whose cache aliases an already-owned graph key fall back to eager decode.
+TP/SP/ring parallelism, HSDP, offload/cache hooks, quantized KV caches,
+dynamic LoRA, padded text masks, and a second in-flight request whose cache
+aliases an already-owned graph key fall back to eager decode.
 Graphs keep separate entries for different image layouts and copy
 request-owned prefix K/V into static buffers before replay; each graph key
 has a single live owner at a time, so concurrent same-shape requests never
@@ -374,7 +380,10 @@ filled in:
 
 - [ ] TP=2/4 output parity vs single-card — pending measurement.
 - [ ] VAE tiling vs non-tiled decode — pending measurement.
-- [ ] CUDA graph decode vs eager — pending measurement.
+- [x] CUDA graph decode vs eager — measured on GB200 (1024×1024, 50 steps,
+  seed 42, BF16): compile+graph 3.04 s/image vs eager 4.3 s/image (~1.4×),
+  44.5–45.6 dB PSNR vs eager; graph-only is bit-identical to eager. See
+  "CUDA Graph decode" above.
 
 ## Known Limitations
 

vllm_omni/diffusion/data.py

@@ -959,9 +959,10 @@ class OmniDiffusionConfig:
     # Currently only Qwen-Image-2.1's transformer implements this; other models
     # ignore the flag. This is a per-model diffusion path, independent of the
     # AR engine's ``compilation_config.cudagraph_mode`` (which does not apply
-    # to diffusion stages) and of ``diffusion_compile_granularity`` — models
-    # with decode graphs skip automatic torch.compile of their DiT blocks, so
-    # the two never stack. ``enforce_eager=True`` forces eager decode and
+    # to diffusion stages). It stacks with ``diffusion_compile_granularity``:
+    # the DiT blocks are still torch.compile'd and graph capture records the
+    # compiled (fused) kernels, while inductor's own cudagraphs stay off.
+    # ``enforce_eager=True`` forces eager decode and
     # disables capture regardless of this flag; unsupported configurations
     # (SP/TP, ring, HSDP, dynamic LoRA, padded masks, quantized prefix KV, or
     # a second in-flight request aliasing the same graph key) log and fall

vllm_omni/diffusion/models/qwen_image_21/decode_graph.py

@@ -24,12 +24,16 @@ Memory note: graph entries own an additional copy of the prefix K/V. Decode
 attention concatenates that prefix with target K/V inside the captured region;
 transient tensors use the graph memory pool.
 
+The model runner regionally torch.compile's the transformer blocks even when
+decode graphs are enabled, so capture records the compiled (fused) kernels;
+inductor's own cudagraphs stay off so the two graph layers never stack.
+
 Fallbacks (all logged, all eager): CUDA unavailable, model-level CPU offload
 without persistent weight staging, sequence/ring/tensor parallelism,
-HSDP/offload hooks, torch.compile'd blocks, KV-cache quantization, dynamic
-LoRA wrappers, padded text masks (the masked attention path branches on
-mask contents, which cannot be captured), unknown graph key at decode, a
-second live request aliasing an owned graph key, and capture failure.
+HSDP/offload hooks, KV-cache quantization, dynamic LoRA wrappers, padded text
+masks (the masked attention path branches on mask contents, which cannot be
+captured), unknown graph key at decode, a second live request aliasing an
+owned graph key, and capture failure.
 """
 
 from __future__ import annotations
@@ -223,9 +227,6 @@ class QwenImage21DecodeGraphManager:
                     reason = "tensor parallelism runs collectives inside the captured region"
             if reason is None:
                 for block in model.transformer_blocks:
-                    if getattr(block.forward, "_torchdynamo_orig_callable", None) is not None:
-                        reason = "transformer blocks are torch.compile'd"
-                        break
                     if hasattr(block, "_hook_registry") or hasattr(block, "_omni_original_forward"):
                         reason = "transformer blocks carry offload/cache hooks"
                         break

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -240,8 +240,21 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             return
 
         if getattr(model, "enable_cuda_graph_decode", False):
-            logger.info("Model runner: %s uses CUDA graph decode instead of torch.compile.", attr_name)
-            return
+            # Decode-graph models still compile their blocks: graph capture
+            # records the compiled (fused) kernels. Inductor's own cudagraphs
+            # must stay off so the two graph layers never stack.
+            try:
+                import torch._inductor.config as inductor_config
+
+                inductor_config.triton.cudagraphs = False
+            except Exception as e:
+                logger.warning(
+                    "Model runner: could not disable inductor cudagraphs for the "
+                    "compile+CUDA-graph combo (%s); capture may fail and fall back "
+                    "to compiled eager decode.",
+                    e,
+                )
+            logger.info("Model runner: %s combines CUDA graph decode with torch.compile.", attr_name)
 
         compile_granularity = self.od_config.diffusion_compile_granularity
         compile_dynamic = self.od_config.diffusion_compile_dynamic