Repository navigation
[MiniMax-M3.1] Mono decode over the MXFP4 KV and index caches - #2812
Draft
ganyi1996ppo wants to merge 8 commits into
Draft
ganyi1996ppo wants to merge 8 commits into
ganyi1996ppo wants to merge 8 commits into
Conversation
…-dtype mxfp4) K1's index-K insert writes the 8704 B MXFP4 page the way aiter's fused writer does (mxfp4ScaleE8m0 + v_cvt_scalef32_pk_fp4_f32), and the scorer reads it: a 16 B key load and a dword of e8m0 scales a tile, widened with v_cvt_scalef32_pk_bf16_fp4 into the same bf16 MFMA (no permlane transpose); q_frag follows the 32-element groups. Mono's dispatch / weights accept an MXFP4 index cache with fp8 KV; MXFP4 / NVFP4 main KV, NVFP4 index and indexer CP stay refused. ATOM_MONO_CHECK also compares the step's index-cache bytes after the mono layer and after the original one. Tests: index-K row store == aiter writer byte for byte, block scores == score_flydsl_fp4, dispatch refusal matrix, build key, contract trace. Co-Authored-By: Claude <noreply@anthropic.com>
K1 writes the main K (post norm / rope) and V rows into the 8704 B MXFP4 pages with the same row store as the index key, so a page is the one aiter's fused writer leaves; the dispatch no longer refuses the MXFP4 main KV (NVFP4 stays refused). The split stage over those pages: QK is one v_mfma_scale_f32_16x16x128 (A = K fp4 with its e8m0 group scale, B = the e4m3 q), softmax per wave with P in bf16 registers, V widened to bf16 into a wave-private LDS tile and read back transposed (ds_read_b64_tr_b16) for 8 bf16 MFMAs per tile, the 8 waves' (m, O) merged through LDS; sp_o / sp_m / sp_l, merge and o_proj are unchanged. The V tiles share the LDS of the PV partials (a barrier between the last tile and the first partial store): no Smem growth, the 16-token limit is unchanged. A selected block is one page-128 page in the sparse table (PAGES_PER_BLOCK is a build parameter). Tests: the split stage alone against a torch softmax over the dequantized pages (NaN-poisoned tail pages, empty partitions, a wrong-table control), contract traces and build keys with kv_mxfp4, dispatch matrix. ATOM_MONO_CHECK compares the K / V / index pages (and FP8 scales) of the mono layer with the original layer's. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
…out convention) With the MXFP4 KV cache, K1's head task writes q as e4m3 (unit scale, saturating, cast from the f32 post-norm/rope elements), as the aiter fused writer's q_out does, and the split stage reads those bytes into LDS as is, instead of a bf16 q rounded and cast in the stage. fp8-KV Mono keeps bf16 q. Co-Authored-By: Claude <noreply@anthropic.com>
Factor the clamp-and-pack into pack_q_e4m3 and test it on a GPU: in-range values are torch's e4m3 cast, beyond +-448 (inf too) saturates to +-448, NaN stays NaN. Without the clamp the hardware converts >448 to NaN, so the test fails on it. Co-Authored-By: Claude <noreply@anthropic.com>
A DSpark block drafter first-compiles at B >= 2 because torch.compile's mark_dynamic cannot make a size-1 batch dim dynamic. With --max-num-seqs 1 the batch never exceeds 1 and the per-seq buffers hold one row, so the B = 2 warmup ran past them; keep the usual >= 1 floor there. This is the recipe's CONC=1 point.
…shuffled GEMM (opt-in) ATOM_MXFP8_1X32_PRESHUFFLE=1 (default 0) holds native FP8 weights with 1x32 E8M0 scales (16, 16)-preshuffled on gfx950, so native_quant_linear's existing AITER call (gemm_a8w8_blockscale, isBpreshuffled=True) runs them on FlyDSL's MXFP8 bmm / mxpsh kernels, picked by AITER's mxscale bpreshuffle tuned table, instead of the Triton group32 GEMM. Weight bytes and scales stay the checkpoint's, so the output is bit-identical. The target is the MiniMax-M3.1 DSpark draft, whose MXFP8 linears may not be re-quantized; AITER already ships MiniMax-M3 1x32 rows for its shapes (ROCm/aiter#6060). weight_is_stored_preshuffled answers for both the load and the weight sync. Tests: truth table, dispatch, and a gfx950 bit-exact comparison against the row-major path.
…ncurrency The agentic recipe's M3.1 section now runs one configuration from c1 to c48: MXFP4 sparse KV and MXFP4 index cache with Mono decode, the draft as shipped (no draft PTPC). Per point only the prefill delayer (on up to c16), --max-num-seqs (CONC at c1/c2, else 2 x CONC) and the capture sizes change; c32 captures every size 1..64 (P90 interactivity 122.7 -> 140.6). The measured table is recomputed from each 3600 s run's aiperf export, with total and output throughput, against SGLang CI run 40110172586. The optional ATOM_MXFP8_1X32_PRESHUFFLE=1 (draft MXFP8 linears on AITER's preshuffled GEMM) is documented with its c32 run and gsm8k.
This branch has not been deployed
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
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
The Mono decode (fused per-layer kernels) now serves MiniMax-M3.1 with MXFP4 caches on gfx950:
--index-cache-dtype mxfp4over an fp8 KV cache: K1 writes the index-K MXFP4 pages and the index scorer reads them.ATOM_M3_SPARSE_KV_MXFP4=1: K1 writes the MXFP4 K / V pages, and the split stage reads them.q_outconvention, which the non-Mono path now uses too ([MiniMax-M3.1] MXFP4 sparse KV cache and index cache #2770).Before this, Mono refused MXFP4 caches, so MXFP4 and Mono could not be combined.
Depends on
ganyi/m31-mxfp4-kv-index); retarget tomainafter it lands. The diff here is only the five Mono commits.Tested
Stack: aiter main 849f711d39 + the four PR heads above, ATOM head f328ca2; gfx950, M3.1-Flash MXFP4, TP4, DSpark K=7, 2026-10-29. Each number below is one run unless it says otherwise.
Unit tests:
tests/test_minimax_m3_*.py364 passed, 1 skipped, fresh FlyDSL cache.FLYDSL_RUNTIME_ENABLE_CACHE=0 ARCH=gfx950 python -m atom.models.minimax_m3.mono.contractpasses.black --checkandruff checkare clean on the changed files. A new GPU test covers K1's e4m3 q store (in-range values equal torch's e4m3 cast, beyond +-448 and inf saturate, NaN stays NaN); it fails if the clamp is removed.Layer check (
ATOM_MONO_CHECK=1, eager, MXFP4 K / V + MXFP4 index, ~38K-token prompts): each Mono layer against the original layer on the same step, 57 sparse layers x 7 steps = 399 layer checks per run.GSM8K (lm-eval, 5-shot, chat template, temperature 0, all 1319, concurrency 2, real acceptance, Mono on): flexible 0.9469 +- 0.0062, strict 0.9098 +- 0.0079. Before this change (MXFP4 + Mono, bf16 main Q, other machine, 2026-10-28): 0.9477 / 0.9030; non-Mono 0.9477 / 0.8999. In three 6 s profiler windows during the run, Mono served 1830 of 1884 decode forwards (97.1%).
Step time (ctx ~100K, B=1 / B=2, chunk gap per decode step, synthetic acceptance 4.04, same 4 cards, 5 repeats in one server session each, median of the repeat means):
The earlier bf16-main-Q version of Mono MXFP4 measured 6.56 / 8.54 ms. A second 3-repeat round on another machine gave 6.70 / 8.41 ms.
Limits
select_longpath (ctx > 65K) is untested.-k monotest run has 3 failures intests/test_deepseek_v41_mono_build_key.py; they fail the same way on [MiniMax-M3.1] MXFP4 sparse KV cache and index cache #2770's head.🤖 Generated with Claude Code