[XPU][MiniMax-M3] Add XPU support for MSA sparse-attention path + tune gemma_rmsnorm for BMG

vllm/models/minimax_m3/__init__.py

@@ -2,23 +2,30 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """MiniMax M3 model — hardware-isolated entry point.
 
-The implementation lives under ``nvidia/`` and ``amd/``; this module picks the
-right one for the current platform and re-exports the public classes used by
-the model registry. (Mirrors ``vllm.models.deepseek_v4``.)
+The implementation lives under ``nvidia/``, ``amd/`` and ``xpu/``; this module
+picks the right one for the current platform and re-exports the public classes
+used by the model registry. (Mirrors ``vllm.models.deepseek_v4``.) The ``xpu``
+package reuses the NVIDIA classes and overrides only the Gemma RMSNorm.
 """
 
 from typing import TYPE_CHECKING
 
 from vllm.platforms import current_platform
 
-# The NVIDIA branch is the static default that type-checkers see; the ROCm
-# branch overrides it at runtime (kept type-compatible via type: ignore).
-if TYPE_CHECKING or not current_platform.is_rocm():
+# The NVIDIA branch is the static default that type-checkers see; the ROCm and
+# XPU branches override it at runtime (kept type-compatible via type: ignore).
+if TYPE_CHECKING or not (current_platform.is_rocm() or current_platform.is_xpu()):
     from .nvidia.model import (
         MiniMaxM3SparseForCausalLM,
         MiniMaxM3SparseForConditionalGeneration,
     )
     from .nvidia.mtp import MiniMaxM3MTP
