[kevin-mii]

Part of the MiniMax-M3.1-Flash AMD series that replaces #45936 (closed with its base branch minimax-m3.1-flash): kernel PRs #46690 (attention), #46691 (MoE / GEMM) and #46692 (quick all-reduce), each on main; then #46693 with the M3.1 call sites and #46694 (indexer CP). It follows the DeepSeek-V4.1 AMD layout (#41018-#41021, then #41308).

Summary

• MoE sorting + quant. moe/moe_sort_quant.py (one launch for decode batches) and moe/moe_sort_quant_large.py (four launches for prefill) sort and MXFP4-quantize the stage-1 activations bit-for-bit like aiter's multi-phase path; moe_sorting_small.py hands them to aiter as runtime patches.
• Shared-expert fold. moe/aiter_moe_addend.py folds the shared-expert output into aiter's stage-2 reduction as a runtime patch, applied only when aiter's signature matches.
• Router. gemm/router_gemv.py raises the gfx950 split-K GEMV to 512 rows with a reduce launch past 16 rows, exposes its split-K partials, and fuses them with the shared expert's fp8 SwiGLU in one launch; moe_fused_gate.py accepts those partials and sums them in order. moe/smallm_router_gfx950/m31_router.* is a one-launch decode router for 1-16 tokens, bitwise equal to router GEMV, gate and aiter's sort.
• GEMM. gemm/aiter_asm_blockscale.py launches aiter's asm block-fp8 split-K GEMM into an output its producer already zeroed (aiter always memsets it first; [HIP] [Feature] a8w8 blockscale asm: out_zeroed skips the split-K output memset ROCm/aiter#6786 adds out_zeroed), and jit/utils/hipcc.py launches the 2-D grid it needs. smallm_fp8_gfx950 gains the M3.1 shared expert's 1536 x 6144 gate_up and exposes its kernel body, which m31_router runs in its fused launch.
• fp8 producers. quantization/aiter_fp8_quant_triton.py holds Triton epilogues that write exactly aiter's per-token and 1x128 fp8 bytes and scales; layernorm/minimax_m3_rmsnorm.py adds the Gemma RMSNorm (+ residual) producers that use them.
• Merged with main's Beam FP32-gate path ([Feat] kernels: Beam model support #45254): FP32 gates keep their CUDA configs and one-row tile; BF16 gates stay gfx950-only.
• The M3.1 call sites (fused routers, fp8 producers, shared-expert fold, pre-zeroed GEMM launch) are in the integration PR.

Serving gain of each kernel

Measured in the integration PR at TP4 on MI350X (NVFP4 KV, DSpark acceptance simulated at 3 tokens per step, 32K-token contexts), each change turned off alone; the share decode TPOT or prefill latency grows without it, mean of two GPU sets:

• Shared-expert fold into aiter's stage-2 reduction: 32K-token prefill +15.9%, decode +1.9-5.0% at 8-64 requests. Not bit-identical to adding the shared output afterwards (its test bounds the relative error at 1e-2).
• Gemma RMSNorm producers writing the next GEMM's fp8: decode +8.5% at 1 request, +1.5-2.9% at 8-64.
• SwiGLU writing the down projection's fp8: decode +4.4% at 1 request, prefill +5.3%.
• Fused decode routers (m31_router, router GEMV partials): decode +4.2% at 1 request, +1.5% at 8. Bitwise equal to the unfused chain.
• The MoE sort + quant kernels were not toggled on their own; they are bitwise equal to aiter's multi-phase path.
• The pre-zeroed asm split-K launch is bit-exact and saves one memset launch per prefill block-fp8 GEMM; end to end it measured within run-to-run noise.

Not in this PR

• moe_sorting_small.sorting_impl_wrapper is ~160 lines; its two one-launch branches could share a helper.
• m31_router keeps module-level kernel caches and global _disabled, like its smallm_moe_gfx950 / smallm_fp8_gfx950 siblings.

Changes to existing kernels

router_gemv, moe_fused_gate, moe_sorting_small and smallm_fp8_gfx950 keep their signatures. On CUDA, router_gemv's FP32-gate path runs main's kernel body and tiles, with a split-K scratch now sized for its own tiles (64 KiB at N = 256, against main's 512 KiB), and moe_fused_gate's 2-D path is unchanged; split-K partials never take its SM120 or radix fast paths.

On gfx950, existing callers change when this merges, with no new call site:

• router_gemv's bf16 tiles change for its callers (MiniMax-M3, Beam, DeepSeek-V4.1 VL routing), which changes the fp32 summation order; the fp64 accuracy and determinism test covers them.
• main's aiter MoE runner applies moe_sorting_small's patch, which now sorts 128-expert top-4 routing at hidden 6144 in one launch with the stage1 MXFP4 activation quant: moe_sort_quant up to 1024 tokens, moe_sort_quant_large above (aiter's sort again past 349,525 tokens, the large sort's int32 offset limit), both bitwise equal to aiter's multi-phase sort. It no longer emits an mxfp8 quant for separated-SwiGLU MXFP4 weights, which never read one, and it stays off, with an info log, when aiter's sort or quant signatures differ from the ones it wraps.
• smallm_fp8_gfx950 gains (1536, 6144), which main's per-token fp8 linear also dispatches up to 16 rows; it stays within one bf16 ulp of aiter's gemm_a8w8_bpreshuffle there.

Verification

• MI350X (same stack), this branch alone on main: test_minimax_m3_rmsnorm 32 (the kernel cases), test_aiter_asm_blockscale 1 (the launch case), test_aiter_moe_addend 3, test_m31_router_gfx950 6, test_moe_sort_quant_large 3, test_moe_sorting_small 7 (192 subtests), test_router_gemv_gfx950 4, test_aiter_fp8_quant_triton 7, main's test_smallm_fp8_proj_gfx950 7 (with the M3.1 shape). Main's FP32-gate path run on gfx950 through the merged kernel: max error 3.9e-7 to 6.2e-7 against fp64 at 1-8 rows (main's bound 2e-6), deterministic. pre-commit clean.
• Not executed: main's CUDA test_router_gemv.py (needs CUDA), end-to-end serving.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ⏳ Run #41356557775
Latest PR Test (Extra): ❌ Run #41356557742
Latest PR Test (AMD ROCm 10): ⏳ Run #41356557741

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[kevin-mii]

/tag-and-rerun-ci