grouped-nf4-gemm

Run MoE expert math directly on 4-bit weights.

One GPU launch for the active experts. No full-size bf16 weight copies.

Triton kernels for grouped NF4 and native MXFP4 matrix multiplication, training gradients, INT4 decode, FP8 attention, and streaming expert weights from RAM or SSD. The packers, reference implementations, and byte-level checks ship alongside the kernels.

Want to train a model? Start with Loggetta. It installs this package through the experts4bit-qlora runtime. Use grouped-nf4-gemm directly to build or tune a kernel integration.

Where it loses. A per-expert baseline can be faster at small shapes and in some graphed decode workloads. With weights already resident in bf16, Unsloth's H100 prefill kernel won by 2.6–5.3×. A faster kernel can leave training time unchanged. Details.

Install · Results · Kernel API · Task guides · Hugging Face

Measured on one RTX 4090 and one H100. The receipts record the GPU, torch and (for the Unsloth comparison) the NVIDIA driver; clock locking, ECC state and the Triton version were not recorded. How these were measured.

Install

pip install grouped-nf4-gemm

GPU kernels need Linux, an NVIDIA sm_80+ GPU, torch ≥ 2.8 and Triton ≥ 3.4. CI tests Python 3.11. CPU pack/decode and provenance tools work without CUDA; macOS and Windows are not exercised by CI. ROCm and XPU are port targets.

New in 0.51.0: with NF4_QLORA_COMPACT_DELTA=1, the compact delta's forward and backward each hold one padded block fewer at their peaks: the forward frees the zero-padded block before the delta is built, and the backward frees the rebuilt block before its grad. Every output and gradient is the same bytes as 0.50.0's. experts4bit-qlora's fit of Qwen3-30B-A3B's field recipe on a 24 GiB RTX 4090 (its TC1 amendment 78) ran on this code. The compact delta stays opt-in, and no memory or speed number is claimed in this package. Release notes

Try it on your GPU

From a checkout of this repository; about one minute, no model download:

pip install grouped-nf4-gemm bitsandbytes
python examples/dequant_tax.py

The example compares packed compute with dequantize-then-GEMM at three workload sizes and prints a repeated-baseline control. On CPU it explains which GPU measurement is unavailable.

Which entry point?

Task	API or guide
Grouped NF4 forward and backward	nf4_grouped.gemm_4bit_grouped, dgrad_4bit_grouped
NF4 single-row decode, bandwidth route (default at Qwen3-30B-A3B's, Granite-3.1-3b-a800m's, OLMoE-1B-7B's and Mixtral-8x7B's expert shapes on >= 160-SM parts; Qwen3's also on the RTX 4090: sm_89, 128 SMs)	GNF4_GEMV_BW (auto; 1 everywhere, 0 off) with gemm_4bit_grouped; NF4 guide
Dense bf16 single-row decode GEMV (lane K35; an API only: it routes nothing, and its consumer's served read is the timing of record)	dense_gemv.gemv_bf16_bw
Int8-weight single-row decode GEMV for an int8 lm_head (lane K36; an API only, as K35's)	dense_gemv.gemv_i8_bw, dense_gemv.quantize_rows_int8
Int8-weight projections at decode, 1 to 16 rows, for Mixtral-8x7B's attention (lane K37; an API only, as K35's)	dense_i8.linear_i8
Native MXFP4 expert math	mxfp4_grouped.gemm_mxfp4_grouped
INT4 decode	INT4 guide
FP8 paged attention	Attention guide
Stream experts from RAM or NVMe	Streaming guide
Check checkpoint bytes

Layouts and signatures · All supported entry points

These examples are executed by CI on Linux. They use pure PyTorch; the GPU kernels remain CUDA-only.

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.

What is measured

Each result below is confirmed by a pre-registered run with public evidence. CI checks the values against claims.json.

Experiment	Result	Tier	Claim ID
OLMoE QLoRA on real prose, fused vs this project's per-expert loop	4.50× on RTX 4090; 4.75× on H100	confirmed	gnf4.kernel.e2e-training-real-prose
Decode vs Unsloth's kernel, with weights stored in 4-bit	1.70× on H100; 2.79× on RTX 4090	confirmed	gnf4.kernel.h2h-unsloth
gpt-oss-120b native-MXFP4 QLoRA experiment, L40S	9.82 GB peak VRAM; 144/144 frozen expert tensors hash-identical after training	confirmed	gnf4.mxfp4.train-9.82gb

OLMoE training · Unsloth kernel comparison · 120B experiment

• Unsloth comparison: same pod and process, three repetitions per GPU. The receipts record the GPU and compute capability, torch 2.8.0+cu128, the NVIDIA driver (570.195.03 on the RTX 4090, 580.159.03 on the H100), bitsandbytes 0.50.0, Unsloth 2026.8.15, and whether Unsloth's TMA path was available (no on the RTX 4090, yes on the H100). Receipts
• OLMoE training: the receipts record the GPU, compute capability, VRAM, torch 2.8.0+cu128 and experts4bit-qlora 0.17.5 from published wheels; the grouped-nf4-gemm version was not recorded. Receipts
• Not recorded in either: clock locking, ECC state, the Triton version.
• Analytic ceiling, not a measurement: at the design stage, batch-1 decode was estimated memory-bound with a ceiling of about 8.1× over the two-pass dequant path (phase 1 notes).
• The chart is generated from docs/claims.json by scripts/build_readme_chart.py; CI fails when it is stale.

The first comparator is this project's per-expert dequantize-to-bf16 loop, without CUDA graphs. The Unsloth row is a kernel comparison, not end-to-end training. The 120B row is an experiment using host memory; it does not mean the whole model fits in 9.82 GB or that Loggetta runs this path.

Where it loses

• Small shapes and some graphed decode workloads: a per-expert baseline can be faster.
• Weights already resident in bf16: Unsloth's H100 prefill kernel won that comparison, by 2.6–5.3×.
• Kernel speed vs whole-model speed: a faster kernel can leave training time unchanged. The decoded training route stays opt-in: its end-to-end training-step read found no measurable saving.

Use the current status to choose a route, and benchmark on real text rather than random token IDs.

Documentation and machine-readable data

Task guides · Capabilities · Claims and evidence · Package compatibility · Reproduce the benchmarks · llms.txt · Contributing

License & attribution

MIT. Developed with AI assistance under the author's direction and review; attribution. For ports, integration, or research work: jordananderson.work.