Add optimized_w3: 3-bit LUT GEMM on a new mainloop

tests/optimized_w3/DESCRIPTION.md

@@ -0,0 +1,153 @@
+# optimized_w3 — a faster 3-bit LUT decode GEMM for FLUTE
+
+A drop-in **3-bit LUT decode GEMM** for FLUTE that keeps FLUTE's learned
+non-uniform quant map and offline-pack model, but replaces the CuTe mainloop with
+a memory-bound-optimal one. Across the **entire M = 1–16 decode band** it is
+**1.1–1.4× faster than FLUTE's best-tuned 3-bit kernel** while being **~2×
+tighter in accuracy** — measured with FLUTE's own `do_bench` on an A100. It lifts
+DRAM utilisation from **~36% → ~68%**, i.e. onto the same bandwidth roofline the
+W4 kernels hit.
+
+Self-contained under `tests/optimized_w3/`: CUDA kernel + PyTorch binding, offline
+packing, learned-qmap prmt-LUT, data-driven per-shape dispatch, an autotuner, a
+correctness+bench test, and docs (README, results, KERNEL_ANATOMY).
+
+---
+
+## Why this exists
+
+Decode-phase GEMM at batch M ≤ 16 is **memory-bound**: it streams far more weight
+bytes than it does math, so the roofline is DRAM bandwidth and *any dequant work
+not hidden under memory traffic is pure overhead on the critical path.* FLUTE's
+3-bit kernel dequantises via a two-plane bit-slice unpack **plus a shared-memory
+paired-LDS** through the learned map, on a thin-tiled loop that hides latency by
+running two CTAs/SM. That unpack + smem round-trip sits *on* the critical path
+and the thin tile feeds each weight to only a few `mma`, so the kernel stalls
+issue-bound at **~36% DRAM** — the dequant is *added to* memory time, not hidden
+under it.
+
+optimized_w3 keeps the two things FLUTE gets right — the **learned non-uniform
+qmap** (its accuracy edge) and the **offline-pack model** — and rebuilds the
+mainloop so dequant is **hidden under memory** (`time ≈ max(mem, compute)`),
+putting the kernel on the DRAM roofline at **~68%**.
+
+---
+
+## Core contributions (what's novel here)
+
+The memory-bound mainloop *structure* is adapted from **Marlin** (Frantar et al.,
+IST-DASLab — see *Citation*), which showed a low-M mixed-precision GEMM can be
+held at near-peak DRAM bandwidth by running **one fat CTA per SM** with a deep
+prefetch pipeline and a large register accumulator, instead of relying on
+multi-CTA occupancy. Marlin's fast path, however, dequantises **uniform** W4
+levels with a `lop3` magic-bias arithmetic trick that **does not apply to a
+learned, non-uniform map**. The novel contributions here carry that learned-LUT
+case onto the memory-bound mainloop:
+
+1. **In-register `prmt`-LUT dequant for a learned map.** The 8-entry learned qmap
+   is held in registers as two fp16 byte-planes; dequant is a `prmt` (byte
+   permute) that uses the fetched code word *itself* as the selector — **8 values
+   = 1 shift + 8 `prmt`, zero shared-memory traffic.** `prmt` *is* the table
+   lookup, so FLUTE's non-uniform accuracy is preserved while FLUTE's smem
+   paired-LDS is eliminated entirely. This is the piece Marlin's arithmetic path
+   can't express.
+
+2. **Nibble + mma-fragment offline pack.** Each 3-bit code is packed into its own
+   nibble, so the word loaded from DRAM *is already* the `prmt` selector — **no
+   unpack ALU on the critical path** (vs FLUTE's cross-plane shift/mask/recombine).
+   The packer also lays weights in `mma.sync` fragment order (`_base_perm`), so
+   the weight tile needs **no `ldmatrix` transpose**: after the `cp.async` stage
+   into shared memory, a plain vectorized 16-byte `ld.shared` drops each lane's
+   bytes **directly into its `mma` B-fragment registers** (only the A tile still
+   pays an `ldmatrix`). The weights still transit the `cp.async` pipeline buffer —
+   that staging *is* the latency-hiding mechanism; what's eliminated is the
+   `ldmatrix` and the unpack ALU. The one smem round-trip actually removed is
+   FLUTE's dequant-table LDS, which contribution 1 serves from registers instead.
+   Trade-off, stated honestly: a nibble spends 4 bits on a 3-bit code, so the
+   weight stream is ~25% larger than FLUTE's 3.2-bit bit-slice — paid back many
+   times over by moving DRAM utilisation 36% → 68% (net ~1.5× on the widest shape).
+
+3. **FLUTE-qmap integration on the fat-tile mainloop.** Wiring the learned map,
+   group scales, and offline pack through Marlin's fat register tile + 4-stage
+   `cp.async` pipeline so the dequant overlaps the weight loads. The **fat
+   register accumulator** is the lever that moves the roofline (36% → 68% DRAM),
+   but the mechanism depends on where you are in the band: as M grows, the M loop
+   is innermost so each dequantised weight is **reused across all `thread_m_blocks`
+   `mma`** (classical arithmetic intensity). At the bottom of the band (M ≤ 16,
+   `thread_m_blocks = 1`) there is **no such reuse** — there the 36% → 68% comes
+   from **memory-level parallelism** (the 4-stage `cp.async` keeps 4 weight tiles
+   in flight, saturating DRAM by request depth instead of by FLUTE's occupancy
+   switching) plus **freeing issue slots**: with the smem-LDS, unpack, and
+   `ldmatrix` all gone (contributions 1–2), the SM spends its issue bandwidth on
+   *loads* rather than dequant bookkeeping — which is exactly what an issue-capped
+   memory-bound kernel needs. Either way DRAM bandwidth, not instruction issue,
+   becomes the bound.
+
+**Adopted from Marlin (attributed, Apache-2.0):** one CTA/SM at 256 threads, the
+4-stage `cp.async` pipeline with fully-unrolled static addressing, striped
+**Stream-K** with in-L2 serial reduction, and L2 evict-first streaming of the
+single-use weights.
+
+### It's a co-design, not one knob
+The same pipeline-depth knob has **opposite sign** on the two kernels
+(28672×8192, M=16): this kernel at `Stages=2` is *slower* than FLUTE (165 µs vs
+122 µs), and FLUTE at `Stages=4` is slower than FLUTE at 2. Only the **fat-tile ×
+deep-pipe product** crosses (→ 88 µs). There is no incremental config-space path
+from FLUTE — the crossing is the whole mainloop rewrite, which is why it ships as
+a separate kernel.
+
+---
+
+## Results — full M = 1–16 decode band (A100-SXM4-80GB)
+
+`triton.testing.do_bench` (L2 flush, rep=100 — FLUTE's own tuning method), vs
+**FLUTE's best 3-bit template per shape** (autotuned, not just tid20). Both
+kernels use identical output allocation (apples-to-apples). The speedup is
+**flat across M = 1–16** because both kernels are weight-bandwidth-bound and the
+`M×K` activation is negligible — so the decode band is won *uniformly*, not just
+at one point.
+
+| shape (N×K) | FLUTE (µs) | optimized_w3 (µs) | **speedup, M=1–16** | rel-err (FLUTE → opt) |
+|---|---|---|---|---|
+| 28672×8192  | 125–126 | 88–90 | **1.39–1.42×** | 4.2e-4 → **2.5e-4** |
+| 8192×8192   |  46–50  | 35–37 | **1.28–1.33×** | 6.3e-4 → **2.8e-4** |
+| 14336×4096  |  41–42  | 32–33 | **1.27–1.29×** | 5.1e-4 → **2.9e-4** |
+| 4096×4096   |  22–22  | 20–20 | **1.09–1.12×** | 6.2e-4 → **3.1e-4** |
+
+Per-M detail (speedup at M ∈ {1, 2, 4, 8, 16}):
+
+| shape | M=1 | M=2 | M=4 | M=8 | M=16 |
+|---|---|---|---|---|---|
+| 28672×8192 | 1.42× | 1.42× | 1.41× | 1.41× | 1.39× |
+| 8192×8192  | 1.33× | 1.30× | 1.28× | 1.29× | 1.29× |
+| 14336×4096 | 1.29× | 1.27× | 1.28× | 1.27× | 1.28× |
+| 4096×4096  | 1.11× | 1.09× | 1.12× | 1.10× | 1.10× |
+
+- **Widest-N wins most.** 28672×8192 is the most memory-bound (largest
+  weight-bandwidth share), so it gets the biggest lift — ~88 µs at ~68% DRAM, the
+  W4 roofline.
+- **~2× tighter accuracy, everywhere.** Same learned qmap on both sides, but the
+  dequant here rounds to fp16 before the accumulate → relative error is roughly
+  **half** FLUTE's own 3-bit output on every shape. Every autotune/test candidate
+  is exactness-gated (`err ≤ 2e-3` and must not regress vs FLUTE) *before* it is
+  allowed to be timed.
+
+Reproduce: `python -m tests.optimized_w3.test_correctness` (A100). Numbers above
+are a fresh `do_bench` sweep; run-to-run variance is ±~0.03× on the ratio.
+
+---
+
+## Citation
+
+The mainloop structure — one CTA/SM, deep `cp.async` pipeline, fat register
+tiles, striped Stream-K, static unrolling — derives from **Marlin**:
+
+> E. Frantar, R. L. Castro, J. Chen, T. Hoefler, D. Alistarh.
+> *MARLIN: Mixed-Precision Auto-Regressive Parallel Inference on Large Language
+> Models.* IST-DASLab. https://github.com/IST-DASLab/marlin (Apache-2.0).
+
+The Apache-2.0 license header is preserved in `csrc/optimized_w3_kernel.cu`. The
+**contributions in this package** are the W3-specific pieces Marlin's uniform-W4
+path does not cover: the in-register `prmt`-LUT dequant for a *learned,
+non-uniform* map, the nibble + `mma`-fragment offline pack, and the FLUTE-qmap
+integration on the fat-tile mainloop.

tests/optimized_w3/KERNEL_ANATOMY.md

@@ -0,0 +1,135 @@
+# optimized_w3 — Kernel Anatomy
+
+How the kernel works, what it changes vs FLUTE's 3-bit production kernel, and why
+those changes pay off. For the measured numbers see `results.md`.
+
+**Problem.** Decode GEMM `C[M,N] = A[M,K] · dequant(Wq[K,N])`, weights 3-bit,
+group-quantised scales, **M ≤ 64** (memory-bound). FLUTE stores weights `(K, N)`.
+
+**One-line idea.** Keep FLUTE's two levers of value — the **learned non-uniform
+8-entry qmap** (its accuracy) and the **offline-pack model** — but replace the
+CuTe mainloop with a **Marlin-derived mainloop** so the dequant is *hidden under*
+memory instead of *added to* it. DRAM utilisation goes ~36% → ~68%.
+
+---
+
+## 1. FLUTE's 3-bit production kernel (the baseline, config tid20)
+
+`SMs_Multiple=2` (216 blocks, **2 CTA/SM**), **128 threads/CTA**, **2 pipeline
+stages**, thin tiles (TileM16/TileK64/TileP32). Dequant = **two-plane bit-slice
+unpack + a shared-memory paired-LDS through `qmap2`**. Hides latency by
+**occupancy** (the 2nd CTA issues while the 1st stalls).
+
+```
+kernel flute_qgemm(A, Bq, S, qmap2, ...):          # 128 threads, 2 CTA/SM
+    for k_tile in streamk_range():
+        cp.async(smemA, smemB <- next tile)         # 2-deep
+        for packed 32b word w in smemB:             # 10 codes / 32b, TWO PLANES
+            code   = unpack_two_plane(w)            # shift+mask+recombine across planes
+            b_frag = LDS qmap2[code]                # shared-mem paired LDS
+            b_frag = hmul2(b_frag, scale)
+        accum += TiledMMA(smemA, b_frag)            # thin N-tile: few mma / weight
+    streamk_reduce(accum)
+```
+
+Result: 36% DRAM. The per-code unpack + smem LDS sits **on the critical path**,
+and the thin N-tile feeds each dequanted weight to only a few `mma` — so the
+kernel is *work/issue-bound*, dequant time is **added** to memory time.
+
+> FLUTE is **not** a fixed 2-stage kernel: it ships 36 3-bit configs spanning
+> Stages {2,3,4,5}, Threads {128,256}, SMs_Multiple {1,2}, and autotunes per
+> shape. tid20 is the config its tuner selects for the decode band (measured:
+> deeper stages are monotonically worse there — a deeper pipe costs smem and
+> adds no MLP in this memory-bound regime). We benchmark against FLUTE's *best*
+> config per shape, not just tid20.
+
+## 2. optimized_w3 (this kernel)
+
+`sms` blocks (**1 CTA/SM**), **256 threads/CTA**, **4 pipeline stages**, fat
+register tiles. Dequant = **in-register prmt-LUT**. Hides latency by **ILP
+inside one CTA** (deep pipe + register tiles + fully-unrolled static addressing).
+
+```
+kernel optimized_w3(A, Bq, S, lut, ...):            # 256 threads, 1 CTA/SM
+    FragC accum[thread_m_blocks][4][2] = 0          # FAT register tile
+    fill 4-stage cp.async pipe
+    for k_tile in striped_streamk_range():          # fully unrolled body
+        cp.async(smemA, smemB <- k+4)               # 4-deep prefetch
+        ldmatrix frag_a <- smemA[k]
+        for w in smemB[k]:                          # 8 codes / 32b, ONE plane (nibbles)
+            frag_b = dequant_w3(w, lut)             # 2 prmt from registers, no smem
+            frag_b = scale(frag_b, S_group)
+        for mb, n: accum[mb][n] += mma(frag_a[mb], frag_b[n])   # many mma / weight
+    serial_in_L2_reduce(accum)                       # lock-stepped across column slices
+```
+
+```
+dequant_w3(sel /*fetched nibble word == prmt selector*/, L /*qmap as byte planes*/):
+    lo = prmt(L.lo0, L.lo1, sel)     # low bytes of qmap[code] for 4 codes
+    hi = prmt(L.hi0, L.hi1, sel)     # high bytes
+    return interleave(lo, hi)        # 8 values = 1 SHR + 8 prmt, zero smem
+```
+
+Result: 68% DRAM. dequant + mma live in one CTA's ILP, overlapped with the
+4-deep cp.async; the fat tile amortises each weight load over many mma. Dequant
+is **hidden under** memory (`time ≈ max(mem, compute)`).
+
+---
+
+## 3. What changed, and why each is necessary
+
+| # | Change | vs FLUTE | Role |
+|---|--------|----------|------|
+| 1 | **Fat register tile** (many mma / weight load) | thin TileP=32 | **the lever** — raises arithmetic intensity so DRAM (not issue) bounds → 36%→68% |
+| 2 | **In-register prmt-LUT dequant** | smem paired-LDS | removes the smem round-trip for the dequant table |
+| 3 | **Nibble pack + `_base_perm` fragment order** | two-plane, flat | removes per-code unpack ALU and the smem round-trip for B |
+| 4 | **Static unroll + 256 thr / 1 CTA/SM** | dynamic, 128 thr / 2 CTA | ILP + register/smem budget that makes #1 runnable |
+
+**Kept, deliberately:** FLUTE's learned qmap. Because it is *non-uniform*, dequant
+must be a **table lookup** — Marlin-W4's lop3 magic-bias arithmetic only works for
+uniform levels. So we keep the table but serve it via **prmt in registers** (#2)
+instead of smem. `prmt` *is* the lookup; qmap is what it reads from.
+
+**Came free with the mainloop** (already present, no work): striped Stream-K +
+in-L2 serial reduction, L2 evict-first streaming of the single-use weights, and
+scale double-buffering.
+
+### Why it's a co-design, not one knob (measured, single variable)
+The same pipeline knob has **opposite sign** on the two kernels (28672×8192, M16):
+
+| Stages | FLUTE | optimized_w3 |
+|---|---|---|
+| 2 | **122µs (best)** | 165µs (*slower than FLUTE*) |
+| 4 | 134µs (worse) | **88µs (best)** |
+
+optimized_w3 at Stages=2 **loses to FLUTE** — the fat-tile structure without the
+deep pipe loses, and the deep pipe without the structure (on FLUTE) also loses.
+Only the **fat-tile × deep-pipe product** crosses. That is why no incremental
+config-space path exists between the two kernels (FLUTE's Threads, MMA-layout and
+pack are static-asserted as one welded unit): the crossing is the mainloop rewrite.
+
+---
+
+## 4. Tuning granularities
+
+| Granularity | Knob | FLUTE | optimized_w3 | Tunable here |
+|---|---|---|---|---|
+| Grid / SM | block count / `sms` | 216 (2/SM) | `sms` (1/SM) | ✅ `sms` in dispatch (≤ device SMs) |
+| CTA | threads/CTA | 128 | 256 | fixed (kernel) |
+| Tile | (thread_k, thread_n) | TileP-locked | (128,128) or (64,256) | ✅ autotuned per shape |
+| Pipeline | stages | 2–5 (tuner) | 4 | fixed (4 is best, measured) |
+| Pack | layout | two-plane 16B | nibble 16B + perm | fixed |
+| Dequant | mechanism | smem LDS | register prmt | fixed |
+
+The autotuned surface for this kernel is **(thread_k, thread_n) × sms per
+(M, N, K)** — see `dispatch.py` / `autotune.py`. Wide-N/deep-K shapes want the fat
+`(64,256)` tile; the rest want `(128,128)`; `sms=108` (or `-1` auto) generally
+best on A100. `sms` must never exceed the device SM count — an oversubscribed grid
+breaks the striped lock-order reduction.
+
+---
+
+## Provenance
+The mainloop derives from Marlin (Frantar et al., IST-DASLab, Apache-2.0; see the
+license header in `csrc/`). The 3-bit prmt-LUT dequant, the nibble+`_base_perm`
+pack, and the FLUTE-qmap integration are the additions here.

tests/optimized_w3/README.md

@@ -0,0 +1,54 @@
+# optimized_w3
+
+A drop-in **3-bit LUT GEMM** for the FLUTE decode path (M ≤ 64). Keeps FLUTE's
+learned quant map and offline-pack model, but runs on a **Marlin-derived
+mainloop** (deep cp.async pipeline, striped Stream-K, fat register tiles,
+in-register prmt-LUT dequant). On the memory-bound decode band this raises DRAM
+utilisation ~36% → ~68% and gives **1.1–1.6× over FLUTE's best 3-bit template**,
+at strictly tighter accuracy.
+
+- **[results.md](results.md)** — optimizations and measured numbers.
+- **[KERNEL_ANATOMY.md](KERNEL_ANATOMY.md)** — how it works and why.
+
+## Layout
+```
+optimized_w3/
+├── __init__.py           public API
+├── qgemm.py              pack_w3 / pack_scales / lut_planes / qgemm_w3 + dispatch
+├── autotune.py           per-shape autotuner -> writes dispatch.py
+├── dispatch.py           tuned (M,N,K) -> (thread_k, thread_n, sms) table (data)
+├── test_correctness.py   exactness vs fp64 + smoke bench (both harnesses)
+├── csrc/
+│   ├── optimized_w3_kernel.cu    Marlin-derived mainloop + W3 prmt-LUT dequant
+│   └── optimized_w3_binding.cpp
+├── README.md · results.md · KERNEL_ANATOMY.md
+```
+
+## Usage
+```python
+import torch
+from tests.optimized_w3 import pack_w3, pack_scales, lut_planes, qgemm_w3
+
+N, K, M, group = 8192, 8192, 16, 128
+Wp  = pack_w3(codes, N, K)                    # codes: (K, N) ints 0..7
+Sp  = pack_scales(scales, N, K)               # scales: (N, K/group) fp16
+lut = lut_planes(qmap)                        # 8-entry fp16 qmap; once per layer
+ws  = torch.zeros(N // 128 * 16, dtype=torch.int32, device="cuda")
+out = qgemm_w3(A, Wp, Sp, lut, ws)            # A: (M, K) fp16 -> (M, N) fp16
+```
+The launch config is dispatched from `dispatch.py` by shape; pass
+`thread_k/thread_n/sms` explicitly to override.
+
+## Build env (A100)
+The extension JIT-compiles on first import and needs a libstdc++ providing
+`GLIBCXX_3.4.29` plus CUDA:
+```bash
+export LD_LIBRARY_PATH=/orcd/software/core/001/spack/pkg/gcc/12.2.0/yt6vabm/lib64:$LD_LIBRARY_PATH
+export CUDA_HOME=/orcd/software/core/001/pkg/cuda/12.9.1
+python -m tests.optimized_w3.test_correctness
+```
+
+## Provenance
+The mainloop derives from Marlin (Frantar et al., IST-DASLab, Apache-2.0 — see
+the license header in `csrc/`). The 3-bit prmt-LUT dequant, the nibble +
+fragment-order pack, and the FLUTE-qmap integration are the additions here.

tests/optimized_w3/__init__.py

@@ -0,0 +1,11 @@
+"""optimized_w3 -- 3-bit LUT GEMM on a Marlin-derived mainloop for FLUTE decode.
+
+Public API:
+    pack_w3(codes, N, K)          offline weight pack (nibble + fragment perm)
+    pack_scales(scales, N, K)     offline scale pack
+    lut_planes(qmap)              learned 8-entry qmap -> prmt byte planes
+    qgemm_w3(A, Wp, Sp, lut, ws)  the matmul (shape-dispatched launch config)
+"""
+from .qgemm import pack_w3, pack_scales, lut_planes, qgemm_w3, best_config
+
+__all__ = ["pack_w3", "pack_scales", "lut_planes", "qgemm_w3", "best_config"]