feat(moe): mxfp4 marlin grouped gemm and resident-ep bf16 seam - #307
Merged
Conversation
Contributor
Author
|
Validation recorded internally; CI is the public record. |
Decode-only companion to the streamed pure-DP path: per-rank stacked shard built once from the copy-engine host tensors; pad -> all_gather -> global router -> fused_moe_bf16 with non-local slots weight-masked -> all_reduce(SUM) -> local slice. Empty ranks participate fully in the collectives (their experts serve other ranks). Sized purely from the worker-synced num_tokens_per_rank scalar.
Repack, dequant, SiTU epilogue, hard-fail contract. Nibble order and E8M0 semantics settled against compressed-tensors 0.17.1 as an INDEPENDENT oracle (ours was bent to match our own kernel, so our pair proved nothing), on a real K3 expert tensor pulled from the checkpoint: verdict frozen into mxfp4_oracle_vector.py with sha256 pins and a clamp-free reference dequant, so it never needs re-deriving. marlin_weight_prep.py: repack_mxfp4_to_marlin_gs32 — pure nibble rearrangement carrying E8M0 uint8 scales through (E2M1 magnitudes are non-uniform, so INT4 conversion would be a second lossy quantization and is refused), exact E8M0->bf16 expansion (<<7; edge bytes 0x00/0xFF raise), the w1||w3 storage-adjacent fused repack, an exact CPU inverse for the round-trip proof, and R1-R8 hard-fail contract checks. marlin_grouped_gemm.cu: dequant_e2m1 + WCodec/Act templating of the grouped kernels, a SiTU epilogue, and two new pybind entries (grouped_marlin_gemm_m16_mxfp4, _m16_s1_mxfp4_situ) with TORCH_CHECK hard-fail seams so the raw bindings cannot bypass the contract. HONESTY NOTE: the existing INT4 kernel bodies are TEMPLATED (dequant_u4b8 -> dequant_w4<CODEC>), behavior-preserving by construction but not byte-identical -- the staged GPU ladder carries INT4 M16 and fused-S1 regression arms (T6/T6b) that MUST pass before any merge to main. K2.5's python surface is append-only. marlin_grouped_moe.py: K3 wrappers with L1-L5 + activation-contract hard-fail checks per the 2026-08-04 POIS ruling -- no operational fallback; the reference path exists only behind an explicit debug opt-in whose server-side consumer is a named follow-up.
A stale installed copy in site-packages shadows the repo's in-tree build whenever the repo root is not on sys.path, and the old message said "rebuild" — the wrong fix for that failure. The L1 error now names the extension file actually loaded so an import-path problem is diagnosed as one. Split out of the Kimi-K3 M2 model commit (kernel-scope file).
Andrewxu313
force-pushed
the
tairanxu/k3-prefill-kernel
branch
from
August 12, 2026 21:46
7ad15d8 to
765971a
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Extends the production marlin grouped-GEMM machinery to MXFP4 (E2M1 nibbles + E8M0
scales) for Kimi-K3's routed experts: an on-device repack from checkpoint order to
marlin tile order, an E2M1 dequant path in the CUDA kernel, a frozen oracle vector that
pins the packing convention, and a hard-fail policy — every contract violation raises,
with the unfused reference path reachable only through the explicit
batchgen_debugparity opt-in. Also adds the resident-EP BF16 decode seam (stacked per-rank expert
shards + fused grouped GEMM + combine) that the kimi-linear decode path selects, and
the CPU repack suite plus the staged GPU parity ladder for both. Companion PRs:
tairanxu/k3-prefill-core(merges first) andtairanxu/k3-prefill-model(mergeslast).
Why
Kimi-K3 ships MXFP4 expert weights, which the existing uniform-INT4 marlin path cannot
represent; routing K3 numerics onto the production-proven marlin machinery avoids
reviving an unvalidated standalone MXFP4 kernel. The resident-EP seam gives the model
family a decode MoE mode that keeps expert shards resident instead of streaming them
per step. Parity is gated by the frozen oracle vector, the CPU repack suite, and the
staged GPU parity ladder in
tests/moe/; validated qualitatively against the BF16reference chain.
Type of Change
modelkernelcorefixinfradocsFile changes
batchgen/moe/fused_moe_bf16_resident.pybatchgen/moe/marlin_grouped_moe.pybatchgen/moe/marlin_weight_prep.pybatchgen/moe/mxfp4_grouped_gemm.pybatchgen/moe/mxfp4_oracle_vector.pybatchgen_kernels/src/moe/marlin_grouped_gemm.cutests/moe/__init__.pytests/moe/_loader.pytests/moe/gpu_parity_mxfp4_marlin.pytests/moe/test_mxfp4_marlin_repack.pyChecklist
model/kernelPR does not touch the scheduling/scaffolding layer (§2.5–§2.6).debug_*/scratch_*/tmp_*scripts,BATCHGEN_*env-guards, strayprint(), committed artifacts, orCo-Authored-By(§1/§4). (bench_*benchmarks are fine.)tests/; touched modules'MODULE.mdupdated if the public API changed (§2).bash .github/workflows/scripts/check-pr-hygiene.sh origin/mainlocally; the CI hygiene check is green.🤖 Generated with Claude Code