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.
Install · Results · Kernel API · Task guides
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.43.0: two defaults, each read in experts4bit-qlora on Qwen3-30B-A3B and one RTX 5090. On GPUs with at least
160 SMs, single-request decode at Qwen3-30B-A3B's expert shapes uses the bandwidth-targeted GEMV: 1.24× as fast at one
request (GNF4_GEMV_BW=0 turns it off). Bucketed LoRA-delta calls run as one compact autograd node with the same bytes:
0.65 GB less training memory on packed rows, and slightly faster steps (NF4_QLORA_COMPACT_BUCKETS=0 turns it off).
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 shapes on >= 160-SM parts) | GNF4_GEMV_BW (auto; 1 everywhere, 0 off) with gemm_4bit_grouped; NF4 guide |
| 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 | Provenance guide |
Layouts and signatures · All supported entry points
Try the packers and byte checks on CPU
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
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.
1 maintainer
Metadata
1 maintainer
Release files for grouped-nf4-gemm 0.43.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| grouped_nf4_gemm-0.43.0.tar.gz | 310.1 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| grouped_nf4_gemm-0.43.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 640.9 kB
Release files / grouped_nf4_gemm-0.43.0.tar.gz
| Download URL | grouped_nf4_gemm-0.43.0.tar.gz |
|---|---|
| Size | 310.1 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
3c7b3cb9cf26c3e379fc3bdf186d05eb9e93db8c520df7d9f5efb7cee0b61f58
|
|
BLAKE2b-256 checksum How to use checksums |
387b775f132b523ff7480a56a5bcbd5b9d87785e4a6f9a79f8c130ebffe01658
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 8, 2026.
Transparency logRelease files / grouped_nf4_gemm-0.43.0-py3-none-any.whl
| Download URL | grouped_nf4_gemm-0.43.0-py3-none-any.whl |
|---|---|
| Size | 330.8 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
e72c63d2300689e37c72dd3d878f266f04be2d596bf2ff815199f1def138f779
|
|
BLAKE2b-256 checksum How to use checksums |
2db05bb8e97502a52085fca50ae7f31f4bac5cd9189b6663bf50c1053dd139bd
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 8, 2026.
Transparency log
