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:
parent
141cb3ae10
commit
ede94176f6
4 changed files with 289 additions and 5 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue