[ganyi1996ppo]

Summary

Adds the MiniMax-M3 index-score prefill scorer that reads an MXFP4 index K cache (e2m1 data plus one e8m0 scale per 32 elements), taking a bf16 or an e4m3 query. M3.1 can keep its index cache in MXFP4, which halves the index-cache bytes per token (and so raises the KV pool capacity); this is the prefill reading side. The fp8 prefill kernel gains an MXFP4-cache switch (k_fp4); its fp8 instantiations generate identical LLVM IR and ISA.

Changed: minimax_m3_index_score_prefill.py (1327 to 1835 lines: the MXFP4 path as a k_fp4 switch of build_prefill_score, optional k_scale argument of score_prefill_flydsl, the cache layout constants and _validate_cache, and the score_prefill_flydsl_fp4 entry; it was already over the usual 1000-line limit) and one lazy export in aiter/ops/flydsl/__init__.py: minimax_m3_index_score_prefill_fp4_flydsl. New: one op_test (683 lines). No C++ is touched. Net diff against main (git diff origin/main...HEAD): 3 files, +1356 -160.

Split from #7058, the earlier combined version that this PR and its decode sibling #7060 replace (whether #7058 gets closed is pending the author's decision). This PR does not depend on the decode PR: the cache layout constants and _validate_cache that the decode file also defines are copied into the prefill module, and the fp4 test imports nothing from another op test (main's op_tests/test_flydsl_minimax_m3_index_score.py imports pytest at module level, so importing from it pulled pytest in indirectly): LOG2E and alloc_score come from the prefill kernel module, and the shape lists and the helpers make_prefill_case, aiter_scores, run_prefill and _prefill_work are local copies with the same logic as the fp8 test's. The copies are about 110 lines and can drift from the fp8 file; a shared pytest-free helper module would remove them but means editing the fp8 test. Neither the fp4 test nor anything it imports needs pytest.

Format

• MXFP4: e2m1 data, one e8m0 scale per 32 elements.
• Page of 128 tokens = [128 x 64 B data][128 x 4 B scale] = 8704 B. The kernels take two views, k_data [pages,128,64] and k_scale [pages,128,4], each with its own page stride, so the layout is the caller's choice (separate tensors or the single-region page above both work; a test checks the two agree bitwise).
• e8m0 bytes 0 and 255 are outside the contract. The cases that are judged by aiter's fp8 scorer (pa_sparse_block_score_prefill, and the fp8 FlyDSL kernel) dequantize to e4m3, which is exact for biased e8m0 in [119, 133], so those stay in that window and walk the byte across it. The kernel applies the scale inside the MFMA and has no such window, so a separate check (scale_bytes_beyond_e4m3_*) builds a cache from e2m1 codes and per-token scale bytes spread over 100..150 (below 119 and above 133), computes the expected scores in fp64 in the test (value = e2m1 x 2^(byte-127), then max over live tokens of the dot with Q), and compares bit for bit. Real M3 tensors were not measured.

Kernels

score_prefill_flydsl_fp4(idx_q, k_data, k_scale, block_table, cu_seqlens_q, seq_lens, prefix_lens, max_query_len, max_seq_len, sm_scale, out=None) is a thin wrapper of score_prefill_flydsl(..., k_scale=...).

• The fp8 prefill kernel itself, with its K side switched by k_fp4; tiling, Q staging, causal mask, folds and stores are shared. The page stays MXFP4 through LDS (8.5 KB per buffer instead of 16 KB) and is fed to the scaled MFMA as the fp4 A operand with the e8m0 byte as block scale. A bf16 Q is rounded to e4m3 in-kernel (as the fp8 prefill kernel does); an e4m3 Q is taken directly.
• The two variants share one builder, so an epilogue or work-map change applies to both. The kernel argument list is its ABI, so the MXFP4 kernel has its own signature (it takes the scale view) and the fp8 one is unchanged; both run one shared score_body.

Performance

