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

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