docs: Update kernel spec for format unification and dispatch

Remove grouped scalar GEMV references (kernel removed in ac7d6ff).
Update five-kernel → four-kernel strategy. Document tiled-only
runtime format and kbit_linear/kbit_expert_linear dispatch functions.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

benchmarking-report.md

@@ -5,16 +5,17 @@ All kernel times are NCU `gpu__time_duration.avg` unless stated otherwise.
 
 ## Kernel dispatch
 
-Five kernels cover the full inference workload. Dispatch selects the fastest
-kernel per (layer_type, M) pair:
+Four kernels cover the full inference workload. `kbit_linear` and
+`kbit_expert_linear` dispatch to the fastest kernel per (layer_type, M):
 
 | Kernel | M range | Layers | Status |
 |--------|---------|--------|--------|
 | Scalar GEMV | 1-4 | Dense + attention | Done (V8), 1.5-1.9x faster than fp16 at M=1 |
 | MMA dequant | 5-16 | Dense + attention | Done, ~1.0-1.3x vs fp16 |
 | Dequant + cuBLAS | 17+ | Dense + attention | Done, ~0.95-1.0x vs fp16 |
-| Grouped scalar GEMV | 1-4 | MoE experts | Done, competitive with fp16 |
-| Grouped MMA | 5+ | MoE experts | Done, competitive with fp16 |
+| Grouped MMA | 1-16 | MoE experts | Done, competitive with fp16 |
+
+All kernels read tiled format (from `repack_kbit`) with E4M4 absmax.
 
 ## Per-shape speedups at M=1 (decode, dominant workload)
 
@@ -166,10 +167,10 @@ kernel launches (dequant + matmul), doubling the dispatch tax.
    end-to-end throughput by up to 1.5x on top of the current kernel
    speedups.
 
-4. **MoE grouped kernels need V8 optimizations.** The grouped scalar GEMV
-   currently matches fp16 but does not beat it. Porting the V8 inner loop
-   (vectorized A loads, 2-warp config, M-dispatch) would bring it closer
-   to the 1.5-1.9x speedups seen on dense layers.
+4. **MoE dispatch is unified.** The grouped MMA handles M<=16 for MoE
+   layers; for larger M, `kbit_expert_linear` falls back to per-expert
+   dequant + cuBLAS matmul. The grouped scalar GEMV was removed (it only
+   won one shape at M=1 by 0.3 us).
 
 5. **Lower k is strictly better for inference speed, not just model size.**
    k=2 is fastest at every M value because it reads the least data. The

deployment-summary.md

@@ -14,18 +14,19 @@ is 80-84% of total GEMM wall-clock time in typical sessions.
 At 16+ concurrent users, the advantage disappears because large prefill
 chunks dominate and the dequant overhead exceeds the bandwidth savings.
 
-The system uses **5 CUDA kernels** dispatched per (layer_type, M):
+The system uses **4 CUDA kernels** dispatched by `kbit_linear` and
+`kbit_expert_linear` per (layer_type, M). All kernels read tiled
+format (from `repack_kbit`) with E4M4 absmax:
 
 | Kernel | M range | Layers | Mechanism |
 |--------|---------|--------|-----------|
 | Scalar GEMV | 1-4 | Dense + attn | 64 threads, shuffle codebook, no tensor cores |
 | MMA dequant | 5-16 | Dense + attn | Tensor core m16n8k16, inline dequant |
 | Dequant + cuBLAS | 17+ | Dense + attn | Separate dequant kernel → cuBLAS GEMM |
-| Grouped scalar GEMV | 1-4 | MoE experts | Same as scalar, batched across experts |
-| Grouped MMA | 1+ | MoE experts | Same as MMA, batched across experts |
+| Grouped MMA | 1-16 | MoE experts | Same as MMA, batched across experts |
 
-For MoE layers at large M (prefill), the grouped MMA kernel loses to
-fp16 BMM, so a hybrid dequant + cuBLAS BMM path is available.
+For MoE layers at max_M > 16 (prefill), `kbit_expert_linear` falls
+back to per-expert dequant + cuBLAS matmul.
 
 ---
 