Measured on gfx950 (MI355X), one GPU, FlyDSL 0.4.0, image rocm/atom-dev:nightly_202610261558. Method: CUDA-graph replay of K calls (K = 2..8 prefill) over disjoint regions of a 300k-page pool to defeat the MALL, median of 25 replays, same card and process, every cell run in two rounds with the arm order reversed (the table reports the median). Random K data, e8m0 in [122,127], uniform context per batch, no padding bound beyond the context. Times are microseconds per call. The fp8 baseline is the existing fp8 kernel with a bf16 Q. Ratios below 1 mean MXFP4 is faster.

Prefill (fp8 = existing fp8 scorer, bf16 Q; native = the MXFP4 path):

cell	fp8 bf16 Q (us)	native bf16 Q (us)	native e4m3 Q (us)	native bf16 / fp8	native e4m3 / fp8
b1 q4096 ctx 8K	12.8	11.6	10.7	0.90	0.83
b1 q4096 ctx 128K	63.8	58.6	58.6	0.92	0.92
b4 q512 ctx 8K	8.8	8.5	7.9	0.96	0.90
b8 q2048 ctx 16K	38.9	34.0	32.3	0.88	0.83
MTP7 b32 q8 ctx 32K	39.6	30.5	29.5	0.77	0.75
chunk 32768, prefix 0	102.3	93.2	90.8	0.91	0.89
chunk 32768, prefix 100K (ctx 132768)	465.5	430.1	414.3	0.92	0.89

The native path is faster than fp8 on all seven measured cells (0.75-0.96). These are kernel-level numbers only. The table was measured on an earlier head that had a separate native kernel; after folding it into the fp8 builder the same cells moved by at most +1.7% per cell (geomean within 0.6%) with identical ISA, i.e. noise; removing the staged route afterwards left the ISA identical (timing not repeated).

Tests

python op_tests/test_flydsl_minimax_m3_index_score_prefill_fp4.py   # 60 correctness cases (PASS/FAIL table, exit 1 on any failure), then the perf table

At 3f4108b (a test-only commit: the kernel files are identical to 7ad68c3) the script ran on one MI355X (gfx950, one GPU), same image and overlay: 60 of 60 PASS, exit 0, every perf err column 0, from the repo root and again from /tmp; and once more with sys.modules["pytest"] = None set before it ran, so that any import pytest raises: 60 of 60 PASS, exit 0.

