[XPU][MiniMax-M3] Add and optimize fused qknorm/rope/kv-insert SYCL kernel for BMG

[yangulei]

Summary

Adds the fused_minimax_m3_qknorm_rope_kv_insert SYCL kernel (horizontally-fused MiniMax-M3 attention pre-processing: Gemma-RMSNorm + partial-NeoX RoPE + KV/index-cache insert) on Intel BMG (Arc Pro B60), and then optimizes it. Two commits:

1. the SYCL port / kernel add, and
2. a profile-guided (unitrace) optimization (~3.34× faster: 100,935 → 30,224 ns/call, pinned @ 2.4 GHz, min-of-7 device timing; MBU 18% → 56% of HBM).

Commit 1 — Add fused qknorm_rope_kv_insert SYCL kernel

SYCL port of the horizontally-fused M3 attention pre-processing kernel from vllm-project/vllm#45381 (csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu), registered in the _C extension as torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert so the unchanged upstream vllm._custom_ops wrapper resolves it on XPU (no vllm-side dispatch shim).

Unlike DeepSeek-V4's qnorm_rope_kv_insert (head_dim 512, GPT-J interleaved rope, 2 standard RMSNorms, MLA layout), M3 uses head_dim 128, partial-NeoX rope (rotary_dim leading dims rotated, rest pass-through), four Gemma-style RMSNorms (x * rsqrt(mean(x^2)+eps) * (1+weight)), and the packed [q|k|v|index_q|index_k] layout. One work-item handles one (token, head-slot) and loops over the 128-dim head:

• q → Gemma-norm + rope → q_out (sparse) / in place (dense)
• k → Gemma-norm + rope → in place + scatter kv_cache[:,0]
• v → raw → scatter kv_cache[:,1]
• index_q → Gemma-norm + rope → index_q_out
• index_k → Gemma-norm + rope → in place + scatter index_cache

Commit 2 — Optimize for BMG (~3.3×)

The original used a per-work-item buf[128] (16 KB private-memory spill) and a partial-sub-group launch geometry (~num_tokens threads) → 34% occupancy, 72% XVE stall. The optimization:

• Streaming two-pass norm+RoPE: sum-of-squares pass, then normalize+RoPE written straight to dst/cache — removes buf[128], private mem 16 KB → 0.
• Cooperative sub-group per head: one M3_SG_SIZE-lane sub-group per (token, head-slot) on a flat 1-D nd_range; lanes split the 128-dim head, sum-of-squares via reduce_over_group, rotary pairs computed lane-locally. Occupancy 34% → 92%, stall 72% → 31%.

A register-cache-of-src variant and a large-GRF build were both tried and measured worse (occupancy-bound), so they are not included.

Testing

• tests/test_fused_minimax_m3_qknorm_rope_kv_insert.py — 20 passed (dense norm+rope parity; sparse-full incl. index branch + KV/index cache inserts) vs an fp32 reference.
• The kernel translation unit compiles cleanly against main headers (csrc/utils.h); the op registration is a self-contained addition to csrc/ops.h / csrc/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.

[Copilot]

Pull request overview

Adds a new SYCL kernel and op registration for MiniMax-M3’s fused attention pre-processing on Intel XPU (BMG), plus XPU-side correctness tests to validate dense and sparse (cache-insert) behavior.

Changes:

• Register torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert for XPU and expose it via the test ops registry.
• Add the SYCL implementation of the fused qk-norm + partial NeoX RoPE + KV/index-cache insert kernel (head_dim=128).
• Add dense + sparse correctness tests comparing against an fp32 reference.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
csrc/fused_minimax_m3_qknorm_rope_kv_insert.cpp	New SYCL kernel + C++ entrypoint for fused M3 qknorm/rope/cache insert
csrc/torch_bindings.cpp	Registers the new op schema and XPU implementation
csrc/ops.h	Adds the C++ declaration for the new op
tests/register_ops.py	Adds Python wrapper used by tests to call the new op
tests/test_fused_minimax_m3_qknorm_rope_kv_insert.py	New correctness tests for dense + sparse/cache-insert paths
CMakeLists.txt	Adds the new source file to the build

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.