TODO.md

# TODO — FLUX.2 klein-4B quantization track

Backlog. Status: [ ] todo · [~] in progress · [x] done. (Block-surgery backlog archived in
`block_surgery_todo.md`.)

## ★ 2026-06-13 — NVFP4 + deployable fused model — DONE
- [x] **NVFP4 in `svdquant.py`** (E2M1 + group-16 + FP8 scales; `w_fmt`/`a_fmt`) + unit test
      (`scripts/test_nvfp4.py`). Beats INT4 ~2× at matched bits.
- [x] **NVFP4 sweep** (was the "NVFP4 sim" TODO below): W4A4 r{32,64,128} = 0.0390/0.0364/0.0303;
      NVFP4-W+FP8-A r{64,128} = 0.0204/**0.0169** (overall champion). `scripts/run_nvfp4_cell.sh`.
- [x] **Real kernel speed** — built Nunchaku from source for sm_120 (`/workspace/build_nunchaku`,
      `do_build.sh`); NVFP4 W4A4 = 2.49–2.75× per-layer; INT4 broken on Blackwell (hardware). Profilers
      `scripts/20–25`.
- [x] **bf16→Nunchaku NVFP4 exporter** (`flux2distill/nunchaku_export.py`, no DeepCompressor; validated
      2.99% in `scripts/26`) → **deployable fused klein-4B** (`scripts/27/28`):
      `outputs/nvfp4/deploy/klein4b_nvfp4_fused.safetensors` (2.9 GB), correct images, 1.74×@512 /
      1.90×@1024, −24% VRAM. Report `report/NVFP4_REPORT.md`; figures `scripts/make_nvfp4_figures.py`.
- [x] Docs: `docs/CUDA_SETUP_RUNBOOK.md` (setup + footguns), `docs/SPEEDUP_IDEAS.md` (next steps).

## ★★ 2026-06-14 — NVFP4 head-to-head (N=512) — DONE
- [x] **Matched head-to-head** A teacher / C ours-NVFP4-r128-real-kernel / D plain-NVFP4-r0 /
      E BFL-official-FP8, N=512 MJHQ-30k. Low-rank branch helps: **LPIPS −19.7%, PSNR +1.27 dB,
      FID-teacher −14.7%**; real kernel ≈ fake-q; no semantic loss. `report/HEADTOHEAD_klein4b_nvfp4.md`,
      `RESULTS.md` (2026-06-14), `outputs/eval/h2h/metrics.json`, `report/figures/h2h_*.png`.
      Pipeline `scripts/run_h2h.sh` → `scripts/34_metrics.py` → `scripts/42_h2h_figures.py`.
- [x] **BFL-official comparison (was DEFERRED in `report/NVFP4_SVDQUANT_EVAL.md`) — RESOLVED:**
      their **FP8** loads faithfully (W8A8, `scripts/41_gen_bfl_fp8.py`) = model E; their **NVFP4 is
      BLOCKED** — cutlass/TensorRT swizzled layout, needs their runtime (`scripts/40_load_bfl.py` +
      HEADTOHEAD §5). Do NOT re-attempt manual dequant. Use our r0 as the "plain NVFP4" stand-in.
- [ ] (open, low priority) Run BFL-official NVFP4 if a `tensorrt_llm` / TRT FLUX path becomes available
      on sm_120 — only path to a *direct* BFL-NVFP4 number.

## ★★ NEXT — kernel/pipeline speedup (full menu in `docs/SPEEDUP_IDEAS.md`)
Lift the ~1.9× toward the per-step 2.3× / per-layer 2.5× by attacking the non-transformer overhead
and the attention. Priority order:
- [ ] **CUDA graphs** over the denoise loop (4B is launch-bound at small sizes) — cheapest big win.
- [ ] **Quantize text-encoder (Qwen3) + VAE decoder** — the fixed ~30% overhead cap.
- [ ] **Fix the fused batch=1 limitation** (packed-rotary broadcast over batch) — serving throughput.
- [ ] **Output-quant fusion** (`qout`: fold layer N+1's act-quant into N's epilogue).
- [ ] **FP8 attention** (QKᵀ/AV in E4M3) — the high-res ceiling; verify the GEMM uses Blackwell
      `tcgen05`/TMEM, not a Hopper `wgmma` path.
- [ ] GEMM tile/warp autotune per shape; TMA loads; rank Pareto + per-layer rank; real `wcscales`; FP6.
- [ ] Measure everything across {quant level × batch × res} vs bf16 (extend `outputs/nvfp4/benchmark.json`).

## ★ Deployment coverage — INT4 export for OLDER GPUs (Ada/Ampere/Turing, RTX 20/30/40)
- [ ] **Extend `flux2distill/nunchaku_export.py` to `precision='int4'`** (int4 codes, fp16/bf16
      per-group-64 scales via `pack_scale` not `pack_micro_scale`, no wtscale/wcscales) and produce
      real deployable checkpoints `klein4b_int4_{w4a8,w4a4}.safetensors` — the artifacts a 4090/30-series
      user actually loads (Blackwell has NO INT4 tensor cores; INT4 is the older-GPU format). Quality is
      verifiable HERE (the INT4 kernel runs on sm_120, just emulated/slow); REAL speed needs an Ada GPU.
      NB: the saved fake-quant `quant_state.pt` champions are bf16-stored quality refs, NOT runnable-fast
      on any GPU — they're the source recipe; the INT4 export is what makes them deployable on older cards.

## Build-out (lead) — DONE
- [x] `flux2distill/svdquant.py` — SVDQuantLinear (smooth + SVD low-rank + 4-bit residual),
      fake-quant primitives, act-stat+Gram hooks, build/load surgery.
- [x] **Whitened (activation-aware) SVD** — min OUTPUT error ‖X̂(Ŵ−L)‖ via eigen-√ of the
      smoothed-act Gram (robust to non-PD; per-layer plain-SVD fallback). Makes rank a real knob.
- [x] **Iterative low-rank refinement** (SVDQuant §4.2) — re-fit L to absorb 4-bit quant error,
      keep best iterate. `REFINE` env (default 3). −13% at r16 vs plain.
- [x] `scripts/11–13` (calib download / build / eval); 2000-img calib cached at `data/monet_calib`
      (100-img subset is plenty, ~3min/build).
- [x] **Initial rank sweep (100-calib, whitened+refine)** — 0.0574/0.0494/0.0454. SUPERSEDED by the
      full 300-calib grid below (which revises the per-knob story — see RESULTS.md). Kept as history.

## ACTIVE campaign (2026-06-01) — full method×rank GRID, fixed calib — ✅ DONE
Closed the "L-shape, not a grid" gap: ran the **complete 2×2 method grid at EVERY rank**, all on
one fixed 300-img calib (`data/monet_cache` latents) so every cell is comparable.

- [x] **Grid: ranks {16,32,64} × {plain, plain+whiten, plain+refine, plain+whiten+refine}** = 12
      builds+evals, fixed 300-calib, α=0.5, group=64, W4A8. Dirs `outputs/abl_c300_r{R}_{variant}`,
      8-probe montages each. Run one-at-a-time w/ per-run logs + Monitor. Full grid in RESULTS.md.
- [x] **plain+refine (the missing cell)** — now run at all 3 ranks. Refine's solo effect: helps
      everywhere EXCEPT r16 (overfits weight error). **r64 plain+refine = 0.0446 is the grid best.**
- [x] Collate into RESULTS.md grid (done) + report assets/montages (in progress).

**Grid takeaways (overturn the old "each upgrade compounds" claim):**
- Refine = reliable workhorse (best variant at every rank uses it; only r16-plain regresses).
- Whitening ALONE is unreliable at 300-calib — non-monotonic in rank (hurts r16 & r64, helps r32);
  it overfits the noisy 300-img Gram. Only earns its keep at r32 or when paired with refine.
- Strong whiten×refine interaction (refine in the whitened metric fixes whitening's overfit).
- At high rank, **skip whitening**: r64 plain+refine (0.0446, no Gram, simpler) ≥ r64 whiten+refine (0.0451).

## Deferred follow-ups (queued — see plan.md §5)
- [ ] **★ NOW THE KEY EXPERIMENT — re-sweep at 2000-image calib.** The grid shows whitening
      misbehaving at 300-calib (overfit Gram). Re-run the grid (or at least the whiten / whiten+refine
      cells) at **2000-img full-res calib** (`scripts/11` → `data/monet_calib`, needs image download)
      to test: does whitening become *reliably* beneficial with richer activation stats? This decides
      whether to keep whitening at all. Reuse the cached-stats idea below to make it cheap.
- [x] **Eval probes 4 → 8** (richer visual comparison): kept the original 4 (storefront text,
      mountain lake, fisherman, neon street) + added chalkboard-text, multi-object breakfast flat-lay,
      hand/face five-fingers, dewy-spiderweb macro. In `scripts/13_eval_svdquant.py`.

## ★★ NOW THE #1 EXPERIMENT — SmoothQuant is HURTING W4A8 (confirmed 2026-06-01)
Mechanism ablation (rank-0, 300-calib) found **removing SmoothQuant IMPROVES eval-loss**:
RTN (no smooth, `SMOOTH=0`, s=1) = **0.0573** vs SmoothQuant rank-0 (α=0.5) = 0.0729 (**−21%**), and
RTN even beats the *smoothed* SVD r16/r32 cells (0.0620/0.0586). The whole α=0.5 grid is mis-tuned —
the SVD branch was partly papering over smoothing damage (worst-layer wrecon 0.15→0.26 with smoothing).
- [x] **RTN floor (SMOOTH=0)** and **SmoothQuant rank-0** baselines built+evaled (RESULTS.md).
      New `SMOOTH` env on `scripts/12` (`from_linear(smooth=False)` → s=1). NB **α=0 ≠ no-smooth**.
- [x] **★ DONE (2026-06-10, new box): SMOOTH=0 at r{16,32,64} plain+refine** — wins at every rank
      (−14.1/−8.6/−8.6%), new best **r64 ns = 0.0297** (new axis; vs smoothed twin re-eval 0.0325).
      `SMOOTH=0` is now the default. NB: box rebuilt → eval axis shifted (0.0446→0.0325 for the
      same checkpoint); compare within-box only. Full table in RESULTS.md.
- [x] **★ DONE (2026-06-10): α sweep {off, 0.1, 0.25, 0.5} × r{32,64}** — replicated U-shape,
      identical ordering both ranks: off < 0.25 < 0.5 ≪ 0.1. No α beats off; α=0.1 (the
      weight-equalizing extreme) is the WORST (+20% vs off-adjacent wrecon). `SMOOTH=0` is the
      permanent W4A8 default. Full table in RESULTS.md.
- [x] **DONE (2026-06-10): W4A4 3-cell ablation** — r64 nosmooth 0.5103 (catastrophic, 17× the A8
      twin) / r64 α=0.5 0.3885 (smoothing FLIPS to +24% helpful at A4) / r128 α=0.5 **0.3060**
      (rank = activation shield, +21%). Naive per-token A4 not viable. Full table in RESULTS.md.
- [x] **DONE (2026-06-10): W4A4 α-up sweep** — r64 a75 0.2819; r128 a75 **0.2080** (W4A4 best);
      r128 a100 0.2397 (turns — weight damage outruns act relief). **α*≈0.75 for per-token A4**;
      the optimal α tracks the bottleneck. Full table in RESULTS.md.
- [x] **DONE (2026-06-10): per-group activation quant (AGROUP=64)** — the W4A4 fix: r64-ns
      **0.0742** (−85% vs per-token), r128-ns **0.0610** = W4A4 best (~2× the A8 champion).
      Smoothing dead weight again at both ranks → recipe converges to plain+refine, no smooth,
      per-group W+A. Implemented as `a_group` through svdquant.py + scripts/12+13 (config-recorded,
      back-compat: old configs load as per-token). Full 2×2 + qualitative in RESULTS.md.
- [ ] **NVFP4 sim** (explicit, user-requested): FP4 E2M1 element grid + group-16 + FP8-quantized
      scales in `fake_quant_act`/`fake_quant_weight` (~30 lines). This is the Blackwell-native
      deployable format (this box's GPU runs it in silicon). Pair with the AGROUP=16 ladder cell.
- [ ] **(corrected 2026-06-10)** the `jasperai/monet` calib data is NOT Monet paintings — it's a
      diverse photographic set (260/400 captions mention text/signs; 262 mention people/hands).
      So calib content-narrowness is NOT a confound; whitening's 300-calib instability remains
      attributable to Gram sample noise (the 2k re-sweep still tests that). A *probe-aligned*
      eval batch (text-heavy crops scored directly by the metric) is still a nice-to-have, but
      the "calib lacks text" premise was wrong.
- [ ] **Act group-size ladder at A4** — unit test shows ~−25% act-error per halving (g64 0.129 →
      g32 0.099 → g16 0.077): run `AGROUP=16` r128 nosmooth as the sim-only upper bound. NB g64
      = Ada-INT4 deployable; **g16 is only deployable as NVFP4 on Blackwell** → pairs with the
      NVFP4 sim (E2M1 grid + FP8 scales, ~30 lines in `fake_quant_act`/`fake_quant_weight`).
- [ ] (deferred) whiten+refine at A4 (act-aware L as act shield; Gram ~18 GB → tight on 32 GB,
      may need CPU-Gram offload).
- [ ] (deferred) the 2000-img calib re-sweep below now ranks BELOW fixing α — do α first.
- [ ] **Weight group size** {32, 64, 128} at rank 32 — finer = better, ~free in fake-quant.
- [ ] **More refine iters** {5, 8} — all layers hit the iter=3 cap, so the knob isn't saturated.
- [ ] **W4A4** — the Nunchaku sweet spot; map the aggressive corner vs W4A8.

## ⚠ Model states NOT saved — rebuild before any use
Due to the ~250 GB volume quota, every cell's 7.4 GB fake-quant `quant_state.pt` was deleted
after its eval during the 2026-06-10 campaign (losses/configs/montages all kept; builds verified
bit-deterministic on this box). **Currently on disk: ONLY the 3 W4A8 SMOOTH=0 states (incl. the
W4A8 champion r64 = 0.0297). The W4A4 champion and all other cells' weights are gone.**
- [ ] **Rebuild the W4A4 champion when needed:** `ABITS=4 AGROUP=64 bash scripts/run_cell.sh 128
      w4a4g64_nosmooth 0 3 0` (~15 min; expect eval_vel_loss=0.0610 to reproduce).
- [ ] Any other deleted cell rebuilds the same way from its recorded knobs (dir name encodes
      them; `quant_config.json` in each dir is the contract). NB determinism is same-box/stack —
      if the box is rebuilt again, regenerated models follow the same recipe but may not be
      bit-identical and the eval axis may move (re-anchor baselines first, as done 2026-06-10).
- [ ] (optional) `hf sync` the champion states to the HF bucket for durability (needs HF_TOKEN).

## Efficiency / infra
- [ ] **One-time stats cache** — absmax+Gram depend only on (teacher, calib), not rank/α/group;
      cache once, reuse across the whole sweep (one calib pass instead of per-build). Biggest sweep
      speedup; makes the ablation grid + α/group sweeps cheap. Also lets the Gram offload to CPU
      between builds → makes a 24 GB card viable for the whiten path.
- [ ] **Batch the SVD across same-shape layers** — currently 100 sequential per-layer SVDs (the
      dominant cost; a single small SVD never saturates the GPU). The 100 Linears have only a few
      distinct shapes (qkv/mlp repeat across blocks) → stack same-shape layers, do batched
      `svd`/`eigh`/matmul (works for whiten + refine too). Est. 3–5× on decompose. Code change +
      re-validate against locked numbers.
- [ ] **GPU: drop the A100-80GB.** Workload peaks ~18 GB (plain/RTN/refine) / ~28 GB (whiten); the
      SVD loop is latency-bound, not throughput-bound → A100 compute/VRAM mostly idle. Move to a
      **RTX 5090 (32 GB, Blackwell)** — cheaper, fits everything incl. the Gram, AND it's the real
      INT4/FP4 kernel target (A100 sm_80 INT4 is weak, can't measure real speed). 4090 (24 GB) ok for
      no-whiten cells; 48 GB (A6000/L40) if you want whiten headroom with zero code changes.
- [ ] **Per-layer rel-err report** — which layers are hardest to quantize (higher-rank candidates).
- [ ] Ablate calib size (100 vs 2k) and σ coverage (uniform vs the 4 schedule sigmas) on eval-loss.

## Deployment (deferred — needs real kernels, not this A100)
- [ ] **Real low-bit number** on Ada/Blackwell (RTX 40/50) via Nunchaku fused INT4/FP4 — our A100
      (sm_80) under-represents it; fake-quant here is quality-only. Expect ~3× latency, ~3–3.5× mem.
- [ ] **Nunchaku-ready export** — align our decomposition layout/packing so a checkpoint loads in
      Nunchaku (we don't build with DeepCompressor). Verify numerics match.
- [ ] **B200 scale-up** + text-encoder caching + larger batch (remaining fixed-overhead wins).

## Knob summary (what each does) — updated from the 2026-06-01 full grid
- **rank**      = quality knob. ~Free vs a bf16 baseline, but NOT free on real 4-bit kernels: the
                  fused Nunchaku branch still runs 16-bit GEMMs whose share of the (now 3–4× faster)
                  layer is ~5–10% latency at r32 (SVDQuant §4.3/E.5, measured), scaling ~linearly →
                  r64 ≈ 10–20%, r128 ≈ 20–40%. Plus MEMORY (bf16 branch). Pick method per rank
                  (r64→plain+refine no-smooth is the W4A8 champion).
- **bit-width** = the SPEED knob (W4A8 vs W4A4). NOT rank.
- **α / SMOOTH** = activation↔weight outlier split (smoothing). **SMOOTH=0 (off) is the validated
                  default at W4A8** (2026-06-10: beats α=0.5 at every rank); small-α sweep still open.
- **group**     = weight-quant granularity. Untuned (64 default).
- **whiten**    = output-error vs weight-error SVD (eigen-√, robust to non-PD). UNRELIABLE at 300-calib
                  — non-monotonic in rank (overfits the Gram). Only use paired w/ refine or at r32.
- **refine**    = error-feedback iters (3, all layers still hit the cap → not saturated). The RELIABLE
                  upgrade — helps everywhere except r16-plain. Default ON.

## Notes
- Continuity: eval uses the SAME held-out first-16 of `data/monet_cache` as `08_train_recover.py`,
  so quant numbers sit on the same axis as the surgery frontier in `RESULTS.md`.
- Quantization is **orthogonal** to the (shelved) block surgery — in principle they stack, but we
  lead with quant alone since it's the bigger, cleaner lever.