diff --git a/scripts/compare_engines.py b/scripts/compare_engines.py new file mode 100644 index 0000000000..48e50ac981 --- /dev/null +++ b/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()) diff --git a/scripts/leverage_probe.py b/scripts/leverage_probe.py new file mode 100644 index 0000000000..287d005242 --- /dev/null +++ b/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()) diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 4350b6aa9a..447825c361 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/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) diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 51705442eb..fb7c0fdb79 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/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):