scripts/20_nunchaku_profile.py

"""Profile Nunchaku FLUX.2-klein-9B low-bit (FP4 / INT4) fused SVDQuant kernels on Blackwell.

Loads the SVDQuant-quantized transformer (svdq-{prec}_r32-...safetensors) into a stock
`Flux2KleinPipeline`, runs a warm-up + N timed text-to-image generations, and prints a
per-step + end-to-end latency / VRAM / throughput breakdown.

Timing is done by monkeypatching `transformer.forward` (one call == one denoise step) and
`vae.decode`, each wrapped in cuda.synchronize() + perf_counter() for clean, async-free
numbers. The warm-up run (kernel autotune / lazy CUDA init / cudnn pick) is excluded.

Usage:
  PYTHONPATH=. python3 scripts/20_nunchaku_profile.py [precision] [steps] [H] [W] [n_timed]
    precision : fp4 | int4   (default: nunchaku.utils.get_precision() -> fp4 on sm_120)
    steps     : denoise steps (default 4 — klein is distilled / CFG-free)
"""
import os, sys, time, json, statistics as st
import torch

PREC  = sys.argv[1] if len(sys.argv) > 1 and sys.argv[1] not in ("-",) else None
STEPS = int(sys.argv[2]) if len(sys.argv) > 2 else 4
H     = int(sys.argv[3]) if len(sys.argv) > 3 else 1024
W     = int(sys.argv[4]) if len(sys.argv) > 4 else 1024
NRUN  = int(sys.argv[5]) if len(sys.argv) > 5 else 3
REPO  = "models/klein-9b-nunchaku"
NAME  = "FLUX.2-klein-9B-Nunchaku"
PROMPT = ("a photorealistic storefront at golden hour with a glowing neon sign that reads "
          "\"NUNCHAKU\", wet pavement reflections, shallow depth of field, 50mm")

from diffusers import Flux2KleinPipeline
from nunchaku.models.transformers.transformer_flux2 import NunchakuFlux2Transformer2DModel
from nunchaku.utils import get_precision

prec = PREC or get_precision()
ckpt = f"{REPO}/svdq-{prec}_r32-{NAME}.safetensors"
print(f"=== Nunchaku FLUX.2-klein-9B | precision={prec.upper()} r32 | {W}x{H} | {STEPS} steps | {NRUN} timed runs ===")
print(f"device={torch.cuda.get_device_name(0)} | torch={torch.__version__} cuda={torch.version.cuda}")

torch.cuda.reset_peak_memory_stats()
t0 = time.perf_counter()
transformer = NunchakuFlux2Transformer2DModel.from_pretrained(ckpt, torch_dtype=torch.bfloat16)
pipe = Flux2KleinPipeline.from_pretrained(REPO, torch_dtype=torch.bfloat16, transformer=transformer)
pipe.to("cuda")
torch.cuda.synchronize()
print(f"[load] {time.perf_counter()-t0:.1f}s | post-load VRAM {torch.cuda.max_memory_allocated()/1e9:.2f} GB")

# ---- instrument per-step transformer + vae decode (sync'd) ----
step_times, dec_times = [], []
_dumped = {"done": False}
DUMP_FWD = os.environ.get("DUMP_FWD")  # path: save the first forward's inputs for the bf16 replay bench
_fwd = transformer.forward
def _to_cpu(x):
    if torch.is_tensor(x): return x.detach().to("cpu")
    if isinstance(x, (list, tuple)): return type(x)(_to_cpu(v) for v in x)
    if isinstance(x, dict): return {k: _to_cpu(v) for k, v in x.items()}
    return x
def timed_fwd(*a, **k):
    if DUMP_FWD and not _dumped["done"]:
        torch.save({"args": _to_cpu(a), "kwargs": _to_cpu(k)}, DUMP_FWD)
        _dumped["done"] = True
        print(f"[dump] saved first transformer.forward inputs -> {DUMP_FWD}")
    torch.cuda.synchronize(); s = time.perf_counter()
    out = _fwd(*a, **k)
    torch.cuda.synchronize(); step_times.append(time.perf_counter() - s)
    return out
