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
[MiniMax-M3.1] FP4 index cache and NVFP4 sparse KV cache - #2744
This PR adds two opt-in caches for MiniMax-M3. Both are off by default and only reach MiniMax-M3.
--index-cache-dtype fp4 stores the lightning-indexer keys as NVFP4: a 64 B e2m1 payload plus 8 B of E4M3 scales per key, against 128 B for FP8.
ATOM_M3_SPARSE_KV_NVFP4=1 stores the sparse layers' K/V the same way, in their own MhaKvPool(kv_format="nvfp4").
At high concurrency the FP8 cache is what limits MiniMax-M3.1. At c48 on TP4 (--gpu-memory-utilization 0.8):
the cache holds 6.06M tokens, against 7.97M in flight at p50;
the token-weighted prefix hit falls to 19.9%;
output is ~94 tok/s/GPU.
With both caches on, the KV block shrinks from 3,179,520 to 1,827,840 bytes. That is 1.74x the tokens at the same memory: 12.7M tokens at --gpu-memory-utilization 0.9.
This PR replaces #2745 and #2746, which were stacked on it.
Design
Kernels.atom/model_ops/minimax_m3/nvfp4_kernels/ holds the gfx950 kernels. The derived files carry a license header.
decode.py: Gluon block-sparse attention over a packed NVFP4 K/V cache, with rows addressed through the block table.
verify_topk.py: fused index score + streaming top-k for speculative-verify rows.
nvfp4.py: the NVFP4 cache view and its e2m1 -> fp16 even/odd row loads, in Triton and Gluon.
Each kernel has a shape guard: can_use_nvfp4_decode and can_fuse_verify_topk.
Write path.
The fused qk-norm/rope/cache kernel only writes BF16 or FP8. It now writes the step's rows BF16 into a shared token-major staging buffer, with an identity slot mapping.
index_fp4_store / kv_nvfp4_store then quantize those rows into the paged planes at the real slots: per 16 elements, scale = e4m3(amax / 6), codes rounded half-to-even and saturating at 6.
The staging buffer is sized for the larger of max_num_batched_tokens and max_num_seqs x (draft + 1) rows.
Index. The prefill and decode scorers read the FP4 plane (IS_FP4) through the gfx950 e2m1 decode, as even/odd fp16 tiles.
Decode score grid (FP8 and FP4 index)._decode_index_score_tiled_kernel now runs on a (chunk, request) grid instead of (request, chunk). Workgroups go to the XCDs round-robin by linear id. With the request on the fastest axis, a batch of 8 put every chunk of request b on XCD b (two XCDs per request at batch 4 and 12), so the longest request set the kernel time. Scores and top-k are unchanged.
Attention.
Decode and prefill rows run nvfp4_sparse_attention over their top-k blocks.
One decode_rows launch per layer maps each row to its request and causal length. The attention and the fused top-k share it.
The shape guard runs once per layer at the first forward. The shapes are fixed per deployment, so one call is enough.
Fused verify top-k. It replaces score + top-k whenever can_fuse_verify_topk holds: NVFP4 KV, FP4 index, at least 3 requests, and within the kernel's shape bounds. It emits top-k ids only, which is all the NVFP4 attention reads.
This branch (main 035a78a + this PR + #2743, image rocm/atom-dev:nightly_202610261558): c48, 1200 s: 397.6 tok/s/GPU, P90 interactivity 71.4, p50 interactivity 105.6, prefix hit 95.58%, TTFT p50 297 ms, 0 errors. The stacked PRs on main with the same image and length gave 396.7 / 70.5 / 104.7, with a 95.96% prefix hit.
Decode score grid, on the FP8 + Mono c16 configuration (FP8 KV + FP8 index, 3600 s, same tree with and without the change):
c16 trace before the change: 85 us per score call at bs 8 against 37 us at bs 7, which is 2.9 ms of a 21.4 ms step.
Graph replay of minimax_m3_index_topk_decode (score + top-k) at bs 8, with one 450K request and seven 150K requests: 82.2 -> 55.3 us. No batch from 4 to 12 is slower.
Where not to use it. At c16, FP8 + Mono stays faster: 181.6 tok/s/GPU and P90 191.2, against 156.4 and 137.9 for NVFP4 over 1200 s. That NVFP4 run predates the per-layer row-map fix. The NVFP4 caches are meant for c >= 32, where the FP8 cache runs out.
Follow-up: MXFP4 caches
The NVFP4 caches in this PR are the first step for c32 and c48. The next step is MXFP4: e2m1 data with one E8M0 scale per 32 elements. It is ahead of NVFP4 on capacity, on the write path and in the readers. The work is in five open PRs:
NVFP4 here writes the step's rows in BF16 to a staging buffer, then quantizes them in separate launches. In a c16 trace of this branch (bs 4, 274K context per request) that is 3 _nvfp4_store_kernel launches per layer, 762 us per step.
With #7057 the quantization happens inside the fused qk-norm/RoPE/cache kernel, so there is no staging buffer and no store launch.
Readers.
gfx950's scaled MFMA (v_mfma_scale_f32_16x16x128_f8f6f4) takes fp4 operands with one E8M0 scale per 32 elements, which is exactly the MXFP4 layout. NVFP4's E4M3 scale per 16 elements does not fit it, so this PR's kernels convert e2m1 to fp16 first.
In kernel benchmarks, the MXFP4 decode scorer takes 0.45-0.80x the time of the fp8 scorer (#7060).
End to end, as reported in #2770. Each point is one run over a 1200 s window, on an older tree than this PR. Compare the rows of this table only with each other, not with the 3600 s results above.
CONC
NVFP4 tok/s/GPU / P90 intvty
MXFP4 tok/s/GPU / P90 intvty
32
344.8 / 85.8
370.3 / 108.3
48
396.6 / 71.3
429.0 / 92.1
At c16, MXFP4 gave 163.8 / 181.0 against 163.6 / 156.5 for fp8 + Mono.
In the kernel tests, the MXFP4 main-KV quantization error is 1.25-1.9x NVFP4's. On GSM8K, fp8, NVFP4 and MXFP4 are within about one standard error of each other.
Mono declines the MXFP4 sparse KV. Mono's FP4 support is planned for a later PR.
Relation to this PR.#2770 is cut from yuhua/m3.1-flash, which carries this PR, and adds its MXFP4 pool next to the NVFP4 one. Startup refuses mixing the two formats. The plan is to move c32 and c48 to MXFP4 once the aiter PRs land. Until then, the NVFP4 path in this PR is the configuration that runs.
Tests
tests/test_minimax_m3_fp4_cache_config.py (CPU): the startup refusal matrix (gfx, TP2/4/8, top-k reuse); the NVFP4 pool's per-entry bytes, views and layer isolation; the geometry name.
the stores match a torch quantizer bit for bit and leave pad rows unwritten;
the K/V store equals the per-head index store;
decode_rows / prefill_rows match torch;
the FP4 scorers select the same blocks as the BF16 scorers fed the dequantized cache, for both decode and prefill;
the NVFP4 attention matches softmax attention over the dequantized cache within 2%;
the attention guard refuses 8 q heads per kv head;
the fused top-k selects what score + top-k selects, at bs 3 and 16.
tests/test_minimax_m3_index_topk.py follows the scorer's new index_kv_scale argument.
The new tests pass on MI355X (gfx950), 24 of 24. The default path is unchanged on this branch: with FP8 KV + FP8 index, gsm8k strict-match is 0.9613 +- 0.0053 and block_bytes is 3,179,520, as on main.
Runs automatically on every eligible PR before approval:
✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests
Heavy model tests:
✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels
Label
Tests
ci:full
Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom
Run native ATOM model accuracy tests
ci:vllm
Run ATOM vLLM OOT model accuracy tests
ci:sglang
Run ATOM SGLang model accuracy tests
Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2744 --add-label <label>
zhuyuhua-v
changed the title
[MiniMax-M3.1] Vendor SGLang gfx950 NVFP4 sparse kernels (1/3)
[MiniMax-M3.1] FP4 index cache and NVFP4 sparse KV cache
Oct 28, 2026
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Critical defects in padded-request handling and NVFP4 cache-reuse validation can produce incorrect attention results, NaNs, and first-forward failures.
Adds opt-in FP4 index and NVFP4 sparse KV caches for MiniMax-M3, reducing cache memory use for high-concurrency inference while preserving the default cache paths.
Changes:
Adds packed cache storage, BF16 staging buffers, and quantization writes.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
FP4 staging dtype mismatches break float16 forwards, the unreleased index-scale view can cause wake-time out-of-memory failures, and incomplete draft groups can trigger out-of-bounds reads in the fused top-k path.
With a float16 MiniMax-M3 model, enabling either FP4 cache creates BF16 staging tensors but passes them to aiter.fused_qknorm_idxrqknorm with the corresponding cache dtype set to "auto". AITER requires each non-FP8 cache tensor to match qkv.dtype, so the first forward fails its dtype check. Allocate staging in the model/QKV dtype so both float16 and BF16 deployments can use these options.
With the FP4 index enabled, KV-cache sleep leaves index_scale referencing the shared backing allocation because _POOL_VIEW_ATTRS in atom/rollout/memory_manager.py includes index_cache but not this new view. The cache memory therefore remains resident, and wake can run out of memory while allocating its replacement. Add index_scale to the release list and extend tests/test_rollout_memory_manager.py to cover it.
The two findings listed as previously missed in the last review are fixed in b629d58:
FP4 staging dtype._bf16_stage became _quant_stage, and it allocates in Config.torch_dtype, so the fused writer's "auto" store matches qkv.dtype for fp16 as well as bf16 models. Covered by test_quant_stage_follows_the_model_dtype[float16|bfloat16].
index_scale on KV-cache sleep.index_scale is added to _POOL_VIEW_ATTRS, so sleep drops it with index_cache and the pool is freed. The binder rebinds both on wake. tests/test_rollout_memory_manager.py now gives the sparse impl an index_scale and asserts that it is released.
The fused top-k row-count finding is answered in its thread.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
The padding fallback passes 32 query rows unchanged to a scorer requiring divisibility by max_query_len=3, causing an AssertionError instead of handling padded rows.
_decode_index_score_tiled_kernel ran on a (request, chunk) grid. Workgroups
are dispatched to the XCDs round-robin by linear id, so with the request on
the fastest axis a batch of 8 sent every chunk of request b to XCD b (two
XCDs per request at batch 4 and 12), and the longest request in the batch
set the kernel time. The grid is now (chunk, request), which spreads every
request over all XCDs. Scores and top-k are unchanged (same per-program
arithmetic).
MiniMax-M3.1-Flash TP4, FP8 index, K=7, MI355X:
- c16 AgentX trace: 85 us per call at bs 8 vs 37 us at bs 7, i.e. 2.9 ms of
a 21.4 ms step.
- Graph replay of minimax_m3_index_topk_decode (score + top-k), bs 8, one
450K and seven 150K requests: 82.2 -> 55.3 us; no shape of bs 4-12 slower.
- c16 AgentX 3600 s: chunk-gap p90 21.24 -> 18.25 ms, request ITL p90
5.37 -> 5.13 ms, 181.6 -> 182.7 tok/s/GPU; gsm8k unchanged.
…le names
- Fused verify top-k ties: blocks of equal score now rank by the higher
block id, as index_topk._pack_score_key does, in the key encoding, the
streaming replacement (it compared scores only) and the merge. The fused
and the score + top-k paths now select the same blocks when the cutoff
falls inside a tie. Test: zero queries (every score 0), 3 and 4 requests.
- Plugin mode: --index-cache-dtype fp4 and ATOM_M3_SPARSE_KV_NVFP4 are
refused when a plugin config is set. The vLLM / SGLang bridges bind
their own caches, so the option would silently keep FP8.
- Top-k reuse guard: m3_reuses_index_topk now evaluates the model's own
schedule over its sparse layers, index_skip_topk_offset and the pattern
bounds included. The schedule moves to models/minimax_m3/index_schedule.py
(config reads only) so atom.config shares it without importing the model.
- NVFP4 KV staging: the fused writer gets only the staging pages the step
writes. AITER checks the index cache against the KV pages it is given,
and a persistent FP8 index pool can be smaller than the whole staging
buffer (max_num_batched_tokens rows). Test: a 1-token write with a 2-page
FP8 index pool and a 4-page staging buffer.
- The decode kernel no longer sets a Triton allocator on every call: none
of these kernels uses global scratch (global_scratch_size 0 for all six),
so nvfp4_kernels/utils.py goes too.
- tests/test_minimax_m3_index_topk.py: the second direct scorer launch in
the gfx1250 launch-policy test passes the scale pointer and strides.
- File names: nvfp4_kernels/decode_gfx950.py -> decode.py,
verify_topk_gfx950.py -> verify_topk.py; nvfp4_gfx950.py (the Gluon
reads) merges into nvfp4.py. can_use_nvfp4_decode_gfx950 and
flash_decode_nvfp4_gfx950 lose the suffix with them.
…r forward
- NVFP4 KV with an FP8 / BF16 index: AITER's page-128 writer checks the
index cache against the KV staging pages it is given, and a captured
graph bucket can run more rows than the pool holds, so slicing the stage
to the step's tokens is not enough. Binding now refuses a persistent
index pool smaller than the staging buffer
(m3_nvfp4_stage_unsupported_reason), with the ways out in the message.
An FP4 index is staged as well and is not coupled.
- The NVFP4 row maps (decode_rows / prefill_rows) are computed by the
forward's first sparse layer and reused by the later ones
(SparseMHAPagedAttentionImpl._row_map) instead of once per layer: 60
launches a decode step become 1. An entry stored at or above the current
layer means a new forward, or a graph capture after a warmup on the same
metadata, so it is recomputed there.
- Test: the pool / staging refusal, including a two-page bucket over a
one-page pool.
- The NVFP4 / FP4 staging buffers hold max_num_batched_tokens rows, the
largest step any forward runs: prefill chunks and decode steps are
scheduled under it, warmup stays within it, and CUDA graph capture drops
decode buckets above max_schedulable_decode_bs, which charges every
sequence mtp_k + 1 rows. The max_num_seqs x (draft + 1) term was never
reachable. The bind check compares the persistent index pool with this
bound, so a pool that fits every reachable step is no longer refused
(max_num_batched_tokens 1024 with max_num_seqs 512 and 7 drafts: 8
staging pages, a 16-page pool now passes).
- NVFP4 prefill and unfused decode call the index top-k with
emit_sparse_block_table=False: the NVFP4 attention reads top-k ids only,
so the page-16 sparse block table and context lengths were allocated and
filled for nothing (8 MiB per sparse layer for a 16K-token prefill at
top-k 16). The indexer-CP merge still emits its table inside the merge
kernel.
- prefill_rows takes searchsorted's int32 output directly.
- Tests: the staging size and the 16-page pool case.
Thanks. Both are fixed in 58764da, which is now rebased on main 456ba2a.
[P2] Bind guard compared against the over-allocation. You're right: the guard should compare the pool against the largest step a forward can run, and the staging buffer itself was over-sized. Every step is bounded by max_num_batched_tokens rows:
prefill chunks and decode steps are scheduled under it;
warmup stays within it (warmup_max_tokens);
capture drops decode buckets above max_schedulable_decode_bs, which charges every sequence mtp_k + 1 rows.
The max_num_seqs × (draft + 1) term in _quant_stage could never be reached. The staging buffers now hold ceil(max_num_batched_tokens / 128) pages, and the bind guard compares the persistent pool with that.
Your case is test_nvfp4_stage_is_the_step_bound_not_the_seq_count: max_num_batched_tokens=1024, max_num_seqs=512, 7 drafts gives 8 staging pages. The largest captured bucket, max_schedulable_decode_bs(512, 1024, 8) × 8 = 1024 rows, fits in them, a 16-page FP8 index pool passes, and a 7-page one is refused.
I kept the guard for eager mode too: warmup runs a max_num_batched_tokens-row prefill whatever the graph mode, and AITER's check is shape-only.
[P2] Unused page-16 sparse tables. NVFP4 prefill and unfused decode now call the index top-k with emit_sparse_block_table=False and keep topk_idx only. The one exception is indexer CP: its merge (merge_candidate_keys) still emits the table inside the merge kernel, as for FP8. Removing it there means changing the CP kernel, so I left that out of this PR.
prefill_rows now takes searchsorted(..., out_int32=True) directly.
Validation, on MI355X with rocm/atom-dev:latest (aiter 6fe1bfc8f, needed by current main):
MiniMax tests plus the offload, rollout and paged-attention dispatch tests: 588 passed, 33 skipped (gfx1250-only).
TP4 gsm8k strict-match:
NVFP4 KV + FP4 index: 0.9621 ± 0.0053 on the rebased head, 0.9651 ± 0.0051 just before;
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The gfx950 inline-assembly/Gluon kernels and cross-layer row-map reuse require specialist hardware review to validate graph replay, TP-dependent cache layouts, and numerical parity unavailable in this source-only environment.
@zhuyuhua-v staging bound looks right: warmup_max_tokens (model_runner.py:1312) and the capture cap (model_runner.py:3898) both sit under max_num_batched_tokens, so _step_stage can't walk past the buffer. emit=False is safe too, sparse_bt/sparse_ctx aren't read once kv_nvfp4 (attention_mha.py:1771, :1974).
any perf number for dropping the sparse table? the 8 MiB is the size, i'd like the prefill latency with and without.
one left: CP + NVFP4 KV is still allowed by the startup guard, so _cp_index_topk_decode keeps emitting that unused table. fine to defer, just leave a short comment there so we know when it can go.
…or NVFP4
With NVFP4 KV, the indexer-CP decode path still builds the page-16 sparse
table inside merge_candidate_keys, and the NVFP4 attention never reads it.
Comment at the call site saying it can go once the merge takes an emit
flag like the TP selector's.
Dropping the sparse table: what it saves. MI355X, TP4 shapes: 1 index head, top-k 16, FP4 index cache, n_valid_column_per_row as the prefill metadata builds it. I timed the prefill index top-k with emit_sparse_block_table=True and False. One call is one sparse layer, and the model has 60. Eager, median of 50 calls:
prefill
emit
no emit
saved / layer
saved / prefill (x60)
16K tokens, prefix 0
156.8 us
146.8 us
10.0 us
0.60 ms
16K tokens, prefix 100K
699.0 us
696.9 us
2.0 us
0.12 ms
32K tokens, prefix 0
325.1 us
318.4 us
6.7 us
0.40 ms
32K tokens, prefix 100K
1454.5 us
1444.6 us
9.9 us
0.59 ms
4 x 4K tokens, prefix 100K
640.8 us
637.9 us
2.9 us
0.18 ms
2K tokens, prefix 150K
202.7 us
201.8 us
0.9 us
0.05 ms
The selected top-k ids are identical with and without the table.
End to end, I measured single-request prefill latency: a fresh random-token prompt, so there is no prefix-cache hit, max_tokens=1, 10 requests per length. NVFP4 KV + FP4 index, TP4, the head before this change (bfa508b) against 58764da:
prompt
with table
without
16K tokens
274.2 ms median (min 269.8)
274.5 ms (min 269.9)
64K tokens
1089.2 ms median (min 1083.0)
1090.3 ms (min 1085.4)
So it does not show up in prefill latency. The saving per prefill is at most about 0.6 ms (16K tokens, prefix 0), against 5-15 ms of run-to-run spread. What the change really buys is the per-layer allocation: about 8 MiB for a 16K-token chunk, never read.
CP + NVFP4 KV.aa4a812 adds a short shortcut: comment at the _cp_index_topk_decode call in _sparse_decode. It says the CP merge still emits the page-16 table under NVFP4 KV, inside merge_candidate_keys, and that the table can go once the merge takes an emit flag like the TP selector's.
@zhuyuhua-v numbers work, 0.6 ms is inside the spread so it's the 8 MiB allocation that carries it, keep emit=False. checked the cp note: merge_candidate_keys hardcodes EMIT=True and _alloc_emit(..., True) (indexer_candidate_exchange.py:299, :326), so it matches attention_mha.py:1941. lgtm, thanks.
The reason will be displayed to describe this comment to others. Learn more.
not ready yet.. ci is red
please re-run pre checkin and paste the log, and any perf number / gsm8k for the new kernels?
ci Pre Checkin / Run non-GPU unit tests failed — not sure yet (low): Cannot be determined: the supplied job log is empty (0 bytes), so there is no failing test name or traceback. The job passes 12/12 on main and 11/11 on other recent runs, and this PR adds two test files to exactly that suite, so a PR-caused collection/assertion failure is plausible but unverified.
next: Re-run the job and fetch the log plus the JUnit XML artifact; if it reproduces, run python -m pytest tests/test_minimax_m3_fp4_cache_config.py tests/test_minimax_m3_index_topk.py tests/test_minimax_m3_nvfp4_kernels.py -rs on a CPU-only runner. I audited the new tests for CPU collectability and found no defect: both triton-based files call pytest.importorskip("triton") before importing the op module, atom.config imports get_gfx lazily inside the new M3 branch, and MhaKvPool/KvGeometry pull in only torch. https://github.com/ROCm/ATOM/actions/runs/41044965417/job/124026693371
ci ATOM Test failed — ci machine/env issue, not the pr (high): Not an independent failure — it is a cascade of the Pre Checkin failure. atom-test.yaml's Check Pre Checkin Signal job runs .github/scripts/check_signal.sh, which exits 78 when the Pre Checkin run for the same SHA is completed with any non-success conclusion, and every GPU job is gated on needs.check-signal.result == 'success' (atom-test.yaml:114-116), so no GPU work ran at all.
echo "Pre Checkin did not pass (conclusion: ${CONCLUSION:-unknown}), skipping workflow."
exit 78 # 78 = neutral/skip
next: Fix or re-run the Pre Checkin unit-test job; this workflow should clear on its own once Pre Checkin concludes success for the same head SHA. https://github.com/ROCm/ATOM/actions/runs/41044965430
The reason will be displayed to describe this comment to others. Learn more.
Not intentional: I introduced it when I folded the Gluon reads into nvfp4.py. 357e1c3 moves them to nvfp4_kernels/nvfp4_gluon.py; nvfp4.py is Triton-only again, so importing it (and index_topk) no longer loads triton.experimental.gluon, and the existing importorskip("triton") guard matches.
…on.py
index_topk imports nvfp4_load_even_odd from nvfp4.py, and folding the
Gluon reads into that module (e8ee35a) made every MiniMax-M3 sparse path
import triton.experimental.gluon, FP4 on or off. nvfp4.py now holds the
cache struct and the Triton loads only; the Gluon reads that the decode
attention and the fused top-k use live in nvfp4_gluon.py. Importing
nvfp4.py no longer loads Gluon.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The gfx950 inline-assembly and Gluon kernels require hardware-specialist review because source-only verification cannot validate their integrated numerical behavior, device addressing, and CUDA-graph replay.
@valar-cop thanks. Both CI and the import are dealt with in 357e1c3.
Pre Checkin. It is green on 357e1c3: Run non-GPU unit tests reports 10624 passed, 1124 skipped, 3 xfailed. Ruff, Black, the offline smoke test and the signal check all pass, and the GPU accuracy jobs are queued.
The earlier failure was tests/test_engine_stats_step_interval.py::test_a_negative_interval_falls_back_to_the_clock[-5], which this PR does not touch. It depends on the clock reading. window_expired returns now - last >= 10.0. With that run's last = 503.482308852, (last + 10.0) - last is 9.999999999999943 in float64, so the >= was false. The [-1] case runs the same code and passed. It passes in the new run.
In the new run's CPU suite, test_minimax_m3_fp4_cache_config.py runs (26 passed). test_minimax_m3_index_topk.py and test_minimax_m3_nvfp4_kernels.py skip at their triton guards, as intended on that runner.
Import (your inline note). This was not intentional: I introduced it when I folded the Gluon reads into nvfp4.py. nvfp4.py is Triton-only again, holding the cache struct and nvfp4_load_even_odd. The Gluon reads that the decode attention and the fused top-k use now live in nvfp4_kernels/nvfp4_gluon.py. Importing nvfp4.py no longer loads triton.experimental.gluon, so index_topk is back to its pre-PR Triton floor, and its existing importorskip("triton") guard matches. AITER itself still imports Gluon, as it did before this PR.
gsm8k / perf for the new kernels. MiniMax-M3.1-Flash TP4, MI355X, K=7.
gsm8k is 5-shot with the chat template, all 1319 questions, real DSpark acceptance. Image rocm/atom-dev:latest (aiter 6fe1bfc8f), measured on this branch at 58764da; the commits after it only add a comment and move the Gluon helpers:
config
strict-match
NVFP4 KV + FP4 index
0.9621 +- 0.0053
NVFP4 KV + FP8 index
0.9598 +- 0.0054
FP8 KV + FP8 index (default)
0.9659 +- 0.0050
Performance is AgentX semianalysis_cc_traces_weka_062126, 3600 s per point, synthetic accepted length 4.04. It was measured on the 10-28 head (a841c57) with image nightly_202610261558, against SGLang CI run 40110172586:
@zhuyuhua-v import fix looks right: nvfp4.py only pulls triton now, index_topk imports just that one, and gluon comes in through nvfp4_gluon only at the nvfp4 branches in attention_mha.py, so that point is closed. pre checkin is green here too, and the clock test is not ours (it is not in the file list). numbers are enough for me, verify top-k within a couple us of SGLang's kernel is fine.
hold before ATOM Test goes green, ping me with c32/c48 on current head.
gsm8k (5-shot, chat template, all 1319 questions, real DSpark acceptance): strict-match 0.9636 +- 0.0052, flexible-extract 0.9629. The server log shows block_bytes=1827840, and all 4 ranks report sparse KV is NVFP4.
AgentX (semianalysis_cc_traces_weka_062126, 3600 s per point, synthetic accepted length 4.04):
The current head is at or slightly above the 10-28 numbers. The two runs used different images and aiter versions (10-28 was nightly_202610261558), so I would not attribute the small differences to the code.
Each run had a few failed requests (3 at c32, 5 at c48). All were client-side Broken pipe / Connection reset / Server disconnected, and the server logs contain no traceback or error, the same as in the earlier runs.
@zhuyuhua-v numbers are enough: c32 325.4 / 111.8 and c48 437.8 / 66.6 are at or above the 10-28 head and level with SGLang at c48, gsm8k 0.9636 sits inside the earlier band, and block_bytes 1827840 matches the design. pre checkin and ATOM Test are green on 357e1c3, so the hold is cleared. the 3/5 drops are client side and you say earlier runs show the same rate, taking that. lgtm, thanks.
The reason will be displayed to describe this comment to others. Learn more.
Thanks — I reviewed HEAD 357e1c3cb. The PR's CPU tests (98) pass locally; I didn't run the gfx950 kernel tests. I found nothing that crashes or goes out of bounds, but there are format and consistency issues worth addressing before merge (16 inline comments).
First and most important (aiter_attention.py:49):--index_cache_dtype fp4 already means MXFP4 (E2M1 + e8m0, used by DSA / DeepSeek V4 / V4.1), and this PR reuses the same "fp4" for M3's NVFP4 (E4M3 scales) with a separate implementation, plus a fp4: bool that duplicates index_cache_dtype. Please either reuse the MXFP4 format or give NVFP4 its own value and remove the bool.
Then: the fused verify top-k can drop the current token's block when init + local > topk and orders ids differently from the non-fused path, so results depend on batch size (verify_topk.py:535/539); live-request-count constexprs cause JIT during eager serving (:602); and plain Q=1 decode also takes the verify kernel (attention_mha.py:1917). The rest are boundaries (env not in Config, M3 rules in shared config), performance (odd-chunk dummy work, 3 store launches per layer) and tests that never run in CI.
Please remove / split out (marked inline where applicable): the duplicate mooncake key selector (keys.py:89), the dead non-FP4 SGLang branches (verify_topk.py:571). Not inline: the _decode_index_score_tiled_kernel grid swap from (request, chunk) to (chunk, request) is an independent FP8-path perf fix unrelated to FP4 — please split it into its own PR; and model.py aliases should_skip_index_topk back to the old private name just to avoid touching call sites — please use the new name directly.
The reason will be displayed to describe this comment to others. Learn more.
fp4: bool duplicates index_cache_dtype, and "fp4" already means a different format._resolve_index_cache_dtype maps "fp4" to torch.uint8 (L82), which drops the format, and this bool is then added to recover it from runner.config.index_cache_dtype again (L620). The spec should carry the format (the index_cache_dtype string or an enum) and _index_cache_fields branch on that — not dtype plus a side flag.
More importantly, --index_cache_dtype fp4 already exists on main and means MXFP4: a packed E2M1 plane plus an e8m0 index_scale plane, one scale per 32 (see the KVCacheTensor comment at atom/config.py:55-68; layout from sparse_indexer_fp4.fp4_index_block_shapes, index_scale declared in mla_kv_pool.py:75, V4.1's v41_pool_geometry.py:204). It's what DSA (GLM-5.2 / V3.2), DeepSeek V4 (the gfx950 default, config.py:2769) and V4.1 (deepseek_v41/backend.py:165) use. This PR reuses the same "fp4" for M3's NVFP4 (E4M3 scales, FP4_GROUP grouping) with a separate implementation (minimax_m3/index_fp4.py: its own fp4_index_row_bytes, its own E2M1 rounding, its own index/index_scale fields). So one CLI value now means two incompatible formats depending on the model, and anything keyed on index_cache_dtype (the offload / mooncake namespaces this PR also touches) can't tell them apart.
Please either (a) use the existing MXFP4 index format for M3 and reuse sparse_indexer_fp4's layout/writer — if M3's scorer can't take e8m0 scales, please say why NVFP4 is required — or (b) give NVFP4 its own value (e.g. "nvfp4"), carry the format string in the spec, and remove this bool.
The reason will be displayed to describe this comment to others. Learn more.
When init_blocks + local_blocks > topk, the fused verify top-k keeps only the KT forced blocks with the smallest ids, which drops the current token's block; the non-fused path keeps the highest local blocks instead. can_fuse_verify_topk doesn't check init + local <= topk (it doesn't even take init_blocks). E.g. topk=16, init=4, local=16 on a long sequence: fused selects 0-3 and n-16..n-5, dropping n-4..n-1 (including the token's own partial block), so attention is wrong. Since fused only kicks in at ≥3 requests, the same request is right at small batch and wrong at large batch. Whether M3's actual sparse config hits this I couldn't confirm (no local config) — plausible.
The reason will be displayed to describe this comment to others. Learn more.
The fused path emits ascending block ids (ids = -tl.topk(-...)), while the non-fused path emits score order with forced blocks first; ±0.0 and NaN order differently too. And the guard looks at the padded graph-bucket row count, so which path a request takes depends on batch size / bucket padding (e.g. real batch 2 in bucket 4 goes fused, eager goes non-fused). NVFP4 attention accumulates in list order, so fused vs non-fused round the logits differently and results aren't batch-invariant. The tests compare only sets (_same_sets), which hides this, and the docstring's "ids equal" (L11) doesn't hold.
The reason will be displayed to describe this comment to others. Learn more.
PROGRAMS, MAX_SPLITS, MIN_BLOCKS, R_P2 = next_power_of_2(num_reqs) and the merge's NC are passed as constexpr but vary with the live request count. Under eager decode (--enforce-eager, or batches above the largest capture size), nearly every new request count ≥17 JITs a fresh gluon + triton pair while serving — multi-second stalls. These should be runtime arguments (or bucketed) so the build count is bounded.
The reason will be displayed to describe this comment to others. Learn more.
_fused_verify_topk doesn't require max_query_len > 1, so plain decode (Q=1) with ≥3 requests also takes the "verify" kernel — the name no longer matches. With Q=1 and one index head per rank, the 16-row MFMA tile uses one row; VERIFY_FUSED_MIN_REQS_NVFP4 = 3 was measured on spec verify, so switching plain decode has no measurement behind it, and it brings in the ordering difference above.
The reason will be displayed to describe this comment to others. Learn more.
_row_map stashes its result on the metadata object with setattr and decides "new forward" by layer_num size — a new "compute once per forward" mechanism. make_sparse_decode/prefill_metadata already computes n_valid_column_per_row once per step for all sparse layers; row_req / row_seq belong there. No stale data today, but any future caller sharing the metadata with the same (rows, Q) and a larger layer_num (e.g. a sparse draft layer) would silently get the target's row map. The prefill pad branch (live = tok < cu[-1]) is dead too — the M3 builder has supports_token_graph=False, so prefill is never padded.
The reason will be displayed to describe this comment to others. Learn more.
Please remove the non-FP4 branches carried over from SGLang: the bf16/fp8 k_ok in the guard, VERIFY_FUSED_MIN_REQS = 16, the launcher's non-NVFP4 path, _load_k_tile's else branch, and the never-passed out= parameter. The only caller always passes Nvfp4Cache, so these several hundred lines have no caller and no test. Related in decode.py: v_scale is always 1.0, % max_slots silently wraps an invalid block id onto a valid page, and the guard page_size % block_size == 0 passes for page_size=256 but then addresses 128-blocks on 256-pages (wrong page) — it should be == 128. The call-site literals 128 / 1.0 also duplicate the guard's SPARSE_BLOCK_SIZE.
The reason will be displayed to describe this comment to others. Learn more.
The _quant_stage staging buffers are allocated at bind, outside the KV memory budget and the profile; they're lazily created via getattr, cached by name only (not by inner shape), and not released on rollout sleep. At max_num_batched_tokens=32K that's ~24 MiB per rank not counted in gpu-memory-utilization, kv_stage / index_stage stay resident after sleep, and a second geometry would silently get a wrong-shaped buffer. Please declare them in __init__ and tie them to the pool's lifecycle.
The reason will be displayed to describe this comment to others. Learn more.
Boundary: m3_reuses_index_topk is just a forwarding wrapper with a local import, and M3-specific rules (m3_fp4_cache_unsupported_reason, m3_nvfp4_stage_unsupported_reason) live in the shared config.py; indexer CP's fp4 refusal is a special case outside the pure indexer_cp_unsupported_reason. index_schedule.py is already a pure config module — move these three there, drop the wrapper, and put the CP refusal inside the reason function. Also sparse_attn_nvfp4.check says "needs 16 heads" though config already enforces the head count (the runtime failure would be alignment/addressing), and docs are stale: mha_kv_pool's kv_views / scale_views docstrings still say SHUFFLE/fp32, attention_mha.py:1372 still says "binder has no M3-specific code", and kv_pools is still annotated tuple[int, int].
The reason will be displayed to describe this comment to others. Learn more.
All ~430 lines of kernel tests are gated on gfx950, so CI (CPU torch) always skips them — store, scorer, attention and fused top-k regressions are only caught by someone running on MI355X by hand. The reference quantizer copies the kernel's thresholds and 1/s step by step, so a quant-formula edge error would appear in both and cancel; the 2% attention tolerance is loose enough that a dropped/duplicated token or an off-by-one causal mask still passes; the prefill row map is never checked against attention output; and the fused test covers only init_blocks=0, so the init+local and ordering cases above aren't covered.
…oncurrency
NVFP4 sparse KV (ATOM_M3_SPARSE_KV_NVFP4) and the FP4 index cache come from
#2744, which is not on main, so the c32 / c48 band and its flags go: every
concurrency runs FP8 KV + FP8 index, with the prefill delayer up to c16.
The FP8 cache's capacity limit at c48 stays documented as a known limit,
with the earlier measurement. The measured rows, all taken with the previous
command (and NVFP4 at c32 / c48), stay marked as to be re-measured.
This branch has not been deployed
No deployments
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.
zhuyuhua-v
changed the title
valar-cop
left a comment
valarLip
left a comment
Summary
This PR adds two opt-in caches for MiniMax-M3. Both are off by default and only reach MiniMax-M3.
--index-cache-dtype fp4stores the lightning-indexer keys as NVFP4: a 64 B e2m1 payload plus 8 B of E4M3 scales per key, against 128 B for FP8.ATOM_M3_SPARSE_KV_NVFP4=1stores the sparse layers' K/V the same way, in their ownMhaKvPool(kv_format="nvfp4").At high concurrency the FP8 cache is what limits MiniMax-M3.1. At c48 on TP4 (
--gpu-memory-utilization 0.8):With both caches on, the KV block shrinks from 3,179,520 to 1,827,840 bytes. That is 1.74x the tokens at the same memory: 12.7M tokens at
--gpu-memory-utilization 0.9.This PR replaces #2745 and #2746, which were stacked on it.
Design
Kernels.
atom/model_ops/minimax_m3/nvfp4_kernels/holds the gfx950 kernels. The derived files carry a license header.decode.py: Gluon block-sparse attention over a packed NVFP4 K/V cache, with rows addressed through the block table.verify_topk.py: fused index score + streaming top-k for speculative-verify rows.nvfp4.py: the NVFP4 cache view and its e2m1 -> fp16 even/odd row loads, in Triton and Gluon.Each kernel has a shape guard:
can_use_nvfp4_decodeandcan_fuse_verify_topk.Write path.
index_fp4_store/kv_nvfp4_storethen quantize those rows into the paged planes at the real slots: per 16 elements,scale = e4m3(amax / 6), codes rounded half-to-even and saturating at 6.max_num_batched_tokensandmax_num_seqsx (draft + 1) rows.Index. The prefill and decode scorers read the FP4 plane (
IS_FP4) through the gfx950 e2m1 decode, as even/odd fp16 tiles.Decode score grid (FP8 and FP4 index).
_decode_index_score_tiled_kernelnow runs on a (chunk, request) grid instead of (request, chunk). Workgroups go to the XCDs round-robin by linear id. With the request on the fastest axis, a batch of 8 put every chunk of request b on XCD b (two XCDs per request at batch 4 and 12), so the longest request set the kernel time. Scores and top-k are unchanged.Attention.
nvfp4_sparse_attentionover their top-k blocks.decode_rowslaunch per layer maps each row to its request and causal length. The attention and the fused top-k share it.Fused verify top-k. It replaces score + top-k whenever
can_fuse_verify_topkholds: NVFP4 KV, FP4 index, at least 3 requests, and within the kernel's shape bounds. It emits top-k ids only, which is all the NVFP4 attention reads.Startup refusals (
m3_fp4_cache_unsupported_reason):Indexer CP refuses the FP4 index; Mono refuses NVFP4 KV.
Code shared with other models.
MhaKvPoolgains akv_formatfield. For the default"", the fields, sizes and views are unchanged.KvGeometrygains a defaultedkv_format, so its name is unchanged for every existing pool.aiter_attention.pyandattention_mha.pyrun only for MiniMax-M3 sparse layers.Results
MiniMax-M3.1-Flash-Quark-MXFP4, 4x MI355X TP4, fixed 7-draft DSpark,
ATOM_M3_DSPARK_DRAFT_PTPC=1(#2743), gpu-memory-utilization 0.9.Throughput is measured with AgentX
semianalysis_cc_traces_weka_062126, prefix cache on, synthetic accepted length 4.04.atom:prefix_cache_cached_tokens / atom:prompt_tokens.3600 s per point, measured on the stacked #2744 + #2745 + #2746 (on #2686):
This branch (main 035a78a + this PR + #2743, image
rocm/atom-dev:nightly_202610261558): c48, 1200 s: 397.6 tok/s/GPU, P90 interactivity 71.4, p50 interactivity 105.6, prefix hit 95.58%, TTFT p50 297 ms, 0 errors. The stacked PRs on main with the same image and length gave 396.7 / 70.5 / 104.7, with a 95.96% prefix hit.Decode score grid, on the FP8 + Mono c16 configuration (FP8 KV + FP8 index, 3600 s, same tree with and without the change):
minimax_m3_index_topk_decode(score + top-k) at bs 8, with one 450K request and seven 150K requests: 82.2 -> 55.3 us. No batch from 4 to 12 is slower.Accuracy. lm-eval gsm8k as InferenceX CI runs it: 5-shot, chat template, temperature 0, all 1319 samples, real acceptance.
Where not to use it. At c16, FP8 + Mono stays faster: 181.6 tok/s/GPU and P90 191.2, against 156.4 and 137.9 for NVFP4 over 1200 s. That NVFP4 run predates the per-layer row-map fix. The NVFP4 caches are meant for c >= 32, where the FP8 cache runs out.
Follow-up: MXFP4 caches
The NVFP4 caches in this PR are the first step for c32 and c48. The next step is MXFP4: e2m1 data with one E8M0 scale per 32 elements. It is ahead of NVFP4 on capacity, on the write path and in the readers. The work is in five open PRs:
fused_qknorm_idxrqknormquantizes K, V and the index key to MXFP4 straight into 128-token pages, and can emit the main Q as e4m3pa_decode_mxfp4: paged-attention decode over an MXFP4 KV cache with an fp8 QATOM_M3_SPARSE_KV_MXFP4=1 --index-cache-dtype mxfp4Why MXFP4 is faster than the NVFP4 path here:
--gpu-memory-utilization 0.9, that is 104,776 blocks per GPU against 99,295 ([MiniMax-M3.1] MXFP4 sparse KV cache and index cache #2770)._nvfp4_store_kernellaunches per layer, 762 us per step.v_mfma_scale_f32_16x16x128_f8f6f4) takes fp4 operands with one E8M0 scale per 32 elements, which is exactly the MXFP4 layout. NVFP4's E4M3 scale per 16 elements does not fit it, so this PR's kernels convert e2m1 to fp16 first.pa_decode_mxfp4(planned path, which [MiniMax-M3.1] MXFP4 sparse KV cache and index cache #2770 uses for decode) takes 0.57-0.64x the time of the fp8pa_decodeat 512-4096 rows, and 1.05-1.16x at 1-64 rows (#7056).End to end, as reported in #2770. Each point is one run over a 1200 s window, on an older tree than this PR. Compare the rows of this table only with each other, not with the 3600 s results above.
At c16, MXFP4 gave 163.8 / 181.0 against 163.6 / 156.5 for fp8 + Mono.
Open points (from #2770):
Relation to this PR. #2770 is cut from
yuhua/m3.1-flash, which carries this PR, and adds its MXFP4 pool next to the NVFP4 one. Startup refuses mixing the two formats. The plan is to move c32 and c48 to MXFP4 once the aiter PRs land. Until then, the NVFP4 path in this PR is the configuration that runs.Tests
tests/test_minimax_m3_fp4_cache_config.py(CPU): the startup refusal matrix (gfx, TP2/4/8, top-k reuse); the NVFP4 pool's per-entry bytes, views and layer isolation; the geometry name.tests/test_minimax_m3_nvfp4_kernels.py(gfx950, skipped elsewhere):decode_rows/prefill_rowsmatch torch;tests/test_minimax_m3_index_topk.pyfollows the scorer's newindex_kv_scaleargument.The new tests pass on MI355X (gfx950), 24 of 24. The default path is unchanged on this branch: with FP8 KV + FP8 index, gsm8k strict-match is 0.9613 +- 0.0053 and
block_bytesis 3,179,520, as on main.🤖 Generated with Claude Code