Fusing RMSNormTP in minimax_m2 (#14416)

Co-authored-by: xuebi <xuebi@minimaxi.com>

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())
+            # q = self.q_norm(q.contiguous())
-            k = self.k_norm(k.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)