ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

20
studio/backend/core/inference/diffusion_speed.py
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)