Optimizing all_reduce in RMSNormTP in minimax_m2 (#16483)

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",