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
@@ -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
|
||||
|
||||
Reference in new issue
Block a user