grouped-nf4-gemm

Triton kernels for 4-bit MoE inference: grouped NF4/MXFP4 GEMM, INT4 GEMV, FP8 paged attention, and CPU/NVMe expert streaming.

Also discoverable as nf4gemm, gnf4, grouped-mxfp4-gemm. These are lookup-only aliases from the system manifest; grouped-nf4-gemm is the canonical project.

Direct answer

Use this package when dequantising each active expert into a high-precision weight before its matmul is the bottleneck. It computes the grouped expert GEMM directly over the packed NF4 (bitsandbytes gemm_4bit layout) or native MXFP4 (e2m1 + e8m0) representations, and ships the int4 decode GEMV, the fp8 paged attention, the decode glue and the host/NVMe primitives a 4-bit MoE serving path needs around it. It is a kernel package: experts4bit-qlora drives it.

Use this when

• A per-expert 4-bit MoE loop pays a launch per active expert per projection, and wherever that loop dequantises (bitsandbytes releases before 0.50.0, cells outside its packed gemm_4bit inference forward, the conventional 4-bit backward) it also writes and re-reads a bf16 copy of each expert; compute the grouped expert GEMM directly on the NF4-packed weights in one launch instead. — gemm_4bit_grouped runs one Triton launch for all active experts over the bitsandbytes gemm_4bit NF4 layout (LUT decode to fp32 in registers, blockwise absmax scaling, fp32 accumulation, bf16 epilogue); dequant_ref is the pure-torch oracle it is asserted against; dgrad_4bit_grouped is the one-launch training backward. solution page
• Serve or train a MoE whose experts are released as MXFP4 (gpt-oss, DeepSeek-V4, Kimi lineage) directly from the checkpoint's own bytes, without re-quantising them. — gemm_mxfp4_grouped is the table-swap of the NF4 kernel over native blocks [E, N, K//2] + e8m0 scales [E, N, K//32]; gemv_mxfp4_b32 is its decode-grade GEMV; mxfp4_loader.to_kernel_shapes views the released tensors in kernel layout with no copy; mxfp4_qlora and mxfp4_native_load train LoRA over the frozen native stacks. solution page
• Single-token (decode) expert and attention projections on 4-bit weights are bandwidth-bound; an int4 GEMV with 32-wide scales and a calibrated packer serves them near the memory ceiling. — quant_x_rows quantises activation rows to int8 with per-32 scales; gemv_int4_b32 runs the split-K int4-b32 GEMV over packed [E, N, K//2] + fp16 scales [E, N, K//32] (reduce_partials folds the split-K sum and bf16 cast into one launch); pack_int4_b32 is the round-to-nearest packer and gptq_pack_int4_b32 + HessianAccumulator the calibrated one; gemm_int4_b32_grouped_captured is the capture-legal batched variant. solution page
• A 4-bit MoE serving path needs a paged decode attention over an fp8 KV cache -- sliding windows, attention sinks, a custom scale, per-layer KV geometry -- computed with fp8 tensor-core dots on the cards that have them. — fp8_kv quantises and appends KV blocks in e4m3 with per-(token, head) or grouped scales; fp8_paged_decode_attention is the paged decode kernel with sliding windows, sinks, custom scale and per-layer KV geometry, asserted against paged_attn_ref. With GNF4_ATTN_COMPUTE unset it selects the fp8 compute path wherever fp8_compute_unsupported returns None, an explicit request is never downgraded, and compute_counts records which path ran. solution page
• The same paged decode attention over the fp8 KV cache computed with f32 dots: the default on sm_80-sm_88 (no fp8 tensor-core dot), the fallback wherever the fp8 constraints fail with GNF4_ATTN_COMPUTE unset, and every explicit f32 request. — fp8_paged_decode_attention with compute='f32' (or selected by the capability-conditional default) runs the split kernel, or the packed kernel with pack_heads=True, dequantising e4m3 in registers and accumulating in f32; it imports and launches on any sm_80+ card and is asserted against the same paged_attn_ref. solution page
• The non-GEMM part of a batch-1 decode step is a tail of small launches (RMSNorm, residual add, rotary, router softmax and top-k, activation, top-k combine, split-K reduce); fold each into one launch. — int4_b32 carries the decode glue: rmsnorm_rows, rmsnorm_resid_rows, scaled_resid_add_rows, rope_norm_heads, rope_heads, router_epilogue (softmax-then-top-k, or top-k on logits with an optional bias), swiglu_rows, combine_rows and reduce_partials; each has an in-tree parity test against a torch reference. solution page
• The expert weights do not fit in VRAM, or not even in host RAM: stream them from pinned host memory or from an on-disk expert arena on NVMe. — nvme_arena.bake / bake_expert_tensors relocate a checkpoint's expert tensors into a page-aligned arena with an index and manifest; nvme_bake_nf4.bake_nf4 quantises at bake; nvme_reader.ArenaReader reads it with O_DIRECT queues; nvme_residency.ColdTier and row_pool.RowPool hold the hot rows in pinned DRAM; arena_experts.ArenaExpertSource fetches per-layer fused stacks; mxfp4_residency serves native MXFP4 arenas. solution page
• Prove that the quantised bytes a kernel serves are the checkpoint's released bytes, unchanged, before quoting any result. — mxfp4_loader.file_tensor_sha256 hashes a safetensors tensor's data-section bytes without loading or dequantising; provenance_table builds the per-tensor hash table for a snapshot; verify_arena_matches raises on any mismatch (a provenance failure is not a tolerance); nvme_arena.verify checks an arena against its source; python -m verify_provenance is the CLI. solution page

Do not use this when

• dense (non-MoE) GEMMs: use cuBLAS, torch, or bitsandbytes' own 4-bit path -- since 0.50.0 its supported ordinary 2-D inference cells compute from the packed weights directly (torch.ops.bitsandbytes.gemm_4bit); there is no grouped routed-MoE contract upstream, which is what this package is
• a serving engine or a vLLM replacement: this is a kernel package driven by experts4bit-qlora
• operating systems other than Linux for the kernels (triton is Linux-only); ROCm/XPU are port targets, not supported
• models that already fit in bf16 with headroom: the grouped kernel loses at small shapes and to a CUDA-graphed per-expert loop at some decode shapes (claims gnf4.kernel.graphed-baseline-decode-loses)
• silent fallbacks: nothing falls back silently on CPU — nf4_grouped.gemm_4bit_grouped and dgrad_4bit_grouped refuse CPU tensors with an error that names dequant_ref; gemm_mxfp4_grouped and the int4_b32 kernels carry no device guard and fail inside the Triton launch, and the fp8_kv appends and the paged attention refuse with a CUDA+Triton message that names no reference

Install

Routes from docs/capabilities.json at the pinned commit; the first entry of the register is the primary route.

pip install grouped-nf4-gemm

Which path?

Smallest verified example

CPU-only, no Triton: the pure-torch pack reference and the oracle every kernel is asserted against. Executed on a triton-less host by the page authors; the GPU form (gemm_4bit_grouped against the same oracle) is in the solution page.

import torch
from nf4_pack_ref import quantize_pack_nf4
from nf4_grouped import dequant_ref

w = torch.randn(64, 128)
packed, absmax = quantize_pack_nf4(w)                 # [64, 64] uint8, [64, 2] fp32
back = dequant_ref(packed, absmax, 64, 128)           # the checkable oracle
assert (back - w).abs().max() < 0.25 * w.abs().max()  # NF4 round-trip inside the codebook's spacing

What it supports

Modes

benchmark, calibration, inference, offload, packing, serving, streaming, training, verification

Formats

bf16, fp8, int4, mxfp4, nf4

Model families

DeepSeek-V4 lineage; Kimi K3 lineage (reference decode); Qwen3-235B-A22B (host-RAM flagship); any checkpoint whose experts the bake can enumerate; any expert or projection weight packed with pack_int4_b32 / gptq_pack_int4_b32 (driven by experts4bit-qlora's serve lanes); any fused-expert MoE whose experts are packed in the bitsandbytes NF4 gemm_4bit layout (Qwen3-MoE, OLMoE, Mixtral, GraniteMoe, Gemma-4 text tower via experts4bit-qlora); any safetensors checkpoint; MXFP4 expert tensors for gpt-oss / DeepSeek-V4 / Kimi lineage; driven by experts4bit-qlora's paged engine for Qwen3-MoE, OLMoE, GraniteMoe, gpt-oss, Gemma-4, Mixtral; gpt-oss (20b, 120b); gpt-oss / DeepSeek-V4 / Kimi lineage MXFP4 arenas

Operating system

Linux

Python

>=3.11 tested in CI (pyproject says >=3.9; 3.9/3.10 are not tested)

Accelerator

NVIDIA CUDA GPU for the serving tiers and for the default NF4 quantise bake (nvme_bake_nf4.bake_nf4 with quantize_fn=None); the relocation bake and verify (nvme_arena.bake / bake_expert_tensors / verify) need no GPU; NVIDIA CUDA GPU, sm_80 or newer (sm_120 is the primary serving target); NVIDIA CUDA GPU, sm_80 or newer (the composition was measured on the RTX 5090, sm_120); NVIDIA CUDA GPU, sm_80 or newer; this is the path sm_80-sm_88 take by default. Its reference tests were run on an RTX 5090 (sm_120) with the mode forced; no registered cell exists on sm_80-sm_88; NVIDIA CUDA GPU, sm_89 or newer (the fp8 tensor-core dot); measured on the RTX 5090 (sm_120) only -- sm_89 and sm_90 meet the requirement but no registered cell was run there; none: pure torch (a CUDA GPU is optional)

Public entry points

arena_experts:ArenaExpertSource, cli:python -m nvme_arena, cli:python -m nvme_bake_nf4, cli:python -m verify_provenance, flag:GNF4_ATTN_COMPUTE, fp8_kv:fp8_kv_append_t1, fp8_kv:quantize_kv_fp8, fp8_paged_attn:compute_counts, fp8_paged_attn:fp8_compute_unsupported, fp8_paged_attn:fp8_paged_decode_attention, fp8_paged_attn:paged_attn_available, fp8_paged_attn:paged_attn_ref, gptq_pack:HessianAccumulator, gptq_pack:gptq_pack_int4_b32, int4_b32:combine_rows, int4_b32:gemm_int4_b32_grouped_captured, int4_b32:gemv_int4_b32, int4_b32:quant_x_rows, int4_b32:reduce_partials, int4_b32:rmsnorm_resid_rows, int4_b32:rmsnorm_rows, int4_b32:rope_heads, int4_b32:rope_norm_heads, int4_b32:router_epilogue, int4_b32:scaled_resid_add_rows, int4_b32:swiglu_rows, int4_pack_ref:dequant_int4_ref, int4_pack_ref:pack_int4_b32, mxfp4_grouped:gemm_mxfp4_grouped, mxfp4_grouped:gemv_mxfp4_b32, mxfp4_loader:file_tensor_sha256, mxfp4_loader:provenance_table, mxfp4_loader:to_kernel_shapes, mxfp4_loader:verify_arena_matches, mxfp4_native_load:build_native_qlora_model, mxfp4_pack_ref:dequant_mxfp4, mxfp4_pack_ref:quantize_pack_mxfp4, mxfp4_qlora, mxfp4_residency:Mxfp4NvmeResidency, nf4_grouped:build_group_tiles_device, nf4_grouped:dequant_ref, nf4_grouped:dgrad_4bit_grouped, nf4_grouped:gemm_4bit_grouped, nf4_grouped:gemm_4bit_grouped_captured, nf4_grouped:repack_from_bnb, nf4_pack_ref:quantize_pack_nf4, nf4_qlora:fused_grouped_lora, nf4_qlora:gemm_4bit_grouped_train, nvme_arena:bake, nvme_arena:bake_expert_tensors, nvme_arena:verify, nvme_bake_nf4:bake_nf4, nvme_reader:ArenaReader, nvme_reader:alloc_landing, nvme_residency:ColdTier, nvme_residency:capacity_for_bytes, row_pool:RowPool, verify_provenance:main

Owns (system manifest)

• grouped packed GEMM (NF4 and native MXFP4) and its dgrad
• decode GEMV (int4-b32, MXFP4) and split-K reduce
• FP8 paged attention and the fp8 KV appends
• decode glue kernels (rope, norm and residual folds, swiglu, combine)
• pack and pure-torch reference operations, GPTQ packing, provenance verification
• host/NVMe primitives (arena, reader, row tier, bake, residency primitives)
• kernel-level measurement and the claims register for kernel-level numbers

Related package

experts4bit-qlora — one system, two independently released packages: the runtime package depends on the kernel package through its optional `fast` extra; the kernel package never depends on the runtime package Dependency direction: experts4bit-qlora -> grouped-nf4-gemm.

Limitations — what it does not promise

• Loses to a CUDA-graphed per-expert baseline at some decode shapes (claim gnf4.kernel.graphed-baseline-decode-loses); shapes below the routing threshold go back to the reference path (docs/STATUS.md)
• CUDA + Triton only; calling the kernel on CPU raises and names dequant_ref
• The fidelity claim is measured against fp32 accumulation on NVIDIA tensor cores; other backends must re-measure (docs/PORTABILITY.md)
• Every speed figure is per card -- the census is sm_86, the head-to-head and the real-finetune cells are RTX 4090 / H100, the torch._grouped_mm cell is RTX 5090 -- and none is an architecture-wide statement; the registered comparator is this repository's own per-expert dequantize-then-GEMM loop, and no registered cell times bitsandbytes' own packed 2-D inference forward (0.50.0 and later)
• A uniform int4 grid cannot represent e2m1 levels: do not quantise-bake an already-MXFP4 checkpoint to NF4 or int4 (README)
• gemv_mxfp4_b32 re-streams weights per row: it is a decode lever for a handful of rows; batched rows use the grouped GEMM or the consumer's NF4 path
• No ROCm/XPU
• The claims here are fidelity and provenance results from single runs per card (a rented H200 pod for serving, an L40S and an A2000 for training, the K3 oracle on CPU); the MXFP4 GEMM has no speed census of its own and gemv_mxfp4_b32 is inside no registered claim
• The int4-b32 GEMV and the calibrated packer are measured-private (claims gnf4.serve.int4-b32-gemv, gnf4.serve.gptq-pack-int4-b32): the receipts live in the consumer's private audit tree
• Split-K on the decode GEMV was refuted once and ships dormant (claim gnf4.retired.splitk-gemv records the refutation); the shipped plan is the measured one
• Whether int4 experts pass a quality gate is the consumer's per-family question (experts4bit-qlora docs/STATUS.md): small-active-parameter models fail it with round-to-nearest packing
• The single-stream decode anchor (claim gnf4.serve.decode-anchor-5090) is a consumer-measured anchor of the serving class on a knob-off basis, not a measurement of this lane, and is not cited here as its evidence
• One card class measured: both claims were run on the RTX 5090 (sm_120). The sm_89+ requirement is the kernel's precondition, not a measured statement about Ada or Hopper parts
• Adds one e4m3 rounding on q and one on p, so its test bound is wider than the f32 path's (kernel/test_fp8_paged_attn.py)
• pack_heads=True falls back to the split fp8 kernel with a one-time RuntimeWarning where the packed tile exceeds the card's shared memory
• Served-path parity (Granite, gpt-oss, Gemma-4) is measured in experts4bit-qlora's private receipts; the evidence here is kernel-level parity against paged_attn_ref
• KV budgets in docs/context-budgets.md are rung-one (A2000) numbers; pending rows must not be promoted
• Open, #319 (claim gnf4.open.f32-compute-modes-triton34, status open, so it backs nothing): on torch 2.8.0+cu128 / triton 3.4.0, on unmodified main, the split and packed f32 modes miss their fp32 reference beyond the registered tolerance in a subset of the paged-attention tests; the claim carries the numbers and the issue. A kernel that imports and launches is not numerically supported; this entry stays unsupported until the issue closes and the register carries a passing cell
• A lane on this path cannot be gated with -x on those torch/Triton pairs; gate the fp8 modes with -k "f8dot or pf8" (README, docs/STATUS.md)
• Explicit GNF4_ATTN_COMPUTE=f32 or compute='f32' selects this path on any card, sm_89+ included: the explicit request is honoured, never downgraded to fp8
• The composition timing and the paired perplexity deltas are measured-private (claim gnf4.serve.decode-glue-kernels): real runs through the consumer on the RTX 5090 whose receipts live in a private audit tree; what is checkable here is each kernel's parity test
• int4_b32 imports triton at module level and is not importable without it
• reduce_partials (0.29.0) and combine_rows (0.30.0) are inside no registered claim: capability without a published measurement
• NVMe is a batch tier: per-token latency is bounded by the device link (claim gnf4.nvme.tier-batch-only)
• Expert prefetch closed negative over four arcs (claim gnf4.flagship.prefetch-closed-negative)
• The cold-engine 'free floor' premise was refuted (claim gnf4.cold-engine.phase0-premise-refuted); read docs/INDEX.md before quoting any cold-engine number
• This entry is the primitive layer (arena, reader, row tier, low-level residency); which bytes live where and the end-to-end run of a MoE larger than VRAM are experts4bit-qlora's pages run-moe-larger-than-vram (decision/router) and offload-moe-experts-to-cpu-or-nvme (model-level integration)
• A quantize-at-bake arena (nvme_bake_nf4) is a documented two-hop transform: its manifest's bake_mode says it is bit-identical to the quantiser's output, not to a bf16 release
• Making the verifier pass on non-matching bytes is the highest-severity class in SECURITY.md; report privately

Current measured position

Representative active claims, selected from docs/claims.json at the source commit; retired and superseded claims never appear here, and the register's tiers are counted above (claims tiered in the public register: 11 confirmed · 9 measured · 3 measured-private · 1 projected · 3 open · 1 superseded · 2 retired). Each card links its receipt at the pinned commit.

Current status and compatibility

v0.30.2 · current status → — the position page is docs/STATUS.md at the pinned commit (latest on main, unpinned; not the source of any fact rendered here); the site's status and compatibility pages are generated from the same commit, and the evidence words are defined once in the system manifest's vocabulary.

Problems this solves

• How do I run paged decode attention over an FP8 KV cache for a quantized MoE serving path, with sliding windows, attention sinks and a custom scale? — fp8_paged_decode_attention runs paged flash-decode over an E4M3 KV cache with windows, sinks and a custom scale; the fp8 compute path is supported on sm_89+, the f32 path is open under #319.
• How do I run a single-token INT4 decode GEMV over routed MoE experts, and pack calibrated (GPTQ) weights for it? — gemv_int4_b32 runs the batch-1 int4-b32 decode GEMV with exact integer accumulation, and gptq_pack_int4_b32 packs calibrated weights onto the same bytes; its numbers are measured-private.
• How do I run MoE expert inference natively on MXFP4 (e2m1 + e8m0) weights, straight from the released checkpoint bytes? — gemm_mxfp4_grouped multiplies gpt-oss and Kimi-class experts on their released e2m1 blocks and e8m0 scales, with no requantization to NF4 and no bf16 materialization.
• How do I run a grouped GEMM directly on NF4 packed MoE expert weights, without dequantizing to bf16 first? — gemm_4bit_grouped runs the routed MoE expert GEMM on bitsandbytes NF4 packed weights in one Triton launch, decoding nibbles in registers with fp32 accumulation and no bf16 expert tensor.
• How do I run a MoE whose experts do not fit in VRAM, streaming them from pinned host RAM or an NVMe arena? — The kernel-side storage primitives for experts that do not fit VRAM: the NVMe arena bake and O_DIRECT reader, the pinned-DRAM row tier and low-level residency, and the GPU-driven host gather.
• How do I verify that the quantized expert bytes I serve or train on are the released checkpoint bytes? — file_tensor_sha256, provenance_table and verify_arena_matches hash safetensors byte ranges against the loaded and baked bytes, and verify_provenance re-derives a run's hash table.

Canonical links

• GitHub: https://github.com/pjordanandrsn/grouped-nf4-gemm · issues · releases
• PyPI: https://pypi.org/project/grouped-nf4-gemm/
• Pinned at e2af4cfb91b2 (the source of every fact on this page): docs/capabilities.json · docs/STATUS.md · docs/claims.json · docs/SOLUTIONS.md · docs/system-manifest.json
• Latest on main (unpinned; not the source of any fact rendered here), beside each pinned document above: docs/capabilities.json · docs/STATUS.md · docs/claims.json · docs/SOLUTIONS.md · pinned counterparts: docs/claims.json@e2af4cfb91b2
• Machine-readable: this page as JSON · as Markdown · packages.json · system-manifest.json · system-manifest.provenance.json · /ml/llms.txt