[None][perf] Fuse MiniMax-M3 MSA per-layer KV-cache writes into one kernel

[zheyuf]

Description

Where this PR targets at:

MiniMax-M3 MSA per-layer KV-cache writes fusion (K, V, index-K):

Take a sparse layer as example:

1. input RMSNorm
2. QKV proj  (+ index-Q/K proj for the indexer)
3. ►► WRITE: scatter new K, V, index-K into the paged caches   ◄◄ our fix
4. indexer:  proxy-FMHA (reads index-K) → max-reduce → SelectBlocks (top-16 of 128-token blocks)
5. main sparse attention (reads K,V from selected blocks) → o_proj
6. post-attn RMSNorm
7. MoE FFN (router → top-4 experts + shared) + allreduce

The problem

Each MSA layer writes its new-token main K, main V, and (sparse layers) index-K. Writing of K, V and index-K takes four tiny kernels: division, remainder, index cast, index_put scatter. So that's 3*4=12 kernel launches per sparse layer. At 60 layers that is ~720 tiny kernels per decode step, all captured into decode CUDA graphs and re-executed on every step. It accumulates to large launch overhead.

The fix

This PR replaces them with one Triton launch per layer (~720 → 60) launches on the write path each iter).

Measured impact (4x B300, InferenceMAX-style serving benchmark)

trtllm-serve + benchmark_serving, NVFP4 + fp8 KV cache, MSA, decode CUDA graphs + overlap scheduler, no spec decode, random 8k/1k, identical seeded prompt sets, JIT/autotuner warmed.

metric	conc 64, TP4/EP4	conc 32, TP2/EP1
output throughput	1,817.1 → 1,942.5 tok/s (+6.9%)	1,082.1 → 1,171.7 tok/s (+8.3%)
median TPOT (≈ decode iteration time)	33.86 → 31.76 ms (−2.1 ms)	28.39 → 26.20 ms (−2.2 ms)
mean TTFT	995 → 954 ms (−4.1%)	715 → 690 ms (−3.5%)

Nsys trace for a layer (look at green parts)

Before fix: 4 tiny kernels-to-be-fused for index-K, then 3 kernels for indexer, then 8 tiny kernels-to-be-fused for K and V.

After fix: 1 fused Triton launch (contains previous 12 kernels for K, V and K-index), then 3 kernels for indexer. So it's fuse + reorder: compute the page/offset once, write K + V + index-K together, before the indexer runs.

Test Coverage

• Spec-decode: TestMiniMaxM3::test_nvfp4_eagle3[tp4-ep4-adp=False-overlap=True-msa=True-graphs=True-eval_mode=inferencemax] PASSED on this branch — strict-match 95.45 (baseline same tip: 95.07), acceptance rate 0.844 vs 0.846, mean acceptance length 3.532 vs 3.537

PR Checklist

• PR title and description
• Test coverage
• CI (/bot run)