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.46.0: a reference backward through frozen GGUF k-quant weights (dX only, scored against an fp64 oracle, with a model-level fp64 gradcheck); an opt-in split tile table that runs a heavy expert's prefill rows on 32- or 64-row K19 tiles; refusals that name AMD on a ROCm build; and an AMD compile-only CI check. No default changes. 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
How these were measured
- 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.jsonbyscripts/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.
1 maintainer
Metadata
1 maintainer
Release files for grouped-nf4-gemm 0.46.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.46.0.tar.gz | 322.5 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| grouped_nf4_gemm-0.46.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 665.0 kB
Release files / grouped_nf4_gemm-0.46.0.tar.gz
| Download URL | grouped_nf4_gemm-0.46.0.tar.gz |
|---|---|
| Size | 322.5 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
a52e703b6d5ac1601be80fbc3354a9da0e734ac419ac972f429691bf72006add
|
|
BLAKE2b-256 checksum How to use checksums |
e21ae08a8f4b95caf076efbd9ef36741f4df0edbc13dc15da86432fdb8cda1cc
|
| 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 11, 2026.
Transparency logRelease files / grouped_nf4_gemm-0.46.0-py3-none-any.whl
| Download URL | grouped_nf4_gemm-0.46.0-py3-none-any.whl |
|---|---|
| Size | 342.5 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
8dbefb18df530bbe093666a52b99ed508b4f3046d7f10f2716539bed876b6fec
|
|
BLAKE2b-256 checksum How to use checksums |
e08720dec106b556514a3e1ef9c736571607b627fd8d163ea760d4ff44cb6670
|
| 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 11, 2026.
Transparency log
