unsloth/scripts/perf_levers_probe.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Measure the next-phase diffusion levers on the real model, vs today's compiled baseline.

Variants (Z-Image dense bf16, regional compile = the shipped "default" speed profile):
  baseline            -- channels_last + compile_repeated_blocks (reference image)
  inductor_flags      -- + the lossless inductor autotune flags (conv_1x1_as_mm,
                         coordinate_descent_tuning(+all_dirs), epilogue_fusion=False)
  attn_cudnn          -- + set_attention_backend("_native_cudnn") (exact)
  attn_flash4         -- + set_attention_backend("flash_4_hub")   (exact, SM100)
  attn_sage           -- + set_attention_backend("sage")          (INT8 QK, quantized)
  fbcache             -- + First-Block-Cache (threshold 0.12)      (few-step headroom test)

Reports median latency, vs-baseline speedup, peak VRAM, and LPIPS vs baseline. One CUDA GPU."""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import numpy as np

BASE = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "perf_levers_images"


_LP = {"fn": None}


def _lpips(ref, arr):
    try:
        import lpips
        import torch

        # Keep the metric model on CPU: caching it on CUDA leaves it resident across
        # variants, and each run resets peak-memory stats, so its VRAM would be charged
        # to (and reduce headroom for) every later variant's measurement.
        if _LP["fn"] is None:
            _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval()

        def t(x):
            return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0

        with torch.no_grad():
            return float(_LP["fn"](t(ref), t(arr)).item())
    except Exception as exc:  # noqa: BLE001
        print(f"    (lpips: {type(exc).__name__})", flush = True)
        return None


def _set_inductor_flags():
    import torch._inductor.config as ic

    ic.conv_1x1_as_mm = True
    ic.coordinate_descent_tuning = True
    ic.coordinate_descent_check_all_directions = True
    ic.epilogue_fusion = False
    try:
        ic.force_fuse_int_mm_with_mul = True
    except Exception:  # noqa: BLE001
        pass


def _reset_inductor_flags():
    import torch._inductor.config as ic

    ic.conv_1x1_as_mm = False
    ic.coordinate_descent_tuning = False
    ic.coordinate_descent_check_all_directions = False
    ic.epilogue_fusion = True
    # Reset the int-mm fusion flag too, or it leaks from the inductor_flags variant into
    # every later compiled row and the attention/fbcache measurements stop being isolated.
    try:
        ic.force_fuse_int_mm_with_mul = False
    except Exception:  # noqa: BLE001
        pass


def _load():
    import diffusers
    import torch

    t = diffusers.ZImageTransformer2DModel.from_pretrained(
        BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
    )
    pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
    pipe.to("cuda")
    try:
        pipe.vae.to(memory_format = torch.channels_last)
    except Exception:  # noqa: BLE001
        pass
    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 _median(xs):
    return sorted(xs)[len(xs) // 2]


def run(
    tag,
    steps,
    seed,
    res,
    iters,
    *,
    attn = None,
    fbcache = None,
    inductor = False,
):
    import torch

    torch.compiler.reset()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    _reset_inductor_flags()
    if inductor:
        _set_inductor_flags()
    pipe = _load()
    note = ""
    if attn is not None:
        try:
            pipe.transformer.set_attention_backend(attn)
        except Exception as exc:  # noqa: BLE001
            note = f"attn({attn})={type(exc).__name__}:{str(exc)[:60]}"
            print(f"    [{tag}] {note}", flush = True)
            del pipe  # free the resident pipe so a skipped variant doesn't leak VRAM
            torch.cuda.empty_cache()
            return None
    else:
        # set_attention_backend pins diffusers' PROCESS-WIDE active backend, and a fresh
        # transformer's processors (backend None) inherit it. Force native for the no-attn
        # variants so they aren't silently measured under a prior variant's kernel (e.g.
        # fbcache running with a leftover sage backend).
        try:
            pipe.transformer.set_attention_backend("native")
        except Exception as exc:  # noqa: BLE001 — best-effort isolation
            print(f"    [{tag}] attn(native-reset)={type(exc).__name__}:{str(exc)[:60]}", flush = True)
    if fbcache is not None:
        try:
            from diffusers.hooks import FirstBlockCacheConfig, apply_first_block_cache
            apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = fbcache))
        except Exception as exc:  # noqa: BLE001
            print(f"    [{tag}] fbcache={type(exc).__name__}:{str(exc)[:60]}", flush = True)
            del pipe
            torch.cuda.empty_cache()
            return None
    try:
        pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True)
    except Exception as exc:  # noqa: BLE001
        print(f"    [{tag}] compile={type(exc).__name__}:{str(exc)[:60]}", flush = True)
    try:
        _gen(pipe, steps, seed, res)  # warmup / compile
    except Exception as exc:  # noqa: BLE001
        import traceback

        traceback.print_exc()
        print(f"    [{tag}] FAILED first gen: {type(exc).__name__}:{str(exc)[:80]}", flush = True)
        del pipe
        torch.cuda.empty_cache()
        return None
    dts, img = [], None
    for _ in range(iters):
        img, dt = _gen(pipe, steps, seed, res)
        dts.append(dt)
    peak = torch.cuda.max_memory_allocated() / 1e9
    arr = np.array(img)
    OUT.mkdir(parents = True, exist_ok = True)
    img.save(OUT / f"{tag}.png")
    del pipe
    torch.cuda.empty_cache()
    return _median(dts), arr, peak


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)
    p.add_argument("--iters", type = int, default = 3)
    args = p.parse_args(argv)
    s, r, seed, it = args.steps, args.res, args.seed, args.iters

    print(f"== perf levers (Z-Image dense, {r}px, {s} steps) ==", flush = True)
    base = run("baseline", s, seed, r, it)
    if base is None:
        print("baseline FAILED", flush = True)
        return 1
    bmed, ref, bpeak = base
    print(f"  baseline {bmed:.3f}s  peak={bpeak:.1f}G", flush = True)
    rows = [("baseline", bmed, bpeak, 0.0)]

    variants = [
        ("inductor_flags", dict(inductor = True)),
        ("attn_cudnn", dict(attn = "_native_cudnn")),
        ("attn_flash4", dict(attn = "flash_4_hub")),
        ("attn_sage", dict(attn = "sage")),
        ("attn_sage_inductor", dict(attn = "sage", inductor = True)),
        ("fbcache_0p12", dict(fbcache = 0.12)),
    ]
    for tag, kw in variants:
        out = run(tag, s, seed, r, it, **kw)
        if out is None:
            rows.append((tag, None, None, None))
            continue
        med, arr, peak = out
        lp = _lpips(ref, arr)
        rows.append((tag, med, peak, lp))
        spd = f"{bmed/med:.2f}x" if med else "-"
        print(f"  {tag:20s} {med:.3f}s ({spd} vs base) peak={peak:.1f}G LPIPS={lp}", flush = True)

    print("\n==== SUMMARY (ref = baseline compile) ====", flush = True)
    for tag, med, peak, lp in rows:
        if med is None:
            print(f"  {tag:20s} FAILED")
            continue
        spd = f"{bmed/med:.2f}x" if med else "-"
        lpv = "ref" if (tag == "baseline") else (f"{lp:.3f}" if lp is not None else "n/a")
        print(f"  {tag:20s} {med:.3f}s  {spd:>6s}  peak={peak:.1f}G  LPIPS={lpv:>6s}", flush = True)
    print("PERF-LEVERS-DONE", flush = True)
    return 0


if __name__ == "__main__":
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
    sys.exit(main())