[zhuyuhua-v]

Summary

ATOM_M3_SPARSE_KV_NVFP4=1 stores the K/V of MiniMax-M3's sparse (indexed) attention layers as NVFP4. Per token and kv head that is a 64 B e2m1 payload plus 8 B of E4M3 block scales, against 128 + 4 B for FP8. With --index-cache-dtype fp4 from the previous PR, decode index scoring and top-k run in SGLang's fused gfx950 verify kernel.

Stacked on #2745 (FP4 index), which is stacked on #2744 (vendored kernels). Review only the top commit.

• Pool: KvGeometry gains kv_format. The sparse layers declare "nvfp4" and get their own MhaKvPool, whose k/v and k_scale/v_scale fields are uint8 planes. Dense layers, the draft and --kv_cache_dtype are unchanged.
• Write: the fused qk-norm / rope / cache kernel writes BF16 K/V into a shared token-major staging buffer (asm_layout=False, identity slots). kv_nvfp4_store (index_fp4.py's store kernel, per head) then quantizes them into the pool at the real slots.
• Attention: decode and prefill rows run through SGLang's gfx950 NVFP4 sparse decode kernel (vendored), addressed by the top-k block ids and the request block table.
  • A one-launch Triton kernel builds each decode row's (request, causal length) on device, so it is valid under CUDA graph capture.
  • An earlier per-forward Python cache of these rows went stale under capture and faulted.
• Fused verify top-k (ATOM_M3_FUSED_VERIFY_TOPK, default 1): with an FP4 index and >= 3 requests, decode score + top-k is one fused kernel. It emits top-k ids only, which is all the NVFP4 attention reads, so it is tied to this path.
• Mono decode refuses this configuration.

Results

MiniMax-M3.1-Flash-Quark-MXFP4, 4x MI355X TP4, AgentX semianalysis_cc_traces_weka_062126 with prefix cache, fixed 7-draft DSpark with synthetic AL 4.04, gpu-memory-utilization 0.9.

• tok/s/GPU = output throughput / 4.
• P90 interactivity = 1000 / p90 ITL.
• SGLang is the InferenceX CI run on the same recipe.

KV capacity (from the server log):

config	block_bytes	blocks	tokens
FP8 KV + FP8 index	3,179,520	57,011	7.3M
NVFP4 KV + FP4 index	1,827,840	99,130	12.7M (+74%)

c48, 3600 s (on #2686 + draft ptpc):

config	tok/s/GPU	P90 intvty	prefix hit	TTFT p50
FP8 KV + FP8 index (gpu-mem 0.8, before)	~94	—	0.17	44 s
NVFP4 KV + FP8 index	380.7	43.3	0.865	462 ms
NVFP4 KV + FP4 index + fused verify (this PR)	434.8	67.8	0.929	439 ms
SGLang CI	436.1	70.9

c32, 3600 s:

config	tok/s/GPU	P90 intvty
NVFP4 KV + FP8 index	315.0	96.5
this PR	324.1	106.6
SGLang CI	335.4	141.6

On this branch (main 481c2bf + this stack + ATOM_M3_DSPARK_DRAFT_PTPC=1, image rocm/atom-dev:nightly_202610261558), c48, 1200 s: 396.7 tok/s/GPU, P90 interactivity 70.5 (SGLang CI 436.1 / 70.9), p50 ITL 9.55 ms, prefix hit 0.905, TTFT p50 375 ms, 0 errors. That run differs from the 3600 s table above in three ways at once: run length, base tree (main without #2686) and image. The 9% lower throughput is not attributed to any one of them.

Where this is not the best config. At c16, FP8 KV + FP8 index with mono is faster: 181.6 tok/s/GPU and P90 interactivity 191.2, against 156.4 and 137.9 for this config (1200 s, measured before the on-device decode-row fix, not re-measured). This config is meant for c >= 32, where the FP8 cache runs out.

Accuracy

lm-eval gsm8k as InferenceX CI runs it: 5-shot, chat template, temperature 0, full 1319 samples, real acceptance, NVFP4 KV + FP4 index + fused verify.

tree	strict-match
#2686 + this stack	0.9659 +- 0.0050
this branch on main	0.9644 +- 0.0051
SGLang CI (MI355X)	0.9598 +- 0.0054

Tests (gfx950)

• kv_nvfp4_store over H heads equals the per-head index store.
• The decode-row kernel matches the torch mapping.
• NVFP4 sparse attention: max relative error 0.0046 against the dequantized reference.
• Fused verify top-k: identical to the Triton score + top-k on 100% of rows.

🤖 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.