CERIN AMROTH · ML systems

How do I run MoE expert inference natively on MXFP4 (e2m1 + e8m0) weights, straight from the released checkpoint bytes?

Solved by grouped-nf4-gemm · this page in the repository (pinned e2af4cfb91b2, the source of this rendering; latest on main, unpinned; not the source of any fact rendered here)

Install routes

From docs/capabilities.json at the pinned commit.

pip install grouped-nf4-gemm

Primary route: kernel package.

Alternatives:

Environment (from the capability register): OS: Linux · Python: >=3.11 tested in CI (pyproject says >=3.9; 3.9/3.10 are not tested) · Accelerator: NVIDIA CUDA GPU, sm_80 or newer (sm_120 is the primary serving target) · Requires: torch>=2.8 (pre-releases accepted); triton>=3.4 (Linux-only distribution)

Use mxfp4_grouped.gemm_mxfp4_grouped from grouped-nf4-gemm: the grouped expert GEMM runs on the checkpoint's own MXFP4 blocks and e8m0 scales, so a gpt-oss or Kimi-K3-class expert is multiplied as shipped, with no requantization to NF4 and no bf16 materialization. mxfp4_pack_ref.dequant_mxfp4 is the pure-torch reference the kernel is gated against.

Symptoms

Why it happens

MXFP4 (OCP MX v1.0) is a different codebook from NF4: sixteen e2m1 values, ±{0, 0.5, 1, 1.5, 2, 3, 4, 6}, one power-of-two e8m0 scale per 32 elements, even element in the low nibble, against NF4's per-64 fp32 absmax and high-nibble-first packing. A kernel built for one layout cannot read the other, so the usual answer is to convert, and every conversion either loses information (MXFP4 to bf16 to NF4) or reintroduces the bf16 round trip the fused kernel exists to delete. docs/mxfp4/PHASE0-seam-map.md records the four places the decode differs; the MXFP4 kernel is the NF4 kernel with exactly those swapped.

Which project solves it

grouped-nf4-gemm owns the native MXFP4 kernels (gemm_mxfp4_grouped; gemv_mxfp4_b32, the packed MXFP4 decode GEMV — optimising that GEMV is a change in this kernel repository, per AGENTS.md section 8, not in the consumer), the pack/decode reference, the loader helpers that map checkpoint shapes to kernel shapes without copying (mxfp4_loader.to_kernel_shapes), the arena source for per-expert releases (arena_experts.ArenaExpertSource), and the training wrapper (mxfp4_qlora). experts4bit-qlora (PyPI) drives them from a model; mxfp4_native_load.build_native_qlora_model(snap, r, alpha) here is the gpt-oss-specific loader that builds a QLoRA model without entering the dequant path.

Install

Kernel package (the minimum route):

pip install grouped-nf4-gemm

Linux, NVIDIA GPU sm_80 or newer (sm_120 is the primary serving target), triton>=3.4 (Linux-only distribution), torch>=2.8 (pre-releases accepted); CI tests Python 3.11. The pack/decode reference and loader hashing are pure torch. Through the model consumer:

pip install "experts4bit-qlora[fast]"

Smallest correct example

CPU-only: pack, decode, and confirm the checkpoint-to-kernel reshape moves no bytes.

# CPU-only (pure torch; no GPU, no triton launch)
import torch
from mxfp4_pack_ref import quantize_pack_mxfp4, dequant_mxfp4
from mxfp4_loader import to_kernel_shapes, tensor_sha256

E, N, K = 2, 64, 256                                  # K % 32 == 0
blocks, scales = quantize_pack_mxfp4(torch.randn(E, N, K))   # [E,N,K//32,16] u8, [E,N,K//32] u8
wq = dequant_mxfp4(blocks, scales)                    # [E, N, K] fp32
assert wq.shape == (E, N, K)
kb, ks = to_kernel_shapes(blocks, scales)             # [E, N, K//2] view; scales unchanged
assert tensor_sha256(kb) == tensor_sha256(blocks)     # a view, not a reorder

GPU: the grouped kernel against the reference decode of the same bytes.

# GPU (sm_80+) + triton
import torch
from mxfp4_pack_ref import quantize_pack_mxfp4, dequant_mxfp4
from mxfp4_loader import to_kernel_shapes
from mxfp4_grouped import gemm_mxfp4_grouped

E, N, K = 4, 128, 256
blocks, scales = quantize_pack_mxfp4(torch.randn(E, N, K) * 0.3)
kb, ks = (t.cuda() for t in to_kernel_shapes(blocks, scales))
sizes, expert_ids = [3, 1, 4], [1, 0, 3]              # group-sorted tokens per expert
a_cat = torch.randn(sum(sizes), K, device="cuda", dtype=torch.bfloat16)

out = gemm_mxfp4_grouped(a_cat, kb, ks, sizes, expert_ids)   # [T, N] bf16

row, refs = 0, []
for m, e in zip(sizes, expert_ids):
    W = dequant_mxfp4(blocks[e].cuda(), scales[e].cuda())     # [N, K] fp32, same bytes
    refs.append(a_cat[row:row + m].float() @ W.t()); row += m
ref = torch.cat(refs)
assert ((out.float() - ref).abs().max() / ref.abs().max()).item() < 2e-2

Expected result

The CPU block completes: the packed shapes are the gpt-oss tensor shapes and the kernel-shaped view hashes identically to its source. The GPU block returns [T, N] bf16 in group order, within the bound kernel/test_mxfp4_grouped.py uses. On a real checkpoint, mxfp4_loader.file_tensor_sha256(path, name) over the shard's byte range equals tensor_sha256 of the loaded view (verify-quantized-checkpoint-provenance.md).

Supported scope

Limitations

STATUS.md · claims.json · K3-PROVENANCE-CHAIN.md · INDEX.md (the docs/mxfp4/ pre-registrations and results) · verify-quantized-checkpoint-provenance.md · int4-decode-gemv.md

Evidence

Register: claims.json. Confirmed: claim gnf4.mxfp4.serve-tax-deleted (gpt-oss-120b served on its released bytes against the shipped-precision reference; its P1 sub-clause missed as stamped and the receipt says why), claim gnf4.mxfp4.train-9.82gb (QLoRA on native bytes; every expert hash identical before and after training). Measured: claim gnf4.k3.oracle-exact (reference decode reproduces Kimi K3's declared reference). Receipts: mxfp4/RESULTS-mxfp4-serve.md, mxfp4/RESULTS-mxfp4-train.md, RESULTS-k3-phase1-oracle.md.

Common wrong approaches

Source and freshness

This page is a rendering of docs/solutions/native-mxfp4-moe-inference.md at commit e2af4cfb91b2 (sha256 7dc1d765eb161c67…). Numbers are never copied here: every measured statement cites a claim ID in docs/claims.json at that commit. Repository-relative links resolve to this site's pages where the document is published and to the pinned commit otherwise.

CURRENT · source: pjordanandrsn/grouped-nf4-gemm@e2af4cfb91b2 · rendered package: 0.30.2 · latest published package: 0.30.2