You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[FlyDSL] MiniMax-M3 index score: MXFP4 index-cache decode scorer - #7060
Adds a MiniMax-M3 index-score decode kernel 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 kernel is the decode reading side. The fp8 index-score builder gains an mxfp4 specialisation; the generated code (LLVM IR and ISA) of the existing fp8 kernels is unchanged.
Files: aiter/ops/flydsl/kernels/minimax_m3_index_score.py (1474 to 1928 lines), op_tests/test_flydsl_minimax_m3_index_score.py (2169 to 2605 lines, the MXFP4 tests added), and two lazy exports in aiter/ops/flydsl/__init__.py: minimax_m3_index_score_fp4_flydsl (score_flydsl_fp4) and minimax_m3_index_score_fp4_flydsl_supported (index_score_supported_fp4). No new files, no C++.
Split from #7058, the earlier combined version that this PR and its prefill sibling #7061 replace (whether #7058 gets closed is pending the author's decision).
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. e2m1 x 2^e is exact in e4m3 for biased e8m0 in [119, 133], so widening an MXFP4 cache to the unit-scale fp8 index cache (the baseline in bench_score_mxfp4) saturates outside that window. The MXFP4 scorers apply the scale themselves. Values outside the window were not measured on real tensors; the tests cover wide scale ranges against the oracle.
Kernel
score_flydsl_fp4 uses the same work map, geometry, Q loading, causal cutoff and score store as the fp8 scorer: MXFP4 is a compile-time specialisation (mxfp4, plus q_fp8 for an e4m3 Q) of build_index_score, and only the K data+scale loads and the dequant/MFMA step are MXFP4-specific. The kernel signature differs (two cache views), so the shared body is a closure called from two thin kernels. Generated LLVM IR and ISA of the fp8 kernels are unchanged, and the MXFP4 kernels match the earlier separate MXFP4 kernel (checked at the fold, outputs bit-equal; unchanged since: 31 specialisations, 13 fp8 and 18 MXFP4, byte-identical before and after the last cleanup). The MXFP4 entry points take no q_scale and no CP shard (no caller needs them); score_flydsl and index_score_supported gain an optional k_scale that selects the MXFP4 cache.
bf16 Q: K is expanded to bf16 with v_cvt_scalef32_pk_bf16_fp4 and fed to the bf16 16x16x32 MFMA, so the numerics contract (K lifted to Q's precision) is unchanged.
e4m3 Q: one v_mfma_scale_f32_16x16x128_f8f6f4 per 16-token tile, A = fp4 K with its e8m0 byte, B = Q as e4m3; products are exact in fp32.
K tiles in flight = 6 (the fp8 scorer uses 3; an MXFP4 tile is half the bytes). In a 2/3/4/6/8 sweep, 6 was 15-35% faster than 3 on small batch x context cells.
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..16 decode) 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 mean of the two rounds). 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.
S = queries per request, H = index heads, B = batch. Range over the 9 cells below: bf16 Q / fp8 = 0.48-0.65, e4m3 Q / fp8 = 0.47-0.59. Kernel-level numbers only. The table was taken with the earlier separate MXFP4 kernel; the folded kernel has identical ISA, and a base-vs-folded re-run of the same protocol (A-B-B-A) was within noise (geomean 1.003 fp8, 1.008 MXFP4 bf16 Q, 1.006 e4m3 Q, worst cell 1.021, run-to-run spread up to 3%). Cells measured with a CP4 shard were dropped with the CP kwarg.
Decode table (9 cells)
shape
B
ctx
fp8 (us)
MXFP4 bf16 Q (us)
MXFP4 e4m3 Q (us)
bf16 Q / fp8
e4m3 Q / fp8
S4 H1
8
32768
9.65
6.24
5.71
0.65
0.59
S4 H1
8
100000
24.09
13.48
12.29
0.56
0.51
S4 H1
8
131072
26.72
14.59
13.47
0.55
0.50
S4 H1
16
32768
15.25
8.88
8.39
0.58
0.55
S4 H1
16
100000
40.33
20.88
19.99
0.52
0.50
S4 H1
16
131072
48.79
24.83
23.55
0.51
0.48
S4 H1
32
32768
26.65
14.38
13.52
0.54
0.51
S4 H1
32
100000
78.49
38.35
36.67
0.49
0.47
S4 H1
32
131072
93.79
45.37
44.20
0.48
0.47
Tests
python op_tests/test_flydsl_minimax_m3_index_score.py # 90 MXFP4 correctness cases (PASS/FAIL table, exit 1 on any failure), then the decode, MXFP4 and prefill perf tables
pytest op_tests/test_flydsl_minimax_m3_index_score.py # main's existing pytest suite, unchanged: 269 passed, 4 skipped
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 pre-existing fp8 tests in that file are unchanged and give 269 passed, 4 skipped (as before).
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. Scores are judged against the fp8 test's own oracle. They cover decode geometries, CUDA-graph buffers, split vs single-region views, alignment and rejection rules (including a non-positive or NaN sm_scale), and e8m0 extremes.
Mutation check: swapping the decode nibble order failed 60 tests when measured on the earlier combined branch; not re-run after the split or the later cleanups.
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 the saturating range of the fp8-cache widening used as the baseline; they have no real-data measurement.
End-to-end numbers are not part of this PR; they belong to the ATOM integration PR.
black 26.10.1 and ruff 0.16.0 are clean on the two changed Python files at this head; the whole-repo ruff check . / black --check . was not re-run at this head.
Adds score_flydsl_fp4 / index_score_supported_fp4 (and an optional k_scale on score_flydsl / index_score_supported): the MiniMax-M3 decode 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 decode 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; the decode kernel takes 0.47-0.65 of the fp8 kernel's time in the table above (measured at an earlier head). Behaviour of the fp8 path is unchanged.
Validation and unverified scope
Measured at this head (827441b), 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 and the aiter JIT cache removed:
python op_tests/test_flydsl_minimax_m3_index_score.py run from the repo root and again from /tmp: both exit 0, 90 of 90 MXFP4 correctness cases PASS, and the decode, MXFP4 and prefill perf tables are printed.
LLVM IR and final ISA, before vs after this head's commit, for 8 fp8 kernels (4 decode, 4 prefill compiled from the same tree) and 8 MXFP4 decode kernels (4 shapes, bf16 and e4m3 Q): all byte-identical.
black 26.10.1 and ruff 0.16.0 are clean on the two changed Python files.
Unverified scope: other GPUs than gfx950; real M3 tensors; end-to-end use (ATOM integration PR); the pytest run of main's existing fp8 suite (269 passed, 4 skipped, quoted above) was done at an earlier head and not repeated at this one; the performance table is from an earlier head; the mutation check quoted under Tests was not re-run; 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_flydsl_fp4 and index_score_supported_fp4 do not exist, so every MXFP4 case fails with an attribute error. The pre-existing pytest cases in this file cover fp8 caches only (they build no MXFP4 cache, never pass k_scale) and are unchanged.
Changed op tests
op_tests/test_flydsl_minimax_m3_index_score.py (the only op test changed): the existing fp8 pytest cases are untouched; this PR adds an MXFP4 section driven from main(), 90 cases, each of which would fail as follows:
matches_oracle_* and every_geometry_*: scores against an independent oracle on the dequantised cache, over the decode matrix and every geometry the selector reaches, bf16 and e4m3 Q. Catches a wrong nibble order, k-axis map or scale byte.
wide_scale_range_*: scale bytes across 2^+-20, so each lane group must pick its own byte.
graph_buffers_*: preallocated score buffers and work map captured once in a CUDA graph and replayed with different lengths and block tables.
sm_scale_must_be_positive_finite (added at this head): sm_scale of -1, 0 and NaN must raise a ValueError, because the epilogue takes the page max before multiplying by the scale and a non-positive scale would pick the wrong block silently. The existing test_scale_fp32_range covers only 1e-50 and 1e39.
default_score_buffer_reads_as_minus_inf: with the default (kernel-allocated) score buffer, slots past a request's end read as -inf and the rest are finite.
bench_score_mxfp4: the perf sweep (us, TFLOPS, TB/s for MXFP4 against the fp8 scorer), each result compared with the oracle.
These are not redundant with the fp8 cases in the same file: none of them builds an MXFP4 cache.
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:
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).
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 7060 --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.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
The benchmark unfairly includes full-cache conversion in every timed fp8 invocation while using precomputed fp4 inputs, inflating fp8 latency with work omitted from its reported byte count.
Adds MXFP4 index-cache decode scoring for MiniMax-M3, extending the existing FP8 scorer with a lower-storage cache format. This PR covers decode; prefill is handled separately in #7061.
Changes:
Specializes the shared scorer for MXFP4 caches with BF16 or E4M3 queries and optional query scaling.
Exposes lazy-loaded MXFP4 scoring and support-check APIs.
Adds correctness, cache-layout, graph-replay, input-validation, and performance tests.
File
Description
op_tests/test_flydsl_minimax_m3_index_score.py
Adds MXFP4 tests and benchmarks; runs correctness checks before script-mode benchmarks.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Negative query scales can bypass validation, score tests can miss infinity-sign mismatches, and FP8 benchmarks include conversion overhead absent from the FP4 path.
With sm_scale=-1.0 and q_scale=-1.0, the folded multiplier is positive, so both the entry point and support predicate accept a query scale that violates the documented contract. Validate q_scale > 0 independently in _score and _supported (line 1733), while retaining the combined multiplier's finite FP32 range check.
Use check_scores to validate negative-infinity entries
This helper accepts +inf where the top-k reference contains -inf: both finite masks are false, so those entries are excluded from the error check. Replace the local comparison with check_scores(a, b) to check the negative-infinity mask and use the required checkAllclose path with standard mismatch diagnostics.
Negative query scale (kernels/minimax_m3_index_score.py): valid, fixed. q_scale <= 0 or NaN is now refused on its own in _score and _supported (the folded-multiplier FP32 range check stays); sm_scale = q_scale = -1 is rejected by both, with a new test.
Infinity-sign mismatch in the top-k check (op_tests): valid, fixed. assert_topk_agrees now calls check_scores, which compares the -inf mask and runs checkAllclose; a new test covers a +inf where the reference has -inf.
fp8 baseline timing: same as the inline thread (fixed there).
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The gfx950 FP4/scaled-MFMA paths and shared fp8 refactor require exact-head GPU correctness, generated-code equivalence, and paired performance validation that this source-only review cannot independently verify.
New kernel `score_flydsl_fp4` scores an MXFP4 index cache (e2m1 data plus one
e8m0 scale per 32 elements) given as two views, k_data [pages,128,64] and
k_scale [pages,128,4], each with its own page stride, so separate tensors and
a single-region 8704 B page both work.
Query dtype selects the K path:
- bf16 Q: K is widened to bf16 with v_cvt_scalef32_pk_bf16_fp4 and fed to the
same bf16 16x16x32 MFMA as the fp8 scorer; the numerics contract (K lifted to
Q's precision) is unchanged.
- e4m3 Q: one v_mfma_scale_f32_16x16x128_f8f6f4 per 16-token tile, A = the
lane's 16 B of fp4 K plus its e8m0 byte, B = Q as e4m3 (unit scale). Products
are exact in fp32, so scores are those of dequantize(K) against Q as stored.
An optional q_scale (default 1.0, positive) is folded into the score
multiplier.
K tiles in flight is 6 rather than the fp8 scorer's 3: an MXFP4 tile is half
the bytes, and 6 was 15-35% faster than 3 on small batch*context cells in a
2/3/4/6/8 sweep.
Work map, geometry and CP shard are imported from the fp8 scorer. The builder
is a copy of `build_index_score` with the K side replaced, not a
parameterisation of it, so the two can drift and an epilogue or work-map change
must be mirrored. e8m0 bytes 0 and 255 are outside the contract.
Only exports are added to aiter/ops/flydsl/__init__.py; the fp8 path is
untouched.
Co-Authored-By: Claude <noreply@anthropic.com>
The sweep used checkAllclose, which only logs, so CI (python3 file)
could not fail on a wrong kernel; use the fp8 test's check_scores.
Import the fp8 test as siblings do and build the case matrix once.
Co-Authored-By: Claude <noreply@anthropic.com>
python3 op_tests/<file> now runs the pytest matrix before the perf sweep
(as test_flydsl_hc_mix.py does), so CI can fail on a wrong kernel. Import
the fp8 test with a script-mode fallback, hoist torch, fix the docstring
about e8m0 scale clamping. ruff 0.16.0 (ruff check .) and black are clean.
Co-Authored-By: Claude <noreply@anthropic.com>
One builder and one test file now cover both index caches: the MXFP4 cache
is an `mxfp4` specialisation of build_index_score, and score_flydsl_fp4 /
index_score_supported_fp4 are thin wrappers over the shared host path.
Generated LLVM IR and ISA are byte-identical to the separate kernels for 39
specialisations (13 fp8, 26 MXFP4); outputs are bit-equal on 54 cases.
python3 op_tests/<file> runs the pytest checks before the perf sweeps.
Co-Authored-By: Claude <noreply@anthropic.com>
Replaces the internal ASTRewriter.transform call, so a FlyDSL bump cannot
break the import. q_scale's docstring now says it scales a bf16 Q too (the
code always did) and a test checks it. LLVM IR and ISA are byte-identical
for 39 specialisations.
Co-Authored-By: Claude <noreply@anthropic.com>
q_scale and the CP kwargs have no caller in ATOM#2770 or vLLM#64198, so the
MXFP4 entry points lose them, and score_flydsl / index_score_supported take an
optional k_scale instead of a parallel private core. Remove the tests that only
re-checked other tests' logic (top-k, fp8 cross-check, bf16-vs-e4m3 gap,
quantizer rule) and the CP and q-dtype multiples of graph and view tests.
IR and ISA are byte-identical for all 31 remaining specialisations.
Co-Authored-By: Claude <noreply@anthropic.com>
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The new oracle and wide-scale tests bypass the required shared checkAllclose path and should use check_scores to retain standard mismatch diagnostics and special-value handling.
The MXFP4 tests this PR added used pytest fixtures, parametrize, skip, raises
and a pytest.main call in __main__. They are now plain check functions driven
from main() by a case table (PASS/FAIL table, exit 1 on any failure), like the
sibling flydsl op_tests, and __main__ is main's again. The same 89 cases run.
They compare with check_scores instead of the local compare. The pytest code
that was already on main is untouched.
Co-Authored-By: Claude <noreply@anthropic.com>
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Human assessment is needed to resolve the mismatch between the documented e4m3 saturation contract and the directly scaled MXFP4 decode path, which can produce different scores.
…tion check
Kernel: use fx.rocdl for the scaled MFMA like the rest of the file, name the
0x7F7F7F7F unit e8m0 scale UNIT_SCALE, say why the MFMA is raw. No change to
the generated code.
Test: add a check that sm_scale <= 0 and NaN are refused by score_flydsl_fp4.
Co-Authored-By: Claude <noreply@anthropic.com>
kernel: fx.rocdl. used consistently (the bare rocdl import is gone); the unit e8m0 scale 0x7F7F7F7F is the named constant UNIT_SCALE; one comment says the scaled MFMA is called raw on dword vectors (I did not try the fx.gemm MFMA_Scale atom here). Generated code unchanged: LLVM IR and ISA byte-identical before/after for 8 fp8 and 8 MXFP4 decode kernels.
op test: new sm_scale_must_be_positive_finite check, score_flydsl_fp4 must raise ValueError for sm_scale of -1, 0 and NaN (my earlier comment of 2026-10-29 02:16 said sm_scale = q_scale = -1 is rejected by both with a new test; q_scale has since been removed and no test of sm_scale <= 0 or NaN was in the tree at the previous head, so this adds it).
Run on gfx950: python op_tests/test_flydsl_minimax_m3_index_score.py from the repo root and from /tmp, exit 0, 90/90 MXFP4 checks PASS, perf tables printed.
PR description updated (stale numbers, the PR Self-Review sections).
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Human assessment is needed to determine how buffer initialization included in benchmark timing but excluded from throughput accounting affects the reported performance results.
…m0 255 doc
- every-geometry check now runs the CP points (cp_world>1 with k_scale) against
the oracle; MXFP4 + CP works, so it is covered rather than rejected.
- q_for keeps q's strides, so the qslice e4m3 case reaches load_q8's strided path.
- bench_score_mxfp4 allocates the score buffer once, outside the timed region.
- Docs: e8m0 byte 255 is +inf as f32 (bf16-Q path) and NaN in the scaled MFMA
(e4m3-Q path); measured, the score is NaN on both.
Co-Authored-By: Claude Code <noreply@anthropic.com>
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The shared-kernel refactor and new gfx950 FP4 scaled-MFMA paths require target-GPU correctness, generated-code equivalence, graph-replay, and paired performance checks that this source-only review cannot perform.
* [FlyDSL] M3 index score: MXFP4 index-cache decode scorer
New kernel `score_flydsl_fp4` scores an MXFP4 index cache (e2m1 data plus one
e8m0 scale per 32 elements) given as two views, k_data [pages,128,64] and
k_scale [pages,128,4], each with its own page stride, so separate tensors and
a single-region 8704 B page both work.
Query dtype selects the K path:
- bf16 Q: K is widened to bf16 with v_cvt_scalef32_pk_bf16_fp4 and fed to the
same bf16 16x16x32 MFMA as the fp8 scorer; the numerics contract (K lifted to
Q's precision) is unchanged.
- e4m3 Q: one v_mfma_scale_f32_16x16x128_f8f6f4 per 16-token tile, A = the
lane's 16 B of fp4 K plus its e8m0 byte, B = Q as e4m3 (unit scale). Products
are exact in fp32, so scores are those of dequantize(K) against Q as stored.
An optional q_scale (default 1.0, positive) is folded into the score
multiplier.
K tiles in flight is 6 rather than the fp8 scorer's 3: an MXFP4 tile is half
the bytes, and 6 was 15-35% faster than 3 on small batch*context cells in a
2/3/4/6/8 sweep.
Work map, geometry and CP shard are imported from the fp8 scorer. The builder
is a copy of `build_index_score` with the K side replaced, not a
parameterisation of it, so the two can drift and an epilogue or work-map change
must be mirrored. e8m0 bytes 0 and 255 are outside the contract.
Only exports are added to aiter/ops/flydsl/__init__.py; the fp8 path is
untouched.
Co-Authored-By: Claude <noreply@anthropic.com>
* Trim MXFP4 decode scorer dead indirection and docstring
Generated LLVM IR and ISA are byte-identical across 14 specialisations.
Co-Authored-By: Claude <noreply@anthropic.com>
* Make MXFP4 index-score sweep assert like the fp8 sweep
The sweep used checkAllclose, which only logs, so CI (python3 file)
could not fail on a wrong kernel; use the fp8 test's check_scores.
Import the fp8 test as siblings do and build the case matrix once.
Co-Authored-By: Claude <noreply@anthropic.com>
* Run MXFP4 index-score checks from main() and fix import/lint
python3 op_tests/<file> now runs the pytest matrix before the perf sweep
(as test_flydsl_hc_mix.py does), so CI can fail on a wrong kernel. Import
the fp8 test with a script-mode fallback, hoist torch, fix the docstring
about e8m0 scale clamping. ruff 0.16.0 (ruff check .) and black are clean.
Co-Authored-By: Claude <noreply@anthropic.com>
* Fold the MXFP4 decode scorer into the fp8 index-score builder
One builder and one test file now cover both index caches: the MXFP4 cache
is an `mxfp4` specialisation of build_index_score, and score_flydsl_fp4 /
index_score_supported_fp4 are thin wrappers over the shared host path.
Generated LLVM IR and ISA are byte-identical to the separate kernels for 39
specialisations (13 fp8, 26 MXFP4); outputs are bit-equal on 54 cases.
python3 op_tests/<file> runs the pytest checks before the perf sweeps.
Co-Authored-By: Claude <noreply@anthropic.com>
* Use public flyc.jit for the shared index-score body
Replaces the internal ASTRewriter.transform call, so a FlyDSL bump cannot
break the import. q_scale's docstring now says it scales a bf16 Q too (the
code always did) and a test checks it. LLVM IR and ISA are byte-identical
for 39 specialisations.
Co-Authored-By: Claude <noreply@anthropic.com>
* Address review: validate q_scale on its own, fair fp8 baseline, sign-aware top-k check
q_scale was only checked through the folded multiplier, so sm_scale = q_scale = -1
passed; reject q_scale <= 0 in the entry point and the predicate. The perf sweep
now converts the fp8 cache once outside the timed callable. assert_topk_agrees
uses check_scores, which also compares the -inf mask (a +inf slot slipped through).
Co-Authored-By: Claude <noreply@anthropic.com>
* Drop MXFP4 scorer routes nothing calls
q_scale and the CP kwargs have no caller in ATOM#2770 or vLLM#64198, so the
MXFP4 entry points lose them, and score_flydsl / index_score_supported take an
optional k_scale instead of a parallel private core. Remove the tests that only
re-checked other tests' logic (top-k, fp8 cross-check, bf16-vs-e4m3 gap,
quantizer rule) and the CP and q-dtype multiples of graph and view tests.
IR and ISA are byte-identical for all 31 remaining specialisations.
Co-Authored-By: Claude <noreply@anthropic.com>
* Run the MXFP4 decode op tests from main() instead of pytest
The MXFP4 tests this PR added used pytest fixtures, parametrize, skip, raises
and a pytest.main call in __main__. They are now plain check functions driven
from main() by a case table (PASS/FAIL table, exit 1 on any failure), like the
sibling flydsl op_tests, and __main__ is main's again. The same 89 cases run.
They compare with check_scores instead of the local compare. The pytest code
that was already on main is untouched.
Co-Authored-By: Claude <noreply@anthropic.com>
* MXFP4 decode: fx.rocdl consistently, named unit scale, sm_scale rejection check
Kernel: use fx.rocdl for the scaled MFMA like the rest of the file, name the
0x7F7F7F7F unit e8m0 scale UNIT_SCALE, say why the MFMA is raw. No change to
the generated code.
Test: add a check that sm_scale <= 0 and NaN are refused by score_flydsl_fp4.
Co-Authored-By: Claude <noreply@anthropic.com>
* MXFP4 decode tests: CP shards, strided e4m3 Q, prealloc bench; fix e8m0 255 doc
- every-geometry check now runs the CP points (cp_world>1 with k_scale) against
the oracle; MXFP4 + CP works, so it is covered rather than rejected.
- q_for keeps q's strides, so the qslice e4m3 case reaches load_q8's strided path.
- bench_score_mxfp4 allocates the score buffer once, outside the timed region.
- Docs: e8m0 byte 255 is +inf as f32 (bf16-Q path) and NaN in the scaled MFMA
(e4m3-Q path); measured, the score is NaN on both.
Co-Authored-By: Claude Code <noreply@anthropic.com>
---------
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: r <r@x>
The MXFP4 CP decode hands aiter's scorer k_scale and the shard (cp_world /
cp_rank), and its work map and capacity calls the shard. A new GPU-free
test runs those three calls and binds each to the installed aiter's
signature, so an aiter older than ROCm/aiter#7060 (no k_scale) fails at
test time rather than at the first MXFP4 CP decode. It skips where
aiter.ops.flydsl does not import.
Co-authored-by: Cursor <cursoragent@cursor.com>
The MXFP4 CP decode hands aiter's scorer k_scale and the shard (cp_world /
cp_rank), and its work map and capacity calls the shard. A new GPU-free
test runs those three calls and binds each to the installed aiter's
signature, so an aiter older than ROCm/aiter#7060 (no k_scale) fails at
test time rather than at the first MXFP4 CP decode. It skips where
aiter.ops.flydsl does not import.
Co-authored-by: Cursor <cursoragent@cursor.com>
The MXFP4 CP decode hands aiter's scorer k_scale and the shard (cp_world /
cp_rank), and its work map and capacity calls the shard. A new GPU-free
test runs those three calls and binds each to the installed aiter's
signature, so an aiter older than ROCm/aiter#7060 (no k_scale) fails at
test time rather than at the first MXFP4 CP decode. It skips where
aiter.ops.flydsl does not import.
Co-authored-by: Cursor <cursoragent@cursor.com>
The MXFP4 CP decode hands aiter's scorer k_scale and the shard (cp_world /
cp_rank), and its work map and capacity calls the shard. A new GPU-free
test runs those three calls and binds each to the installed aiter's
signature, so an aiter older than ROCm/aiter#7060 (no k_scale) fails at
test time rather than at the first MXFP4 CP decode. It skips where
aiter.ops.flydsl does not import.
Co-authored-by: Cursor <cursoragent@cursor.com>
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
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.
valar-cop
left a comment
Summary
Adds a MiniMax-M3 index-score decode kernel 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 kernel is the decode reading side. The fp8 index-score builder gains an
mxfp4specialisation; the generated code (LLVM IR and ISA) of the existing fp8 kernels is unchanged.Files:
aiter/ops/flydsl/kernels/minimax_m3_index_score.py(1474 to 1928 lines),op_tests/test_flydsl_minimax_m3_index_score.py(2169 to 2605 lines, the MXFP4 tests added), and two lazy exports inaiter/ops/flydsl/__init__.py:minimax_m3_index_score_fp4_flydsl(score_flydsl_fp4) andminimax_m3_index_score_fp4_flydsl_supported(index_score_supported_fp4). No new files, no C++.Split from #7058, the earlier combined version that this PR and its prefill sibling #7061 replace (whether #7058 gets closed is pending the author's decision).
Format
[128 x 64 B data][128 x 4 B scale]= 8704 B. The kernels take two views,k_data [pages,128,64]andk_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).bench_score_mxfp4) saturates outside that window. The MXFP4 scorers apply the scale themselves. Values outside the window were not measured on real tensors; the tests cover wide scale ranges against the oracle.Kernel
score_flydsl_fp4uses the same work map, geometry, Q loading, causal cutoff and score store as the fp8 scorer: MXFP4 is a compile-time specialisation (mxfp4, plusq_fp8for an e4m3 Q) ofbuild_index_score, and only the K data+scale loads and the dequant/MFMA step are MXFP4-specific. The kernel signature differs (two cache views), so the shared body is a closure called from two thin kernels. Generated LLVM IR and ISA of the fp8 kernels are unchanged, and the MXFP4 kernels match the earlier separate MXFP4 kernel (checked at the fold, outputs bit-equal; unchanged since: 31 specialisations, 13 fp8 and 18 MXFP4, byte-identical before and after the last cleanup). The MXFP4 entry points take noq_scaleand no CP shard (no caller needs them);score_flydslandindex_score_supportedgain an optionalk_scalethat selects the MXFP4 cache.v_cvt_scalef32_pk_bf16_fp4and fed to the bf16 16x16x32 MFMA, so the numerics contract (K lifted to Q's precision) is unchanged.v_mfma_scale_f32_16x16x128_f8f6f4per 16-token tile, A = fp4 K with its e8m0 byte, B = Q as e4m3; products are exact in fp32.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..16 decode) 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 mean of the two rounds). 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.S = queries per request, H = index heads, B = batch. Range over the 9 cells below: bf16 Q / fp8 = 0.48-0.65, e4m3 Q / fp8 = 0.47-0.59. Kernel-level numbers only. The table was taken with the earlier separate MXFP4 kernel; the folded kernel has identical ISA, and a base-vs-folded re-run of the same protocol (A-B-B-A) was within noise (geomean 1.003 fp8, 1.008 MXFP4 bf16 Q, 1.006 e4m3 Q, worst cell 1.021, run-to-run spread up to 3%). Cells measured with a CP4 shard were dropped with the CP kwarg.
Decode table (9 cells)
Tests
Run on one MI355X (gfx950), image
rocm/atom-dev:nightly_202610261558, with this branch'saiter/ops/flydsl/andop_tests/overlaid on the image's aiter (41d4d8d), because main's full Python tree does not import against the image's csrc. The pre-existing fp8 tests in that file are unchanged and give 269 passed, 4 skipped (as before).The tests are standalone: the cache is built with aiter's
per_1x32_f4_quantand dequantized withfp4_utils, with no dependency on any writer kernel. Scores are judged against the fp8 test's own oracle. They cover decode geometries, CUDA-graph buffers, split vs single-region views, alignment and rejection rules (including a non-positive or NaNsm_scale), and e8m0 extremes.Mutation check: swapping the decode nibble order failed 60 tests when measured on the earlier combined branch; not re-run after the split or the later cleanups.
Known limits / not measured
black26.10.1 andruff0.16.0 are clean on the two changed Python files at this head; the whole-reporuff check ./black --check .was not re-run at this head.PR Self-Review
Reviewed revision
827441b
Behavior and intent
Adds
score_flydsl_fp4/index_score_supported_fp4(and an optionalk_scaleonscore_flydsl/index_score_supported): the MiniMax-M3 decode 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 decode 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; the decode kernel takes 0.47-0.65 of the fp8 kernel's time in the table above (measured at an earlier head). Behaviour of the fp8 path is unchanged.Validation and unverified scope
Measured at this head (827441b), on one MI355X (gfx950), image
rocm/atom-dev:nightly_202610261558with this branch'saiter/ops/flydsl/andop_tests/overlaid on the image's aiter (see Tests), FlyDSL cache disabled and the aiter JIT cache removed:python op_tests/test_flydsl_minimax_m3_index_score.pyrun from the repo root and again from/tmp: both exit 0, 90 of 90 MXFP4 correctness cases PASS, and the decode, MXFP4 and prefill perf tables are printed.black26.10.1 andruff0.16.0 are clean on the two changed Python files.Unverified scope: other GPUs than gfx950; real M3 tensors; end-to-end use (ATOM integration PR); the pytest run of main's existing fp8 suite (269 passed, 4 skipped, quoted above) was done at an earlier head and not repeated at this one; the performance table is from an earlier head; the mutation check quoted under Tests was not re-run; 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_flydsl_fp4andindex_score_supported_fp4do not exist, so every MXFP4 case fails with an attribute error. The pre-existing pytest cases in this file cover fp8 caches only (they build no MXFP4 cache, never passk_scale) and are unchanged.Changed op tests
op_tests/test_flydsl_minimax_m3_index_score.py(the only op test changed): the existing fp8 pytest cases are untouched; this PR adds an MXFP4 section driven frommain(), 90 cases, each of which would fail as follows:matches_oracle_*andevery_geometry_*: scores against an independent oracle on the dequantised cache, over the decode matrix and every geometry the selector reaches, bf16 and e4m3 Q. Catches a wrong nibble order, k-axis map or scale byte.wide_scale_range_*: scale bytes across 2^+-20, so each lane group must pick its own byte.graph_buffers_*: preallocated score buffers and work map captured once in a CUDA graph and replayed with different lengths and block tables.split_and_single_region_views_agree_bitwise,e4m3_q_rows_must_be_16_byte_aligned,invalid_cache_rejected_*,metadata_rejects_cpu,supported_accepts_valid_inputs: layout and rejection contract (dtypes, shapes, strides, overlap, misalignment).sm_scale_must_be_positive_finite(added at this head):sm_scaleof -1, 0 and NaN must raise a ValueError, because the epilogue takes the page max before multiplying by the scale and a non-positive scale would pick the wrong block silently. The existingtest_scale_fp32_rangecovers only 1e-50 and 1e39.default_score_buffer_reads_as_minus_inf: with the default (kernel-allocated) score buffer, slots past a request's end read as -inf and the rest are finite.bench_score_mxfp4: the perf sweep (us, TFLOPS, TB/s for MXFP4 against the fp8 scorer), each result compared with the oracle.These are not redundant with the fp8 cases in the same file: none of them builds an MXFP4 cache.
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:
fused_qknorm_idxrqknorm).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