grouped-nf4-gemm 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- grouped_nf4_gemm-0.1.0/LICENSE +21 -0
- grouped_nf4_gemm-0.1.0/PKG-INFO +268 -0
- grouped_nf4_gemm-0.1.0/README.md +249 -0
- grouped_nf4_gemm-0.1.0/kernel/grouped_nf4_gemm.egg-info/PKG-INFO +268 -0
- grouped_nf4_gemm-0.1.0/kernel/grouped_nf4_gemm.egg-info/SOURCES.txt +11 -0
- grouped_nf4_gemm-0.1.0/kernel/grouped_nf4_gemm.egg-info/dependency_links.txt +1 -0
- grouped_nf4_gemm-0.1.0/kernel/grouped_nf4_gemm.egg-info/requires.txt +2 -0
- grouped_nf4_gemm-0.1.0/kernel/grouped_nf4_gemm.egg-info/top_level.txt +3 -0
- grouped_nf4_gemm-0.1.0/kernel/host_gather.py +89 -0
- grouped_nf4_gemm-0.1.0/kernel/nf4_grouped.py +557 -0
- grouped_nf4_gemm-0.1.0/kernel/nf4_pack_ref.py +43 -0
- grouped_nf4_gemm-0.1.0/pyproject.toml +30 -0
- grouped_nf4_gemm-0.1.0/setup.cfg +4 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Cerin Amroth LLC
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,268 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: grouped-nf4-gemm
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Fused grouped GEMM on bitsandbytes NF4-packed weights with host streaming, on consumer GPUs.
|
|
5
|
+
Author: Jordan Anderson
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://cerinamroth.com/ml/grouped-nf4-gemm/
|
|
8
|
+
Project-URL: Source, https://github.com/pjordanandrsn/grouped-nf4-gemm
|
|
9
|
+
Project-URL: Receipts, https://github.com/pjordanandrsn/grouped-nf4-gemm/tree/main/kernel
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
13
|
+
Requires-Python: >=3.9
|
|
14
|
+
Description-Content-Type: text/markdown
|
|
15
|
+
License-File: LICENSE
|
|
16
|
+
Requires-Dist: torch>=2.8
|
|
17
|
+
Requires-Dist: triton>=3.4
|
|
18
|
+
Dynamic: license-file
|
|
19
|
+
|
|
20
|
+
# grouped-nf4-gemm — single-launch W4A16 GEMM over fused NF4 MoE expert stacks
|
|
21
|
+
|
|
22
|
+
A Triton kernel that runs the grouped expert GEMM **directly on NF4-packed
|
|
23
|
+
weights** — one launch for all active experts, LUT decode to fp32 in
|
|
24
|
+
registers, blockwise fp32 absmax, fp32 accumulation, bf16 epilogue. No
|
|
25
|
+
per-expert dequantize-then-`bmm` round trip, no bf16 weight materialization.
|
|
26
|
+
It consumes the canonical packed layout and conventions of the bitsandbytes
|
|
27
|
+
`gemm_4bit` family ([#1949](https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1949)):
|
|
28
|
+
`[E, N, K/2]` uint8 + fp32 blockwise absmax, with `(sizes, expert_ids)`
|
|
29
|
+
supplied after the usual token→expert sort.
|
|
30
|
+
|
|
31
|
+
**Why:** for frozen 4-bit MoE experts, the standard path pays to decode the
|
|
32
|
+
weights into bf16 and then reads them again — at batch-1 decode that round
|
|
33
|
+
trip (plus ~3 kernel launches per active expert) dominates. Fusing the decode
|
|
34
|
+
into the GEMM deletes it. The measured side effect worth stating plainly:
|
|
35
|
+
**fp32 accumulation makes the fused path *more accurate* than the
|
|
36
|
+
materialize-to-bf16 baseline, in every cell ever measured here.**
|
|
37
|
+
|
|
38
|
+
## The claim (blind-confirmed, receipts in-repo)
|
|
39
|
+
|
|
40
|
+
Everything below is from **pre-registered, OpenTimestamps-stamped blind
|
|
41
|
+
confirmatory runs** (protocol + pass/fail criteria stamped before data; two
|
|
42
|
+
devices; n=3 fresh-process reps; worst/median-rep reduction; failures
|
|
43
|
+
reported at full volume). On sm_86 at batch-1 decode, versus the
|
|
44
|
+
dequantize-then-matmul baseline on the same stacks:
|
|
45
|
+
|
|
46
|
+
- **Fidelity:** property suite green on every device, every run (35 → 44
|
|
47
|
+
tests as the kernel grew); fused output error **below the baseline's in
|
|
48
|
+
every cell ever measured** (fp32 accumulate).
|
|
49
|
+
- **Energy:** fused J/token **below the baseline in 104 of 112
|
|
50
|
+
confirmatory-grade cells across v1–v3**. Six of the eight misses are the
|
|
51
|
+
`top_k=1`/tiny class (named below); the other two are parity-margin
|
|
52
|
+
readings (1.005, 1.010) on a single instance. On bandwidth-bound cells the
|
|
53
|
+
energy win has never failed to replicate.
|
|
54
|
+
- **Speed:** census MoE shapes (OLMoE, Qwen3-30B, Gemma-4, GPT-OSS-120B,
|
|
55
|
+
gate_up + down) run **1.16–2.73× at median** (one census cell —
|
|
56
|
+
gpt-oss `down`, 2880×2880 — is instance-sensitive: 0.7–2.0× across five
|
|
57
|
+
instances). Fresh off-census shapes with `top_k ≥ 6` (DeepSeek-V3,
|
|
58
|
+
granite-3.1, Qwen3-Next) run **1.0–1.8× at median**; `k=2`-large shapes
|
|
59
|
+
(Grok-1, Mixtral-8x22B) 1.0–1.24×, never slower.
|
|
60
|
+
- **Versus the other execution classes** (same-run census on the v6
|
|
61
|
+
kernel, [receipts](bench/phase1/results/comparators_v6/RESULTS-comparators-v6.md)):
|
|
62
|
+
the grouped-bf16-GEMM class that unsloth's MoE backend rides
|
|
63
|
+
(`grouped_gemm.ops.gmm`, dequant inside the timed path as 4-bit storage
|
|
64
|
+
requires) loses to the fused kernel on **every census cell — decode
|
|
65
|
+
median 4.67×, prefill median 3.02×**. Axolotl/PEFT stacks have no kernel
|
|
66
|
+
of their own (their QLoRA forward is bitsandbytes `Linear4bit` — see the
|
|
67
|
+
flagship bnb baseline). GPTQ-Marlin is fidelity-excellent but per-expert
|
|
68
|
+
(launch-storm at MoE decode) and format-incompatible with NF4 checkpoints.
|
|
69
|
+
- **Known losers:** `top_k=1` cells are **instance-unstable in both
|
|
70
|
+
directions** (Scout `down` measured 0.47–1.12 across six contexts on
|
|
71
|
+
identical code — split-K helps paired but can't stabilize the class), and
|
|
72
|
+
**tiny shapes (≲5 M weight elements) lose outright** (0.24–0.35× speed,
|
|
73
|
+
4–7× energy). v4 adds a dispatch floor that routes tiny cells back to the
|
|
74
|
+
dequant path.
|
|
75
|
+
- **Prefill** (compute-bound M): the v6 register-LUT mainloop rewrite
|
|
76
|
+
(blind-CONFIRMED) runs **1.39–1.54× the prior mainloop on every census
|
|
77
|
+
prefill cell**; against the dequant path the census reads 1.14–2.78× with
|
|
78
|
+
all three large gate_ups above 1.15 — gate_up is no longer a loser class.
|
|
79
|
+
One caveat carried at full volume: the dequant *baseline* itself swings
|
|
80
|
+
~25% between cloud instances (the fused kernel holds within 0.2 ms), and
|
|
81
|
+
OLMoE gate_up (the smallest-expert shape) remains below parity at ~0.6×.
|
|
82
|
+
|
|
83
|
+
Six blind confirmatories have run; the first five **did not fully pass as
|
|
84
|
+
registered**, each results doc says exactly what failed and why, and the
|
|
85
|
+
sixth passed clean:
|
|
86
|
+
[v1](kernel/RESULTS-gate2-confirmatory.md) (caught the original per-shape
|
|
87
|
+
config table overfitting its census), [v2](kernel/RESULTS-v2-confirmatory.md)
|
|
88
|
+
(validated the replacement single-constant config on 64-SM parts and the
|
|
89
|
+
off-census `k≥6` wins), [v3](kernel/RESULTS-v3-confirmatory.md) (found the
|
|
90
|
+
v2-era SM-conditional premise was measurement noise, quantified the
|
|
91
|
+
`top_k=1` and tiny-shape loss classes, and established the methodology rule
|
|
92
|
+
that latency-bound cells only support paired claims),
|
|
93
|
+
[v4](kernel/RESULTS-v4-confirmatory.md) (dispatch floor + split-K work floor
|
|
94
|
+
+ prefill config; caught its own dispatch-point regression),
|
|
95
|
+
[v5](kernel/RESULTS-v5-confirmatory.md) (the load-time dispatch fix, clean on
|
|
96
|
+
the A5000 11/11 with energy 8/8 on both devices; one contended-A2000 noise
|
|
97
|
+
cell kept it from a full pass — the dispatch line is closed),
|
|
98
|
+
[v6](kernel/RESULTS-v6-confirmatory.md) (**CONFIRMED**, all five criteria:
|
|
99
|
+
the register-LUT M-tile mainloop, adjudicated on the instance-robust paired
|
|
100
|
+
rewrite ratio after the dress rehearsal exposed the dequant baseline's
|
|
101
|
+
host lottery). The preregs,
|
|
102
|
+
amendments, evidence JSONs, sweeps, and mechanical reducers are all
|
|
103
|
+
committed; `.ots` files anchor the protocols to Bitcoin before the data
|
|
104
|
+
existed.
|
|
105
|
+
|
|
106
|
+
## Flagship: a 235B MoE decoding at the PCIe physical limit on ≤16 GB of VRAM
|
|
107
|
+
|
|
108
|
+
`bench/phase3/` runs Qwen3-235B-A22B with **all expert weights NF4-packed in
|
|
109
|
+
host pinned RAM (~128 GB)** and streamed per-token over PCIe, with this
|
|
110
|
+
kernel as the sole MoE compute. Same discipline (prereg + OTS, receipts
|
|
111
|
+
in-repo):
|
|
112
|
+
|
|
113
|
+
- **[Phase A](bench/phase3/flagship/RESULTS-flagship-offload.md)** (synthetic
|
|
114
|
+
weights, real GQA attention + router): **5.57 tok/s = 102–103% of the
|
|
115
|
+
measured 44.3 GB/s link's waterfall ceiling** — the stream fully hides
|
|
116
|
+
compute — on a **13.6 GB** working set. The dequantize-then-matmul path on
|
|
117
|
+
the identical pipeline: 1.81 tok/s (34% of ceiling). ALL PASS. (Fractions
|
|
118
|
+
marginally above 100% are microbench conservatism: the 1 GiB×10 ceiling
|
|
119
|
+
measurement brackets every copy with a host sync, paying launch +
|
|
120
|
+
sync-return latency the pipeline's continuously-queued copy stream never
|
|
121
|
+
pays.)
|
|
122
|
+
- **[The gap is architectural](bench/phase3/flagship/RESULTS-flagship-bnb-baseline.md)** —
|
|
123
|
+
we registered the prediction that bnb's own CUDA dequant kernel would
|
|
124
|
+
also hide under the copy shadow (which would have narrowed our claim),
|
|
125
|
+
and it was **refuted**: the standard path reaches **40% of waterfall**
|
|
126
|
+
(per-expert dequant+GEMM compute outlasts the shadow), versus 93–94%
|
|
127
|
+
fused on the same pod. Against the strongest standard comparator the
|
|
128
|
+
fused path is **2.33× tokens/s and 2.21× J/token**.
|
|
129
|
+
- **[Phase B](bench/phase3/flagship/RESULTS-flagship-phaseB.md)** (the real
|
|
130
|
+
438 GB checkpoint, stream-quantized to NF4 in place): **coherent greedy
|
|
131
|
+
text at 4.3–4.4 tok/s on 15.2 GB VRAM**, replicated across five pods.
|
|
132
|
+
- **Expert prefetch is measured CLOSED, negative** — four registered arcs
|
|
133
|
+
([B2](bench/phase3/flagship/RESULTS-flagship-phaseB2.md) speculation:
|
|
134
|
+
token-to-token expert stickiness is only 0.44;
|
|
135
|
+
[B3](bench/phase3/flagship/RESULTS-flagship-phaseB3.md) early routing: the
|
|
136
|
+
pre-attention router predicts the post-attention top-8 at **0.93** but the
|
|
137
|
+
CPU sync tax eats the win;
|
|
138
|
+
[B4](bench/phase3/flagship/RESULTS-flagship-phaseB4.md) threaded issuance:
|
|
139
|
+
GIL tax, 0.57×;
|
|
140
|
+
[B5](bench/phase3/flagship/RESULTS-flagship-phaseB5.md) GPU-driven
|
|
141
|
+
zero-copy gather: hit rate H makes speculation move (2−H)× the bytes, and
|
|
142
|
+
the observed loss matches that law to ~1% — break-even needs H ≳ 0.95,
|
|
143
|
+
above this model's 0.93 predictor ceiling).
|
|
144
|
+
- **Recommended configuration: `--prefetch-mode gpu`** — expert ids stay
|
|
145
|
+
GPU-resident and a triton kernel ([`kernel/host_gather.py`](kernel/host_gather.py))
|
|
146
|
+
gathers expert rows straight from pinned host RAM over UVA (zero-copy),
|
|
147
|
+
with no per-layer memcpy launches and no GPU→CPU syncs. It is the fastest
|
|
148
|
+
measured arm (4.39–4.41 tok/s, +1.5% over serialized memcpy, byte-identical
|
|
149
|
+
greedy output 6/6) and validates SM-issued UVA reads at ≥ copy-engine
|
|
150
|
+
throughput at 7.98 GB/token.
|
|
151
|
+
|
|
152
|
+
## Reproduce
|
|
153
|
+
|
|
154
|
+
See [REPRO.md](REPRO.md) — suite, benchmark, and verdict reduction are each
|
|
155
|
+
one command from a frozen tree. Requires an sm_86 GPU, `torch ≥ 2.8`,
|
|
156
|
+
`bitsandbytes`, and a C compiler on PATH (triton builds launcher stubs at
|
|
157
|
+
runtime).
|
|
158
|
+
|
|
159
|
+
```
|
|
160
|
+
python -m pytest kernel/test_nf4_grouped.py -q # 44 tests, ~2.5 min
|
|
161
|
+
python bench/phase1/harness.py --models OLMoE --regimes decode_bs1 \
|
|
162
|
+
--backends dequant_grouped fused_nf4 --out receipts.json
|
|
163
|
+
```
|
|
164
|
+
|
|
165
|
+
## Layout
|
|
166
|
+
|
|
167
|
+
- `kernel/nf4_grouped.py` — the kernel (decode gemv path + M-tile path),
|
|
168
|
+
packing helpers, torch reference decode
|
|
169
|
+
- `kernel/test_nf4_grouped.py` — property suite (bnb decode exactness at
|
|
170
|
+
bf16 output precision, fidelity ordering, adversarial absmax, boundaries)
|
|
171
|
+
- `kernel/prereg_*.json` + `.ots` — pre-registered protocols, stamped
|
|
172
|
+
- `kernel/RESULTS-*.md` — results, including the failures
|
|
173
|
+
- `bench/phase1/` — backend-registry harness (dequant/gemv/grouped-mm/
|
|
174
|
+
unsloth/marlin/fused), confirmatory evidence, reducers
|
|
175
|
+
- `bench/phase2/` — decode config sweeps (both devices); `arch/` —
|
|
176
|
+
cross-architecture census (sm_86/89/90)
|
|
177
|
+
- `bench/phase3/` — the 235B offload flagship: `offload_decode_235b.py`
|
|
178
|
+
(Phase A, synthetic), `offload_generate_235b.py` (real checkpoint,
|
|
179
|
+
generation, prefetch arms), `flagship/` — results + receipts
|
|
180
|
+
- `kernel/host_gather.py` — GPU-driven zero-copy gather from pinned host
|
|
181
|
+
memory (UVA), the recommended offload copy path
|
|
182
|
+
- `docs/KERNEL_CONTRACT.md`, `docs/TOLERANCE_CONTRACT.md` — op contract and
|
|
183
|
+
fidelity spec; `census/`, `roofline/` — shape census + ceilings
|
|
184
|
+
|
|
185
|
+
Regenerate the machine-generated artifacts:
|
|
186
|
+
|
|
187
|
+
```
|
|
188
|
+
python3 census/make_census.py # census/shape_census.json
|
|
189
|
+
python3 roofline/roofline.py # roofline/ceilings.json
|
|
190
|
+
```
|
|
191
|
+
|
|
192
|
+
## Status / roadmap
|
|
193
|
+
|
|
194
|
+
Landed through v6: universal decode constant (the dense-sweep result),
|
|
195
|
+
split-K for starved grids (with a per-split work floor), a load-time
|
|
196
|
+
min-bytes dispatch floor (tiny cells route to the dequant path via
|
|
197
|
+
`decode_dispatch()`), the register-LUT prefill mainloop (v6, confirmed),
|
|
198
|
+
and the flagship offload pipeline (Phase A/B + the closed prefetch program
|
|
199
|
+
+ the UVA gather path + the bnb-CUDA-dequant baseline, whose registered
|
|
200
|
+
prediction was refuted — see the flagship section). Pending: the v6 A2000
|
|
201
|
+
report-only addendum; a bare-metal gen4 replication when stock returns.
|
|
202
|
+
Parked: sm_120 (three consecutive cloud provisioning failures on 5090s —
|
|
203
|
+
availability, not code). Ecosystem landing is calendar-gated on the
|
|
204
|
+
bitsandbytes v0.50.0 release; see the coordination note on #1949.
|
|
205
|
+
|
|
206
|
+
## Cross-vendor projections (stamped, PROJECTED tier — help us confirm them)
|
|
207
|
+
|
|
208
|
+
The waterfall arithmetic doesn't care which vendor's bus you're on, so we've
|
|
209
|
+
extended it — under the same receipts discipline — into a stamped, pre-silicon
|
|
210
|
+
projection table for AMD, Intel, and NVIDIA unified-memory parts:
|
|
211
|
+
[`PROJECTIONS-multiarch.md`](PROJECTIONS-multiarch.md) (protocol:
|
|
212
|
+
[`PROTOCOL-multiarch.md`](PROTOCOL-multiarch.md); model + R1 anchor gate:
|
|
213
|
+
[`projections/`](projections/)). Both docs are OpenTimestamps-anchored (`.ots`)
|
|
214
|
+
**before any of this silicon was run** — the projections are a falsifiable
|
|
215
|
+
prediction, not a marketing table.
|
|
216
|
+
|
|
217
|
+
**Every row is `PROJECTED` — none is confirmed on its hardware.** Streaming rows
|
|
218
|
+
(discrete GPU) are anchored to the measured flagship numbers to <2%; unified-memory
|
|
219
|
+
rows are *ceilings only* and real decode sits below them. Headlines, all
|
|
220
|
+
projected: 235B-A22B at **3.0–3.5 tok/s** on a gen4-×16 desktop, **6.0–6.9** on
|
|
221
|
+
gen5 (5090 / RDNA4), and a **17–22 tok/s ceiling** on 128 GB unified boxes
|
|
222
|
+
(Strix Halo / DGX Spark / Jetson Thor). NF4-vs-bf16 is a **3.56×** byte reduction
|
|
223
|
+
(absmax-inclusive), not the round 4×.
|
|
224
|
+
|
|
225
|
+
**Call for confirmatories.** If you own any listed part, run
|
|
226
|
+
`PROTOCOL-multiarch.md` and file the result — **pass or fail** — as an issue.
|
|
227
|
+
A refuting measurement is as welcome as a confirming one; that's the point.
|
|
228
|
+
Template:
|
|
229
|
+
|
|
230
|
+
```
|
|
231
|
+
Title: [confirmatory] <platform> — <model>
|
|
232
|
+
Environment: vendor / device / driver / runtime / triton / torch / bnb;
|
|
233
|
+
link measured via lspci + on-box microbench (streaming) OR mem-band spec
|
|
234
|
+
(unified)
|
|
235
|
+
Correctness gate: max rel-err vs dequant_ref = <value> (pass < 1e-2)
|
|
236
|
+
Measured decode: <tok/s> per census cell Projected band: <from table>
|
|
237
|
+
Verdict: within band? / refutes row? Attach: results JSONL
|
|
238
|
+
```
|
|
239
|
+
|
|
240
|
+
## License & attribution
|
|
241
|
+
|
|
242
|
+
MIT ([LICENSE](LICENSE)). Portions developed with Claude Code as an AI
|
|
243
|
+
assistant under the author's direction and review — see
|
|
244
|
+
[ATTRIBUTION.md](ATTRIBUTION.md). All claims are the author's responsibility.
|
|
245
|
+
|
|
246
|
+
## Portability program
|
|
247
|
+
|
|
248
|
+
The kernel is single-source Triton; everything that must differ per vendor
|
|
249
|
+
is being pulled into `backends/` — device detection, warp/wavefront/sub-group
|
|
250
|
+
width, per-arch autotune search spaces. `bench/hw_contract.py` validates
|
|
251
|
+
kernel correctness on any torch device **without a bitsandbytes build**; if
|
|
252
|
+
you have ROCm or XPU silicon, that is the entry point. `docs/PORTABILITY.md`
|
|
253
|
+
is the pre-port hazard register. Per the repo's tier language, every
|
|
254
|
+
non-CUDA row is `port target` until a confirmatory passes on that silicon.
|
|
255
|
+
|
|
256
|
+
## Router-predictability probe
|
|
257
|
+
|
|
258
|
+
`router_probe/` asks whether the measured H = 0.93 one-layer-lead prediction
|
|
259
|
+
ceiling is the router's conditional entropy or the probe's capacity limit.
|
|
260
|
+
The charter and procedure were OTS-stamped before any real-model capture;
|
|
261
|
+
the Phase-0 instrument gate passed 4/4 on planted fixtures (see
|
|
262
|
+
`router_probe/RESULTS.md`, exploratory tier). Phase 1 is pending.
|
|
263
|
+
|
|
264
|
+
## Contact
|
|
265
|
+
|
|
266
|
+
Cerin Amroth Research takes contract and pilot engagements on this work —
|
|
267
|
+
kernel ports, offload integration, and sponsored research lanes with
|
|
268
|
+
stamped receipts. Contact **jordan@cerinamroth.com**.
|
|
@@ -0,0 +1,249 @@
|
|
|
1
|
+
# grouped-nf4-gemm — single-launch W4A16 GEMM over fused NF4 MoE expert stacks
|
|
2
|
+
|
|
3
|
+
A Triton kernel that runs the grouped expert GEMM **directly on NF4-packed
|
|
4
|
+
weights** — one launch for all active experts, LUT decode to fp32 in
|
|
5
|
+
registers, blockwise fp32 absmax, fp32 accumulation, bf16 epilogue. No
|
|
6
|
+
per-expert dequantize-then-`bmm` round trip, no bf16 weight materialization.
|
|
7
|
+
It consumes the canonical packed layout and conventions of the bitsandbytes
|
|
8
|
+
`gemm_4bit` family ([#1949](https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1949)):
|
|
9
|
+
`[E, N, K/2]` uint8 + fp32 blockwise absmax, with `(sizes, expert_ids)`
|
|
10
|
+
supplied after the usual token→expert sort.
|
|
11
|
+
|
|
12
|
+
**Why:** for frozen 4-bit MoE experts, the standard path pays to decode the
|
|
13
|
+
weights into bf16 and then reads them again — at batch-1 decode that round
|
|
14
|
+
trip (plus ~3 kernel launches per active expert) dominates. Fusing the decode
|
|
15
|
+
into the GEMM deletes it. The measured side effect worth stating plainly:
|
|
16
|
+
**fp32 accumulation makes the fused path *more accurate* than the
|
|
17
|
+
materialize-to-bf16 baseline, in every cell ever measured here.**
|
|
18
|
+
|
|
19
|
+
## The claim (blind-confirmed, receipts in-repo)
|
|
20
|
+
|
|
21
|
+
Everything below is from **pre-registered, OpenTimestamps-stamped blind
|
|
22
|
+
confirmatory runs** (protocol + pass/fail criteria stamped before data; two
|
|
23
|
+
devices; n=3 fresh-process reps; worst/median-rep reduction; failures
|
|
24
|
+
reported at full volume). On sm_86 at batch-1 decode, versus the
|
|
25
|
+
dequantize-then-matmul baseline on the same stacks:
|
|
26
|
+
|
|
27
|
+
- **Fidelity:** property suite green on every device, every run (35 → 44
|
|
28
|
+
tests as the kernel grew); fused output error **below the baseline's in
|
|
29
|
+
every cell ever measured** (fp32 accumulate).
|
|
30
|
+
- **Energy:** fused J/token **below the baseline in 104 of 112
|
|
31
|
+
confirmatory-grade cells across v1–v3**. Six of the eight misses are the
|
|
32
|
+
`top_k=1`/tiny class (named below); the other two are parity-margin
|
|
33
|
+
readings (1.005, 1.010) on a single instance. On bandwidth-bound cells the
|
|
34
|
+
energy win has never failed to replicate.
|
|
35
|
+
- **Speed:** census MoE shapes (OLMoE, Qwen3-30B, Gemma-4, GPT-OSS-120B,
|
|
36
|
+
gate_up + down) run **1.16–2.73× at median** (one census cell —
|
|
37
|
+
gpt-oss `down`, 2880×2880 — is instance-sensitive: 0.7–2.0× across five
|
|
38
|
+
instances). Fresh off-census shapes with `top_k ≥ 6` (DeepSeek-V3,
|
|
39
|
+
granite-3.1, Qwen3-Next) run **1.0–1.8× at median**; `k=2`-large shapes
|
|
40
|
+
(Grok-1, Mixtral-8x22B) 1.0–1.24×, never slower.
|
|
41
|
+
- **Versus the other execution classes** (same-run census on the v6
|
|
42
|
+
kernel, [receipts](bench/phase1/results/comparators_v6/RESULTS-comparators-v6.md)):
|
|
43
|
+
the grouped-bf16-GEMM class that unsloth's MoE backend rides
|
|
44
|
+
(`grouped_gemm.ops.gmm`, dequant inside the timed path as 4-bit storage
|
|
45
|
+
requires) loses to the fused kernel on **every census cell — decode
|
|
46
|
+
median 4.67×, prefill median 3.02×**. Axolotl/PEFT stacks have no kernel
|
|
47
|
+
of their own (their QLoRA forward is bitsandbytes `Linear4bit` — see the
|
|
48
|
+
flagship bnb baseline). GPTQ-Marlin is fidelity-excellent but per-expert
|
|
49
|
+
(launch-storm at MoE decode) and format-incompatible with NF4 checkpoints.
|
|
50
|
+
- **Known losers:** `top_k=1` cells are **instance-unstable in both
|
|
51
|
+
directions** (Scout `down` measured 0.47–1.12 across six contexts on
|
|
52
|
+
identical code — split-K helps paired but can't stabilize the class), and
|
|
53
|
+
**tiny shapes (≲5 M weight elements) lose outright** (0.24–0.35× speed,
|
|
54
|
+
4–7× energy). v4 adds a dispatch floor that routes tiny cells back to the
|
|
55
|
+
dequant path.
|
|
56
|
+
- **Prefill** (compute-bound M): the v6 register-LUT mainloop rewrite
|
|
57
|
+
(blind-CONFIRMED) runs **1.39–1.54× the prior mainloop on every census
|
|
58
|
+
prefill cell**; against the dequant path the census reads 1.14–2.78× with
|
|
59
|
+
all three large gate_ups above 1.15 — gate_up is no longer a loser class.
|
|
60
|
+
One caveat carried at full volume: the dequant *baseline* itself swings
|
|
61
|
+
~25% between cloud instances (the fused kernel holds within 0.2 ms), and
|
|
62
|
+
OLMoE gate_up (the smallest-expert shape) remains below parity at ~0.6×.
|
|
63
|
+
|
|
64
|
+
Six blind confirmatories have run; the first five **did not fully pass as
|
|
65
|
+
registered**, each results doc says exactly what failed and why, and the
|
|
66
|
+
sixth passed clean:
|
|
67
|
+
[v1](kernel/RESULTS-gate2-confirmatory.md) (caught the original per-shape
|
|
68
|
+
config table overfitting its census), [v2](kernel/RESULTS-v2-confirmatory.md)
|
|
69
|
+
(validated the replacement single-constant config on 64-SM parts and the
|
|
70
|
+
off-census `k≥6` wins), [v3](kernel/RESULTS-v3-confirmatory.md) (found the
|
|
71
|
+
v2-era SM-conditional premise was measurement noise, quantified the
|
|
72
|
+
`top_k=1` and tiny-shape loss classes, and established the methodology rule
|
|
73
|
+
that latency-bound cells only support paired claims),
|
|
74
|
+
[v4](kernel/RESULTS-v4-confirmatory.md) (dispatch floor + split-K work floor
|
|
75
|
+
+ prefill config; caught its own dispatch-point regression),
|
|
76
|
+
[v5](kernel/RESULTS-v5-confirmatory.md) (the load-time dispatch fix, clean on
|
|
77
|
+
the A5000 11/11 with energy 8/8 on both devices; one contended-A2000 noise
|
|
78
|
+
cell kept it from a full pass — the dispatch line is closed),
|
|
79
|
+
[v6](kernel/RESULTS-v6-confirmatory.md) (**CONFIRMED**, all five criteria:
|
|
80
|
+
the register-LUT M-tile mainloop, adjudicated on the instance-robust paired
|
|
81
|
+
rewrite ratio after the dress rehearsal exposed the dequant baseline's
|
|
82
|
+
host lottery). The preregs,
|
|
83
|
+
amendments, evidence JSONs, sweeps, and mechanical reducers are all
|
|
84
|
+
committed; `.ots` files anchor the protocols to Bitcoin before the data
|
|
85
|
+
existed.
|
|
86
|
+
|
|
87
|
+
## Flagship: a 235B MoE decoding at the PCIe physical limit on ≤16 GB of VRAM
|
|
88
|
+
|
|
89
|
+
`bench/phase3/` runs Qwen3-235B-A22B with **all expert weights NF4-packed in
|
|
90
|
+
host pinned RAM (~128 GB)** and streamed per-token over PCIe, with this
|
|
91
|
+
kernel as the sole MoE compute. Same discipline (prereg + OTS, receipts
|
|
92
|
+
in-repo):
|
|
93
|
+
|
|
94
|
+
- **[Phase A](bench/phase3/flagship/RESULTS-flagship-offload.md)** (synthetic
|
|
95
|
+
weights, real GQA attention + router): **5.57 tok/s = 102–103% of the
|
|
96
|
+
measured 44.3 GB/s link's waterfall ceiling** — the stream fully hides
|
|
97
|
+
compute — on a **13.6 GB** working set. The dequantize-then-matmul path on
|
|
98
|
+
the identical pipeline: 1.81 tok/s (34% of ceiling). ALL PASS. (Fractions
|
|
99
|
+
marginally above 100% are microbench conservatism: the 1 GiB×10 ceiling
|
|
100
|
+
measurement brackets every copy with a host sync, paying launch +
|
|
101
|
+
sync-return latency the pipeline's continuously-queued copy stream never
|
|
102
|
+
pays.)
|
|
103
|
+
- **[The gap is architectural](bench/phase3/flagship/RESULTS-flagship-bnb-baseline.md)** —
|
|
104
|
+
we registered the prediction that bnb's own CUDA dequant kernel would
|
|
105
|
+
also hide under the copy shadow (which would have narrowed our claim),
|
|
106
|
+
and it was **refuted**: the standard path reaches **40% of waterfall**
|
|
107
|
+
(per-expert dequant+GEMM compute outlasts the shadow), versus 93–94%
|
|
108
|
+
fused on the same pod. Against the strongest standard comparator the
|
|
109
|
+
fused path is **2.33× tokens/s and 2.21× J/token**.
|
|
110
|
+
- **[Phase B](bench/phase3/flagship/RESULTS-flagship-phaseB.md)** (the real
|
|
111
|
+
438 GB checkpoint, stream-quantized to NF4 in place): **coherent greedy
|
|
112
|
+
text at 4.3–4.4 tok/s on 15.2 GB VRAM**, replicated across five pods.
|
|
113
|
+
- **Expert prefetch is measured CLOSED, negative** — four registered arcs
|
|
114
|
+
([B2](bench/phase3/flagship/RESULTS-flagship-phaseB2.md) speculation:
|
|
115
|
+
token-to-token expert stickiness is only 0.44;
|
|
116
|
+
[B3](bench/phase3/flagship/RESULTS-flagship-phaseB3.md) early routing: the
|
|
117
|
+
pre-attention router predicts the post-attention top-8 at **0.93** but the
|
|
118
|
+
CPU sync tax eats the win;
|
|
119
|
+
[B4](bench/phase3/flagship/RESULTS-flagship-phaseB4.md) threaded issuance:
|
|
120
|
+
GIL tax, 0.57×;
|
|
121
|
+
[B5](bench/phase3/flagship/RESULTS-flagship-phaseB5.md) GPU-driven
|
|
122
|
+
zero-copy gather: hit rate H makes speculation move (2−H)× the bytes, and
|
|
123
|
+
the observed loss matches that law to ~1% — break-even needs H ≳ 0.95,
|
|
124
|
+
above this model's 0.93 predictor ceiling).
|
|
125
|
+
- **Recommended configuration: `--prefetch-mode gpu`** — expert ids stay
|
|
126
|
+
GPU-resident and a triton kernel ([`kernel/host_gather.py`](kernel/host_gather.py))
|
|
127
|
+
gathers expert rows straight from pinned host RAM over UVA (zero-copy),
|
|
128
|
+
with no per-layer memcpy launches and no GPU→CPU syncs. It is the fastest
|
|
129
|
+
measured arm (4.39–4.41 tok/s, +1.5% over serialized memcpy, byte-identical
|
|
130
|
+
greedy output 6/6) and validates SM-issued UVA reads at ≥ copy-engine
|
|
131
|
+
throughput at 7.98 GB/token.
|
|
132
|
+
|
|
133
|
+
## Reproduce
|
|
134
|
+
|
|
135
|
+
See [REPRO.md](REPRO.md) — suite, benchmark, and verdict reduction are each
|
|
136
|
+
one command from a frozen tree. Requires an sm_86 GPU, `torch ≥ 2.8`,
|
|
137
|
+
`bitsandbytes`, and a C compiler on PATH (triton builds launcher stubs at
|
|
138
|
+
runtime).
|
|
139
|
+
|
|
140
|
+
```
|
|
141
|
+
python -m pytest kernel/test_nf4_grouped.py -q # 44 tests, ~2.5 min
|
|
142
|
+
python bench/phase1/harness.py --models OLMoE --regimes decode_bs1 \
|
|
143
|
+
--backends dequant_grouped fused_nf4 --out receipts.json
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
## Layout
|
|
147
|
+
|
|
148
|
+
- `kernel/nf4_grouped.py` — the kernel (decode gemv path + M-tile path),
|
|
149
|
+
packing helpers, torch reference decode
|
|
150
|
+
- `kernel/test_nf4_grouped.py` — property suite (bnb decode exactness at
|
|
151
|
+
bf16 output precision, fidelity ordering, adversarial absmax, boundaries)
|
|
152
|
+
- `kernel/prereg_*.json` + `.ots` — pre-registered protocols, stamped
|
|
153
|
+
- `kernel/RESULTS-*.md` — results, including the failures
|
|
154
|
+
- `bench/phase1/` — backend-registry harness (dequant/gemv/grouped-mm/
|
|
155
|
+
unsloth/marlin/fused), confirmatory evidence, reducers
|
|
156
|
+
- `bench/phase2/` — decode config sweeps (both devices); `arch/` —
|
|
157
|
+
cross-architecture census (sm_86/89/90)
|
|
158
|
+
- `bench/phase3/` — the 235B offload flagship: `offload_decode_235b.py`
|
|
159
|
+
(Phase A, synthetic), `offload_generate_235b.py` (real checkpoint,
|
|
160
|
+
generation, prefetch arms), `flagship/` — results + receipts
|
|
161
|
+
- `kernel/host_gather.py` — GPU-driven zero-copy gather from pinned host
|
|
162
|
+
memory (UVA), the recommended offload copy path
|
|
163
|
+
- `docs/KERNEL_CONTRACT.md`, `docs/TOLERANCE_CONTRACT.md` — op contract and
|
|
164
|
+
fidelity spec; `census/`, `roofline/` — shape census + ceilings
|
|
165
|
+
|
|
166
|
+
Regenerate the machine-generated artifacts:
|
|
167
|
+
|
|
168
|
+
```
|
|
169
|
+
python3 census/make_census.py # census/shape_census.json
|
|
170
|
+
python3 roofline/roofline.py # roofline/ceilings.json
|
|
171
|
+
```
|
|
172
|
+
|
|
173
|
+
## Status / roadmap
|
|
174
|
+
|
|
175
|
+
Landed through v6: universal decode constant (the dense-sweep result),
|
|
176
|
+
split-K for starved grids (with a per-split work floor), a load-time
|
|
177
|
+
min-bytes dispatch floor (tiny cells route to the dequant path via
|
|
178
|
+
`decode_dispatch()`), the register-LUT prefill mainloop (v6, confirmed),
|
|
179
|
+
and the flagship offload pipeline (Phase A/B + the closed prefetch program
|
|
180
|
+
+ the UVA gather path + the bnb-CUDA-dequant baseline, whose registered
|
|
181
|
+
prediction was refuted — see the flagship section). Pending: the v6 A2000
|
|
182
|
+
report-only addendum; a bare-metal gen4 replication when stock returns.
|
|
183
|
+
Parked: sm_120 (three consecutive cloud provisioning failures on 5090s —
|
|
184
|
+
availability, not code). Ecosystem landing is calendar-gated on the
|
|
185
|
+
bitsandbytes v0.50.0 release; see the coordination note on #1949.
|
|
186
|
+
|
|
187
|
+
## Cross-vendor projections (stamped, PROJECTED tier — help us confirm them)
|
|
188
|
+
|
|
189
|
+
The waterfall arithmetic doesn't care which vendor's bus you're on, so we've
|
|
190
|
+
extended it — under the same receipts discipline — into a stamped, pre-silicon
|
|
191
|
+
projection table for AMD, Intel, and NVIDIA unified-memory parts:
|
|
192
|
+
[`PROJECTIONS-multiarch.md`](PROJECTIONS-multiarch.md) (protocol:
|
|
193
|
+
[`PROTOCOL-multiarch.md`](PROTOCOL-multiarch.md); model + R1 anchor gate:
|
|
194
|
+
[`projections/`](projections/)). Both docs are OpenTimestamps-anchored (`.ots`)
|
|
195
|
+
**before any of this silicon was run** — the projections are a falsifiable
|
|
196
|
+
prediction, not a marketing table.
|
|
197
|
+
|
|
198
|
+
**Every row is `PROJECTED` — none is confirmed on its hardware.** Streaming rows
|
|
199
|
+
(discrete GPU) are anchored to the measured flagship numbers to <2%; unified-memory
|
|
200
|
+
rows are *ceilings only* and real decode sits below them. Headlines, all
|
|
201
|
+
projected: 235B-A22B at **3.0–3.5 tok/s** on a gen4-×16 desktop, **6.0–6.9** on
|
|
202
|
+
gen5 (5090 / RDNA4), and a **17–22 tok/s ceiling** on 128 GB unified boxes
|
|
203
|
+
(Strix Halo / DGX Spark / Jetson Thor). NF4-vs-bf16 is a **3.56×** byte reduction
|
|
204
|
+
(absmax-inclusive), not the round 4×.
|
|
205
|
+
|
|
206
|
+
**Call for confirmatories.** If you own any listed part, run
|
|
207
|
+
`PROTOCOL-multiarch.md` and file the result — **pass or fail** — as an issue.
|
|
208
|
+
A refuting measurement is as welcome as a confirming one; that's the point.
|
|
209
|
+
Template:
|
|
210
|
+
|
|
211
|
+
```
|
|
212
|
+
Title: [confirmatory] <platform> — <model>
|
|
213
|
+
Environment: vendor / device / driver / runtime / triton / torch / bnb;
|
|
214
|
+
link measured via lspci + on-box microbench (streaming) OR mem-band spec
|
|
215
|
+
(unified)
|
|
216
|
+
Correctness gate: max rel-err vs dequant_ref = <value> (pass < 1e-2)
|
|
217
|
+
Measured decode: <tok/s> per census cell Projected band: <from table>
|
|
218
|
+
Verdict: within band? / refutes row? Attach: results JSONL
|
|
219
|
+
```
|
|
220
|
+
|
|
221
|
+
## License & attribution
|
|
222
|
+
|
|
223
|
+
MIT ([LICENSE](LICENSE)). Portions developed with Claude Code as an AI
|
|
224
|
+
assistant under the author's direction and review — see
|
|
225
|
+
[ATTRIBUTION.md](ATTRIBUTION.md). All claims are the author's responsibility.
|
|
226
|
+
|
|
227
|
+
## Portability program
|
|
228
|
+
|
|
229
|
+
The kernel is single-source Triton; everything that must differ per vendor
|
|
230
|
+
is being pulled into `backends/` — device detection, warp/wavefront/sub-group
|
|
231
|
+
width, per-arch autotune search spaces. `bench/hw_contract.py` validates
|
|
232
|
+
kernel correctness on any torch device **without a bitsandbytes build**; if
|
|
233
|
+
you have ROCm or XPU silicon, that is the entry point. `docs/PORTABILITY.md`
|
|
234
|
+
is the pre-port hazard register. Per the repo's tier language, every
|
|
235
|
+
non-CUDA row is `port target` until a confirmatory passes on that silicon.
|
|
236
|
+
|
|
237
|
+
## Router-predictability probe
|
|
238
|
+
|
|
239
|
+
`router_probe/` asks whether the measured H = 0.93 one-layer-lead prediction
|
|
240
|
+
ceiling is the router's conditional entropy or the probe's capacity limit.
|
|
241
|
+
The charter and procedure were OTS-stamped before any real-model capture;
|
|
242
|
+
the Phase-0 instrument gate passed 4/4 on planted fixtures (see
|
|
243
|
+
`router_probe/RESULTS.md`, exploratory tier). Phase 1 is pending.
|
|
244
|
+
|
|
245
|
+
## Contact
|
|
246
|
+
|
|
247
|
+
Cerin Amroth Research takes contract and pilot engagements on this work —
|
|
248
|
+
kernel ports, offload integration, and sponsored research lanes with
|
|
249
|
+
stamped receipts. Contact **jordan@cerinamroth.com**.
|