Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9

scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer.
Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with
use_triton_kernel=False (the default triton path needs the missing MSLK library), but
only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than
using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's
CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto
on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast.

scripts/nvfp4_probe.py

@@ -0,0 +1,145 @@
+# 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 NVFP4 via torchao with use_triton_kernel=False (no MSLK) on the real dense
+Z-Image transformer: is it a genuine FP4-tensor-core speedup over fp8, and is quality
+in-bar? Reference for LPIPS is dense bf16 eager. Run on one CUDA (Blackwell) 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("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/nvfp4_images")
+
+
+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))
+
+
+_LP = {"fn": None}
+
+
+def _lpips(ref, arr):
+    try:
+        import torch, lpips
+        if _LP["fn"] is None:
+            _LP["fn"] = lpips.LPIPS(net="alex", verbose=False).cuda().eval()
+        def t(x):
+            return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda()
+        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 _load_dense():
+    import torch, diffusers
+    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")
+    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 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)
+    p.add_argument("--min-feat", type=int, default=512)
+    args = p.parse_args(argv)
+    steps, res, seed, mf = args.steps, args.res, args.seed, args.min_feat
+    import torch
+    import torch.nn as nn
+    OUT.mkdir(parents=True, exist_ok=True)
+
+    def filt(mod, fqn=""):
+        return (isinstance(mod, nn.Linear)
+                and mod.in_features >= mf and mod.out_features >= mf)
+
+    def run(tag, *, cfg=None, compile=True):
+        torch.compiler.reset(); torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
+        pipe = _load_dense()
+        if cfg is not None:
+            from torchao.quantization import quantize_
+            quantize_(pipe.transformer, cfg, filter_fn=filt)
+        if compile:
+            try:
+                pipe.transformer.compile_repeated_blocks(fullgraph=True, dynamic=True)
+            except Exception as exc:  # noqa: BLE001
+                print(f"    [{tag}] compile failed: {type(exc).__name__}: {str(exc)[:90]}", flush=True)
+        _gen(pipe, steps, seed, res)  # warmup / compile
+        dts, img = [], None
+        for _ in range(args.iters):
+            img, dt = _gen(pipe, steps, seed, res); dts.append(dt)
+        gp = torch.cuda.max_memory_allocated() / 1e9
+        arr = np.array(img); img.save(OUT / f"{tag}.png")
+        del pipe; torch.cuda.empty_cache()
+        return _median(dts), arr, gp
+
+    from torchao.quantization import (
+        Float8DynamicActivationFloat8WeightConfig as FP8)
+    from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV
+
+    print(f"== nvfp4 probe (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==", flush=True)
+    bref, ref, _ = run("bf16_eager", cfg=None, compile=False)
+    print(f"  bf16 eager ref: {bref:.3f}s", flush=True)
+    rows = [("bf16_eager", bref, float("inf"), 0.0, None)]
+
+    specs = [
+        ("bf16_compile", None, True),
+        ("fp8_compile", FP8(), True),
+        ("nvfp4_notriton_compile", NV(use_triton_kernel=False), True),
+        ("nvfp4_notriton_eager", NV(use_triton_kernel=False), False),
+    ]
+    for tag, cfg, comp in specs:
+        try:
+            med, arr, gp = run(tag, cfg=cfg, compile=comp)
+            ps, lp = _psnr(ref, arr), _lpips(ref, arr)
+            rows.append((tag, med, ps, lp, gp))
+            print(f"  {tag:24s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G", flush=True)
+        except Exception as exc:  # noqa: BLE001
+            import traceback; traceback.print_exc()
+            print(f"  {tag:24s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush=True)
+            rows.append((tag, None, None, None, None))
+
+    fp8 = next((r[1] for r in rows if r[0] == "fp8_compile" and r[1]), None)
+    print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush=True)
+    for tag, med, ps, lp, gp in rows:
+        if med is None:
+            print(f"  {tag:24s} FAILED"); continue
+        vs_fp8 = f"{fp8/med:.2f}x" if fp8 else "-"
+        psv = "inf" if ps == float("inf") else f"{ps:.1f}"
+        lpv = "ref" if (lp == 0.0 and tag == 'bf16_eager') else (f"{lp:.3f}" if lp is not None else "n/a")
+        print(f"  {tag:24s} {med:.3f}s  vs_fp8:{vs_fp8:>6s}  PSNR={psv:>5s}  LPIPS={lpv:>6s}", flush=True)
+    print("NVFP4-PROBE-DONE", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
+    sys.exit(main())