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).

scripts/compare_engines.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Head-to-head: PyTorch (diffusers GGUF) vs native stable-diffusion.cpp.
+
+Same Z-Image GGUF transformer, same VAE + text encoder, same resolution / steps /
+seed, both resident (no CPU offload) on the same GPU. Reports per-engine compute
+latency (model already loaded) so the denoise + VAE + TE work is compared fairly;
+for sd.cpp it also reports the one-shot wall time (compute + the per-call model
+reload, which a persistent sd-server would remove).
+
+PyTorch runs first (load / warmup / median), is unloaded, then sd.cpp runs.
+"""
+
+from __future__ import annotations
+
+import argparse
+import re
+import sys
+import time
+from pathlib import Path
+
+_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend"
+if str(_BACKEND_ROOT) not in sys.path:
+    sys.path.insert(0, str(_BACKEND_ROOT))
+
+PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
+_DONE_RE = re.compile(r"generate_image completed in ([0-9.]+)s")
+
+
+def _median(xs):
+    return sorted(xs)[len(xs) // 2]
+
+
+def bench_pytorch(repo, gguf, resolutions, steps, seed, iters):
+    import torch
+    from core.inference.diffusion import DiffusionBackend
+
+    rows = []
+    backend = DiffusionBackend()
+    for speed in ("off", "default"):
+        backend.begin_load(repo, gguf_filename = gguf, speed_mode = speed)
+        while backend.load_progress().get("phase") != "ready":
+            if backend.load_progress().get("phase") == "error":
+                raise RuntimeError(backend.load_progress())
+            time.sleep(0.5)
+        for res in resolutions:
+            def gen():
+                torch.cuda.synchronize()
+                t0 = time.time()
+                backend.generate(prompt = PROMPT, width = res, height = res,
+                                 steps = steps, guidance = 0.0, seed = seed, batch_size = 1)
+                torch.cuda.synchronize()
+                return time.time() - t0
+            gen()  # warmup (compiles for `default`)
+            med = _median([gen() for _ in range(iters)])
+            rows.append(("pytorch", speed, res, med, None))
+            print(f"  pytorch  speed={speed:7s} {res}px  compute={med:.3f}s", flush = True)
+        backend.unload()
+    return rows
+
+
+def bench_sdcpp(binary, gguf, vae, llm, resolutions, steps, seed, iters):
+    from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles
+    from core.inference.sd_cpp_engine import SdCppEngine
+
+    engine = SdCppEngine(binary = binary)
+    if not engine.is_available():
+        print("  sd.cpp binary not available; skipping", flush = True)
+        return []
+    files = SdCppModelFiles(diffusion_model = gguf, vae = vae, llm = llm)
+    rows = []
+    out_dir = Path("outputs/compare_engines")
+    out_dir.mkdir(parents = True, exist_ok = True)
+    for native in (None, "default"):  # resident-no-fa vs resident+--diffusion-fa
+        for res in resolutions:
+            params = SdCppGenParams(prompt = PROMPT, width = res, height = res,
+                                    steps = steps, cfg_scale = 1.0, seed = seed)
+            computes, walls = [], []
+            for _ in range(iters):
+                captured = {"c": None}
+                def _log(ln):
+                    m = _DONE_RE.search(ln)
+                    if m:
+                        captured["c"] = float(m.group(1))
+                t0 = time.time()
+                engine.generate(files, params, output_path = str(out_dir / f"sd_{native}_{res}.png"),
+                                offload = [], native_speed = native, on_log = _log)
+                walls.append(time.time() - t0)
+                if captured["c"] is not None:
+                    computes.append(captured["c"])
+            med_c = _median(computes) if computes else None
+            med_w = _median(walls)
+            tag = "default(+fa)" if native == "default" else "off"
+            rows.append(("sdcpp", tag, res, med_c, med_w))
+            print(f"  sdcpp    speed={tag:12s} {res}px  compute={med_c}s  wall={med_w:.3f}s", flush = True)
+    return rows
+
+
+def main(argv = None) -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--repo", default = "unsloth/Z-Image-Turbo-GGUF")
+    p.add_argument("--gguf-name", default = "z-image-turbo-Q4_K_M.gguf")
+    p.add_argument("--sd-binary", default = None)
+    p.add_argument("--sd-gguf", default = None, help = "local gguf for sd.cpp (default: same as pytorch via cache)")
+    p.add_argument("--vae", default = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/flux_vae/ae.safetensors")
+    p.add_argument("--llm", default = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/qwen3_te/Qwen3-4B-Instruct-2507-Q4_K_M.gguf")
+    p.add_argument("--resolutions", default = "512,1024")
+    p.add_argument("--steps", type = int, default = 8)
+    p.add_argument("--seed", type = int, default = 42)
+    p.add_argument("--iters", type = int, default = 3)
+    args = p.parse_args(argv)
+
+    from huggingface_hub import hf_hub_download
+    from core.inference.sd_cpp_engine import find_sd_cpp_binary
+
+    resolutions = [int(x) for x in args.resolutions.split(",")]
+    sd_gguf = args.sd_gguf or hf_hub_download(args.repo, args.gguf_name)
+    binary = args.sd_binary or find_sd_cpp_binary()
+
+    print("== PyTorch (diffusers GGUF) ==", flush = True)
+    pt = bench_pytorch(args.repo, args.gguf_name, resolutions, args.steps, args.seed, args.iters)
+    print("== stable-diffusion.cpp (native) ==", flush = True)
+    sd = bench_sdcpp(binary, sd_gguf, args.vae, args.llm, resolutions, args.steps, args.seed, args.iters)
+
+    print("\n==== COMPARISON (Z-Image-Turbo Q4, fixed seed, resident) ====", flush = True)
+    print(f"{'engine':9s} {'config':13s} {'res':>5s} {'compute_s':>10s} {'wall_s':>8s}", flush = True)
+    for eng, cfg, res, c, w in pt + sd:
+        cs = f"{c:.3f}" if c is not None else "n/a"
+        ws = f"{w:.3f}" if w is not None else "-"
+        print(f"{eng:9s} {cfg:13s} {res:5d} {cs:>10s} {ws:>8s}", flush = True)
+    print("COMPARE-DONE", flush = True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/leverage_probe.py

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Probe two candidate levers from the optimization research, on the real GGUF path:
+
+  * `coordinate_descent_tuning` (Inductor) -- lossless extra kernel autotuning.
+  * FirstBlockCache (diffusers `apply_first_block_cache`) -- step-skip cache, lossy,
+    evaluated at a low (8) step count where its ceiling is lower.
+
+Each config is a fresh pipeline load (so Inductor config / compile artifacts don't
+cross-contaminate). Reports latency + PSNR vs the eager reference. Run on one CUDA GPU.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+REPO = "unsloth/Z-Image-Turbo-GGUF"
+GGUF = "z-image-turbo-Q4_K_M.gguf"
+BASE = "Tongyi-MAI/Z-Image-Turbo"
+PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
+
+
+def _psnr(a, b):
+    mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
+    return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse))
+
+
+def _load():
+    import torch
+    import diffusers
+    from huggingface_hub import hf_hub_download
+
+    t = diffusers.ZImageTransformer2DModel.from_single_file(
+        hf_hub_download(REPO, GGUF),
+        quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = torch.bfloat16),
+        torch_dtype = torch.bfloat16, config = BASE, subfolder = "transformer",
+    )
+    pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
+    pipe.to("cuda")
+    return pipe
+
+
+def _gen(pipe, steps, seed, res):
+    import torch
+
+    g = torch.Generator(device = "cuda").manual_seed(seed)
+    torch.cuda.synchronize()
+    t0 = time.time()
+    img = pipe(prompt = PROMPT, width = res, height = res, num_inference_steps = steps,
+               guidance_scale = 0.0, generator = g).images[0]
+    torch.cuda.synchronize()
+    return img, time.time() - t0
+
+
+def main(argv = None) -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--steps", type = int, default = 8)
+    p.add_argument("--res", type = int, default = 1024)
+    p.add_argument("--seed", type = int, default = 42)
+    args = p.parse_args(argv)
+    steps, res, seed = args.steps, args.res, args.seed
+
+    import torch
+
+    def compile_blocks(pipe, *, cdt = False):
+        if cdt:
+            import torch._inductor.config as ic
+            ic.coordinate_descent_tuning = True
+        pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True)
+
+    def run(tag, *, compile = False, cdt = False, fbc = None):
+        # reset inductor config between runs
+        import torch._inductor.config as ic
+        ic.coordinate_descent_tuning = False
+        torch.compiler.reset()
+        pipe = _load()
+        if fbc is not None:
+            from diffusers.hooks import FirstBlockCacheConfig, apply_first_block_cache
+            apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = fbc))
+        if compile:
+            compile_blocks(pipe, cdt = cdt)
+            _gen(pipe, steps, seed, res)  # warmup / compilation
+        else:
+            _gen(pipe, steps, seed, res)  # allocator warmup
+        img, dt = _gen(pipe, steps, seed, res)
+        del pipe
+        torch.cuda.empty_cache()
+        return tag, np.array(img), dt
+
+    results = []
+    print(f"== leverage probe (Z-Image Q4_K_M, {res}px, {steps} steps) ==", flush = True)
+    _, eager, eager_t = run("eager")
+    print(f"  eager: {eager_t:.3f}s", flush = True)
+    results.append(("eager", eager_t, 0.0))
+
+    for tag, kw in [
+        ("compile(default)", dict(compile = True)),
+        ("compile+coord_desc", dict(compile = True, cdt = True)),
+        ("fbc0.12+compile", dict(compile = True, fbc = 0.12)),
+        ("fbc0.20+compile", dict(compile = True, fbc = 0.20)),
+        ("fbc0.20(no compile)", dict(fbc = 0.20)),
+    ]:
+        try:
+            t, img, dt = run(tag, **kw)
+            ps = _psnr(eager, img)
+            results.append((tag, dt, ps))
+            print(f"  {tag:22s} {dt:.3f}s  ({(eager_t-dt)/eager_t*100:+.0f}% vs eager)  PSNR={ps:.1f} dB", flush = True)
+        except Exception as exc:  # noqa: BLE001
+            print(f"  {tag:22s} FAILED: {type(exc).__name__}: {str(exc)[:140]}", flush = True)
+
+    print("\n==== SUMMARY ====", flush = True)
+    for tag, dt, ps in results:
+        sp = f"{(results[0][1]-dt)/results[0][1]*100:+.0f}%" if tag != "eager" else "ref"
+        pss = f"{ps:.1f}dB" if ps else "ref"
+        print(f"  {tag:24s} {dt:.3f}s  {sp:>6s}  {pss:>8s}", flush = True)
+    print("LEVERAGE-PROBE-DONE", flush = True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
+    sys.exit(main())

studio/backend/core/inference/diffusion_speed.py

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

studio/backend/tests/test_diffusion_speed.py

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