@@ -191,40 +192,6 @@ fp16, the same model requires 140 GB (two H100s or four 4090s).
 
 ---
 
-## Grouped scalar GEMV: where it fits
-
-The grouped scalar GEMV (`kbit_grouped_scalar_gemv`) is a specialized
-kernel for MoE expert layers at M=1-4. It uses the same flat data format
-and shuffle codebook as the dense scalar GEMV.
-
-### When it wins
-
-Only for **moe_gu (K=2048, N=512) at M=1** — and barely:
-
-| Shape | M | Grouped scalar | Grp MMA | fp16 BMM | Winner |
-|-------|---|---------------|---------|----------|--------|
-| moe_gu | 1 | **11.3** | 11.6 | 11.7 | Grouped (by 0.3 us) |
-| moe_gu | 2 | 12.9 | **11.8** | 12.7 | Grp MMA |
-| moe_gu | 4 | 17.1 | **11.9** | 18.9 | Grp MMA |
-| moe_dn | 1 | 24.9 | **12.1** | 13.1 | Grp MMA |
-| moe_dn | 4 | 38.3 | **12.1** | 12.1 | Grp MMA |
-
-The grouped scalar is terrible on moe_dn (K=512): with only 512/64=8
-quant blocks per thread and C=1 (one column per block), the kernel is
-launch-overhead-dominated. The grouped MMA wins everywhere except that
-one moe_gu M=1 case.
-
-### Why it still exists
-
-1. It uses the flat data format (from `quantize_kbit` directly), no
-   repack step. If you only store weights in flat format, the grouped
-   scalar is the only MoE option at M=1-4.
-2. The moe_gu M=1 win is small but real in the most common workload
-   (single-user decode). Over thousands of layers, 0.3 us adds up.
-3. It provides a correctness cross-check against the grouped MMA.
-
----
-
 ## Remaining optimization opportunities
 
 ### 1. CUDA Graphs for hybrid path (medium impact, low effort)

kbit-kernel-spec.md

@@ -122,20 +122,24 @@ than fp16 is at ~16 concurrent users.
 
 ---
 
-## Five-kernel strategy
+## Four-kernel strategy
 
 Each kernel covers a range of M where it has a structural advantage.
-The dispatch logic selects the best kernel per (layer_type, M) pair.
+The dispatch logic (`kbit_linear`, `kbit_expert_linear`) selects the
+best kernel per (layer_type, M) pair. All kernels read tiled format
+(from repack_kbit) with E4M4 absmax.
 
-| Kernel | M range | Layer types | Data format |
-|--------|---------|-------------|-------------|
-| 1. Scalar GEMV | 1-4 | Dense, attention | Flat (quantize_kbit), float32 absmax |
-| 2. MMA dequant | 5-16 | Dense, attention | Tiled (repack_kbit), E4M4 absmax |
-| 3. Dequant + cuBLAS | 17+ | Dense, attention | Flat -> fp16 |
-| 4. Grouped scalar GEMV | 1-4 | MoE experts | Flat (quantize_kbit), float32 absmax |
-| 5. Grouped MMA | 1+ | MoE experts | Tiled (repack_kbit), E4M4 absmax |
+| Kernel | M range | Layer types | Dispatch function |
+|--------|---------|-------------|-------------------|
+| 1. Scalar GEMV | 1-4 | Dense, attention | `kbit_linear` |
+| 2. MMA dequant | 5-16 | Dense, attention | `kbit_linear` |
+| 3. Dequant + cuBLAS | 17+ | Dense, attention | `kbit_linear` |
+| 4. Grouped MMA | 1-16 | MoE experts | `kbit_expert_linear` |
 
-Why five kernels instead of one:
+For MoE at max_M > 16, `kbit_expert_linear` falls back to per-expert
+dequant + cuBLAS matmul (no dedicated kernel needed).
+
+Why four kernels instead of one:
 - At M=1, tensor cores waste 94% of their compute (m16n8k16 pads 15
   zero rows). A scalar kernel that avoids MMA entirely wins by 3-5x.
 - At M=5-16, MMA utilization rises to 31-100%. The 3.2x data
