[MiniMax-M3] Decode: paged K/V tile loads and a tiny GEMM for the router projection

[thanhhao98]

Two independent decode-path changes for MiniMax-M3, both measured on GB300 (sm_103), TP4, MiniMax-M3-NVFP4, ISL 8192 / OSL 1024, --context-length 1048576, backends auto-resolved.

1. Sparse-decode K/V tiles by page

_gqa_share_sparse_decode_kernel resolved each selected block by loading 128 slot ids from req_to_token and addressing K/V by those ids — two dependent global round trips per block, scattered addressing.

A selected block starts at a multiple of block_size, and PagedTokenToKVPoolAllocator lays slots out as page * page_size + offset (mem_cache/allocator/paged.py). So when page_size >= block_size and page_size % block_size == 0 the block lies inside one page and its slots are contiguous: load the block's first slot, walk it, and K/V addressing becomes affine.

Gated on page_size, which the backend already tracks. page_size == 0 keeps the gather, as do the existing callers, so nothing outside this configuration changes behaviour. The invariant is already load-bearing elsewhere — metadata.py does req_to_token[pool_idx, ::page_size] // page_size, and the M3 page-table top-k kernel does req_to_token[...] / page_size.

nsys, cc=1 decode: _gqa_share_sparse_decode_kernel 8.61 → 5.27 µs/launch (−39%), 57 launches/step, unchanged grid.

2. Router projection on the in-tree tiny GEMM

The router projection is N=128 (num_local_experts) × K=6144 (hidden). cuBLAS splits K on that shape and runs it as two launches. The in-tree tiny_gemm_bf16 (python/sglang/kernels/ops/gemm/tiny_gemm.py) does it in one, and DeepSeek-V2 and Dots3 already dispatch to it through the same tiny_router_gemm_max_tokens helper.

Dispatch order is gfx95 router_gemv (#36557) → NPU → tiny GEMM → torch.mm; the paths are device-disjoint. Guarded off under enable_deterministic_inference, since the row-count threshold makes a token's logits depend on the batch it lands in.

nsys, cc=1 decode: router projection 2.00 → 1.00 launches/layer, 6.20 → 1.81 µs/layer.

Accuracy

GSM8K 1319 q against a matched unpatched control on the same image: 0.856 vs 0.856. bs=1 canary and an EOS/termination check pass on both. The sparse-decode change is additionally verified bitwise-equal to the gather path on the allocator's layout.

Benchmarking

Median TPOT (ms), cumulative over the two changes plus #38615:

cc	1	2	4	8	16	32
main	4.94	5.60	6.26	7.26	9.35	12.84
this PR	4.28	4.84	5.56	6.57	8.82	12.68

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.

---

CI States

Latest PR Test (Base): ❌ Run #34471493086
Latest PR Test (Extra): ❌ Run #34471492937
Latest PR Test (AMD ROCm 10): ❌ Run #34471493321