Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks

The opt-in `max` speed tier now compiles the repeated block with
mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode:
Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer
cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are
deliberately avoided -- both crash on the regionally-compiled block (its static
output buffer is overwritten across denoise steps), measured.

Adds two reproducible benchmarks used to validate the optimization research:
- scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head.
- scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes.

Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen;
coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune);
FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo).
This commit is contained in:
Daniel Han 2026-06-26 05:03:55 +00:00
commit ede94176f6
4 changed files with 289 additions and 5 deletions

View file

@ -152,9 +152,12 @@ def apply_speed_optims(
applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger)
# Near-lossless and the largest win: regional compile of the repeated denoiser
# block, where eligible (now incl. the GGUF transformer).
# block, where eligible (now incl. the GGUF transformer). `max` opts into
# max-autotune (longer compile, autotuned kernels).
if compile_eligible(target, is_gguf = is_gguf, family = family):
applied["compiled"] = _compile_repeated_blocks(pipe, logger)
applied["compiled"] = _compile_repeated_blocks(
pipe, logger, max_autotune = mode == SPEED_MAX
)
if mode == SPEED_MAX:
# Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed.
@ -178,13 +181,22 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool:
return False
def _compile_repeated_blocks(pipe: Any, logger: Any) -> bool:
def _compile_repeated_blocks(pipe: Any, logger: Any, *, max_autotune: bool = False) -> bool:
transformer = getattr(pipe, "transformer", None)
fn = getattr(transformer, "compile_repeated_blocks", None)
if not callable(fn):
return False
# default: mode="default" + dynamic=True -- fast cold start, robust to resolution
# changes (no recompile). max: mode="max-autotune-no-cudagraphs" + dynamic=False --
# Triton autotuning for a few % more on GEMM/conv-heavy models, at a much longer
# compile and a recompile per new resolution. The CUDA-graph modes (reduce-overhead
# / max-autotune) are deliberately NOT used: they crash on the regionally-compiled
# block because its static output buffer is overwritten across denoise steps.
kwargs: dict[str, Any] = {"fullgraph": True, "dynamic": not max_autotune}
if max_autotune:
kwargs["mode"] = "max-autotune-no-cudagraphs"
try:
fn(fullgraph = True, dynamic = True)
fn(**kwargs)
return True
except Exception as exc: # noqa: BLE001 — optimisation only
_warn(logger, "compile_repeated_blocks", exc)

View file

@ -138,8 +138,9 @@ class _Pipe:
def _vae_to(self, *, memory_format):
self.vae.mem_format = memory_format
def _compile(self, *, fullgraph, dynamic):
def _compile(self, **kwargs):
self.compiled = True
self.compile_kwargs = kwargs
def _fuse(self):
self.fused = True
@ -171,6 +172,9 @@ def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch):
)
assert applied["channels_last"] is True and pipe.vae.mem_format == torch.channels_last
assert applied["compiled"] is True and pipe.compiled is True
# default compiles with dynamic=True and no autotune mode (fast cold start,
# resolution-robust, sidesteps the CUDA-graph crash).
assert pipe.compile_kwargs == {"fullgraph": True, "dynamic": True}
# default also autotunes the VAE convs but does NOT flip TF32 or fuse QKV.
assert applied["cudnn_benchmark"] is True and torch.backends.cudnn.benchmark is True
assert applied["tf32"] is False and applied["fused_qkv"] is False
@ -208,6 +212,9 @@ def test_speed_max_enables_tf32_and_fused_qkv(monkeypatch):
)
assert applied["tf32"] is True and torch.backends.cuda.matmul.allow_tf32 is True
assert applied["fused_qkv"] is True and pipe.fused is True
# max opts into autotuned kernels (static shapes); CUDA-graph modes are avoided.
assert pipe.compile_kwargs["mode"] == "max-autotune-no-cudagraphs"
assert pipe.compile_kwargs["dynamic"] is False
def test_speed_max_tf32_only_on_cuda(monkeypatch):