MiniMax-M3: Triton split-K router GEMV with in-kernel fixup

[zcnrex]

MiniMax-M3: Triton split-K router GEMV with in-kernel split-K fixup

The MiniMax-M3 MoE router gate is a skinny GEMV — [M, K] bf16 x [N, K] bf16 -> [M, N] fp32, with N ~ 128 experts and M <= 64 at decode — and the vendor BLAS solutions picked for that shape run far off the memory roofline on gfx950 (~0.1 TB/s for a ~1.6MB weight read at M<=8, N=128). This PR adds python/sglang/kernels/ops/gemm/router_gemv.py, a Triton split-K kernel that reads the gate weight once at near-bandwidth, and routes MiniMaxM3MoE._compute_router_logits through it when router_gemv_supported() says the shape qualifies. Configs are per-M-bucket (BLOCK_K 512 throughout, BLOCK_N and SPLIT_K chosen by batch size, with a dedicated BLOCK_M == 1 GEMV path that uses tl.sum instead of tl.dot) covering M <= 8 / 16 / 32 / 64. The split-K reduction is done inside the same kernel: the last CTA to finish a (n-block, row-tile) pair reduces all SPLIT_K partials and resets its counter slot to zero on the way out, so the scratch buffer stays reusable across launches without either a zero-init launch or a separate reduce launch — each extra launch costs ~2us on gfx950, which is a meaningful fraction of the kernel itself at these sizes. Anything the kernel does not cover — wrong dtype, non-contiguous last dim, N not divisible by BLOCK_N, K not divisible by BLOCK_K, or M > 64 — falls through router_gemv_supported() to the existing torch.mm path unchanged.

Performance

Measured on 8x MI350X (gfx950), MiniMax-M3-MXFP8, TP8, 80k input / 600 output.
Two bench reps per boot; the warm (second) rep is reported and the two agreed
within 0.1%. Noise floor is +/-0.4%. Tables are the harness output verbatim,
trimmed to the first five columns. Baseline is upstream/main.
Baseline (upstream/main)

Input lens: [80000]. Output lens: [600]. Cache hit rate: 90.0%.
|   batch size |   input len |   latency (s) |   input throughput (tok/s) |   output throughput (tok/s) |
|--------------|-------------|---------------|----------------------------|-----------------------------|
|            1 |       80000 |         11.26 |                     187984 |                       55.35 |
|            2 |       80000 |         11.88 |                     191222 |                      108.68 |
|            4 |       80000 |         13.39 |                     192965 |                      204.62 |
|            8 |       80000 |         16.27 |                     193687 |                      370.13 |
|           16 |       80000 |         22.2  |                     193188 |                      616.28 |
|           32 |       80000 |         33.78 |                     193245 |                      934.92 |

With this PR

Input lens: [80000]. Output lens: [600]. Cache hit rate: 90.0%.
|   batch size |   input len |   latency (s) |   input throughput (tok/s) |   output throughput (tok/s) |
|--------------|-------------|---------------|----------------------------|-----------------------------|
|            1 |       80000 |         10.92 |                     189401 |                       57.16 |
|            2 |       80000 |         11.68 |                     191857 |                      110.63 |
|            4 |       80000 |         13.14 |                     194266 |                      208.74 |
|            8 |       80000 |         16.1  |                     195607 |                      374.08 |
|           16 |       80000 |         21.94 |                     194140 |                      625.66 |
|           32 |       80000 |         33.42 |                     194545 |                      947.49 |

Output throughput +1.07%..+3.27%; input throughput +0.33%..+0.99%.

Accuracy

gsm8k, 512 examples, max_tokens 2048, temperature 0, seed 0, against the
same server boot as the perf run.
Baseline

== gsm8k ==
512 examples (single-shot)  |  109.0s  |  1338 tok/s  |  146K tokens

* score           =  97.07%
  stop_rate       =  98.44%
  truncated_rate  =  1.56%  [warn: hitting max_tokens]
  error_rate      =  0.00%

With this PR

== gsm8k ==
512 examples (single-shot)  |  109.5s  |  1334 tok/s  |  146K tokens

* score           =  97.27%
  stop_rate       =  98.05%
  truncated_rate  =  1.95%  [warn: hitting max_tokens]
  error_rate      =  0.00%

Correctness

Not bit-identical. Split-K reassociates the fp32 accumulation in the router GEMV, so the router logits differ in the last bits from the torch.mm path. That is why this carries an accuracy gate rather than a bit-identity claim.

gsm8k is 97.27% against a 97.07% baseline (+0.20 pp), truncated_rate 1.95% vs 1.56%, error_rate 0.00% on both. On a fixed 512-question greedy run +0.20 pp is one answer flipping -- read this as "no measurable accuracy cost", not as an improvement.

Scope is wider than ROCm. SGLANG_OPT_USE_TRITON_ROUTER_GEMV no longer exists; the path is selected by router_gemv_supported(), which tests x.device.type == "cuda" -- true for ROCm torch builds as well as NVIDIA. SGLANG_OPT_USE_BF16_ROUTER_GEMM already defaults true on main. So this is on by default on NVIDIA too, not an opt-in ROCm path, and reviewers should judge it on that basis. If you would rather it ship behind a flag, that is a code change and I am happy to make it.

The kernel fires roughly 57 launches per decode step, and its configs are per-M-bucket, so the win is batch-size dependent by construction -- hence the full bs 1-32 sweep rather than a single point.

Reproduction

Server (identical on both sides):

SGLANG_USE_AITER=1 \
AITER_CONFIG_FMOE=/path/to/minimax_m3_mxfp8_tuned_fmoe.csv \
SGLANG_FORCE_MXFP8_BLOCK_CONVERT_DENSE=1 \
python3 -m sglang.launch_server \
  --model-path /path/to/MiniMax-M3-MXFP8 \
  --trust-remote-code --tp 8 --quantization mxfp8 \
  --moe-runner-backend aiter --fp8-gemm-backend auto \
  --enable-aiter-allreduce-fusion --attention-backend triton \
  --dtype bfloat16 --chunked-prefill-size 8192 --mem-fraction-static 0.80 \
  --host 0.0.0.0 --port 30000

Benchmark:

python3 -m sglang.benchmark.one_batch_server \
  --model None --base-url http://127.0.0.1:30000 \
  --output-len 600 --skip-warmup \
  --batch-size 1 2 4 8 16 32 \
  --input-len 80000 --show-report --cache-hit-rate 0.9

Accuracy gate:

sgl-eval run gsm8k --base-url http://127.0.0.1:30000/v1 \
  --num-examples 512 --num-threads 64 \
  --max-tokens 2048 --temperature 0 --seed 0

Two bench reps per boot, warm rep reported; noise floor +/-0.4%, established by a separate PR measuring flat under a config where its code cannot fire.

Note on the launch command

--moe-runner-backend aiter is not honoured on current main for mxfp8. Arg resolution logs:

mxfp8 quantization supports only cutlass, deep_gemm, flashinfer_trtllm,
flashinfer_trtllm_routed, triton backends. Overriding 'aiter'.

and falls back to triton. Both sides of this A/B therefore ran the Triton MoE runner, so the comparison is apples-to-apples and the deltas stand — but AITER_CONFIG_FMOE is inert under this configuration and the aiter MoE path is not what was measured. The flag is kept in the command above only because it matches the invocation actually used; anyone reproducing will get triton and should see the same numbers.

---

CI States

Latest PR Test (Base): 🚫 Run #34190762630
Latest PR Test (Extra): ❌ Run #34190762501
Latest PR Test (AMD ROCm 7.2): ❌ Run #34190762684

[HaiShaw]

LGTM