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.
This commit is contained in:
Isotr0py committed 2026-09-22 08:58:32 +00:00
1 parent 45aa5b16d9
commit be8a72f0a6
4 files changed
+43 -19

No files matched your search

+4 -3
View File
@@ -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