Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source)

Add an opt-in transformer_quant mode that loads the dense bf16 transformer and
torchao-quantises it onto the low-precision tensor cores, instead of the GGUF
transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a
B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's
0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor.

GGUF+compile stays the low-memory default and the fallback. The mode is gated on
CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB);
any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged
reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe
(Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips
the tiny projections that crash int8's torch._int_mm. New module mirrors
diffusion_precision.py; quant runs before compile before placement.

184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route
coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py
is the standalone torchao lever probe.

scripts/quant_probe.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Empirical quant probe: torchao int8/fp8/fp4 dynamic quant vs GGUF+compile.
+
+Question this answers: GGUF stores the Z-Image DiT at 4-bit but dequantizes to bf16
+per matmul, so it runs at bf16 tensor-core rate. Can a low-precision *tensor-core*
+path (int8dq on any Ampere+, fp8dq on Ada+, NVFP4/MXFP8 on Blackwell), loaded from
+the dense bf16 transformer, beat GGUF+compile on speed while staying inside the
+quality bar -- and how does its quality compare to GGUF's own 4-bit loss?
+
+Reference for all quality numbers is the DENSE bf16 EAGER image (the best this model
+can do). Each config is a fresh pipeline (no compile/quant cross-contamination).
+Reports median latency, PSNR + LPIPS vs reference, and peak VRAM. 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"
+OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/probe_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))
+
+
+_LPIPS = {"fn": None}
+
+
+def _lpips(ref_arr, arr):
+    """Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable."""
+    try:
+        import torch
+        import lpips
+        if _LPIPS["fn"] is None:
+            _LPIPS["fn"] = lpips.LPIPS(net="alex", verbose=False).cuda().eval()
+
+        def t(x):
+            t = torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
+            return t.cuda()
+
+        with torch.no_grad():
+            return float(_LPIPS["fn"](t(ref_arr), t(arr)).item())
+    except Exception as exc:  # noqa: BLE001
+        print(f"    (lpips unavailable: {type(exc).__name__}: {str(exc)[:80]})", flush=True)
+        return None
+
+
+def _load_dense():
+    import torch
+    import 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 _load_gguf():
+    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 _quant_config(name):
+    """Return a torchao config instance for `name`, or raise to mark FAILED."""
+    from torchao.quantization import (
+        Int8WeightOnlyConfig,
+        Int8DynamicActivationInt8WeightConfig,
+        Float8DynamicActivationFloat8WeightConfig,
+    )
+    if name == "int8wo":
+        return Int8WeightOnlyConfig()
+    if name == "int8dq":
+        return Int8DynamicActivationInt8WeightConfig()
+    if name == "fp8dq":
+        return Float8DynamicActivationFloat8WeightConfig()
+    if name == "nvfp4":
+        from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig
+        return NVFP4DynamicActivationNVFP4WeightConfig()
+    if name == "mxfp8":
+        from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig
+        try:
+            import torch
+            return MXDynamicActivationMXWeightConfig(
+                activation_dtype=torch.float8_e4m3fn, weight_dtype=torch.float8_e4m3fn
+            )
+        except TypeError:
+            return MXDynamicActivationMXWeightConfig()
+    raise ValueError(name)
+
+
+def _make_filter_fn(min_features):
+    """Keep only the FLOP-heavy linears: nn.Linear with both in/out >= min_features.
+    The int8 dynamic path uses torch._int_mm (needs activation M>16), and the tiny
+    timestep/pooled projections (in_features=256) run at M=1 and crash it -- skip them."""
+    import torch.nn as nn
+
+    def filter_fn(module, fqn=""):
+        return (
+            isinstance(module, nn.Linear)
+            and getattr(module, "in_features", 0) >= min_features
+            and getattr(module, "out_features", 0) >= min_features
+        )
+
+    return filter_fn
+
+
+def _apply_quant(pipe, name, log, min_features):
+    import torch.nn as nn
+    from torchao.quantization import quantize_
+    cfg = _quant_config(name)
+    total = sum(1 for m in pipe.transformer.modules() if isinstance(m, nn.Linear))
+    filt = _make_filter_fn(min_features)
+    q = sum(1 for n, m in pipe.transformer.named_modules() if filt(m, n))
+    quantize_(pipe.transformer, cfg, filter_fn=filt)
+    log(f"    quantized transformer with {name} ({q}/{total} linears >= {min_features} feat)")
+
+
+def _compile(pipe, log):
+    fn = getattr(pipe.transformer, "compile_repeated_blocks", None)
+    if not callable(fn):
+        return False
+    for kw in ({"fullgraph": True, "dynamic": True}, {"dynamic": True}, {}):
+        try:
+            fn(**kw)
+            log(f"    compiled repeated blocks {kw}")
+            return True
+        except Exception as exc:  # noqa: BLE001
+            log(f"    compile {kw} failed: {type(exc).__name__}: {str(exc)[:90]}")
+    return False
+
+
+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,
+                   help="only quantize Linear with in&out features >= this (int8 _int_mm needs M>16)")
+    p.add_argument("--configs", default="bf16,bf16_c,gguf_c,int8dq_c,fp8dq_c,nvfp4_c,mxfp8_c,int8wo_c",
+                   help="comma list; suffix _c = +compile")
+    args = p.parse_args(argv)
+    steps, res, seed, iters = args.steps, args.res, args.seed, args.iters
+
+    import torch
+    OUT.mkdir(parents=True, exist_ok=True)
+
+    def run(tag, *, source, quant=None, compile=False):
+        torch.compiler.reset()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+        pipe = _load_dense() if source == "dense" else _load_gguf()
+        load_peak = torch.cuda.max_memory_allocated() / 1e9
+        if quant is not None:
+            _apply_quant(pipe, quant, print_, args.min_feat)
+        if compile:
+            _compile(pipe, print_)
+            _gen(pipe, steps, seed, res)  # warmup / compilation
+        else:
+            _gen(pipe, steps, seed, res)  # allocator warmup
+        torch.cuda.reset_peak_memory_stats()
+        dts, img = [], None
+        for _ in range(iters):
+            img, dt = _gen(pipe, steps, seed, res)
+            dts.append(dt)
+        gen_peak = torch.cuda.max_memory_allocated() / 1e9
+        arr = np.array(img)
+        img.save(OUT / f"{tag}.png")
+        del pipe
+        torch.cuda.empty_cache()
+        return tag, _median(dts), arr, load_peak, gen_peak
+
+    print_ = lambda s: print(s, flush=True)  # noqa: E731
+
+    # config table: tag -> (source, quant, compile)
+    table = {
+        "bf16":      ("dense", None,     False),
+        "bf16_c":    ("dense", None,     True),
+        "gguf_c":    ("gguf",  None,     True),
+        "int8wo_c":  ("dense", "int8wo", True),
+        "int8dq_c":  ("dense", "int8dq", True),
+        "fp8dq_c":   ("dense", "fp8dq",  True),
+        "nvfp4_c":   ("dense", "nvfp4",  True),
+        "mxfp8_c":   ("dense", "mxfp8",  True),
+    }
+    want = [c.strip() for c in args.configs.split(",") if c.strip()]
+
+    print(f"== quant probe (Z-Image-Turbo, {res}px, {steps} steps, seed {seed}) ==", flush=True)
+    ref_arr = None
+    rows = []
+    for tag in want:
+        if tag not in table:
+            print(f"  {tag}: unknown config, skipping", flush=True)
+            continue
+        source, quant, compile = table[tag]
+        print(f"-- {tag} (source={source} quant={quant} compile={compile}) --", flush=True)
+        try:
+            _, med, arr, lp, gp = run(tag, source=source, quant=quant, compile=compile)
+        except Exception as exc:  # noqa: BLE001
+            import traceback
+            print(f"  {tag:10s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush=True)
+            traceback.print_exc()
+            rows.append((tag, None, None, None, None, None))
+            continue
+        if ref_arr is None and tag == "bf16":
+            ref_arr = arr
+        psnr = _psnr(ref_arr, arr) if ref_arr is not None else None
+        lpips_v = _lpips(ref_arr, arr) if (ref_arr is not None and tag != "bf16") else (0.0 if tag == "bf16" else None)
+        rows.append((tag, med, psnr, lpips_v, lp, gp))
+        ps = f"{psnr:.1f}dB" if psnr is not None else "n/a"
+        lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a"
+        print(f"  {tag:10s} {med:.3f}s  PSNR={ps:>7s}  LPIPS={lps:>6s}  loadVRAM={lp:.1f}G genVRAM={gp:.1f}G", flush=True)
+
+    base = next((r[1] for r in rows if r[0] == "bf16" and r[1]), None)
+    gguf = next((r[1] for r in rows if r[0] == "gguf_c" and r[1]), None)
+    print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush=True)
+    print(f"{'config':10s} {'sec':>7s} {'vs_bf16':>8s} {'vs_gguf':>8s} {'PSNR':>8s} {'LPIPS':>7s} {'loadG':>6s} {'genG':>6s}", flush=True)
+    for tag, med, psnr, lpips_v, lp, gp in rows:
+        if med is None:
+            print(f"{tag:10s} {'FAILED':>7s}", flush=True)
+            continue
+        vb = f"{base/med:.2f}x" if base else "-"
+        vg = f"{gguf/med:.2f}x" if gguf else "-"
+        ps = f"{psnr:.1f}" if psnr is not None and psnr != float('inf') else ("inf" if psnr == float('inf') else "n/a")
+        lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a"
+        print(f"{tag:10s} {med:>7.3f} {vb:>8s} {vg:>8s} {ps:>8s} {lps:>7s} {lp:>6.1f} {gp:>6.1f}", flush=True)
+    print("QUANT-PROBE-DONE", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
+    sys.exit(main())