docs/SPEEDUP_IDEAS.md

# Speedup ideation — NVFP4-fused klein-4B on Blackwell (sm_120)

Where we are: the fused NVFP4 model runs at **~1.9× end-to-end / 2.3× per-step** vs bf16, −26% VRAM.
The per-step transformer (2.3×) is already near the per-layer GEMM ceiling (2.5×). The gap to "more"
is NOT the GEMM — it's everything else. This doc is a prioritized menu for the next session. We are
happy to depart from Nunchaku's literature where it pays.

## 0. Where the time actually goes (measure first, always)
From the data: at 512px/4-step the bf16 step is ~92 ms, fused ~42 ms. End-to-end fused 0.264 s of which
4×42=168 ms is denoise → **~96 ms (36%) is NON-transformer: VAE decode + text-encode + latent prep +
launch overhead.** That fixed overhead is the single biggest cap on end-to-end speedup for the small 4B.
**Action 0:** instrument a per-stage timeline (text-encode | per-step{qkv|attn|ff} | vae-decode | python
glue) with CUDA events. Every idea below should be justified against this breakdown.

## Tier 1 — high impact, low risk (do these first)
1. **CUDA Graphs over the denoise loop.** The 4B's GEMMs are small → the path is partly *launch-bound*
   (dozens of tiny kernel launches × 4 steps). Capture the per-step transformer (or the whole 4-step
   loop) into a CUDA graph → eliminate launch latency + CPU dispatch. Biggest, cheapest win at small
   batch/res. Expect 10–30% end-to-end. Risk: dynamic shapes (fixed at capture — fine, res is fixed).
2. **Quantize the text encoder (Qwen3).** It's bf16 now and a real chunk of the fixed overhead. Nunchaku
   ships `svdq-int4-Qwen3-text` (→ NVFP4 on Blackwell). Or cache prompt embeddings across a batch/session
   (encode-once). Cuts the ~text-encode slice to near-zero on repeat. Expect 5–15% end-to-end.
3. **Quantize / fuse the VAE decoder.** Conv-heavy, bf16 now. FP8 convs on Blackwell, or tile+fuse the
   decoder. The VAE is a fixed per-image cost that dominates more as the transformer gets faster.
4. **Fix the fused path's batch=1 limitation.** The packed-rotary asserts `rot.shape[0]*shape[1]==M`;
   make the rotary broadcast over batch (replicate per-batch, or index by token-mod-N). Unlocks batched
   throughput — critical for serving (right now bf16 b4 beats fused b1 on img/s at some sizes). Expect
   large throughput gains at batch>1.

## Tier 2 — kernel-level (medium effort, we build Nunchaku from source now)
5. **Output-quantization fusion (act-quant of layer N+1 in layer N's epilogue).** The kernel already has
   a `qout` output ("packed quantized input for the next layer"). If the FLUX2 path doesn't use it, each
   Linear pays a separate `quantize_w4a4_act_fuse_lora` pass. Fusing the next layer's act-quant into the
   current GEMM's epilogue removes a full activation read/write per layer. Expect 5–15% per-step.
6. **GEMM tile / warp-config autotune for sm_120.** The kernel hardcodes `warp_n=128`, comp tiles, etc.
   GB20x has a specific SM count / register file / L2; sweep `warp_n ∈ {64,128,256}`, K-unroll, and the
   CTA tiling per layer shape (our 5 shapes: 3072², 3072→27648, 12288→3072, 3072→18432, 9216→3072).
   A small-shape vs large-shape split likely wins. Expect 5–20% on the GEMM.
7. **TMA (Tensor Memory Accelerator) loads.** If the kernel still uses cp.async/manual loads, switch the
   weight/activation/scale loads to TMA bulk async copies with multicast — overlaps load with MMA, frees
   registers. Blackwell-native. Expect 5–15% on memory-bound (small-M) layers.
8. **Rank Pareto + mixed rank.** r128 costs ~11% over r64 per-layer for 0.0303 vs 0.0364 quality. Map the
   quality/speed frontier; use **per-layer rank** (high rank only on the few hard layers — we have per-
   layer wrecon). Cuts the low-rank tax where it doesn't buy quality. Expect 3–8% per-step at iso-quality.

## Tier 3 — Blackwell-architecture-specific (higher effort, biggest ceiling)
9. **`tcgen05` 5th-gen MMA + Tensor Memory (TMEM).** Confirm the FP4 GEMM uses Blackwell's `tcgen05.mma`
   with TMEM accumulation and **hardware block-scaling** (`mma.kind::mxf4nvf4` / NVFP4 microscale in the
   MMA itself), not a Hopper-style `wgmma` path. If it's on the old path, porting to tcgen05 is the
   single biggest kernel lever (the 5th-gen cores roughly double FP4 throughput vs emulated scaling).
10. **FP8 attention (the real attention lever).** Attention is `attention_fp16` today and dominates at
    high res (O(N²), bf16/fp16). Move QKᵀ and the AV matmul to **FP8** (E4M3) à la FlashAttention-3 /
    SageAttention on Blackwell → ~2× the attention. At 1024px this is likely the #1 end-to-end win. Keep
    softmax in fp32. Risk: quality on the score matrix — validate on the probes.
11. **FP4 for the AV / V-projection.** Softmax-weighted V is more robust than QKᵀ; FP4 the V path while
    keeping scores FP8. Squeezes attention further. Quality-gated.
12. **FP6 (E3M2) as a quality/speed knob.** Blackwell does FP6 natively. For the hardest layers (or the
    activations), FP6 weights/acts give bf16-ish quality at ~1.5× — a middle rung between NVFP4 and bf16.

## Tier 4 — numerical / format (cheap to try, quality-gated)
13. **Real per-channel `wcscales`.** We set wcscales=1. A true per-output-channel scale tightens the
    NVFP4 weight quant → same quality at **lower rank** → less low-rank tax (compounds with #8).
14. **Activation group size.** NVFP4 acts are group-16. The micro-scale loads have overhead; test g32
    (MXFP4-style) for the activation path where quality allows — fewer scale loads.
15. **Smoothing/rotation for the residual.** We use no-smooth (right for W4A8). For W4A4-NVFP4 a light
    Hadamard rotation of the residual can shrink outliers → smaller residual → more energy into the
    cheap low-rank branch, or lower wrecon at fixed rank.
16. **Fewer steps.** Orthogonal but huge: the model is 4-step distilled; test 3-step / 2-step quality.
    Each step removed is −25%/−50% of denoise. Pure pipeline lever, no kernel work.

## Measurement plan (record for the report)
For each idea: (a) the per-stage CUDA-event breakdown before/after, (b) per-step & end-to-end at
512/1024 × batch{1,2,4}, (c) VRAM, (d) image-quality check on the 4 probes (text + hand are the
sensitive ones), (e) the velocity-loss if it changes the math. Always vs the bf16 baseline AND the
current NVFP4-fused baseline. Keep the matrix in `outputs/nvfp4/benchmark.json`.

## Rough priority (impact × 1/effort)
CUDA graphs (1) · TE/VAE quant (2,3) · batch fix (4) · output-fusion (5) — first.
Then FP8 attention (10) and tcgen05 verification (9) for the high-res ceiling.
Then tile autotune (6), TMA (7), rank Pareto (8), wcscales (13).