diff --git a/vllm/models/minimax_m3/__init__.py b/vllm/models/minimax_m3/__init__.py index f9ddb2a9d218..9d0c8fa16e26 100644 --- a/vllm/models/minimax_m3/__init__.py +++ b/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, diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index 208c2d690067..452336f18434 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/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, ) diff --git a/vllm/models/minimax_m3/xpu/__init__.py b/vllm/models/minimax_m3/xpu/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/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 diff --git a/vllm/models/minimax_m3/xpu/model.py b/vllm/models/minimax_m3/xpu/model.py new file mode 100644 index 000000000000..e3a81b2492c7 --- /dev/null +++ b/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", +] diff --git a/vllm/models/minimax_m3/xpu/mtp.py b/vllm/models/minimax_m3/xpu/mtp.py new file mode 100644 index 000000000000..8df2dccda353 --- /dev/null +++ b/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"] diff --git a/vllm/models/minimax_m3/xpu/ops/__init__.py b/vllm/models/minimax_m3/xpu/ops/__init__.py new file mode 100644 index 000000000000..388901877bb0 --- /dev/null +++ b/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", +] diff --git a/vllm/models/minimax_m3/xpu/ops/gemma_rmsnorm.py b/vllm/models/minimax_m3/xpu/ops/gemma_rmsnorm.py new file mode 100644 index 000000000000..43a13eb52d6d --- /dev/null +++ b/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)