[zhuyuhua-v]

Summary

This PR adds --index-cache-dtype fp4 for MiniMax-M3: the lightning-indexer keys are stored as NVFP4. It is stacked on the vendored-kernels PR (#2744); review only the top commit.

Per cached key (D = 128):

• a 64 B e2m1 payload, two codes per byte, even element in the low nibble;
• 8 B of E4M3 block scales, one per 16 elements, with a unit outer scale.

That is 72 B, against 128 B for FP8. It is the format SGLang's M3 NVFP4 index pool stores.

• Store (index_fp4.py): a Triton kernel quantizes per 16 elements with scale = e4m3(amax / 6) and codes rounded half-to-even, saturating at 6. A torch dequant reference sits beside it.
• Pool and binding: two uint8 entry fields, index and index_scale, live in the same entry as the KV. Each sparse layer gets its payload and scale views plus a shared token-major BF16 staging buffer.
• Write: aiter's fused qk-norm / rope / cache kernel only writes BF16 or FP8 index keys. It now writes the step's keys into the staging buffer (identity slot mapping), and index_fp4_store quantizes them into the paged planes at the real slots.
• Scorers: _index_block_score_kernel (prefill) and _decode_index_score_tiled_kernel (decode) take the scale plane (IS_FP4). They decode rows with the gfx950 hardware e2m1 -> fp16 conversion into separate even/odd fp16 tiles and chain the two dots, so no interleave shuffle is needed.
• Indexer CP: refused with an FP4 index, because its context scorer reads FP8/BF16 keys only.
• Mono: already refuses a non-FP8 index cache, so nothing changes there.

Why

At high concurrency the FP8 KV + FP8 index cache runs out. On ATOM at c48 (AgentX semianalysis_cc_traces_weka_062126, gpu-memory-utilization 0.8):

• the 6.06M-token cache was full, against 7.97M tokens in flight at p50;
• the prefix-cache hit rate fell to 17% (the traces allow 96.6%);
• output was ~94 tok/s/GPU.

This PR is the index half of the fix. The KV half is the next PR in the stack. KV block bytes for MiniMax-M3.1-Flash TP4 (from the server log):

KV	index	block_bytes
FP8	FP8	3,179,520
FP8	FP4	2,749,440 (-13.5%)
NVFP4 (next PR)	FP4	1,827,840 (-42.5%)

Validation

MiniMax-M3.1-Flash-Quark-MXFP4, TP4, MI355X.

• Kernels (gfx950): stored values match the torch reference exactly. Top-k from the FP4 scorers is identical on 100% of rows, for both decode and prefill, to the BF16 scorers fed the dequantized cache (e2m1 x E4M3 values are exact in bf16).
• Scorer cost (even/odd fp16 dots), FP4 vs FP8 index, per layer:
  • decode bs64 at ~150K context: 228 vs 167 us; bs6: x0.92 of FP8;
  • prefill: ~2x FP8.
  • With the next PR, decode at bs >= 3 instead goes through the fused verify kernel, which is on par with FP8 at bs16 (60.9 vs 61.2 us) and faster at bs64 (141.5 vs 198.8 us).
• gsm8k (InferenceX CI method: 5-shot, chat template, temperature 0, full 1319, real acceptance), FP8 KV + FP4 index on [MiniMax-M3.1] Enable Mono decode with fixed DSpark verify #2686: strict 0.9629 +- 0.0052. SGLang CI on MI355X: 0.9598 +- 0.0054. On main with the whole stack ([MiniMax-M3.1] NVFP4 sparse KV cache + fused verify top-k (3/3) #2746, NVFP4 KV + FP4 index + fused verify): strict 0.9644 +- 0.0051.
• CPU unit tests touched by the change pass. tests/test_minimax_m3_index_topk.py follows the new scorer signature.

On its own, FP4 index does not speed decode up. At c32 (1200 s, NVFP4 KV), FP4 index gave 337.5 tok/s/GPU and P90 interactivity 79.2, against 340.1 and 87.1 with FP8 index. Its value is capacity, which matters at c48 (see the next PR).

🤖 Generated with Claude Code

[zhuyuhua-v]

Folded into #2744, which now carries the FP4 index cache, the NVFP4 sparse KV cache and the fused verify top-k as one change (with review fixes and tests). Closing in favour of #2744.