Skip to content

Releases: pjordanandrsn/grouped-nf4-gemm

v0.4.0

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 31 Jul 14:09
6f95721

Changelog

0.4.0 — 2026-07-31

Version-number correction. No code change from 0.3.1.

0.3.1 shipped #26 (prefill — many tokens per call) alongside a packaging fix and
described itself as "nothing else changed". A new capability went out under a patch
label, so anyone reading versions rather than diffs had no signal it existed. 0.4.0
is the same tree under the number semver says that feature warranted, and 0.3.1's
entry now describes what it actually contained.

Nothing to migrate: if you are on 0.3.1 you already have prefill.

0.3.1 — 2026-07-30

0.3.0 announced two modules it did not ship. mxfp4_residency and
nvme_residency were absent from pyproject's py-modules allowlist, so they
were never in any wheel — while 0.3.0's release notes described
K3_RESIDENCY_KINDS, fuse_gate_up_segments and Mxfp4NvmeResidency as
shipped. Anyone who followed those notes got ModuleNotFoundError.

py-modules is an explicit allowlist: adding a file under kernel/ does not
package it, and nothing warned. kernel/test_packaging_covers_kernel.py now
diffs the directory against the allowlist and fails naming the missing modules,
so the next gap lands on the pull request instead of on a user. A module that
genuinely should not ship goes in _DELIBERATELY_UNPACKAGED with a reason,
which keeps that decision visible rather than silent.

Correction (added after release): 0.3.1 also shipped a new capability, and its
notes said it did not.
#26 — prefill: the engine takes many tokens per call —
merged to main before the packaging fix and was swept into this tag. The line
"nothing else changed" was written from the packaging work alone rather than from
the full v0.3.0..main delta, and it is wrong.

