fix(minimax-m3): unblock fp8 KV cache and EAGLE3 spec decode on the vLLM plugin

[PerryZhang01]

Two fixes that together let MiniMax-M3 serve under --kv-cache-dtype fp8 with an EAGLE3 draft on the vLLM plugin backend. Neither one alone is enough: the first unblocks prefill, the second unblocks decode.

1. The indexer cache must carry a real fp8 dtype

Serving MiniMax-M3 through the vLLM plugin backend with --kv-cache-dtype fp8 kills EngineCore on the first request:

triton.compiler.errors.CompilationError: at 73:17
    qk = tl.dot(q, k, out_dtype=tl.float32) * sm_scale_log2e
Both operands must be same dtype. Got bf16 and uint8

MiniMaxM3SparseIndexerCache takes its KV-cache spec dtype from vLLM's kv_cache_dtype_str_to_dtype(), and vLLM maps every fp8 kv-cache-dtype to torch.uint8 -- a byte buffer its own kernels reinterpret. ATOM's _index_block_score_kernel dispatches on k.dtype.is_fp8() instead, so a uint8-labelled index cache sends it down the bf16 branch, where it dots bf16 against uint8 and fails to compile.

The bytes in that buffer are already fp8: aiter's fused_qknorm_idxrqknorm writes them under kv_cache_dtype="fp8", and the native ATOM server allocates the same cache as dtypes.d_dtypes["fp8"]. Only the torch dtype label was wrong, so label it aiter.dtypes.fp8 (torch.float8_e4m3fn). The element stays one byte, so vLLM's page-size accounting is unchanged, and vLLM already allows a true fp8 dtype in this position -- its own "fp8_inc" entry maps to torch.float8_e4m3fn.

Scoped to the indexer cache deliberately: the main sparse and dense KV caches keep their uint8 label because the aiter paged kernels read them as raw bytes.

2. A decode segment must be uniform before a decode kernel may have it

With the above fixed, M3 + EAGLE3 still dies part way into a long run:

atom/model_ops/minimax_m3/index_topk.py minimax_m3_index_topk_decode
AssertionError: total_q 121 not divisible by max_query_len 4

Three ATOM decode paths recover a request from a flat query row by dividing:

• the M3 index-topk kernels use row // max_query_len, and the row's causal cutoff seq_len - max_query_len + tok + 1
• aiter's gluon paged decode reshapes q to [q.shape[0] // max_query_len, max_query_len, ...]

All three only hold when every decode request contributes exactly max_query_len rows, and none of them checked. Speculative decode breaks the assumption on both models. On the target, a request that joins the batch without draft tokens contributes one row while its neighbours contribute num_spec + 1, and vLLM keeps all of them in the decode segment because each query length is still within the reorder threshold (121 = 30*4 + 1). On the draft, a request contributes however many tokens the previous step accepted, so the segment is ragged by construction -- which surfaces as

aiter/ops/triton/gluon/pa_decode_gluon.py
RuntimeError: shape '[27, 4, 1, 16, 128]' is invalid for input of size 225280

(110 rows where 27 requests * 4 = 108 were assumed).

_uniform_decode_query_len() measures the segment, and all three builders consult it:

• MinimaxM3SparseAttentionMetadataBuilder.build() takes the uniform decode fast path only when the segment really is uniform; otherwise those requests go to the prefill kernel, which derives causality from cu_seqlens_q/context_lens and accepts variable query lengths.
• AiterMhaMetadataBuilderForVllm.build() routes a ragged decode segment to the extend path, which is varlen. Decode requests sort before extends, so widening the extend segment covers them without reordering.
• AiterMhaMetadataBuilderForVllm.build_for_drafting() -- the one drafting actually goes through. Its docstring claimed "during EAGLE/MTP drafting all requests are uniform decodes" and it only tested for prefills, so it never reached build(). It now falls back to build() when the batch is ragged.

Both mixed-batch branches also stopped reporting a max_query_len they had not measured: the M3 one passed reorder_batch_threshold (always num_spec + 1) and the MHA one a max over the segment. A plain-decode segment under a spec-decode threshold takes 4 there, and with a row count that happens to divide by 4 the kernels mis-map every row silently instead of asserting. tests/plugin/test_minimax_m3_decode_uniformity.py pins that case ([4, 4, 6, 2] sums to 16, divides by 4, and is still ragged).

The kernel-side assert stays as a backstop, now spelling out the invariant.

3. CI

The M3 accuracy cell ran with --kv-cache-dtype auto and no speculative config, so it exercised neither fix -- the index cache never reached the fp8 label path and the decode segment was uniform by construction. It now runs fp8 KV plus an EAGLE3 draft (Inferact/MiniMax-M3-EAGLE3-GQA, num_speculative_tokens=3); nothing else about the cell changed. The threshold stays 0.93, and the native accuracy catalog lists 0.9469 for MXFP4 + EAGLE3 at that same threshold.

Test Result

Fix 1, on 4x MI355X (gfx950) TP4, amd/MiniMax-M3-MXFP4, vLLM 0.27.2.dev0+g6e448d0ea, --kv-cache-dtype fp8: before, the first chat request kills EngineCore; after, the index cache is allocated as (num_blocks, 128, 128) torch.float8_e4m3fn, chat completions answer normally, and a 90K-token 90%-prefix-hit aiperf run completes clean (32/32 requests, TTFT p50 2.28 s, TPOT p50 14.25 ms at concurrency 8).

Fix 2, on 4x MI355X TP4 with MiniMax-M3-MXFP8 + Inferact/MiniMax-M3-EAGLE3-GQA (num_speculative_tokens=3), gsm8k 5-shot chat, full 1319 questions:

build	flexible-extract	strict-match	outcome
fp8 KV + EAGLE3, before this PR	--	--	total_q 121 not divisible by max_query_len 4
fp8 KV + EAGLE3, after	0.9530 ±0.0058	0.9538	no errors, mean acceptance length 3.27
fp8 KV, no spec decode (regression)	0.9484 ±0.0061	0.9492	unchanged from 0.9477 before this PR
auto KV, no spec decode	0.9447 ±0.0063	0.9454	--

The no-spec numbers confirm the uniform fast path still behaves; the spec-decode run keeps its acceptance length, so drafting is doing real work rather than degenerating to plain decode.

tests/plugin/test_minimax_m3_decode_uniformity.py: 13 cases, CPU only.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

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 2122 --add-label <label>

[ganyi1996ppo]

Why not handle all of this in kv_cache_dtyoe_str_to_dtype

[PerryZhang01]

vllm kv_cache_dtype_str_to_dtype fucn will transfer str 'fp8' to torch.uint8