[MiniMax-M3.1] MXFP4 sparse KV cache and index cache by ganyi1996ppo · Pull Request #2770 · ROCm/ATOM

ganyi1996ppo · 2026-10-28T02:23:29Z

Draft. Depends on aiter PRs #7056, #7057 and #7058 (below). The base is yuhua/m3.1-flash, the integration branch of yuhua's unmerged MiniMax-M3.1 PRs, not main; once those merge this needs to be retargeted to main and rebased. Not for merge before then.

Summary

For MiniMax-M3.1-Flash, store the sparse-attention layers' K/V cache and the index-K cache as MXFP4 (e2m1 data, e8m0 scale per 32 values) and run the whole sparse path on it: one fused kernel writes K, V and the index key, aiter's FlyDSL scorers score the index cache, and aiter's pa_decode_mxfp4 does the attention.

The point is capacity: the pool holds about 1.83x the blocks of fp8 (104,776 vs 57,114 blocks per GPU at TP4, --gpu-memory-utilization 0.9, i.e. about 13.4M vs 7.3M tokens). NVFP4 gets 99,295 blocks in the same setup. At AgentX c32 and c48 on the 1200 s windows below it is ahead of both fp8 and NVFP4 on throughput per GPU and on P90 interactivity; at c16 it is level on throughput with fp8 + Mono and ahead on P90 interactivity (one run per point).

Dependencies

aiter (this PR calls the new interfaces; it does not import without them). All three PRs are open drafts, and this PR depends on them being merged:

[FlyDSL] MiniMax-M3 index score: MXFP4 index-cache decode and prefill scorers aiter#7058: score_flydsl_fp4, score_prefill_flydsl_fp4 and the fp8-Q variants (MXFP4 index scoring).
[HIP] fused_qknorm_idxrqknorm: write MXFP4 index-K and main KV caches, emit an fp8 main Q aiter#7057: fused_qknorm_idxrqknorm with kv_cache_dtype="mxfp4", and an e4m3 q_out for the main Q (the write kernel).
[FlyDSL] pa_decode: MXFP4 K/V in the shared PA tile kernel (gfx950, query_length 1) aiter#7056: pa_decode_mxfp4, including the plan-free path and the direct write.

ATOM base: unmerged PRs #2686 (Mono, author wuhuikx, branch hattie/minimax-m31-mono), #2743, #2744, #2745 and #2746, all collected on yuhua/m3.1-flash (head 3e7d1fd55 when this branch was cut). This PR's diff against that branch is only the MXFP4 work. Mono's own FP4 support is not in this PR.

Needs gfx950.

How it works

Enable with:

ATOM_M3_SPARSE_KV_MXFP4=1 --kv_cache_dtype fp8 --index-cache-dtype mxfp4 --block-size 128

ATOM_M3_INDEX_Q_FP8=1 (default off) additionally has the writer emit the indexer query as e4m3 and the scorer read it as such. It has no end-to-end measurement; see "Not measured".

