diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index ca9059f89..da590a01e 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -19,6 +19,8 @@ import logging from typing import Iterable, Optional, Set, Tuple, Union import torch +import triton +import triton.language as tl from torch import nn from transformers import PretrainedConfig @@ -73,6 +75,164 @@ from sglang.srt.utils import ( logger = logging.getLogger(__name__) +@triton.jit +def rmsnorm_sumsq_kernel_serial( + x1_ptr, # T* [B, D] + x2_ptr, # T* [B, D] + stride_x1, # int + stride_x2, # int + sum_sq_ptr, # float* [B] + B, # int + D1, # int + D2, # int + BLOCK_SIZE1: tl.constexpr, + BLOCK_SIZE2: tl.constexpr, +): + row_id = tl.program_id(0) + x1_row = x1_ptr + row_id * stride_x1 + x2_row = x2_ptr + row_id * stride_x2 + + offsets1 = tl.arange(0, BLOCK_SIZE1) + mask1 = offsets1 < D1 + offsets2 = tl.arange(0, BLOCK_SIZE2) + mask2 = offsets2 < D2 + + x1 = tl.load(x1_row + offsets1, mask=mask1, other=0.0) + x2 = tl.load(x2_row + offsets2, mask=mask2, other=0.0) + + x1_f32 = x1.to(tl.float32) + sum_sq1 = tl.sum(x1_f32 * x1_f32, axis=0) + + x2_f32 = x2.to(tl.float32) + sum_sq2 = tl.sum(x2_f32 * x2_f32, axis=0) + + tl.store(sum_sq_ptr + row_id, sum_sq1) + tl.store(sum_sq_ptr + row_id + B, sum_sq2) + + +@triton.jit +def rmsnorm_apply_kernel_serial( + x1_ptr, # T* [B, D] + x2_ptr, # T* [B, D] + w1_ptr, # T* [D] + w2_ptr, # T* [D] + sum_sq_ptr, # float* [B] + out1_ptr, # T* [B, D] + out2_ptr, # T* [B, D] + B, # int + D1, # int + D2, # int + stride_x1, # int + stride_x2, # int + tp_world, # int + eps, # float + BLOCK_SIZE1: tl.constexpr, + BLOCK_SIZE2: tl.constexpr, +): + row_id = tl.program_id(0) + x1_row = x1_ptr + row_id * stride_x1 + x2_row = x2_ptr + row_id * stride_x2 + out1_row = out1_ptr + row_id * stride_x1 + out2_row = out2_ptr + row_id * stride_x2 + + sum_sq1 = tl.load(sum_sq_ptr + row_id) + sum_sq2 = tl.load(sum_sq_ptr + row_id + B) + inv_rms1 = tl.rsqrt(sum_sq1 / D1 / tp_world + eps) + inv_rms2 = tl.rsqrt(sum_sq2 / D2 / tp_world + eps) + + offsets1 = tl.arange(0, BLOCK_SIZE1) + offsets2 = tl.arange(0, BLOCK_SIZE2) + + mask1 = offsets1 < D1 + mask2 = offsets2 < D2 + + x1 = tl.load(x1_row + offsets1, mask=mask1, other=0.0) + w1 = tl.load(w1_ptr + offsets1, mask=mask1, other=1.0) + x2 = tl.load(x2_row + offsets2, mask=mask2, other=0.0) + w2 = tl.load(w2_ptr + offsets2, mask=mask2, other=1.0) + + out1 = (x1.to(tl.float32) * inv_rms1 * w1.to(tl.float32)).to(x1.dtype) + out2 = (x2.to(tl.float32) * inv_rms2 * w2.to(tl.float32)).to(x2.dtype) + tl.store(out1_row + offsets1, out1, mask=mask1) + tl.store(out2_row + offsets2, out2, mask=mask2) + + +def rms_sumsq_serial(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: + assert x1.is_cuda and x2.is_cuda + B, D1 = x1.shape + B2, D2 = x2.shape + assert B == B2 + + stride_x1 = x1.stride(0) + stride_x2 = x2.stride(0) + + sum_sq = torch.empty(B + B2, device=x1.device, dtype=torch.float32) + + BLOCK_SIZE1 = triton.next_power_of_2(D1) + BLOCK_SIZE2 = triton.next_power_of_2(D2) + + grid = (B,) + + rmsnorm_sumsq_kernel_serial[grid]( + x1, + x2, + stride_x1, + stride_x2, + sum_sq, + B, + D1, + D2, + BLOCK_SIZE1, + BLOCK_SIZE2, + ) + return sum_sq + + +def rms_apply_serial( + x1: torch.Tensor, + x2: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + sum_sq: torch.Tensor, + tp_world: int = 1, + eps: float = 1e-5, +) -> torch.Tensor: + assert x1.is_cuda and x2.is_cuda and w1.is_cuda and w2.is_cuda and sum_sq.is_cuda + B, D1 = x1.shape + B2, D2 = x2.shape + assert B == B2 + + stride_x1 = x1.stride(0) + stride_x2 = x2.stride(0) + out1 = torch.empty(B, D1, device=x1.device, dtype=x1.dtype) + out2 = torch.empty(B, D2, device=x2.device, dtype=x2.dtype) + + BLOCK_SIZE1 = triton.next_power_of_2(D1) + BLOCK_SIZE2 = triton.next_power_of_2(D2) + + grid = (B,) + + rmsnorm_apply_kernel_serial[grid]( + x1, + x2, + w1, + w2, + sum_sq, + out1, + out2, + B, + D1, + D2, + stride_x1, + stride_x2, + tp_world, + eps, + BLOCK_SIZE1, + BLOCK_SIZE2, + ) + return out1, out2 + + class MiniMaxM2RMSNormTP(nn.Module): """RMSNorm with Tensor Parallel support for QK normalization.""" @@ -124,6 +284,30 @@ class MiniMaxM2RMSNormTP(nn.Module): return x + @staticmethod + @torch.compile(dynamic=True, backend=get_compiler_backend()) + def forward_qk( + q_norm: "MiniMaxM2RMSNormTP", + k_norm: "MiniMaxM2RMSNormTP", + q: torch.Tensor, + k: torch.Tensor, + ) -> torch.Tensor: + sum_sq = rms_sumsq_serial(q, k) + if q_norm.tp_world > 1: + sum_sq = tensor_model_parallel_all_reduce(sum_sq) + + q, k = rms_apply_serial( + q, + k, + q_norm.weight, + k_norm.weight, + sum_sq, + q_norm.tp_world, + q_norm.variance_epsilon, + ) + + return q, k + class MiniMaxM2MoE(nn.Module): """MiniMax MoE implementation using DeepEP for Expert Parallel support.""" @@ -437,8 +621,11 @@ class MiniMaxM2Attention(nn.Module): qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) if self.use_qk_norm: - q = self.q_norm(q.contiguous()) - k = self.k_norm(k.contiguous()) + # q = self.q_norm(q.contiguous()) + # k = self.k_norm(k.contiguous()) + q, k = MiniMaxM2RMSNormTP.forward_qk( + self.q_norm, self.k_norm, q.contiguous(), k.contiguous() + ) else: q, k = q.contiguous(), k.contiguous() q, k = self.rotary_emb(positions, q, k)