[jiacao-amd]

Purpose

The ROCm BF16x3 router tier (bf16x3_router_gemm_rocm, fp32 gate weights as in MiniMax-M2/M3) only serves batches of 2048+ tokens. Decode batches above the gfx950 fp32 low-M kernel (M ≤ 32) and below 2048 fall back to F.linear(x.float(), w) on the fp32 pipe, which is both slower and less accurate.

This PR adds a split-K Triton kernel for that range:

• Each program loads one BLOCK_M × K/SPLIT_K slice of the bf16 activations and the matching slices of the three bf16 weight terms, and accumulates the three terms in separate fp32 accumulators (small terms are summed among themselves before meeting the large one).
• A second kernel sums the K partitions in a fixed order (deterministic, no atomics).
• SPLIT_K is chosen so the grid stays within 768 programs while each split keeps ≥192 of K (within 3% of the best split on average, 16% worst case, over M ∈ [64, 2047] and six router shapes).

Dispatch (rocm_bf16x3_router_gemm_dispatch): M ≤ 32 → existing fp32 low-M kernel (still fastest there); 32 < M < 2048 → split-K; M ≥ 2048 → existing stacked hipBLASLt GEMM. Arbitrary M/N/K are supported via masks (N=17/384, K=1000/2880 and padded activation rows are tested).

Microbenchmark (MI355X, µs per call, CUDA graph over 57 distinct weights)

(K, E)	M	split-K	F.linear fp32 (before)	stacked GEMM
(6144, 128) MiniMax-M3	33	6.8	12.0	11.1
	96	7.6	14.2	11.9
	192	8.7	17.5	15.0
	512	10.9	22.1	18.1
	1024	14.0	30.3	23.1
	2047	22.1	43.9	28.2
(3072, 256) MiniMax-M2	33	5.9	11.0	10.2
	96	6.6	11.6	12.0
	512	10.8	18.7	16.1
	2047	21.5	35.8	23.9

Accuracy (rel. L2 vs fp64): split-K 1.1e-7 – 2.4e-7, F.linear 2e-7 – 1.4e-6.

E2E (MI355X, MiniMax-M3-MXFP4, TP4, EAGLE3 K=3)

vLLM nightly 8e2beaa3 (AITER v0.1.26) plus #61889, which that nightly needs to start MiniMax-M3; the baseline is the same image without this PR. AgenticX: InferenceX recipe minimaxm3/vllm/mi355x-fp4-mtp/agentic.yaml, 3600 s per point, one run each (run-to-run noise is about ±3%). Cells are baseline → this PR.

AgenticX (selected points; also measured: C1, C4, C8, C15, C20)

Conc	Throughput/GPU (tok/s)	P50 interactivity	P90 interactivity	Mean TPOT	P90 TTFT
24	40,256 → 40,344 (+0.2%)	169.9 → 172.0 (+1.2%)	97.5 → 100.2 (+2.7%)	6.80 → 6.74 ms (−0.9%)	1.19 → 1.25 s (+5.3%)
32	45,022 → 45,234 (+0.5%)	123.4 → 126.9 (+2.8%)	67.9 → 70.6 (+4.0%)	9.59 → 9.70 ms (+1.1%)	1.59 → 1.63 s (+2.7%)

The new kernel serves decode batches above 32 tokens (more than 8 running requests with K=3): about 29% of the time at C24 and 63% at C32. At C15 and below the dispatch is unchanged.

8k/1k C24 (240 requests, every decode batch is in the split-K range; measured on nightly 1b2a666 plus ROCm/aiter#6336, 2 interleaved rounds vs a same-round baseline):

	Output throughput	Mean TPOT
baseline → this PR	+1.73% (+1.88 / +1.57)	−3.3%

Accuracy (GSM8K), InferenceX gsm8k.yaml, 5-shot, all 1,319 questions, real EAGLE3 acceptance, thinking on:

	strict-match	flexible-extract
baseline	96.21%	96.13%
this PR	96.06%	96.06%

Standard error ±0.53%.

The kernel is also more accurate than the F.linear fp32 path it replaces (relative L2 vs fp64 above), and a real-acceptance greedy smoke (24 concurrent ≥8k-token prompts + 4 short) generated text identical to the baseline (24/24, 4/4).

Test Plan

pytest tests/kernels/test_bf16x3_router_gemm_rocm.py tests/kernels/test_gate_linear_rocm_dispatch.py on gfx950 — 220 passed. New cases: fp64 reference across M ∈ {1, 33, 96, 257, 1000, 2047, …}, non-power-of-two experts and ragged K, strided activations, determinism, config bounds, more-accurate-than-F.linear and expert-selection checks at decode M, and dispatch routing by M.

Draft: E2E covers MiniMax-M3 only; MiniMax-M2 / HunYuan-V3 E2E not yet measured.

🤖 Generated with Claude Code

[jiacao-amd]

Closing: on MiniMax-M3.1-Flash AgenticX (c16, TP4) decode steps rarely exceed 32 rows, so the split-K path is not reached often enough to show up end to end.