What that feature does: the engine was decode-only (a_buf.copy_(x.expand(k, -1))
broadcasts ONE token's hidden state across the k slots). The kernel never was —
gemm_mxfp4_grouped's sizes is a per-group token count and already switches to
the tiled path above one row — so this is engine plumbing. Prefill is not decode in
a loop, and the difference is I/O: stepping T tokens re-reads the whole dense side T
times and every routed row T times; entering each layer once for the prompt reads
each distinct expert once. Measured on Kimi K3 at full depth, 7 tokens: dense
108.76 GB once vs 761 GB; expert rows 7,080 vs 10,304 (31 % deduped by route
overlap); 233 GB vs 942 GB of I/O, 187.4 s vs 643 s. VRAM peak unchanged at
3.59 GB.

By semver that warranted a minor bump, not a patch. See 0.4.0.

0.3.0 — 2026-07-30

If you are on triton 3.2, the pipelined MXFP4 engine did not work at all.
Both kernel factories imported triton.language as tl into their locals. With
from __future__ import annotations, BLOCK: tl.constexpr is the string
"tl.constexpr", which triton resolves against the jitted function's
__globals__ — triton 3.4 tolerates it, 3.2 raises NameError('tl is not defined') from inside the compiler. Moving tl to module globals takes
test_mxfp4_residency's companion suite from 7 failed to 7 passed on a
triton-3.2 box.

Experts now serve from NVMe, and the arena you already baked is readable
whatever order it was baked in.
Mxfp4NvmeResidency reads gate_up at one
computed offset, so it needs the two blocks segments adjacent and the two scales
segments adjacent — while arena_experts.K3_KINDS, the released-K3 spelling,
interleaves per projection. Both orders are legitimate; they are for different
consumers. Rather than force a re-bake of 1.45 TB, the gather takes a per-piece
(src, dst, len) table and lands segments where the engine expects them: no
extra bandwidth, nothing on disk touched, any bake order readable. The
identity case keeps the original contiguous kernel, so the previously measured
path is untouched.

K3_RESIDENCY_KINDS carries the real tensor names in the order this engine
wants, and the two constants cross-reference each other so the trap is visible
from either file. A mis-ordered arena is now refused with a message naming the
order that works, instead of "trailing dims differ".

The k slots are shared across layers. They were per-layer, so VRAM scaled
with depth — on a 92-layer model that is the difference between fitting and not.

Also:

  • nvme_reader: a pinned tensor is not reliably page-aligned. O_DIRECT
    needs the buffer address aligned, and assuming pin_memory() delivers that
    is how a good checkpoint reads as a corrupt one.
  • K3's SiTU activation is registered from the release source — none of the
    guesses matched.
  • moe_layer_forward passed model-global expert ids into the kernel's
    stack index, which reads out of bounds silently. Every toy fixture used ids
    smaller than the group count, so they indexed validly by coincidence; on
    896-expert K3 with top-16 it would have corrupted essentially every layer.
  • The equivalence fixture quantized random nibbles against random e8m0 scales,
    whose magnitudes a GLU squares into overflow — and torch.equal is False for
    identical NaNs, so byte-identical outputs compared unequal. It now quantizes
    realistic weights and compares bitwise.

Receipts (docs/): the Phase-1 oracle passes — decode bit-identical to
compressed-tensors across 33,030,144 elements, max delta 0. A real-bytes arena
round-trip on released Kimi-K3 matched 48/48 segments against the shipped
safetensors, with a byte-flip negative control. Each carries a note on what its
OpenTimestamps anchor does not prove: the stamps were applied after the runs,
so they establish the text has not changed since, not that the protocol predated
the data.

v0.3.1

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 31 Jul 01:18
0d97711

Changelog

0.3.1 — 2026-07-30

0.3.0 announced two modules it did not ship. mxfp4_residency and
nvme_residency were absent from pyproject's py-modules allowlist, so they
were never in any wheel — while 0.3.0's release notes described
K3_RESIDENCY_KINDS, fuse_gate_up_segments and Mxfp4NvmeResidency as
shipped. Anyone who followed those notes got ModuleNotFoundError.

py-modules is an explicit allowlist: adding a file under kernel/ does not
package it, and nothing warned. kernel/test_packaging_covers_kernel.py now
diffs the directory against the allowlist and fails naming the missing modules,
so the next gap lands on the pull request instead of on a user. A module that
genuinely should not ship goes in _DELIBERATELY_UNPACKAGED with a reason,
which keeps that decision visible rather than silent.

Nothing else changed. Everything 0.3.0 describes is accurate; two of its modules
were simply unreachable from an installed wheel.

0.3.0 — 2026-07-30

If you are on triton 3.2, the pipelined MXFP4 engine did not work at all.
Both kernel factories imported triton.language as tl into their locals. With
from __future__ import annotations, BLOCK: tl.constexpr is the string
"tl.constexpr", which triton resolves against the jitted function's
__globals__ — triton 3.4 tolerates it, 3.2 raises NameError('tl is not defined') from inside the compiler. Moving tl to module globals takes
test_mxfp4_residency's companion suite from 7 failed to 7 passed on a
triton-3.2 box.

Experts now serve from NVMe, and the arena you already baked is readable
whatever order it was baked in.
Mxfp4NvmeResidency reads gate_up at one
computed offset, so it needs the two blocks segments adjacent and the two scales
segments adjacent — while arena_experts.K3_KINDS, the released-K3 spelling,
interleaves per projection. Both orders are legitimate; they are for different
consumers. Rather than force a re-bake of 1.45 TB, the gather takes a per-piece
(src, dst, len) table and lands segments where the engine expects them: no
extra bandwidth, nothing on disk touched, any bake order readable. The
identity case keeps the original contiguous kernel, so the previously measured
path is untouched.

K3_RESIDENCY_KINDS carries the real tensor names in the order this engine
wants, and the two constants cross-reference each other so the trap is visible
from either file. A mis-ordered arena is now refused with a message naming the
order that works, instead of "trailing dims differ".

The k slots are shared across layers. They were per-layer, so VRAM scaled
with depth — on a 92-layer model that is the difference between fitting and not.

Also:

  • nvme_reader: a pinned tensor is not reliably page-aligned. O_DIRECT
    needs the buffer address aligned, and assuming pin_memory() delivers that
    is how a good checkpoint reads as a corrupt one.
  • K3's SiTU activation is registered from the release source — none of the
    guesses matched.
  • moe_layer_forward passed model-global expert ids into the kernel's
    stack index, which reads out of bounds silently. Every toy fixture used ids
    smaller than the group count, so they indexed validly by coincidence; on
    896-expert K3 with top-16 it would have corrupted essentially every layer.
  • The equivalence fixture quantized random nibbles against random e8m0 scales,
    whose magnitudes a GLU squares into overflow — and torch.equal is False for
    identical NaNs, so byte-identical outputs compared unequal. It now quantizes
    realistic weights and compares bitwise.

Receipts (docs/): the Phase-1 oracle passes — decode bit-identical to
compressed-tensors across 33,030,144 elements, max delta 0. A real-bytes arena
round-trip on released Kimi-K3 matched 48/48 segments against the shipped
safetensors, with a byte-flip negative control. Each carries a note on what its
OpenTimestamps anchor does not prove: the stamps were applied after the runs,
so they establish the text has not changed since, not that the protocol predated
the data.

v0.3.0

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 30 Jul 22:13
487f5e7

Changelog

0.3.0 — 2026-07-30

If you are on triton 3.2, the pipelined MXFP4 engine did not work at all.
Both kernel factories imported triton.language as tl into their locals. With
from __future__ import annotations, BLOCK: tl.constexpr is the string
"tl.constexpr", which triton resolves against the jitted function's
__globals__ — triton 3.4 tolerates it, 3.2 raises NameError('tl is not defined') from inside the compiler. Moving tl to module globals takes
test_mxfp4_residency's companion suite from 7 failed to 7 passed on a
triton-3.2 box.

Experts now serve from NVMe, and the arena you already baked is readable
whatever order it was baked in.
Mxfp4NvmeResidency reads gate_up at one
computed offset, so it needs the two blocks segments adjacent and the two scales
segments adjacent — while arena_experts.K3_KINDS, the released-K3 spelling,
interleaves per projection. Both orders are legitimate; they are for different
consumers. Rather than force a re-bake of 1.45 TB, the gather takes a per-piece
(src, dst, len) table and lands segments where the engine expects them: no
extra bandwidth, nothing on disk touched, any bake order readable. The
identity case keeps the original contiguous kernel, so the previously measured
path is untouched.

K3_RESIDENCY_KINDS carries the real tensor names in the order this engine
wants, and the two constants cross-reference each other so the trap is visible
from either file. A mis-ordered arena is now refused with a message naming the
order that works, instead of "trailing dims differ".

The k slots are shared across layers. They were per-layer, so VRAM scaled
with depth — on a 92-layer model that is the difference between fitting and not.

Also:

  • nvme_reader: a pinned tensor is not reliably page-aligned. O_DIRECT
    needs the buffer address aligned, and assuming pin_memory() delivers that
    is how a good checkpoint reads as a corrupt one.
  • K3's SiTU activation is registered from the release source — none of the
    guesses matched.
  • moe_layer_forward passed model-global expert ids into the kernel's
    stack index, which reads out of bounds silently. Every toy fixture used ids
    smaller than the group count, so they indexed validly by coincidence; on
    896-expert K3 with top-16 it would have corrupted essentially every layer.
  • The equivalence fixture quantized random nibbles against random e8m0 scales,
    whose magnitudes a GLU squares into overflow — and torch.equal is False for
    identical NaNs, so byte-identical outputs compared unequal. It now quantizes
    realistic weights and compares bitwise.

Receipts (docs/): the Phase-1 oracle passes — decode bit-identical to
compressed-tensors across 33,030,144 elements, max delta 0. A real-bytes arena
round-trip on released Kimi-K3 matched 48/48 segments against the shipped
safetensors, with a byte-flip negative control. Each carries a note on what its
OpenTimestamps anchor does not prove: the stamps were applied after the runs,
so they establish the text has not changed since, not that the protocol predated
the data.

v0.2.6 — a K3 MoE block computes from the arena

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 30 Jul 05:16
v0.2.6

0.2.5 gave the NVMe arena a compute consumer. Nothing called it — moe_layer_forward's only callers were its own tests. This is the model-side link.

The patch point

KimiSparseMoeBlock.moe_infer turned out to be a better fit than expected: upstream already builds exactly what the grouped kernel wants.

idxs = topk_ids.view(-1).argsort()
sorted_tokens = x[idxs // topk_ids.shape[1]]   # group-sorted [T, K]  == a_cat
tokens_per_expert = cnts.sum(dim=0)            #                      == sizes
for i, num_tokens in enumerate(tokens_per_expert):
    expert_out = expert(tokens_for_this_expert)   # one matmul PER expert

That per-expert loop — each iteration first materialising an expert's weights — is precisely what gemm_mxfp4_grouped collapses into one launch on packed bytes.

So the patch swaps the loop and nothing else. Sorting, weighting, unsorting and the shared-expert path remain upstream's; patching those would change the model rather than accelerate it. use_latent_moe's projections and the shared experts live in forward(), outside moe_infer, and are untouched.

Honest instrumentation

arena_call_stats() reports patched and calls separately, because a patch count is not a call count — both inert fixes in this project's history patched cleanly and never executed. patched=3, calls=0 is an expressible state, and a gate pins that zero before a forward, with a positive control that moves it.

Layer 0 (dense first_k_dense_replace) is skipped. Layers missing from a partial bake stay on resident experts rather than being silently routed to nothing. Patching nothing raises rather than returning 0.

Still required for K3 inference

A baked arena (1446.5 GB relocation — measured, not estimated) and the always-active weights resident. This release makes the model able to use an arena; it does not create one.

v0.2.5 — the arena gets a compute consumer

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 30 Jul 04:35
v0.2.5

The NVMe tier had a storage half and no compute half. This closes it.

arena_experts — arena → kernel

ArenaExpertSource slices a baked arena row by the bake's own recorded segment geometry and returns the fused [E, N, K//2] blocks and [E, N, K//32] e8m0 scales that gemm_mxfp4_grouped already consumes. moe_layer_forward() runs gate → GLU → down straight off it.

Those shapes are not a coincidence: a DeepSeek-V3-lineage MXFP4 release ships each expert as exactly weight_packed [N, K//2] + weight_scale [N, K//32], which is the kernel's input contract. So the bytes travel disk → arena → GEMM with no dequantize round trip and no requantization — the bytes multiplied are the bytes shipped, which is what makes per-tensor hash provenance mean anything at the point of compute.

Routing, attention and the residual stay with the caller. This is the expert block, not a model.

moonshot_gather is name-scheme pluggable (#16)

Released Kimi K3 renamed three parts at once — container mlp.expertsblock_sparse_moe.experts, projections gate/up/down_projw1/w3/w2, and weightweight_packed + weight_scale. Under the K2 default a real K3 index discovered zero experts, silently, because a dense layer legitimately reports 0 too. detect_scheme() now identifies the family or returns None rather than gathering nothing.

Validated against the real released checkpoint: the optional gate ran with K3_SNAPSHOT pointed at moonshotai/Kimi-K3's own 497,220-entry index.

Fix: build_group_tiles mutated its sizes argument

enumerate over a tensor yields 0-dim views, so left = m; left -= take subtracted in place and zeroed the caller's array. Every prior caller passed a list or called once, so it never surfaced; the first to run two GEMMs off one sizes tensor — gate/up/down in a single layer — got [0, 0] on the second. Affects the NF4 and MXFP4 grouped paths alike. Host-side tile indices only, so no computed value moves and no confirmatory result is disturbed.

Verification

19/19 green on an RTX A5000, no skips, including both end-to-end gates — arena-fed GEMM bit-identical to memory-fed, and the full MoE layer forward. Every byte-identity gate is paired with a control that must fail (flip one arena byte, and it is detected).

v0.2.4 — differentiable grouped NF4 GEMM (training) + NVMe arena surface

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 30 Jul 00:23
v0.2.4

v0.2.4 — the differentiable grouped NF4 GEMM (training half) + NVMe arena surface

nf4_qlora makes the fused kernel trainable: an autograd.Function whose backward
re-decodes one expert at a time, plus fused_grouped_lora composing the frozen
4-bit projection with the trainable B(Ax) delta at the pre-activation point.

Verified on real hardware before this tag:

  • gradient correctness -- 5/5 on sm_86 and sm_89, gradient relative error
    exactly 0.0 against the dequantize-then-matmul reference, forward 0.0027
  • offload safety at 48 layers (Qwen3-30B-A3B, RTX 4090): 0.756x the reference
    peak, no OOM, where holding weights on the ctx OOMed at 22.41 GB
  • frozen experts bit-identical after training -- 192 tensors, 16.3 GB hashed,
    zero placeholders, with a positive control proving a single flipped byte IS
    detected
  • loss parity with the reference on varied data and a held-out eval:
    0.337073 vs 0.340271 (0.95%), train median delta 0.0136 (band 0.05)

Also: triton is now a Linux-only dependency, so the pure-torch provenance and
NVMe-arena surface installs on macOS/Windows; nvme_arena, nvme_bake_nf4 and
nvme_reader ship in the wheel with the smoke exercising the CLI and the O_DIRECT
alignment guard.

0.2.3 was committed but never tagged or published; PyPI goes 0.2.2 -> 0.2.4.

Co-Authored-By: Claude Opus 5 noreply@anthropic.com
-----BEGIN PGP SIGNATURE-----

iJEEABYKADkWIQQ7n/u5ZVSf3n/PS1XqrS1lj4+/ngUCamqZMRsUgAAAAAAEAA5t
YW51MiwyLjUrMS4xMiwwLDMACgkQ6q0tZY+Pv54NXAD9FrbYfM6cdN3SNnF2kVep
CuWRNZjY02/kcwG8andZxU4A/2BTNbZs6adxX9nNjtfhSRXKMDp+/R5Z4hwZdODC
hzEG
=n67z
-----END PGP SIGNATURE-----

v0.2.2 — CPU-runnable quickstart (CI-executed) + teaching asserts

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 21 Jul 01:51

docs/tests only, no kernel-math changes.

  • CPU-runnable README quickstart (NF4 roundtrip, MXFP4 roundtrip, provenance-in-4-lines), extracted and executed by CI (test_readme_cpu_block.py) so examples can't drift from the API.
  • Taught CPU refusal: gemm_4bit_grouped on CPU now names the CUDA requirement and points at dequant_ref (pinned by test_cpu_refusal.py; exempts TRITON_INTERPRET=1).
  • Docstring examples on the five surface functions; quantize_pack_nf4 teaches on non-2-D input.
  • Two courtesy clauses (Unsloth bf16-training context; Axolotl/PEFT factual phrasing).

🤖 Generated with Claude Code

0.2.1 — front door repair

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 21 Jul 00:18

Docs + CI only, no kernel changes: every relative link in the PyPI long_description was dead (rendered as pypi.org 404s) — all 25 receipts links are now absolute GitHub URLs pinned to this tag. Permanent built-artifact smoke added to CI and the publish gate (clean-venv wheel install, README-surface import battery, CPU micros, link + ref-consistency check).

Packaging note: the reported '0.2.0 wheel missing its modules' did not reproduce — the 0.2.0 wheel on PyPI contains all 13 modules and a clean-venv battery imports 13/13. The ImportError experience reproduces via a stale package-index edge (or the grouped-mxfp4-gemm stub resolving its dependency against one), not via wheel contents. No packaging change ships in 0.2.1.

0.2.0 — the MXFP4 native-byte lane

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 20 Jul 20:30

The e2m1 codebook is a table swap of the NF4 mainloop, so the MXFP4 lane ships in this package: grouped GEMM on a checkpoint's exact released bytes with per-tensor sha256 provenance, the pipelined hot/cold decode engine, QLoRA with recompute-in-backward over frozen native bytes, and the executable verify_provenance receipt.

Stamped receipts (in docs/mxfp4/): gpt-oss-120b fused-native serve at reference perplexity (the +9.4% NF4-requant tax deleted); 120b QLoRA at 9.82 GB peak VRAM with 144/144 file == loaded == post-train hashes. Kimi K3 gather scaffolding included, K2-verified — every K3-specific number waits for the 2026-07-27 weights drop and the per-model oracle STOP gate.

Wheel ships: nf4_grouped, nf4_pack_ref, host_gather, mxfp4_pack_ref, mxfp4_grouped, mxfp4_loader, mxfp4_pipelined, mxfp4_qlora, mxfp4_native_load, moonshot_gather, verify_provenance, plus the run/gate harnesses.

Hardened over five review rounds before merge (10 findings: wheel self-consistency, padding safety on all three forward paths, sharded-checkpoint provenance, snapshot resolution, marker-guard fixed-string matching).

v0.1.1

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 18 Jul 08:16

Doc-only release: the PyPI project page now carries the badged README with the Install section (aliases, trusted publishing, PEP 740 note). Kernel modules are byte-identical to v0.1.0.

This release was prepared with AI assistance (Claude).