@@ -147,10 +151,6 @@ Why five kernels instead of one:
   its compute pipeline.
 - MoE experts launched individually waste 88-97% of SMs. Grouping
   all active experts into one kernel launch solves this.
-- The grouped scalar GEMV and grouped MMA serve complementary roles:
-  scalar wins at M=1-4 for moe_gu (K=2048, N=512) where its C=1
-  grid gives better parallelism; grouped MMA wins at all M for
-  moe_dn (K=512, N=2048) and at M>4 for moe_gu.
 
 **Practical importance (from workload analysis in `token_analysis.md`):**
 
@@ -194,17 +194,18 @@ range falls in the gap between these modes.
 - No shared memory for B data, no cp.async, no split-K
 
 **Data format:**
-- B_packed: flat from `quantize_kbit` — `[N * num_k_blocks * k]` uint32
-- B_absmax: flat float32 — `[N * num_k_blocks]`
-- No repack step needed
+- B_packed: tiled from `repack_kbit` — tiles of `[TILE_N × KB_PER_TILE × k]` uint32
+- B_absmax: tiled uint8 E4M4 — tiles of `[TILE_N × KB_PER_TILE]`
+- Supports both flat and tiled layouts via `TILED` template bool
+- Flat layout preserved for standalone use; tiled layout used by dispatch
 
 **Inner loop (V8):**
 
 Each thread strides through quantization blocks along K:
 ```
 for each quant block (stride 64):
     load k bit-plane words (vectorized: int2 for k=2, int4 for k=4)
-    load float32 absmax
+    load absmax (E4M4 → float decode)
 
     for sub = 0..3:            // 4 groups of 8 elements
         load A[m, k_base + sub*8 .. +7] via int4 (8 fp16 values)
@@ -329,12 +330,13 @@ the MMA dequant kernel takes ~68 us (instruction-limited, only 1.3%
 of execution is MMA). A fused dequant kernel would take ~5 us for
 this shape, so dequant + cuBLAS ~27 us would beat 68 us.
 
-**Dequant kernel** (`kDequantizeBlockwise_kbit_vec`): a single CUDA
-kernel that reads k-bit packed data + absmax and writes fp16 output.
-Templated on absmax type: float32 (from `quantize_kbit` directly),
-uint8 E4M4, or fp16. The float32 absmax path was added to eliminate
-a previous Python-side E4M4 conversion that launched ~15 PyTorch
-elementwise kernels (~800 us). Now it is a single kernel launch.
+**Dequant kernel:** Two variants:
+- `kDequantizeBlockwise_kbit_vec`: reads flat layout (from quantize_kbit)
+- `kDequantizeBlockwise_kbit_tiled`: reads tiled layout (from repack_kbit)
+
+Both are templated on absmax type (uint8 E4M4, fp16, float32) and
+output type (fp16, bf16, float32). The tiled variant is used by
+`kbit_linear` dispatch for the M>16 dequant+cuBLAS path.
 
 Dequant GPU kernel times (ncu-measured, k=4):
 
@@ -352,51 +354,15 @@ to the matmul. At M>=64, dequant+cuBLAS wins because cuBLAS scales
 efficiently while MMA is instruction-limited. The crossover is
 M=32-64 depending on shape.
 
-**Data format:** Uses flat layout (same as scalar GEMV). The
-`dequantize_kbit` launcher handles float32, uint8 E4M4, and fp16
-absmax via the `_KBIT_ABSMAX_SUFFIX` dispatch map.
-
----
-
-## 4. Grouped scalar GEMV (`kbit_grouped_scalar_gemv`)
-
-**Location:** `ops.cu` (search for `kbit_grouped_scalar_gemv`)
-
-**Operation:** For each expert e: C_e[M_e, N] = A_e[M_e, K] * W_e^T,
-all experts in one kernel launch.
-
-**Architecture (V8):**
-- 64 threads (2 warps), one output column per block (C=1)
-- Grid = (N, num_experts) — Y-dimension indexes experts
-- `__launch_bounds__(64, 24)` for M<=2, `__launch_bounds__(64, 16)` for M>2
-- M_VAL dispatch (1/2/3/4 templates)
-
-**Data format:**
-- B_packed_all: flat from `quantize_kbit` — concatenated per-expert,
-  each `[N * num_k_blocks * k]` uint32 (truncated to exact size)
-- B_absmax_all: flat float32 — concatenated per-expert,
-  each `[N * num_k_blocks]` float32 (truncated to exact size)
-- No repack step needed. Uses same flat layout as the dense scalar GEMV.
-
-**Inner loop:** Identical to the dense scalar GEMV (V8): vectorized
-int4 A loads, 4-group sub-loop of 8 elements, shuffle codebook lookup.
-The only difference is per-expert pointer arithmetic using
-`expert_offsets[expert_id]` to find each expert's A, B, and C regions.
-
-**Why grouped scalar wins for moe_gu (K=2048, N=512) at M<=4:**
-With C=1, the grid is N × num_experts = 512 × 8 = 4096 blocks. This
-gives full SM utilization (32 blocks/SM). The grouped MMA at this shape
-has far fewer blocks due to tiling overhead.
-
-**Quantize_kbit padding:** `quantize_kbit` appends a small padding
-(4 packed words + 1 absmax) to each expert's output. The test and
-benchmark helpers truncate each expert's data to the exact expected
-size before concatenation, so the kernel's arithmetic indexing
-(`expert_id * N * num_k_blocks * K_BITS`) works correctly.
+**Data format:** The flat variant (`dequantize_kbit`) reads flat layout
+from `quantize_kbit`. The tiled variant (`dequantize_kbit_tiled`) reads
+tiled layout from `repack_kbit`, used by `kbit_linear` dispatch.
+Both handle float32, uint8 E4M4, and fp16 absmax via the
+`_KBIT_ABSMAX_SUFFIX` dispatch map.
 
 ---
 
-## 5. Grouped MMA (`kbit_grouped_gemm_prod`)
+## 4. Grouped MMA (`kbit_grouped_gemm_prod`)
 
 **Location:** `ops.cu` (search for `kbit_grouped_gemm_prod`)
 
@@ -486,33 +452,31 @@ targets < 10 us.
 
 ## Data formats
 
-Two formats exist, and which kernel uses which matters:
+All inference kernels read **tiled format** (from `repack_kbit`).
+Flat format exists only as the intermediate output of `quantize_kbit`
+before repacking.
 
-**Flat (from `quantize_kbit`):**
+**Flat (from `quantize_kbit`) — intermediate only:**
 - B_packed: `[N * num_k_blocks * k]` uint32, row-major per column
-- B_absmax: `[N * num_k_blocks]` float32
-- No preprocessing. Used by: scalar GEMV, grouped scalar GEMV,
-  dequant kernel.
+- B_absmax: `[N * num_k_blocks]` float32 or uint8 E4M4
+- Used only during quantization. Converted to tiled by `repack_kbit`
+  at model load time, then discarded.
 
-**Tiled (from `repack_kbit`):**
+**Tiled (from `repack_kbit`) — runtime format:**
 - B_packed: reorganized into `[k_tiles * n_tiles * TILE_N * B_COL_WORDS]`
   for coalesced cp.async loads per tile
 - B_absmax: E4M4-encoded uint8, same tiled layout
-- Requires a one-time repack pass. Used by: MMA kernel, grouped MMA
-  kernel.
+- Used by: scalar GEMV, MMA kernel, grouped MMA kernel,
+  tiled dequant kernel (for dequant+cuBLAS path).
 
 E4M4 encodes each float32 absmax as a single byte (4-bit exponent +
 4-bit mantissa). Decode is branchless: `ldexp(mantissa, exponent-bias)`.
 This saves 4x bandwidth for absmax reads but adds a decode step in
 the inner loop.
 
-**Note:** The grouped scalar GEMV and grouped MMA use different data
-formats. The grouped scalar GEMV uses flat layout with float32 absmax
-(same as the dense scalar GEMV), while the grouped MMA uses tiled
-layout with E4M4 absmax (same as the dense MMA). This means MoE
-expert weights must be stored in both formats if both kernels are used
-in the dispatch, or a runtime conversion must happen. Currently the
-benchmark prepares each format separately.
+The flat dequant kernel (`kDequantizeBlockwise_kbit_vec`) is still
+available for standalone use (e.g., debugging), but `kbit_linear`
+dispatch uses the tiled dequant (`kDequantizeBlockwise_kbit_tiled`).
 
 ---
 
@@ -539,7 +503,7 @@ per element; for k=2, ~8 ops.
 
 | GPU | SM | MMA instruction | Async MMA? | Kernel strategy |
 |-----|-----|-----------------|------------|----------------|
-| RTX 4090 | sm_89 | mma.sync | No | All 5 kernels as described |
+| RTX 4090 | sm_89 | mma.sync | No | All 4 kernels as described |
 | RTX 5090 | sm_120 | mma.sync (ext) | No | Same strategy, more SMs (192) |
 | H100/H200 | sm_90a | wgmma.mma_async | Yes | Could overlap dequant + MMA |
 | B200/GB200 | sm_100a | tcgen05.mma | Yes | Could overlap dequant + MMA |

summary.md

@@ -6,8 +6,8 @@ Base: `23f92e5` (feature/kbit-gemv-v8)
 ## What changed
 
 All kbit kernels now use uint8 E4M4 absmax by default, replacing float32.
-A float16 absmax alternative path is available for scalar GEMV and grouped
-scalar GEMV if higher absmax precision is needed.
+A float16 absmax alternative path is available for scalar GEMV if higher
+absmax precision is needed.
 
 ### Kernel changes
 
@@ -17,19 +17,18 @@ scalar GEMV if higher absmax precision is needed.
   re-encoding from float32.
 - **Scalar GEMV** (dense): `unsigned char*` absmax with `load_absmax<T>`
   decode. Templated on `ABSMAX_T` for uint8 (default) and float16.
-- **Grouped scalar GEMV** (MoE): Same treatment as dense scalar GEMV.
 - **MMA kernels** (dense + grouped): Already used uint8 E4M4 — no change.
 - **Dequantize**: Already supported uint8 — no change.
 
 ### Files modified (8)
 
 - `csrc/ops.cu` — E4M4 encode/decode moved before quantize kernel,
   quantize writes uint8, repack accepts uint8, fp16abs template
-  instantiations for scalar/grouped GEMV
+  instantiations for scalar GEMV
 - `csrc/pythonInterface.cpp` — All wrappers updated for `unsigned char*`;
-  added 16 extern C symbols for fp16abs scalar/grouped GEMV
+  added extern C symbols for fp16abs scalar GEMV
 - `bitsandbytes/backends/cuda/ops.py` — uint8 allocation in quantize,
-  absmax dtype routing in scalar/grouped GEMV dispatch
+  absmax dtype routing in scalar GEMV dispatch
 - `bitsandbytes/_ops.py` — quantize_kbit fake op returns uint8
 - `bitsandbytes/functional.py` — Removed redundant Python-side E4M4 encode
 - `tests/test_scalar_gemv.py` — E4M4 decode in reference functions
@@ -62,10 +61,7 @@ RTX 4090, CUDA events timing, fp16.
 consistent direction. Run-to-run variance dominates. No measurable
 regression.
 
-**Grouped scalar GEMV / MoE** (8 configs, 8 experts):
-- M≥2: within noise (±3%)
-- M=1: possible ~5% overhead from E4M4 decode cost being a larger fraction
-  of the small per-warp workload. One outlier at +22% is likely noise.
+**MoE grouped MMA** (8 configs, 8 experts): No change (already uint8 E4M4).
 
 **MMA kernels**: No change (already uint8 E4M4).