From 486c7de39f5c6ec8d651c5c0c662e940fccb63e0 Mon Sep 17 00:00:00 2001
From: Roger Young <42564206+rogeryoungh@users.noreply.github.com>
Date: Sun, 1 Feb 2026 15:39:37 +0800
Subject: [PATCH] Optimizing all_reduce in RMSNormTP in minimax_m2 (#16483)

---
 python/sglang/srt/models/minimax_m2.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py
index da590a01e..d52d9217d 100644
--- a/python/sglang/srt/models/minimax_m2.py
+++ b/python/sglang/srt/models/minimax_m2.py
@@ -166,7 +166,14 @@ def rms_sumsq_serial(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
     stride_x1 = x1.stride(0)
     stride_x2 = x2.stride(0)
 
-    sum_sq = torch.empty(B + B2, device=x1.device, dtype=torch.float32)
+    # We found that custom all-reduce `sglang::cross_device_reduce_1stage`
+    # is much faster than the nccl all-reduce in torch.
+    # However, `should_custom_ar` checks if the reduced buffer is 16-byte aligned.
+    # RMSNormTP reduces a [B, 2] fp32 tensor, so we pad the total element count to
+    # satisfy the alignment requirement.
+    B_padded = (B + B2 + 3) // 4 * 4
+
+    sum_sq = torch.empty(B_padded, device=x1.device, dtype=torch.float32)
 
     BLOCK_SIZE1 = triton.next_power_of_2(D1)
     BLOCK_SIZE2 = triton.next_power_of_2(D2)
@@ -285,7 +292,6 @@ class MiniMaxM2RMSNormTP(nn.Module):
         return x
 
     @staticmethod
-    @torch.compile(dynamic=True, backend=get_compiler_backend())
     def forward_qk(
         q_norm: "MiniMaxM2RMSNormTP",
         k_norm: "MiniMaxM2RMSNormTP",