Run on one MI355X (gfx950), image rocm/atom-dev:nightly_202610261558, with this branch's aiter/ops/flydsl/ and op_tests/ overlaid on the image's aiter (41d4d8d), because main's full Python tree does not import against the image's csrc. The existing fp8 test file (main's) was re-run on the same overlay, unchanged: 269 passed, 4 skipped. LLVM IR and final ISA of the fp8 and the MXFP4 kernels are identical to the pre-change head for 27 variants (tile_q 64/128/256, L1/2/4, scalar and vector stores, bf16 and e4m3 Q), and the MXFP4 outputs were bit-equal to the earlier separate native kernel on 40 matrix outputs.

The tests are standalone: the cache is built with aiter's per_1x32_f4_quant and dequantized with fp4_utils, with no dependency on any writer kernel or on the decode scorer. Scores are judged against aiter's pa_sparse_block_score_prefill and the fp8 FlyDSL prefill kernel on the same dequantized values, and the e8m0 byte is walked across the exact-e4m3 window, and bytes outside it are checked against an fp64 dequantisation computed in the test. They cover ragged and mixed-length batches, trailing empty requests, garbage past each request's end, strided views, split vs single-region views, and rejection rules (including a strided seq_lens).

Mutation check (applied to the kernel source on the running overlay, then reverted): reading the wrong e8m0 byte gives 18 failures, and disabling the contiguous-vector check on seq_lens fails the rejection test.

Known limits / not measured

• Only gfx950 was tested.
• Kernel numerics on real M3 tensors were not measured; the tests use synthetic caches quantized with aiter's quantizer.
• e8m0 values outside [119, 133] are checked on synthetic caches against an fp64 oracle (100..150) but have no real-data measurement; bytes 1..99 and 151..254 are not generated.
• End-to-end numbers are not part of this PR; they belong to the ATOM integration PR.
• The cache layout constants and _validate_cache duplicate a few lines of the decode PR (see Summary).
• black and ruff==0.16.0 are clean on the files of this PR.

PR Self-Review

Reviewed revision

3f4108b

Behavior and intent

Adds score_prefill_flydsl_fp4 (and the optional k_scale argument of score_prefill_flydsl): the MiniMax-M3 prefill block scorer reading an MXFP4 index K cache (e2m1 data plus one e8m0 byte per 32 elements, 68 B per token against 128 B for fp8), with a bf16 or an e4m3 query. The existing prefill scorer takes only an fp8 index cache, so an MXFP4 cache could not be scored without first expanding it to fp8; M3.1 can keep the cache in MXFP4, which halves its bytes. Behaviour of the fp8 path is unchanged. At this head the fp8 kernels' LLVM IR and ISA are byte-identical to the pre-change head (see below).

Validation and unverified scope

Measured on one MI355X (gfx950), image rocm/atom-dev:nightly_202610261558 with this branch's aiter/ops/flydsl/ and op_tests/ overlaid on the image's aiter (see Tests), FlyDSL cache disabled. This head (3f4108b) differs from 7ad68c3 only in the op test file; the items below say which head each was measured at:

• At 3f4108b: python op_tests/test_flydsl_minimax_m3_index_score_prefill_fp4.py run from the repo root and again from /tmp: both exit 0, 60 of 60 correctness cases PASS, the perf table is printed, and every err column is 0.
• At 3f4108b: the same script with sys.modules["pytest"] = None set first (any import pytest raises): exit 0, 60 of 60 PASS, err columns 0.
• At 7ad68c3 (the kernel files are unchanged in 3f4108b): LLVM IR and final ISA, before vs after that commit, for 4 fp8 prefill kernels (the ones 6 prefill shapes reach: tile_q 64, 128 at L1 and L2, 256), 4 fp8 decode kernels (compiled from the same tree) and 8 MXFP4 prefill kernels (tile_q 64/128/256, L1/L2, bf16 and e4m3 Q): the fp8 ones are byte-identical; the MXFP4 ones differ only in the kernel-name tag (fp4d to fp4).
• At 7ad68c3: mutation of the new scale_bytes_beyond_e4m3 check: shifting every scale byte by +1 or by -1 in the call under test gives mismatch fraction 1.0; the unshifted control passes.
• black 26.10.1 and ruff 0.16.0 are clean on the changed files, rechecked on the op test at 3f4108b.
  Unverified scope: other GPUs than gfx950; real M3 tensors; end-to-end use (ATOM integration PR); main's existing fp8 test file was not re-run at this head (unchanged by this PR); the IR/ISA comparison and the mutation checks were not repeated at 3f4108b; the performance table above was measured at an earlier head and not repeated; CI results at this head.

PR-specific question

The self-review bot lists only the four generic questions on this PR, no PR-specific one. Its question 3, what each changed op test would fail on before this PR and why existing canonical coverage does not already provide it: before this PR score_prefill_flydsl_fp4 and the k_scale argument do not exist, so every case in the new test fails with an import error. The canonical op_tests/test_flydsl_minimax_m3_index_score.py on main covers fp8 caches only and contains no MXFP4 cache, no k_scale rejection rules and no e8m0 scale handling, so none of this is covered there.

Changed op tests

op_tests/test_flydsl_minimax_m3_index_score_prefill_fp4.py (new, the only op test this PR changes). Each group of cases and the failure it catches:

• matches_aiter_and_fp8_kernel_*, every_geometry_*, two_heads_*, ragged_tails_*: scores bit-equal to aiter's pa_sparse_block_score_prefill and to the fp8 FlyDSL kernel on the same dequantised values, over the prefill shapes, every launch geometry the selector reaches, two heads and ragged lengths. Catches a wrong k-axis or nibble order, a wrong scale byte, a wrong geometry/store path.
• scale_extremes_*: per-token magnitudes across 2^-3..2^3, so the scale is applied per lane group and per token and not once.
• scale_bytes_beyond_e4m3_*: e8m0 bytes below 119 and above 133, against an fp64 dequantisation computed in the test, since aiter's fp8 scorer cannot represent those values. Catches a scale applied wrongly outside the e4m3 window (the other cases cannot see it).
• ragged_batch_and_garbage_tails_*, mixed_query_lengths_trailing_empty_*: bytes past a request's end must not reach a live score; a trailing empty request must not be scored.
• e4m3_q_strided_and_single_region: a strided e4m3 query and single-region pages score like the plain inputs.
• rejections: wrong dtypes, narrow views, a narrow block table, an fp16 query, misaligned query rows and a strided seq_lens are refused with a ValueError.
• test_prefill_fp4 (the @benchmark sweep): perf table in us, TFLOPS and TB/s for the fp8 scorer and both MXFP4 query types, each compared against aiter with checkAllclose(rtol=0, atol=0).
  These are not redundant with existing tests: no other op test builds an MXFP4 cache or calls the prefill MXFP4 entry point.

Related work

Independent of the other MXFP4 index/KV PRs: this one builds and passes its tests on main alone. Siblings, each reviewable on its own:

• Decode half of this work: the MXFP4 index decode scorer, split out of [FlyDSL] MiniMax-M3 index score: MXFP4 index-cache decode and prefill scorers #7058 (the earlier combined version; closing it is pending the author's decision): [FlyDSL] MiniMax-M3 index score: MXFP4 index-cache decode scorer #7060.
• [FlyDSL] pa_decode: MXFP4 K/V in the shared PA tile kernel (gfx950, query_length 1) #7056: MXFP4 paged-attention decode.
• [HIP] fused_qknorm_idxrqknorm: write MXFP4 index-K and main KV caches, emit an fp8 main Q #7057: MXFP4 index K / KV write kernel (fused_qknorm_idxrqknorm).
• [MiniMax-M3.1] MXFP4 sparse KV cache and index cache ATOM#2770: ATOM integration draft that consumes all of these.

Both index-score halves add their own export lines to aiter/ops/flydsl/__init__.py; they merge in either order without conflict (checked by merging both branches in both orders).

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 7061 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[Copilot]

🟡 Changes recommended

Noncontiguous seq_lens can cause out-of-pool reads before validation, and timed sentinel fills distort the reported latency and throughput metrics.

2 open issues

• Reject noncontiguous seq_lens before staging · New
• Exclude score buffer initialization from performance timing · New

What changed in this PR

Adds MXFP4 index-cache prefill scoring for MiniMax-M3, extending the existing FlyDSL scorer with a native path and a staged fallback independently of the decode PR.

Changes:

• Shares the prefill kernel builder between fp8 and native MXFP4 scoring, supporting bf16/e4m3 queries and query scaling.
• Adds cache-layout validation, Triton dequantization for staging, and a lazy public export.
• Adds correctness tests and performance sweeps covering both routes.

File	Description
op_tests/​test_flydsl_minimax_m3_index_score_prefill_fp4.py	Adds reference comparisons, layout and rejection tests, and benchmarks.
aiter/​ops/​flydsl/​kernels/​minimax_m3_index_score_prefill.py	Extends the shared builder and launcher with native MXFP4 support.
aiter/​ops/​flydsl/​kernels/​minimax_m3_index_score_prefill_fp4.py	Adds route selection and staged fp8 dequantization.
aiter/​ops/​flydsl/​kernels/​minimax_m3_index_score_fp4_cache.py	Defines cache-layout constants and view validation.
aiter/​ops/​flydsl/​__init__.py	Exposes the MXFP4 prefill scorer through a lazy export.

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The staged path launches dequantization before rejecting noncontiguous seq_lens, allowing stride-blind reads to consume inactive block-table entries and trigger invalid cache accesses before validation raises an error.

2 open issues

• Reject noncontiguous seq_lens before staging
• Exclude score buffer initialization from performance timing

1 previously missed issue

In code that hasn't changed since last review

Use checkAllclose for numerical comparisons

op_tests/​test_flydsl_minimax_m3_index_score_prefill_fp4.py:438

The op-test contract in .claude/skills/aiter-op-test/SKILL.md requires numerical comparisons to use checkAllclose, but this case uses a local torch.testing.assert_close check. Please use the shared checker with the same tolerances and assert a zero returned mismatch ratio, preserving strictness while keeping the standard mismatch diagnostics.

_{Review effort: Balanced}

[ganyi1996ppo]

Review items handled in baac508: (1) noncontiguous seq_lens before staging: rejected up front, with a negative test (thread resolved); (2) timed score-buffer fills: perf sweep now preallocates outputs outside the timed call (thread resolved); (3) review-body note on torch.testing.assert_close in test_native_q_scale: replaced by checkAllclose(rtol=1e-6, atol=0) with a zero-mismatch assert. Tests: 110 passed, 2 skipped; fp8 prefill tests unchanged and passing.

[Copilot]

🟡 Changes recommended

The FP8-constrained test oracle leaves the native path’s wider e8m0 scale range unverified, requiring deterministic boundary cases checked against an independent FP32 reference.

1 open issue

• Cover wider e8m0 scale range with independent FP32 oracle · New

1 resolved since last review

• Prevent cache eviction from invalidating captured launch tables

_{Review effort: Balanced}

[aiter-bot]

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision e86e346b8229 · mode shadow

Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision.

Questions to answer in the PR description

1. What behavior does this PR change, and why is the existing behavior insufficient?
2. What did you validate, and what remains unverified? Include commands, hardware/configuration and results.
3. What behavior would each changed op test fail on before this PR, and why is that coverage not already provided by an existing canonical test?
4. For every changed op_tests/ file, what behavior does it cover and why is it not redundant? Write None if no op test changed.

Edit the PR Self-Review section in the PR description. Include the current revision SHA; this comment refreshes automatically.

Missing or invalid answers

• FAIL — Reviewed revision must contain at least the first 7 characters of the current head SHA

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

[Copilot]

🟡 Changes recommended

The tests leave valid e8m0 scales outside the fp8 baseline’s exact range unverified, requiring a direct FP32 reference to cover the scorer’s broader supported scale range.

2 open issues

• Cover wider e8m0 scale range with independent FP32 oracle
• Use checkAllclose for detailed score comparison failures · New

_{Review effort: Balanced}

[ganyi1996ppo]

Current head: 7ad68c3. Any other SHA quoted in my older replies on this PR belongs to history that was rewritten by a rebase; read this head.

What changed since 54bd603:

• op test: removed --no-perf; only AssertionError from a check is recorded as a FAIL (a HIP RuntimeError or an unexpected ValueError now stops the run instead of continuing on a possibly poisoned context); the benchmark is named test_prefill_fp4 like its siblings; argparse uses RawTextHelpFormatter with an e.g.: line; bit-exact comparisons go through checkAllclose(rtol=0, atol=0).
• op test: new scale_bytes_beyond_e4m3_{bf16,e4m3} case, e8m0 bytes 100..150 (below 119 and above 133) against an fp64 dequantisation computed in the test (answers the Copilot thread on the e4m3-only oracle).
• kernel: name tag fp4d -> fp4 (as the decode kernel); the comment about the kernel ABI states the contract. Generated code: fp8 kernels byte-identical in LLVM IR and ISA, MXFP4 kernels differ only in the kernel name.
• Run on gfx950: python op_tests/test_flydsl_minimax_m3_index_score_prefill_fp4.py from the repo root and from /tmp, exit 0, 60/60 PASS, perf table printed.
• PR title and description updated (stale numbers, one scorer, the PR Self-Review sections).

[Copilot]

🔵 Needs a closer look

The shared FP8/MXFP4 kernel refactor and new gfx950 scaled-MFMA path require exact-head GPU numerical, ISA, and paired latency validation that this source-only review cannot independently perform.

No open issues 🎉

2 resolved since last review

• Cover wider e8m0 scale range with independent FP32 oracle
• Use checkAllclose for detailed score comparison failures

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The gfx950 mixed FP4/FP8 scaled-MFMA path and shared FP8 refactor require revision-matched GPU correctness, ISA-equivalence, and paired performance validation that this source-only review cannot independently verify.

No open issues 🎉

_{Review effort: Balanced}

[yhl-amd]

@valar-cop review this pr

[valar-cop]

LGTM, but the fp4 path only has your own gfx950 run..

any perf number? and a non-author run on gfx950? ci can't reach the fp4 path

ci green.

[Copilot]

🔵 Needs a closer look

The mixed FP4/FP8 MFMA path and shared kernel refactor require exact-head gfx950 correctness and performance validation that this source-only review cannot independently verify.

No open issues 🎉

_{Review effort: Balanced}