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:
1 parent
45aa5b16d9
commit
be8a72f0a6
4 files changed
+43
-19
No files matched your search
@@ -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
|
||||
|
||||
Reference in new issue
Block a user