grouped-nf4-gemm — one-launch 4-bit GEMM over fused MoE expert stacks (NF4 + native MXFP4)

A Triton kernel that runs the grouped expert GEMM directly on 4-bit-packed weights: one launch for all active experts, LUT decode to fp32 in registers, blockwise scaling, fp32 accumulation, bf16 epilogue. No per-expert dequantise-then-bmm, no bf16 weight materialisation. Two codebooks: NF4 on the bitsandbytes gemm_4bit layout, and MXFP4 (OCP e2m1 + e8m0) on a checkpoint's exact released bytes.

Beside the GEMM, the package carries what a 4-bit MoE serving path needs around it: an fp8 paged decode attention (sliding windows, attention sinks, custom scale, per-layer KV geometry), the int4-b32 decode GEMV and its calibrated packer, the decode glue kernels (fused RMSNorm, residual-fold, rotary, router epilogue), and the host-streaming and NVMe tiers for models that do not fit.

Current position, one page: docs/STATUS.md. Every number, with its evidence and tier: docs/claims.json.

See it on your own hardware first

pip install grouped-nf4-gemm bitsandbytes
python examples/dequant_tax.py          # ~1 min, one GPU, no model download

examples/dequant_tax.py times the dequantise-then-GEMM round trip against computing on the packed bytes at three points on the M axis, prints a self-pair beside every ratio (a ratio inside the instrument's own spread is not a measurement), and says what the run does not show.

Install

pip install grouped-nf4-gemm          # nf4gemm and gnf4 are aliases

Trusted publishing; every wheel carries a PEP 740 attestation. The fused GEMM is CUDA-only (triton>=3.4, Linux); the reference decode and the provenance hashing are pure torch and run anywhere.

Try it on CPU right now

No GPU needed for the pack/decode/provenance surface — the fused GEMM is CUDA-only, but the reference decode and the provenance hashing are pure torch.

On Linux this works from a bare pip install; on macOS and Windows it does not, today. nf4_pack_ref imports nf4_grouped, which does a module-level import triton — and triton is declared triton>=3.4; platform_system == 'Linux', so it is simply absent elsewhere and these blocks raise ModuleNotFoundError. The math is pure torch; the import graph is not. CI executes these blocks on Linux, where triton is present, so it validates the code without validating this sentence. Tracked as a real defect — the reference decode should not need the kernel's dependency. These three blocks are extracted and executed by CI (test_readme_cpu_block.py), so they cannot drift from the API.

1. NF4 round-trip — pack a weight, decode it back, check the error:

import torch
from nf4_pack_ref import quantize_pack_nf4
from nf4_grouped import dequant_ref

w = torch.randn(256, 512)                      # a per-expert weight [N, K]
packed, absmax = quantize_pack_nf4(w)          # [256, 256] uint8, [256, 8] fp32
wq = dequant_ref(packed, absmax, 256, 512)     # decode back to [N, K]
print("nf4 rel-err:", round(((wq - w).norm() / w.norm()).item(), 3))     # ~0.09
print("nf4 re-pack idempotent:", torch.equal(quantize_pack_nf4(wq)[0], packed))  # True

2. MXFP4 round-trip — the gpt-oss expert format, same shape story:

import torch
from mxfp4_pack_ref import quantize_pack_mxfp4, dequant_mxfp4

w = torch.randn(128, 256)                      # [.., K], K a multiple of 32
blocks, scales = quantize_pack_mxfp4(w)        # [128, 8, 16] u8, [128, 8] u8 (e8m0)
wq = dequant_mxfp4(blocks, scales)             # [128, 256]
print("mxfp4 rel-err:", round(((wq - w).norm() / w.norm()).item(), 3))   # ~0.12

3. Provenance in four lines — hash on-disk bytes, catch a tampered one:

import torch, json, struct, tempfile, os
from mxfp4_loader import file_tensor_sha256, tensor_sha256

t = torch.arange(64, dtype=torch.uint8)        # stand-in for an expert's packed bytes
hdr = json.dumps({"w": {"dtype": "U8", "shape": [64], "data_offsets": [0, 64]}}).encode()
path = tempfile.mktemp(suffix=".safetensors")
with open(path, "wb") as f:
    f.write(struct.pack("<Q", len(hdr))); f.write(hdr); f.write(t.numpy().tobytes())
print("prov bytes match:", file_tensor_sha256(path, "w") == tensor_sha256(t))    # True
b = bytearray(open(path, "rb").read()); b[-1] ^= 0xFF; open(path, "wb").write(bytes(b))
print("prov tamper detected:", file_tensor_sha256(path, "w") != tensor_sha256(t))  # True
os.remove(path)

That's the same instrument the 144/144 training receipt used.

Which entry point? Pick by where the weights live

Do not quantise-bake a checkpoint that is already MXFP4. Relocation keeps the bytes and hands packed nibbles to the kernel; re-quantising to NF4 costs a dequant per read — measured ~4× slower per request on DeepSeek-V4-Flash.

Training goes through nf4_qlora / mxfp4_qlora, which is what experts4bit-qlora drives (enable_fast(), enable_fast_train()). This package makes one expert-stack matmul cheap; e4b decides which bytes are where.

What is measured

Tiers, used strictly: confirmed = pre-registered, OpenTimestamps- stamped blind confirmatory run; measured = a run with a committed receipt here; measured-private = a real run whose receipt lives in a private audit tree, so you cannot check it from this repository.

Three limits, stated here rather than in a footnote:

1. Against a CUDA-graphed baseline the fused path loses at decode (0.949× on a 4090, 0.858× on an H100). What graphing cannot touch is the memory-traffic win at training shape on bandwidth-limited cards (1.489× on the 4090; parity on the H100). The position is narrow on purpose: competitive at equal VRAM, wins when VRAM binds.
2. Unsloth wins its own regime. Against their bf16-resident kernel they run 2.6–5.3× faster at prefill on an H100.
3. Known losers: top_k=1 cells are instance-unstable in both directions; shapes under ~5 M weight elements lose outright and are routed back to the dequant path.

And one about your benchmark: random token ids understate this kernel by ~1.6×. Prose routes to 98.4% of experts, random ids to 87.5%, and fewer hit experts means fewer iterations of the loop this replaces. Benchmark on real text.

The receipts

Six blind confirmatories (v1–v6); the first five did not fully pass as registered, each results doc says what failed, and the sixth passed clean. Pre-registrations, amendments, evidence JSONs and reducers are committed; .ots files anchor the protocols. The Unsloth head-to-head has its own stamped protocol. All under kernel/ (RESULTS-*.md, prereg_*.json). The 235B flagship and the closed prefetch programme are under bench/phase3/flagship/. The MXFP4 lane and Kimi K3 provenance chain are under docs/mxfp4/ and docs/. What each of the 22 docs is, and whether it is current: docs/INDEX.md.

What was retired

Kept findable in docs/STATUS.md and as retired entries in docs/claims.json: the "sm_120 parked" roadmap line (sm_120 has been the primary serving target since 0.15.0); split-K on the decode GEMV (refuted, ships dormant as the evidence); a fixed fraction-of-waterfall as the offload law; the cold-engine "free floor" premise; expert prefetch (closed, negative, four arcs). The "4.67× vs the grouped-bf16 execution class" number is superseded by the head-to-head — that backend never ran Unsloth's own kernel.

What is open

#319 the f32 paged compute modes miss their reference on torch 2.8 / triton 3.4 (the fp8 modes, the sm_120 default, pass); #87 int32 offset overflow at large max(expert_ids); #73, #60, #58 arena/NVMe efficiency; #71 pinned-row factor on cgroup v2. Every non-CUDA row is a port target — PROJECTIONS-multiarch.md is stamped arithmetic that invites refutation, and docs/PORTABILITY.md is the hazard register.

Reproduce

REPRO.md: suite, benchmark and verdict reduction are each one command from a frozen tree.

python -m pytest kernel/test_nf4_grouped.py -q
python -m pytest kernel/test_fp8_paged_attn.py -q -k "f8dot or pf8"   # the sm_120 serving modes

Layout

kernel/ — the kernels, packers, reference decodes, property suites, pre-registrations and results · bench/phase1..3/ — the confirmatory harnesses, the census, the flagship · bench/sm120-census/ — the 5090 census · bench/cold-engine/ — the cold-tier research record · docs/ — contracts, tolerance spec, MXFP4 and K3 receipts, STATUS.md, claims.json, INDEX.md · census/, roofline/ — shape census and ceilings · router_probe/ — the router-predictability probe.

License & attribution

MIT. Portions developed with Claude Code as an AI assistant under the author's direction and review — see ATTRIBUTION.md. All claims are the author's responsibility.

Contact

Cerin Amroth Research takes contract and pilot engagements on this work — kernel ports, offload integration, sponsored research lanes with stamped receipts. jordan@cerinamroth.com.