LegumMagister/nf4-grouped-gemm

NF4 grouped GEMM for MoE experts

A Triton kernel that multiplies tokens by the experts of a mixture-of-experts layer, with the expert weights stored in bitsandbytes' NF4 format. All experts run in one launch, as torch._grouped_mm does for bf16.

Why

There is no grouped NF4 kernel for MoE experts. bf16 grouped GEMM is native and heavily optimized, and NVFP4 is native on Blackwell, but on A100 and H100 a 4-bit MoE has nothing in between. With load_in_4bit, transformers leaves fused expert weights in bf16, and looping over experts with bitsandbytes is slow. This kernel stores experts at ~4.13 bits per weight (4× less memory than bf16) and runs them in a single grouped launch.

Usage

from kernels import get_kernel

nf4 = get_kernel("LegumMagister/nf4-grouped-gemm")
w = nf4.NF4Weights(experts)                  # [E, N, K] bf16 -> NF4 (bitsandbytes' quantize_4bit, double-quantized)
out = nf4.nf4_grouped_mm(x, offsets, w)      # x sorted by expert; offsets = cumulative rows per expert

out[rows of expert e] = x[rows of e] @ W[e].T. The weights are bitsandbytes' own bytes, so w.dequantize() equals bitsandbytes.functional.dequantize_4bit exactly.

Benchmarks

One NVIDIA H100 80GB HBM3, Qwen3-30B-A3B expert shapes (128 experts, top-8), milliseconds per projection:

projection	tokens	bf16 grouped	bnb, one expert at a time	NF4 grouped	vs bf16	vs bnb
gate_up 128 × 1536 × 2048	1	0.033	0.389	0.061	0.54×	6.4×
	8	0.110	1.867	0.245	0.45×	7.6×
	64	0.230	4.512	0.521	0.44×	8.7×
	512	0.300	5.859	1.112	0.27×	5.3×
	4096	0.530	5.362	2.007	0.26×	2.7×
down 128 × 2048 × 768	1	0.024	0.385	0.042	0.57×	9.2×
	8	0.067	1.634	0.121	0.55×	13.5×
	64	0.141	4.198	0.303	0.46×	13.9×
	512	0.164	5.138	0.578	0.28×	8.9×
	4096	0.297	5.273	1.053	0.28×	5.0×

It is 3–14× faster than running bitsandbytes NF4 expert by expert, and about half the speed of bf16 at small batches, with a quarter of the memory. It does not beat bf16 yet: it reads weights at 110–330 GB/s, against the H100's 3.35 TB/s, because its loads are not pipelined deeply enough. A Marlin/FLUTE-style version (asynchronous copies, weights laid out offline for the tensor cores) is the next step.

Correctness

tests/test_and_benchmark.py compares against bitsandbytes' dequantize_4bit followed by a bf16 matmul, for uneven and empty experts, N and K not multiples of the tile, bf16 and fp16 inputs, with every case run twice. The worst relative error is 3.5e-3 (bf16 rounding), and the dequantized values are bit-identical to bitsandbytes'.

How it works

• Weights: bitsandbytes NF4. A 4-bit index into 16 normal-distribution quantiles, an 8-bit absmax per 64 values, and an FP32 scale per 256 absmax. The per-block scales are precomputed once in FP32 (exactly as bitsandbytes computes them).
• Decode (few rows per expert): one program per expert, N tile and K slice (split-K), with fp32 atomic accumulation; no tile map.
• Prefill: tiles of 16–128 rows of one expert, found in-kernel by a binary search over per-expert tile counts; no host synchronization.
• In both, each 64-wide K tile is unpacked and looked up in registers and fed to bf16 tensor cores (tl.dot).