[XPU][MiniMax-M3] Add XPU support for MSA sparse-attention path + tune gemma_rmsnorm for BMG#45881
[XPU][MiniMax-M3] Add XPU support for MSA sparse-attention path + tune gemma_rmsnorm for BMG#45881yangulei wants to merge 2 commits into
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
This pull request has merge conflicts that must be resolved before it can be |
|
@mergify Rebase done. |
❌ Sorry but I didn't understand the command. Please consult the commands documentation 📚. |
MiniMax-M3 (vllm-project#45381) runs its MSA path on the platform-neutral Triton kernels (indexer + block-sparse GQA attend); the NVIDIA fmha_sm100 MSA path is Blackwell-only. select_main_impl_cls / select_indexer_impl_cls already fall back to the Triton impls on non-CUDA platforms, so the common backend works on XPU. XPU-specific changes (the vendored nvidia/ and amd/ models are untouched): - New vllm/models/minimax_m3/xpu/ package mirroring the nvidia/amd split. It reuses the entire NVIDIA model and overrides only MiniMAXGemmaRMSNorm, which on XPU uses portable Triton Gemma RMSNorm kernels (xpu/ops/gemma_rmsnorm.py, mirroring the ROCm path) instead of the CUDA-only FlashInfer kernels -- no dependency on vllm-xpu-kernels. __init__.py dispatches XPU -> xpu/. - common/ops/index_topk.py _decode_index_score_kernel: the per-block score dot has output dim == num_index_heads (4 for M3); Intel Triton's tl.dot requires >= 16. A USE_DOT constexpr keeps tl.dot on CUDA while XPU uses an equivalent fp32 reduction over the head dim. The fused qknorm/rope/kv-insert kernel is provided by vllm-xpu-kernels under torch.ops._C, so vllm/_custom_ops.py needs no change. is_arch_support_pdl() already defaults to False on XPU (inherited from the base Platform). AI assistance (GitHub Copilot) was used for this change. Co-authored-by: GitHub Copilot Signed-off-by: Youlei Yang <youlei.yang@intel.com>
Mirror of the vllm-xpu change: the _num_warps heuristic picked 16 warps for
the M3 hidden sizes (BLOCK_N up to 4096), which over-subscribes a single row
on BMG/Xe2 and regresses the bandwidth-bound rmsnorm/fused-add kernels. Use a
gentler ramp (16 only for BLOCK_N >= 8192, else 8/4/2). Kernel math unchanged.
Parity vs the fp32 reference verified across n in {128, 2048, 3072, 4096, 6144}.
AI assistance (GitHub Copilot CLI) was used for this change.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Youlei Yang <youlei.yang@intel.com>
|
@mergify Refresh |
❌ Sorry but I didn't understand the command. Please consult the commands documentation 📚. |
|
@Mergifyio Refresh |
❌ Sorry but I didn't understand the command. Please consult the commands documentation 📚. |
|
@Mergifyio refresh |
✅ Pull request refreshed |
|
This pull request has merge conflicts that must be resolved before it can be |
1 participant
Summary
Adds XPU (Intel GPU) support for MiniMax-M3 (#45381), plus a BMG-specific tuning of the Gemma RMSNorm Triton kernel. The vendored
nvidia/andamd/models are untouched; all changes are additive under a newxpu/package.Commits
Add XPU support for the MSA sparse-attention path (
40aa215)vllm/models/minimax_m3/xpu/package mirroring the nvidia/amd split. Reuses the entire NVIDIA model and overrides onlyMiniMAXGemmaRMSNorm, which on XPU uses portable Triton Gemma RMSNorm kernels (xpu/ops/gemma_rmsnorm.py, mirroring the ROCm path) instead of CUDA-only FlashInfer kernels — no dependency on vllm-xpu-kernels for this op.__init__.pydispatches XPU →xpu/.common/ops/index_topk.py_decode_index_score_kernel: the per-block score dot has output dim ==num_index_heads(4 for M3); Intel Triton'stl.dotrequires ≥ 16. AUSE_DOTconstexpr keepstl.doton CUDA while XPU uses an equivalent fp32 reduction over the head dim.select_main_impl_cls/select_indexer_impl_clsalready fall back to the Triton impls on non-CUDA platforms.Tune gemma_rmsnorm Triton num_warps for BMG (
1e224de)_num_warpsheuristic picked 16 warps for the M3 hidden sizes (BLOCK_N up to 4096), over-subscribing a single row on BMG/Xe2 and regressing the bandwidth-bound rmsnorm/fused-add kernels. Uses a gentler ramp (16 only for BLOCK_N ≥ 8192, else 8/4/2). Kernel math unchanged.Not a duplicate
This supersedes draft PR #45797 (now closed), which contained only the gemma_rmsnorm tuning. This PR combines that tuning with its parent MSA sparse-attention XPU support. No other open PR addresses the MiniMax-M3 XPU path.
Testing
n ∈ {128, 2048, 3072, 4096, 6144}.AI assistance
AI assistance (GitHub Copilot CLI) was used for this change. A human submitter has reviewed and defends the change end-to-end.