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

@@ -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