+elif current_platform.is_xpu():
+    from .xpu.model import (  # type: ignore[assignment]
+        MiniMaxM3SparseForCausalLM,
+        MiniMaxM3SparseForConditionalGeneration,
+    )
+    from .xpu.mtp import MiniMaxM3MTP  # type: ignore[assignment]
 else:
     from .amd.model import (  # type: ignore[assignment]
         MiniMaxM3SparseForCausalLM,

vllm/models/minimax_m3/common/ops/index_topk.py

@@ -317,6 +317,7 @@ def _decode_index_score_kernel(
     BLOCK_SIZE_Q: tl.constexpr,
     num_kv_chunks,
     USE_PDL: tl.constexpr,
+    USE_DOT: tl.constexpr,
 ):
     BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q
     pid_r = tl.program_id(0)
@@ -373,7 +374,14 @@ def _decode_index_score_kernel(
             + off_k[:, None] * stride_ik_pos
             + off_d * stride_ik_d,
         )  # [N,D]
-        kq = tl.dot(k, q)  # [N,HQ]
+        if USE_DOT:
+            kq = tl.dot(k, q)  # [N,HQ]
+        else:
+            # Intel Triton's tl.dot requires the output dim >= 16, but it equals
+            # num_idx_heads (4 for MiniMax-M3); reduce over the head dim instead.
+            kq = tl.sum(
+                k[:, :, None].to(tl.float32) * q[None, :, :].to(tl.float32), axis=1
+            )  # [N,HQ]
         kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf"))
         score = tl.max(kq, axis=0)  # [HQ]
         is_visible_block = blk < num_blocks_q
@@ -831,6 +839,7 @@ def minimax_m3_index_decode(
         BLOCK_SIZE_Q=BLOCK_SIZE_Q,
         num_kv_chunks=num_kv_chunks,
         USE_PDL=use_pdl,
+        USE_DOT=not current_platform.is_xpu(),
         **score_kwargs,
     )
 

vllm/models/minimax_m3/xpu/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/minimax_m3/xpu/model.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax M3 (text backbone) — Intel XPU implementation.
+
+This reuses the entire NVIDIA model implementation and overrides only
+``MiniMAXGemmaRMSNorm``: FlashInfer's Gemma RMSNorm kernels are CUDA-only, so
+the XPU norm uses the portable Triton kernels in ``xpu/ops/gemma_rmsnorm.py``
+(no dependency on vllm-xpu-kernels). Mirrors the ``amd`` per-platform override,
+but reuses the NVIDIA classes verbatim instead of copying them (only the norm
+differs).
+
+The NVIDIA ``model`` and ``mtp`` modules bind ``MiniMAXGemmaRMSNorm`` in their
+own namespaces, so the XPU norm is rebound in both before any layer is built.
+"""
+
+import torch
+from torch import nn
+
+from vllm.models.minimax_m3.xpu.ops import gemma_fused_add_rmsnorm, gemma_rmsnorm
+
+
+class MiniMAXGemmaRMSNorm(nn.Module):
+    """Gemma-style RMSNorm backed by the XPU Triton kernels.
+
+    Numerically equivalent to the NVIDIA FlashInfer ``gemma_rmsnorm`` /
+    ``gemma_fused_add_rmsnorm``: ``x * rsqrt(mean(x^2)+eps) * (1+weight)``.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            return gemma_rmsnorm(x, self.weight, self.variance_epsilon)
+        return gemma_fused_add_rmsnorm(x, residual, self.weight, self.variance_epsilon)
+
+
+def _install_xpu_rmsnorm() -> None:
+    from vllm.models.minimax_m3.nvidia import model as _nv_model
+    from vllm.models.minimax_m3.nvidia import mtp as _nv_mtp
+
+    _nv_model.MiniMAXGemmaRMSNorm = MiniMAXGemmaRMSNorm  # type: ignore[misc]
+    _nv_mtp.MiniMAXGemmaRMSNorm = MiniMAXGemmaRMSNorm  # type: ignore[misc]
+
+
+_install_xpu_rmsnorm()
+
+from vllm.models.minimax_m3.nvidia.model import (  # noqa: E402
+    MiniMaxM3SparseForCausalLM,
+    MiniMaxM3SparseForConditionalGeneration,
+)
+
+__all__ = [
+    "MiniMaxM3SparseForCausalLM",
+    "MiniMaxM3SparseForConditionalGeneration",
+]

vllm/models/minimax_m3/xpu/mtp.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax M3 MTP — Intel XPU entry point.
+
+Importing ``xpu.model`` installs the XPU ``MiniMAXGemmaRMSNorm`` into the NVIDIA
+namespaces, so the reused NVIDIA MTP builds with the XPU norm.
+"""
+
+import vllm.models.minimax_m3.xpu.model  # noqa: F401  (installs XPU RMSNorm)
+from vllm.models.minimax_m3.nvidia.mtp import MiniMaxM3MTP
+
+__all__ = ["MiniMaxM3MTP"]

vllm/models/minimax_m3/xpu/ops/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Intel XPU fused ops for MiniMax-M3."""
+
+from vllm.models.minimax_m3.xpu.ops.gemma_rmsnorm import (
+    gemma_fused_add_rmsnorm,
+    gemma_rmsnorm,
+)
+
+__all__ = [
+    "gemma_rmsnorm",
+    "gemma_fused_add_rmsnorm",
+]

vllm/models/minimax_m3/xpu/ops/gemma_rmsnorm.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused Gemma-style RMSNorm for Intel XPU via Triton.
+
+Gemma RMSNorm = ``normalize(x) * (1 + weight)``, computed in fp32. FlashInfer's
+``gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm`` CUDA kernels are unavailable on
+XPU, so this single-pass Triton implementation is used instead (no dependency on
+vllm-xpu-kernels). Mirrors the ROCm Triton path.
+
+Two entry points:
+  * ``gemma_rmsnorm(x, w, eps)``                 -> normalized tensor
+  * ``gemma_fused_add_rmsnorm(x, res, w, eps)``  -> (normalized, x + res)
+
+Both normalize over the last dim and broadcast ``weight`` (shape [N]) over it,
+so they serve both the full-hidden norms (input/post-attn/final) and the
+per-head q_norm/k_norm (N == head_dim). Inputs may be non-contiguous views
+(e.g. ``qkv.split`` slices); strides are passed through and outputs are written
+contiguous.
+"""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _gemma_rmsnorm_kernel(
+    x_ptr,
+    w_ptr,
+    out_ptr,
+    n_cols,
+    stride_row,
+    stride_col,
+    eps,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < n_cols
+    x = tl.load(x_ptr + row * stride_row + cols * stride_col, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    var = tl.sum(x * x, axis=0) / n_cols
+    rstd = 1.0 / tl.sqrt(var + eps)
+    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    out = x * rstd * (1.0 + w)
+    tl.store(
+        out_ptr + row * n_cols + cols,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+@triton.jit
+def _gemma_fused_add_rmsnorm_kernel(
+    x_ptr,
+    res_ptr,
+    w_ptr,
+    out_ptr,
+    res_out_ptr,
+    n_cols,
+    stride_xrow,
+    stride_xcol,
+    stride_rrow,
+    stride_rcol,
+    eps,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < n_cols
+    x = tl.load(
+        x_ptr + row * stride_xrow + cols * stride_xcol, mask=mask, other=0.0
+    ).to(tl.float32)
+    r = tl.load(
+        res_ptr + row * stride_rrow + cols * stride_rcol, mask=mask, other=0.0
+    ).to(tl.float32)
+    s = x + r
+    # residual_out is the pre-norm sum (consumed by the next layer's add).
+    tl.store(
+        res_out_ptr + row * n_cols + cols,
+        s.to(res_out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+    var = tl.sum(s * s, axis=0) / n_cols
+    rstd = 1.0 / tl.sqrt(var + eps)
+    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    out = s * rstd * (1.0 + w)
+    tl.store(
+        out_ptr + row * n_cols + cols,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def _num_warps(block_n: int) -> int:
+    # Tuned on BMG/Xe2: for the M3 hidden sizes (block_n up to 4096) 16 warps
+    # over-subscribes a single row and regresses vs 8; only the very widest
+    # rows benefit from 16.
+    if block_n >= 8192:
+        return 16
+    if block_n >= 2048:
+        return 8
+    if block_n >= 512:
+        return 4
+    return 2
+
+
+def gemma_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    """Return ``x * rsqrt(mean(x^2)+eps) * (1+weight)`` (normalized over -1)."""
+    orig_shape = x.shape
+    n = orig_shape[-1]
+    x2 = x.reshape(-1, n)
+    m = x2.shape[0]
+    out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    block_n = triton.next_power_of_2(n)
+    _gemma_rmsnorm_kernel[(m,)](
+        x2,
+        weight,
+        out,
+        n,
+        x2.stride(0),
+        x2.stride(1),
+        eps,
+        BLOCK_N=block_n,
+        num_warps=_num_warps(block_n),
+    )
+    return out.reshape(orig_shape)
+
+
+def gemma_fused_add_rmsnorm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused add + Gemma RMSNorm.
+
+    Returns ``(gemma_rmsnorm(x + residual), x + residual)``, matching the
+    FlashInfer ``gemma_fused_add_rmsnorm`` contract used by the NVIDIA model.
+    """
+    orig_shape = x.shape
+    n = orig_shape[-1]
+    x2 = x.reshape(-1, n)
+    r2 = residual.reshape(-1, n)
+    m = x2.shape[0]
+    out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    res_out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    block_n = triton.next_power_of_2(n)
+    _gemma_fused_add_rmsnorm_kernel[(m,)](
+        x2,
+        r2,
+        weight,
+        out,
+        res_out,
+        n,
+        x2.stride(0),
+        x2.stride(1),
+        r2.stride(0),
+        r2.stride(1),
+        eps,
+        BLOCK_N=block_n,
+        num_warps=_num_warps(block_n),
+    )
+    return out.reshape(orig_shape), res_out.reshape(orig_shape)