[XPU][MiniMax-M3] Add and optimize fp32_router_gemm SYCL kernel for BMG

[yangulei]

Summary

Adds the fp32_router_gemm SYCL kernel for the MiniMax-M3 MoE router on Intel BMG (Arc Pro B60), and then optimizes it. Two commits, reviewable in isolation:

1. the SYCL port / kernel add, and
2. a profile-guided (unitrace) optimization (~20.3× faster: 491,695 → 24,208 ns/call, pinned @ 2.4 GHz, min-of-7 device timing).

Commit 1 — Add fp32_router_gemm SYCL kernel

SYCL port of csrc/libtorch_stable/fp32_router_gemm.cu from vllm-project/vllm#45381 (m3_release). Computes out[m,n] = sum_k mat_a[m,k] * mat_b[n,k] in fp32, where mat_a (activation) is bf16 or fp32 and mat_b (weight) is fp32. Targets the small MoE-router shape (num_tokens ≤ 32, e.g. 256 experts × 3072 hidden): one sub-group per output column, each lane strides over K accumulating one partial sum per token, then a sub-group reduction writes the column.

Registered as torch.ops._xpu_C.fp32_router_gemm.

Commit 2 — Optimize for BMG (~20×)

The original launched only num_experts hardware threads (one SIMD16 sub-group per expert) → 20% occupancy, 93% XVE stall, memory idle. The optimization:

• K-slicing: split the K reduction across ROUTER_GEMM_NSG sub-groups per expert with an SLM partial-sum reduction → occupancy 20% → 79%.
• Vectorized loads: aligned_vec<T, VEC> K loads (guarded by hidden % VEC == 0, scalar tail otherwise).
• Dual-accumulator ILP: 2-way unrolled inner loop with independent partial sums, to break the acc[m] dependency chain (ALUWR stall) and overlap two loads (SBID latency).
• Tuned NSG=4, VEC=16.

Testing

• tests/test_fp32_router_gemm.py — 9 passed (bf16 + fp32 activation, vs fp32 torch.matmul). Correctness rel err 1.2e-7.
• The kernel translation unit compiles cleanly against main headers (csrc/utils.h); the op registration is a self-contained addition to csrc/xpu/ops.h / csrc/xpu/torch_bindings.cpp / tests/register_ops.py and one CMakeLists.txt source line. Builds standalone on main.

Not a duplicate

Checked gh pr list --repo vllm-project/vllm-xpu-kernels --state open — no open PR touches this kernel.

Notes

• AI assistance (GitHub Copilot CLI) was used. Requires human review of every line before un-drafting.

[jikunshang]

can you add some benchmark to compare torch.matmul and this new kernel?

[yangulei]

Thanks @jikunshang. Added benchmark/benchmark_fp32_router_gemm.py (commit b58de2e) comparing this kernel against torch.matmul(mat_a.float(), mat_b.t()) (the fp32 reference it replaces). Numbers on Arc Pro B60 (BMG), freq pinned @ 2.4 GHz, MiniMax-M3 router shape N=256, K=3072.

Wall-clock per call (triton.testing.do_bench, includes dispatch + the .float() cast for torch)

act	M	torch.matmul (us)	this kernel (us)	kernel/torch
bf16	1	26.8	33.6	0.80x
bf16	8	31.9	41.3	0.77x
bf16	16	31.4	50.3	0.62x
bf16	32	31.4	69.4	0.45x
fp32	1	24.1	38.1	0.63x
fp32	16	29.1	58.5	0.50x
fp32	32	30.5	85.8	0.36x

(max_abs vs the fp32 ref stays ~1e-4 for bf16, exact for fp32.)

Device-time per call (unitrace -d, kernel-only, the fairer apples-to-apples), M=32

path	bf16	fp32
torch.matmul	gemm 12.97us + bf16->fp32 cast 3.11us = 16.1us	gemm 12.8us
this kernel	24.2us	45.6us

Honest takeaway

torch.matmul (oneDNN GEMM) is faster than this kernel on XPU for the router shapes — ~1.5x for bf16, ~3.6x for fp32 in device time — even though this PR already makes the custom kernel ~20x faster than the original port (491us -> 24us). At M=32 the kernel hits ~17% FP32 MFU vs oneDNN's ~31%.

So this optimization is a big win relative to the existing custom kernel, but it does not beat torch.matmul here. Options I'd like your guidance on:

1. Keep the kernel (parity with the upstream CUDA path vllm-project/vllm#45381) with this optimization, or
2. Switch the XPU router to torch.matmul and drop the custom kernel.

Happy to do either. The original CUDA kernel presumably beat cuBLAS for this shape on NVIDIA; on XPU oneDNN's small-GEMM path is clearly stronger.

[yangulei]

Close as the the GateLinear already dispatched to oneDNN fp32 GEMM by default which is optimal.