scripts/24_nunchaku_e2e_speed.py

"""Real END-TO-END speed of our klein-4B W4A4-NVFP4 champion, on this Blackwell card.

Our SVDQuant models are fake-quant (quality already validated: r128 W4A4 = 0.0303, teacher-
indistinguishable). GEMM timing is value-independent, so the real speed of the champion depends
only on its CONFIG (shapes + rank + nvfp4), not the weight values. This swaps every one of
klein-4B's 100 target Linears for a REAL nunchaku SVDQW4A4Linear(precision='nvfp4') at the
champion's config and times the full 4-step pipeline vs bf16 — the true deployable end-to-end
speedup for our architecture (Linears on FP4 kernels; attention/norm/VAE/text-encode stay bf16).

Usage: PYTHONPATH=. python3 scripts/24_nunchaku_e2e_speed.py [H] [W] [steps]
"""
import os, sys, time, json, statistics as st
import torch
import torch.nn as nn
from flux2distill.model_utils import load_pipeline
from flux2distill import svdquant as sq
from nunchaku.models.linear import SVDQW4A4Linear

H = int(sys.argv[1]) if len(sys.argv) > 1 else 512
W = int(sys.argv[2]) if len(sys.argv) > 2 else 512
STEPS = int(sys.argv[3]) if len(sys.argv) > 3 else 4
PROMPT = "a photorealistic storefront at golden hour, neon sign, wet pavement reflections, 50mm"
print(f"=== klein-4B end-to-end speed | {W}x{H} | {STEPS} steps | {torch.cuda.get_device_name(0)} ===")

pipe = load_pipeline(device="cuda")
tf = pipe.transformer
tf.eval().requires_grad_(False)
targets = sq.target_linear_names(tf)
print(f"target Linears: {len(targets)}")

# per-step timing via transformer.forward monkeypatch (sync'd)
_step = []
_fwd = tf.forward
def timed(*a, **k):
    torch.cuda.synchronize(); s = time.perf_counter()
    o = _fwd(*a, **k)
    torch.cuda.synchronize(); _step.append(time.perf_counter() - s)
    return o
tf.forward = timed

def run(seed=0):
    _step.clear(); torch.cuda.reset_peak_memory_stats()
    g = torch.Generator("cpu").manual_seed(seed)
    torch.cuda.synchronize(); t = time.perf_counter()
    with torch.autocast("cuda", dtype=torch.bfloat16):
        pipe(prompt=PROMPT, num_inference_steps=STEPS, guidance_scale=1.0,
             height=H, width=W, generator=g)
    torch.cuda.synchronize()
    return dict(total=time.perf_counter() - t, peak=torch.cuda.max_memory_allocated()/1e9,
                step=st.median(_step[1:]) if len(_step) > 1 else _step[0])

class Wrap(nn.Module):
    """Adapt SVDQW4A4Linear (needs 3D B,S,C) to arbitrary leading dims."""
    def __init__(self, m): super().__init__(); self.m = m
    def forward(self, x):
        s = x.shape
        # SVDQW4A4Linear is a custom CUDA op -> not covered by autocast; the FP4 kernel asserts
        # unless input is fp16/bf16, so cast explicitly (real deployment runs bf16 acts anyway).
        y = self.m(x.reshape(1, -1, s[-1]).to(torch.bfloat16).contiguous())
        return y.reshape(*s[:-1], -1).to(x.dtype)

@torch.no_grad()
def swap_nvfp4(rank):
    for name in targets:
        parent, leaf = sq._get_parent(tf, name)
        lin = getattr(parent, leaf) if not leaf.isdigit() else parent[int(leaf)]
        if isinstance(lin, Wrap):
            lin = lin.m  # already swapped once; rebuild from a fresh dummy of same shape
            in_f, out_f, has_b = lin.in_features, lin.out_features, lin.bias is not None
        else:
            in_f, out_f, has_b = lin.in_features, lin.out_features, lin.bias is not None
        m = SVDQW4A4Linear(in_f, out_f, rank=rank, bias=has_b, precision="nvfp4",
                           torch_dtype=torch.bfloat16, device="cuda")
        m.qweight.random_(-128, 127); m.wscales.copy_(torch.ones_like(m.wscales))
        m.smooth_factor.fill_(1.0); m.smooth_factor_orig.fill_(1.0)
        m.proj_down.normal_(0, 0.02); m.proj_up.normal_(0, 0.02); m.wcscales.fill_(1.0)
        if has_b: m.bias.zero_()
        sq._set_module(tf, name, Wrap(m))
    torch.cuda.empty_cache()

# warmup + measure bf16 baseline
run(seed=99)
bf = st.median([run(i)['total'] for i in range(3)]); bf_r = run(0)
print(f"\nbf16        : {bf:.3f}s/img  step={bf_r['step']*1000:.0f}ms  peakVRAM={bf_r['peak']:.1f}GB")

results = {"bf16": {"total_s": round(bf, 3), "step_ms": round(bf_r['step']*1000, 1), "vram_gb": round(bf_r['peak'], 1)}}
for rank in (128, 64):
    swap_nvfp4(rank); run(seed=99)  # warmup after swap
    tot = st.median([run(i)['total'] for i in range(3)]); r = run(0)
    sp = bf / tot
    print(f"nvfp4 r{rank:<3}: {tot:.3f}s/img  step={r['step']*1000:.0f}ms  peakVRAM={r['peak']:.1f}GB  -> {sp:.2f}x end-to-end vs bf16")
    results[f"nvfp4_r{rank}"] = {"total_s": round(tot, 3), "step_ms": round(r['step']*1000, 1),
                                 "vram_gb": round(r['peak'], 1), "speedup": round(sp, 3)}

os.makedirs("outputs/nvfp4", exist_ok=True)
json.dump({"res": f"{W}x{H}", "steps": STEPS, "results": results},
          open("outputs/nvfp4/e2e_speed.json", "w"), indent=2)
print("\nsaved -> outputs/nvfp4/e2e_speed.json")