transformer.forward = timed_fwd

_dec = pipe.vae.decode
def timed_dec(*a, **k):
    torch.cuda.synchronize(); s = time.perf_counter()
    out = _dec(*a, **k)
    torch.cuda.synchronize(); dec_times.append(time.perf_counter() - s)
    return out
pipe.vae.decode = timed_dec

def run(seed=0):
    step_times.clear(); dec_times.clear()
    torch.cuda.reset_peak_memory_stats()
    g = torch.Generator("cpu").manual_seed(seed)
    torch.cuda.synchronize(); t = time.perf_counter()
    img = pipe(prompt=PROMPT, guidance_scale=1.0, num_inference_steps=STEPS,
               height=H, width=W, generator=g).images[0]
    torch.cuda.synchronize(); total = time.perf_counter() - t
    return dict(img=img, total=total, peak=torch.cuda.max_memory_allocated()/1e9,
                steps=list(step_times), denoise=sum(step_times), decode=sum(dec_times))

print("[warmup] running (excluded) ...")
w = run(seed=999)
print(f"[warmup] total {w['total']:.2f}s | steps {[f'{x*1000:.0f}' for x in w['steps']]}ms")

runs = [run(seed=i) for i in range(NRUN)]
med = st.median
ns = len(runs[0]['steps'])
per_step = [med([r['steps'][i] for r in runs]) for i in range(ns)]
tot, den = med([r['total'] for r in runs]), med([r['denoise'] for r in runs])
dec = med([r['decode'] for r in runs])
ovh = med([r['total'] - r['denoise'] - r['decode'] for r in runs])
peak = med([r['peak'] for r in runs])
stepN = med([med([r['steps'][i] for r in runs]) for i in range(1, ns)]) if ns > 1 else per_step[0]

print("\n================  RESULTS  (median of %d runs)  ================" % NRUN)
print(f"precision            : {prec.upper()}  (Nunchaku SVDQuant W4A4, low-rank r32)")
print(f"resolution / steps   : {W}x{H} / {STEPS}")
print(f"per-step transformer : " + ", ".join(f"{t*1000:.0f}ms" for t in per_step))
print(f"  step1+ mean        : {stepN*1000:.0f}ms   (step0 warms caches; later steps are the true kernel cost)")
print(f"denoise (sum steps)  : {den*1000:.0f}ms   -> {STEPS/den:.2f} steps/s (transformer-only)")
print(f"vae decode           : {dec*1000:.0f}ms")
print(f"text-encode+overhead : {ovh*1000:.0f}ms")
print(f"END-TO-END           : {tot:.2f}s   -> {1/tot:.2f} img/s")
print(f"peak VRAM            : {peak:.2f} GB")

os.makedirs("outputs/nunchaku", exist_ok=True)
out = f"outputs/nunchaku/klein9b_{prec}_{W}x{H}_{STEPS}s.png"
runs[-1]['img'].save(out)
res = dict(model="klein-9b", engine="nunchaku", precision=prec, rank=32, H=H, W=W, steps=STEPS,
           per_step_ms=[round(x*1000,1) for x in per_step], denoise_ms=round(den*1000,1),
           steps_per_s=round(STEPS/den,3), decode_ms=round(dec*1000,1), overhead_ms=round(ovh*1000,1),
           total_s=round(tot,3), img_per_s=round(1/tot,3), peak_vram_gb=round(peak,2),
           device=torch.cuda.get_device_name(0))
with open(f"outputs/nunchaku/klein9b_{prec}_{W}x{H}_{STEPS}s.json", "w") as f:
    json.dump(res, f, indent=2)
print(f"saved sample -> {out}")
print(f"saved metrics -> outputs/nunchaku/klein9b_{prec}_{W}x{H}_{STEPS}s.json")