Pool (mha_kv_pool.py, aiter_attention.py, minimax_m3/kv_mxfp4.py). One 8704 B page (128 x 64 B e2m1 data + 128 x 4 B e8m0 scales) per (page, kv head), K and V each [pages, nkv, 8704]; the index cache is [pages, 8704] (136 + 68 B a token). The sparse layers get their own pool, as with NVFP4.
Write (attention_mha.py). fused_qknorm_idxrqknorm quantizes K, V and the index key straight into the pool pages and emits the main Q as e4m3; there is no bf16 stage and no separate store kernel. A layer that reuses another layer's top-k writes K/V only.
Index scoring (index_topk.py, index_mxfp4.py). The scores come from aiter's FlyDSL scorers over the data and scale views of the pool page, in the layout the Triton scorers produce, so selection and the emit after it are unchanged. The block table is built at page-128 granularity (pages_per_block is a parameter of the emit now: 8 for the page-16 cache, 1 for page-128). The decode scorer's work map is built once per step into a persistent buffer, outside the graph.
Attention (sparse_attn_mxfp4.py). pa_decode_mxfp4 per selected row. Decode uses the planned path (the captured shapes are fixed). Prefill is plan-free: the planned path JIT-compiles a kernel for every row count it has not seen (about 200 ms cold, about 19 ms warm), a prefill chunk has a new row count almost every forward, and in c32 runs that stalled every TP rank by about 300 ms. Plan-free has one JIT key whatever the rows.
Score buffer. Without an out, the aiter decode scorer runs torch.full(-inf) in every layer: about 250 us of FillFunctor a step (60 layers, B=4). The builder now keeps one persistent buffer, refilled once per step outside the graph; a B=4 decode step is about 2.3% faster. Where the scorer lays scores out transposed (S*H >= 16, not TP4) the buffer is not used and nothing changes.
Refused / declined. Startup refuses (ValueError) mixing with the NVFP4 formats, plugin mode (vLLM/SGLang bridges), a block size other than 128, non-gfx950, and ATOM_M3_INDEX_Q_FP8 without --index-cache-dtype mxfp4. The sparse attention path needs sparse_local_block >= 1. Mono declines the MXFP4 sparse KV (Mono's FP4 support is a later PR). Indexer CP falls back with a warning. TBO is not supported (not tested, no startup check).

Commits

Pool and index-cache declaration, the switches and the startup refusals.
Write path: fused kernel writes K, V and the index key; top-k-reusing layers.
Index scoring with FlyDSL scorers, page-128 block table, work map and persistent score buffer.
Attention via pa_decode_mxfp4, planned decode and plan-free prefill.
ATOM_M3_INDEX_Q_FP8.

Results

MiniMax-M3.1-Flash, TP4 on four MI355X (gfx950), DSpark K=7, AgentX client with synthetic acceptance 4.04, 1200 s window, --max-model-len 1048576 --block-size 128, prefix caching on, fp8 and NVFP4 run with --kv_cache_dtype fp8 and their own index dtype. Tok/s/GPU = output throughput / 4; P90 interactivity = 1000 / ITL p90.

What is compared. One run per point. fp8 (Mono enabled), NVFP4 and the older MXFP4 were measured on the older tree (42bf5e247 for c32/c48 and NVFP4 c16; the unmodified yuhua/m3.1-flash for fp8 c16) with the aiter that ships in rocm/atom-dev:nightly_202610261558. "MXFP4" below is this PR's v2 (ATOM 40dafc923) on aiter ganyi/m3-mxfp4-all (the code in PRs A, B, C). The window is 1200 s against 3600 s in the recipe, so none of these are recipe numbers.

conc	config	tok/s/GPU	P90 interactivity	ITL p50 / p90 / p99 (ms)	TTFT p50 / p90 / p99 (ms)
c16	fp8 + Mono	163.6	156.5	3.33 / 6.39 / 9.91	289 / 724 / 3116
c16	NVFP4	158.3	146.4	4.05 / 6.83 / 14.66	318 / 711 / 2070
c16	MXFP4 (this PR)	163.8	181.0	3.09 / 5.52 / 11.36	285 / 578 / 1824
c32	fp8	333.7	78.0	7.87 / 12.82 / 50.72	354 / 1275 / 7908
c32	NVFP4	344.8	85.8	6.70 / 11.66 / 23.53	355 / 717 / 3271
c32	MXFP4 (this PR)	370.3	108.3	5.41 / 9.24 / 18.41	341 / 686 / 2918
c48	fp8	117.5	5.2	7.79 / 192.63 / 566.28	1579 / 29190 / 116149
c48	NVFP4	396.6	71.3	9.50 / 14.03 / 37.21	365 / 1049 / 6765
c48	MXFP4 (this PR)	429.0	92.1	7.30 / 10.85 / 27.45	352 / 1035 / 5379

KV blocks per GPU (c32 runs): fp8 57,114, NVFP4 99,295, MXFP4 104,776. At c48, fp8 runs out of cache capacity (engine prefix-cache hit ratio 0.22) and falls over, which is a property of the 57k-block pool, not of the sparse path. One errored request each in the fp8 c16 and NVFP4 c48 runs; none in the others. Re-running NVFP4 c48 on the v2 tree gave 396.9 tok/s/GPU and 69.6 interactivity (against 396.6 and 71.3), so the tree is not what separates the NVFP4 and MXFP4 rows.

For the MXFP4 path itself, before the persistent score buffer and plan-free prefill (older aiter as well) it measured 158.8 / 343.4 / 372.3 tok/s/GPU and 110.2 / 74.4 / 53.4 P90 interactivity at c16 / c32 / c48. At c32, plan-free prefill alone (the score-buffer-only tree 9da593f59, warm cache, against the v2 run above) gave 339.2 -> 370.3 tok/s/GPU (+9.2%) and 66.7 -> 108.3 P90 interactivity (+62%); that intermediate run is not in the table.

GSM8K (lm-eval 5-shot, chat template, temperature 0, all 1319 questions, concurrency 32, real DSpark acceptance, ATOM 40dafc923, one run each):

config	flexible-extract	strict-match	accepted tokens / forward
fp8	0.9378 +/- 0.0067	0.8923 +/- 0.0085	4.587
NVFP4	0.9401 +/- 0.0065	0.8863 +/- 0.0087	4.576
MXFP4 (this PR)	0.9424 +/- 0.0064	0.8779 +/- 0.0090	4.588

The differences are within about one standard error.

Tests

On a MI355X node (gfx950), image rocm/atom-dev:nightly_202610261558 with aiter from ganyi/m3-mxfp4-all (JIT rebuilt):

New MXFP4 tests, pytest tests/test_minimax_m3_mxfp4_cache.py tests/test_minimax_m3_mxfp4_write.py tests/test_minimax_m3_mxfp4_index.py tests/test_minimax_m3_mxfp4_attention.py: 48 passed. They need a GPU and the new aiter. Each commit passes the tests it contains (10, 12, 25, 46 and 48 passed over the five commits).
Related existing tests (the other MiniMax-M3 tests, test_mha_kv_pool, test_kv_pool_contract, test_config_*, test_quant_config, test_sparse_*, test_indexer_*, test_mtp_index_share, test_block_pool, test_minimax_m31_shared_fusion and others; 830 tests): 772 passed, 49 skipped, 9 failed. The same 9 fail identically on the unmodified yuhua/m3.1-flash head: the four in test_minimax_m31_shared_fusion.py and the five in test_sparse_attn_prefill_composition.py. They are not touched by this PR.
tests/test_envs.py: 105 passed. black --check and ruff check clean on the changed files.

Known limits / Not measured

Hkv > 1 is covered by unit tests only; the end-to-end runs are TP4 with one kv head per rank.
The top-k-reusing layers path is covered by unit tests only.
ATOM_M3_INDEX_Q_FP8 has no end-to-end measurement and is off by default.
Accuracy: the MXFP4 main-KV error is about 3-5x fp8's and 1.25-1.9x NVFP4's in the kernel tests. GSM8K only detects larger regressions (one run per config, standard error about 0.007); no other accuracy evaluation.
Every performance point is a single 1200 s run, so no run-to-run variance is available; the older-tree and new-tree rows are on different ATOM trees and aiter builds, as described above.
Indexer CP and TBO are not supported; Mono declines MXFP4.

🤖 Generated with Claude Code

ATOM_M3_SPARSE_KV_MXFP4=1 with --kv_cache_dtype fp8 --index-cache-dtype mxfp4 declares the MXFP4 caches the aiter kernels read and write: one 8704 B field (128 x 64 B e2m1 data, 128 x 4 B e8m0 scales) per (page, kv head) for K and for V, and a single 8704 B field per block for the index key (136 + 68 B a token). Startup refuses mixing with the NVFP4 formats, plugin mode, a block size other than 128, and non-gfx950; Mono declines it, indexer CP falls back with a warning. This commit only declares the pool and the switches; the write and read paths follow in the next commits. Co-Authored-By: Claude <noreply@anthropic.com>

… index key rope_cache with the MXFP4 switches calls fused_qknorm_idxrqknorm with kv_cache_dtype="mxfp4" straight onto the pool pages (no bf16 stage and no separate store kernel, as the NVFP4 path has), asks for the main q as e4m3 and drops the bf16 copy. A layer that reuses another layer's top-k writes K and V only and takes the kernel's skip_index_branch (the kernel refuses index_cache_dtype="mxfp4" together with it, so the layer does not ask for the index branch). Adds dequantize_pages, the oracle the write test compares the cache against. Needs aiter's fused_qknorm_idxrqknorm with kv_cache_dtype="mxfp4". Co-Authored-By: Claude <noreply@anthropic.com>

…s; page-128 block table index_topk's decode and prefill take index_mxfp4: the scores come from aiter's score_flydsl_fp4 / score_prefill_flydsl_fp4 over the two views (data and scale) of the pool page, in the layout the Triton scorers produce, so the selection and the emit after them are unchanged. pages_per_block is a parameter of the emit now (8 for the page-16 SHUFFLE cache, 1 for page-128 caches), which makes the block table one page a block for the MXFP4 KV. The decode scorer's work map is built once per step into a persistent buffer (capacity over every step width, bound = max_model_len / 128) by make_sparse_decode_metadata, outside the graph. The decode scorer's fp32 score buffer is persistent too: without an `out` the aiter scorer runs a torch.full(-inf) in every layer (a FillFunctor, ~4 us a layer, ~250 us a step at B=4 over 60 layers). The metadata refills the buffer's view once per step, outside the graph, and every layer passes it as `out`. Where the scorer lays scores out transposed (S*H >= 16, not TP4) the view is None and nothing changes. Tests: scorers against a dequantised oracle; replay of a captured decode against a refilled work map; the persistent score buffer equals a fresh -inf buffer bit for bit (scores and selection) over permuted, shortened and smaller steps, with a control that stale scores are really seen without the refill, also through a captured graph. Needs aiter's score_flydsl_fp4 and score_prefill_flydsl_fp4. Nothing in attention_mha calls the scorers yet. Co-Authored-By: Claude <noreply@anthropic.com>

… decode and plan-free for prefill sparse_pa_mxfp4 folds (token, kv head) into rows, hands the e4m3 q and the pool's K and V pages (data and scale views of one 8704 B page) to aiter's pa_decode_mxfp4 and bypasses run_pa_decode's fp8 and ASM gates. attention_mha routes the sparse prefill and decode through it, and asks the selection for a page-128 block table (one page a block) and the MXFP4 scorer, with the step's work map and score buffer on decode. Decode keeps the planned path (the captured shapes are fixed): rows past the planner's batch limit run in segments, each with a plan of its own, refreshed once a step (and recorded in the graph when one is captured) as the tail page is always selected. Prefill is plan-free (work_plan=None, max_partitions=1): plan_pa_decode JIT-compiles a kernel specialised on the batch size, ~200 ms cold (~19 ms warm) for every row count not seen before, and a prefill chunk has a new row count almost every forward, so planned prefill stalled every TP rank about 300 ms in c32 runs. One task and one direct write per row (a row reads at most topk pages) means one JIT specialisation whatever the rows, no plan, no scratch, no segmenting. Tests: decode and prefill against the dequantised oracle (Hkv 1 and 2); any row count 1..700; plan-free bit for bit the planned kernel with a max_partitions=1 plan, and the production split agrees to bf16 rounding; no new planner or PA kernel build and no plan-cache entry across 12 new row counts. Needs aiter's pa_decode_mxfp4 (plan-free and direct write). Co-Authored-By: Claude <noreply@anthropic.com>

…m3 query With the MXFP4 index cache the fused writer can emit the indexer query as e4m3 (cast as stored, with its scale) and the scorer reads it as such. Off by default (bf16 q); refused at startup without --index-cache-dtype mxfp4. Needs aiter's fused_qknorm_idxrqknorm q_fp8_out / q_scale and the fp8-Q scorer variants. Co-Authored-By: Claude <noreply@anthropic.com>

follows the final aiter fused_qknorm_idxrqknorm API (ROCm/aiter#7057): q_fp8_out and q_scale are gone, and an e4m3 q_out makes the kernel write the main q as unit-scale saturating e4m3, with no bf16 q row and no write-back into the packed qkv. The MXFP4 path now allocates one float8_e4m3fn q_out and passes it as q_out. The bytes are those of the old q_fp8_out with q_scale=None. Nothing after the call reads the bf16 q or the normalised qkv q region: the sparse prefill/decode read q only for its shape and hand it to pa_decode_mxfp4, the index scorer reads the separate index_q, and k/v are the pre-call views, ignored by the sparse backends. Co-Authored-By: Claude <noreply@anthropic.com>

aiter ROCm/aiter#7056 merges the standalone pa_decode_mxfp4 entry into pa_decode, so call pa_decode(out, q, k_data, v_data, ctx, bt, scale, 1, key_scale=, value_scale=, work_plan=) instead. max_partitions=1 is dropped: a plan-free call already runs one task per row. The plan-free build-count test reads the shared tile-kernel cache of pa_decode_kernel. Depends on ROCm/aiter#7056. Co-Authored-By: Claude Code <noreply@anthropic.com>

ganyi1996ppo · 2026-11-02T09:28:51Z

Pushed 6ae17ec (one fast-forward commit on top of 6ab524a): the MXFP4 sparse path now calls aiter's pa_decode(out, q, k_data, v_data, ctx, bt, scale, 1, key_scale=, value_scale=, work_plan=) instead of aiter.ops.flydsl.pa_decode_mxfp4. Files: sparse_attn_mxfp4.py (both calls; max_partitions=1 dropped, a plan-free call already runs one task per row), the plan-free build-count test (now reads _PA_DECODE_TILE_CACHE, the tile-kernel cache that replaced build_pa_decode_mxfp4), and docstring/comment mentions.

Why: ROCm/aiter#7056 (head f017c23591) merges the standalone pa_decode_mxfp4 entry into pa_decode at the review's request. Without this commit the MXFP4 sparse path raises ImportError once #7056 lands. Depends on ROCm/aiter#7056: it needs an aiter that has the MXFP4 views in pa_decode, which is #7056 itself.

Tested on gfx950 against aiter#7056 f017c23591 (same PA files): tests/test_minimax_m3_mxfp4_attention.py + tests/test_minimax_m3_mxfp4_cache.py, 32 passed, run from an archive of this commit; black/ruff clean. Not tested: model serving / accuracy runs, ATOM CI's pinned aiter (it predates #7056, so these tests need the #7056 build).

Re-add the stored e4m3 query path removed in 491b929 (accepted by ROCm/ATOM#2770's sparse_pa_mxfp4), selected by a Q_IN_DTYPE trait rather than a bool. The in-kernel BF16 query to FP8 quantization and the BF16 MXFP4 path are unchanged. Co-Authored-By: Claude Code <noreply@anthropic.com>

@benchmark

…uery_length 1) (#7056) * [FlyDSL] Add an MXFP4 KV-cache paged-attention decode (gfx950, query_length 1) pa_decode_mxfp4 reads K and V as e2m1 data + e8m0 scales (one per 32 elements, along head_dim), each as two uint8 views with their own page/head strides, so a caller may use two tensors or one single-region page [block_size x 64 B data][block_size x 4 B scale]. It shares pa_decode's work plan and partition reducer and leaves pa_decode itself untouched. Per 16-token tile and wave: QK is one mixed fp4(A) x fp8(B) v_mfma_scale_f32_16x16x128_f8f6f4 with the K block's e8m0 byte as scaleA and Q used as stored (query_scale, a [1] fp32, multiplies the logits); a bf16 Q takes a reference path that widens K with v_cvt_scalef32_pk_bf16_fp4. V is widened to bf16 (exact) into a wave-private LDS tile and read back transposed with ds_read_b64_tr_b16 as the PV operand; P stays bf16. Invalid tokens get a zero V scale and a -inf score, so tail garbage (e8m0 0/255 included) cannot leak. The fp8 B operand of the 16x16x128 MFMA is not contiguous per lane: lane (g, s) holds Q head s elements 16g..16g+15 in bytes 0..15 and 64+16g.. in bytes 16..31 (measured on gfx950; A lane group g holds K elements 32g..32g+31). tile_depth defaults to 2: with independent blocks (2048 tokens/row, 16 q heads) it is 0.2-2.2 us faster than 4 for 1-512 rows and equal at 4096. Arithmetic is unchanged (a test pins default == depth 4 bitwise). Tests: rel-L2 ~1e-3 vs an fp32 reference over the dequantized K/V (the bf16 output rounding floor), ctx 1..2048, block_size 16/64/128, groups 1..16, several KV heads, plan tasks of 1..N tiles, garbage tails, query_scale. Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode_mxfp4: write the output directly when every row is one task With a plan of max_partitions == 1 each sequence is a single task, so the kernel normalizes and stores the output row itself and the reduce kernel, its scratch and the pmax/psum/pout partials disappear (direct_output, auto on for that plan). Bitwise equal to split + reduce (tested for rows 1/8/64/512, bf16/f16, strided output, graph replay). A context-0 row has no task and is left unwritten on this path; the reduce path zero-fills it. Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode_mxfp4: plan-free form (work_plan=None) Tasks follow from the block index (CTA slot = sequence * P + part, P tiles-per-task chosen on the host from max_partitions, default the planner's 2 * CU task budget), so no plan_pa_decode launch and no work_info read: the context length, block-table entries and Q are loaded in parallel. The kernel writes the reduce info itself and zero-fills context-0 rows on both the reduce and the direct (P == 1) path. Bitwise equal to the planned path wherever the two splits coincide; tested for ragged contexts (1..5000, tasks of several tiles), block sizes 16/64/128, empty rows, the default-split switch at 255/256/257 rows, graph replay. Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode_mxfp4: reuse shared constants and ceildiv, fold view checks Generated ISA is unchanged (82 specializations, identical hash). Co-Authored-By: Claude <noreply@anthropic.com> * [Test] pa_decode_mxfp4: follow op_test conventions Module-level gfx950 skip, checkAllclose + TFLOPS in the benchmark, and run the pytest checks from main() so `python3 op_tests/...` (what CI runs) tests it. Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] Export pa_decode_mxfp4 lazily; derive default-split test rows from the CU count Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode_mxfp4: check output device, return early on zero sequences The test no longer queries the GPU at import. Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode_mxfp4: drop query_scale (the fp8 query is used as stored) Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode_mxfp4: follow main's reducer (partition_bound removed) launch_pa_decode_ps_reduce now takes the bound as max_partitions; same value as before: min(parts, capacity, tiles). Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode_mxfp4: drop the routes no caller uses Removes the bf16-query reference path, the tile_depth and direct_output arguments, multi-partition plan-free and the fp8 comparison in the test benchmark. Generated ISA and outputs are unchanged for what remains. Co-Authored-By: Claude <noreply@anthropic.com> * [Test] pa_decode_mxfp4: rewrite the op_test in aiter script style, no pytest Plain check functions with asserts, a case table and a pandas summary, try/except for the rejection cases (expect_error), one PA_DECODE name for the implementation under test. Same 39 correctness cases and the same timing sweep. Co-Authored-By: Claude <noreply@anthropic.com> * [Test] pa_decode_mxfp4: assertAllclose instead of rel-L2, call pa_decode_mxfp4 directly Replace the custom rel_l2/assert_rel with assertAllclose (rtol=atol=1e-2, at most 1% of elements outside; the measured worst case is 0.1%), drop the PA_DECODE alias, and make the timing section header say what is timed. Same 39 cases and timing sweep. Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode: MXFP4 K/V in the shared tile kernel, replacing the dedicated kernel Fold MXFP4 K/V (e2m1 data + e8m0 scale per 32) into the fp8/bf16 compile_pa_decode_tile (kv_dtype="mxfp4"); pa_decode dispatches a uint8 key cache to it and pa_decode_mxfp4 stays as a thin entry. Delete pa_decode_mxfp4_kernel.py; the MXFP4 cases move into test_flydsl_pa_decode.py. Co-Authored-By: Claude <noreply@anthropic.com> * [Test] pa_decode: drop an unused unpacked name (ruff RUF059) Co-Authored-By: Claude <noreply@anthropic.com> * [Test] pa_decode: rewrite the UT as an aiter script, no pytest Follow the aiter-op-test style (test_quant.py, test_flydsl_dcp_topk_merge.py): @benchmark tests swept from main(), run_perftest for us/TFLOPS/TB/s, an FP32 reference compared with assertAllclose, markdown summary tables, a __main__ guard, and exit 1 listing every failed case. Arch and shape skips filter the sweep up front instead of pytest.skip. The same 242 cases are enumerated (226 run on gfx950, 16 MXFP4 unowned-row cases are filtered by design). Rejected inputs are checked with a small expect-error helper. CI runs the file with plain `python3 "$file"`: drop the pytest special case in aiter_test.sh and the pytest invocation in ffm-bringup.yaml, which would otherwise collect nothing from the script. Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode: MXFP4 validates the operands the generic path already checks The MXFP4 branch reached the kernel with a non-contiguous block table, a context_lengths/block_tables batch that differs from the query, operands on another device, a query or output with a strided head_dim axis, key_scale=None, or a windowed work_plan. Those gave wrong output, illegal memory accesses or an AttributeError. Refuse them with the generic path's checks and messages, and add one regression case per refusal to the UT's MXFP4 rejected-input table. Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode: MXFP4 checks the device of the output and user scratch The generic path refuses an output or scratch tensor on another device; the MXFP4 branch checked only dtype, shape and contiguity, so a CPU output or exp_sums reached the kernel as a device pointer. Refuse both, and add one regression case each to the refused-input table. Co-Authored-By: Claude Code <noreply@anthropic.com> * [Test] pa_decode: collectable zero-arg test_* for CI pytest, restore .github CI dispatches this file with `python3 -m pytest -q`. Rename the arg-taking @benchmark sweeps to run_* and add zero-arg test_pa_decode_{plan,fp8,bf16,mxfp4} that run the same sweeps and assert no failures (unittest.SkipTest when the arch/FlyDSL cannot run). The script route keeps its CLI and summary tables. No pytest import or API. Restore .github/ to the PR base so the PR no longer changes CI dispatch. Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode MXFP4: validate scale page count and base alignment, fix test reference - Require key/value cache and scale views to share the key cache's page count, so a block-table page cannot index past a shorter scale allocation. - Reject data views that are not 16 B aligned and scale views that are not 4 B aligned, the contract mx_init declares for its loads. - Test: add scale-pages-mismatch, misaligned-key and misaligned-scale reject cases; compare an e4m3 query as stored and re-round only a bf16 query. - Test: check torch.cuda.is_available() before probing the GPU arch so a CPU-only host skips instead of raising. Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode MXFP4: keep tiny bf16 queries from flushing to zero quantize_row floored q_scale at 1e-20 before taking its reciprocal, so a bf16 query below ~4e-18 was scaled into the e4m3 subnormal range and rounded to zero even when large e8m0 K scales make its logits significant. For MXFP4 the floor is now the smallest normal f32; it only guards the all-zero row, and the published q_scale is unchanged so the restoration still matches. The FP8 KV path keeps 1e-20. Test: run_pa_decode_mxfp4_tiny_query (Q = 2**-80, K = +-2**80, V = +-1) gave 0 instead of 1 before the change. Co-Authored-By: Claude <noreply@anthropic.com> * [FlyDSL] pa_decode MXFP4: keep the QK accumulator at logit scale The query is normalized to e4m3 (|q| up to 448) and its scale restored after the scaled MFMA, so a tiny query against a large e8m0 K scale (Q = 2**-80, K = 2**120) put ~448 * the logit in the f32 accumulator: inf, then inf - inf in the softmax, although the logit is finite. The exponent of the published q_scale now rides on the MFMA's e8m0 scale of Q (otherwise 1.0) and only its mantissa in [1, 2) is applied afterwards, so the accumulator never exceeds the logit. q_scale is floored at the smallest normal f32 and published floored, which also keeps a bf16 subnormal row (2**-133) in e4m3 instead of mis-scaling it. FP8 KV and e4m3 queries (q_scale 1.0) are unchanged. Test: the tiny-query case is now a sweep of extreme query/K-scale pairs against the FP32 reference: Q = 2**-80 with K = 2**120 and Q = 1 with K = 2**120 gave inf, Q = 2**-133 with K = 2**127 gave 0 instead of 0.175 before the change; the rest pass before and after. Full pa_decode suite and the MXFP4 timings are unchanged on gfx950. Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode MXFP4: block-table fallback past the staged pages, name the rejected argument - mx_pages read the LDS page ids with no bound, while the generic path falls back to the block table past the staged window. A planned task longer than the 256 staged pages (max_context_length below the context) read LDS past the ids and used the garbage as a page. It now takes the block-table entry there, as _staged_page does; the in-window read is unchanged (MXFP4 timings within 1%). - pa_decode: the MXFP4 NotImplementedError was one message for four conditions and never named compute_type. Each condition now says which one failed and its value. - pa_decode_mxfp4: the removed tile_depth and direct_output knobs gave pa_decode's "unexpected keyword argument"; they now raise NotImplementedError saying so. Test: run_pa_decode_mxfp4_short_bound (313 pages, bound 128, one-partition plan with a strided output) gave 0 instead of 0.18 before the change; the reject cases cover compute_type and the two removed knobs. Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode_mxfp4: one-line reason for the removed direct_output knob (ruff ISC004) Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode MXFP4: bound PV accumulation before normalization Scale probabilities for the PV MFMA so a large finite output cannot overflow the intermediate FP32 sum, then restore the scale in the epilogue. Cover the low-value and maximum-context bounds with exact-output tests. Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode MXFP4: address review (merge entry, split flags, gate, tests) - Delete the standalone pa_decode_mxfp4 entry; pa_decode takes the MXFP4 views (key_scale / value_scale) and accepts dtypes.fp4x2 / dtypes.fp8_e8m0 views. - Remove the direct-path clamp to the staged page window: page ids past it come from the block table, so a stale max_context_length no longer drops tokens. - Split is_bf16_kv into is_bf16_kv / PV_BF16 / MXFP4; MXFP4 requires gfx950 in select(), has its own LDS check, QP_ELEM_BYTES is 1 (Q is e4m3 in LDS), and an fp8 query is refused unless KV is MXFP4. - Rename mx_* to mxfp4_* in the kernel modules. - Tests: fold the MXFP4 PV-range, extreme-scale, short-bound, one-task-per-row and rejection cases into run_pa_decode (no new run_* functions), add 128k peaked-softmax cases against an fp64 reference, use dtypes.fp4x2/fp8_e8m0. Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode MXFP4: do not bound the reducer by max_context_length The reducer's partition count was capped by ceil(max_context_length / 256), but a task may own more tokens than the bound (max_context_length only sizes the staged page ids). A two-partition plan over 1024 tokens with max_context_length=128 then merged only the first partial. Bound it by the plan and its capacity only. Add mxfp4-short-bound-np2 (fails before: 22% of elements off). Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode MXFP4: document the V range, inline the select rejection checks - State in the pa_decode docstring and the traits shortcut note that V is staged as bf16, so an MXFP4 V above the bf16 maximum (e8m0 byte 254 with e2m1 >= 2, byte 253 with >= 4) widens to inf even when the normalized output is finite. - Drop the _check_unsupported_select helper: the three PaDecodeSchedule.select rejection cases are dispatched inline in run_pa_decode. Co-Authored-By: Claude Code <noreply@anthropic.com> * [FlyDSL] pa_decode MXFP4: convert test_flydsl_pa_decode.py to pytest Replace the custom sweep runner with pytest parametrization (test_pa_decode_plan, test_pa_decode, test_pa_decode_unowned_kv_rows). Co-Authored-By: Claude Code <noreply@anthropic.com> * test(flydsl_pa_decode): use valid V e8m0 scale byte 1 in MXFP4 pool specs Three MXFP4 pool specs set the V scale e8m0 byte to 0, which the kernel does not support. Use byte 1 (2**-127) and rename the cases and comment to match. Co-Authored-By: Claude Code <noreply@anthropic.com> * Simplify MXFP4 PA comments Co-Authored-By: Claude Code <noreply@anthropic.com> * flydsl pa_decode: trim comments in traits.py Comment-only change; AST is identical to the previous commit. Co-Authored-By: Claude Code <noreply@anthropic.com> * flydsl pa_decode: use explicit dtypes for FP8/MXFP4 pointers Label Q FP8, gfx1250 KV, and MXFP4 K/V data and scale pointers with their real element types instead of fx.Uint8. Byte-addressed fx.Uint8 bases in op_lds.py are intentionally unchanged. Co-Authored-By: Claude Code <noreply@anthropic.com> * flydsl pa_decode: clean up FP8 input-query handling Co-Authored-By: Claude Code <noreply@anthropic.com> * flydsl pa_decode: restore FP8 input-query support for MXFP4 KV Re-add the stored e4m3 query path removed in 491b929 (accepted by ROCm/ATOM#2770's sparse_pa_mxfp4), selected by a Q_IN_DTYPE trait rather than a bool. The in-kernel BF16 query to FP8 quantization and the BF16 MXFP4 path are unchanged. Co-Authored-By: Claude Code <noreply@anthropic.com> * flydsl pa_decode: refuse MXFP4 query/output spans over 2**31 - 1 bytes The kernel adds query row/head offsets in i32, so a query whose last head_dim chunk ends past 2 GiB wrapped to an earlier row, and a direct output store did the same. Refuse such query and output views before launch; two refused-input cases in test_pa_decode cover them. Co-Authored-By: Claude <noreply@anthropic.com> --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: r <r@x>

ganyi1996ppo and others added 5 commits October 27, 2026 21:10

zhuyuhua-v mentioned this pull request Oct 29, 2026

[MiniMax-M3.1] FP4 index cache and NVFP4 sparse KV cache #2744

Open

ganyi1996ppo mentioned this pull request Oct 29, 2026

[MiniMax-M3.1] Mono decode over the MXFP4 KV and index caches #2812

Draft

ganyi1996ppo marked this pull request as ready for review November 4, 2026 01:40

ganyi1996ppo mentioned this pull request Nov 4, 2026

[MiniMax-M3.1] MXFP4 index cache; reject M3 --kv-cache-dtype fp4 #2962

Open

Navigation Menu

Repository navigation

[MiniMax-M3.1] MXFP4 sparse KV cache and index cache - #2770

[MiniMax-M3.1] MXFP4 sparse KV cache and index cache#2770
ganyi1996ppo wants to merge 7 commits into
yuhua/m3.1-flashfrom
ganyi/m31-mxfp4-kv-index

ganyi1996ppo commented Oct 28, 2026 •

edited

Loading

Uh oh!

ganyi1996ppo commented Nov 2, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

ganyi1996ppo commented Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Dependencies

How it works

Commits

Results

Tests

Known limits / Not measured

Uh oh!

ganyi1996ppo commented Nov 2, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

ganyi1996ppo commented Oct 28, 2026 •

edited

Loading