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