Optimize MiniMax-M2.7 on CPU (#31956)

2026-08-13 15:04:16 +08:00
parent 889c2f31aa
commit 3f6ef01322
8 changed files with 687 additions and 72 deletions
@@ -234,6 +234,87 @@ class TestFusedRMSNormGated:
        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)


+class TestFusedQKRMSNorm:
+
+    @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS)
+    @pytest.mark.parametrize(
+        "batch_size,q_size,k_size,v_size",
+        [(1, 256, 64, 64), (17, 512, 128, 128)],
+    )
+    def test_fused_qk_rmsnorm(
+        self, batch_size: int, q_size: int, k_size: int, v_size: int, dtype
+    ):
+        """Q and K split views must be normalized over their distinct full widths."""
+        qkv = torch.randn([batch_size, q_size + k_size + v_size], dtype=dtype)
+        q, k, _ = qkv.split([q_size, k_size, v_size], dim=-1)
+        q_weight = torch.randn(q_size, dtype=dtype)
+        k_weight = torch.randn(k_size, dtype=dtype)
+
+        q_out, k_out = torch.ops.sgl_kernel.fused_qk_rmsnorm_cpu(
+            q, k, q_weight, k_weight, eps
+        )
+        ref_q_out = TestNorm()._forward_native(q, q_weight, eps)
+        ref_k_out = TestNorm()._forward_native(k, k_weight, eps)
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(q_out, ref_q_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(k_out, ref_k_out, atol=atol, rtol=rtol)
+
+    @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS)
+    @pytest.mark.parametrize(
+        "batch_size,q_size,k_size,tp_world_size",
+        [(1, 256, 64, 2), (17, 512, 128, 4)],
+    )
+    def test_fused_qk_rmsnorm_tp(
+        self,
+        batch_size: int,
+        q_size: int,
+        k_size: int,
+        tp_world_size: int,
+        dtype,
+    ):
+        q = torch.randn([batch_size, q_size], dtype=dtype)
+        k = torch.randn([batch_size, k_size], dtype=dtype)
+        q_weight = torch.randn(q_size, dtype=dtype)
+        k_weight = torch.randn(k_size, dtype=dtype)
+
+        q_shards = q.chunk(tp_world_size, dim=-1)
+        k_shards = k.chunk(tp_world_size, dim=-1)
+        q_weight_shards = q_weight.chunk(tp_world_size)
+        k_weight_shards = k_weight.chunk(tp_world_size)
+        local_sum_sq = [
+            torch.ops.sgl_kernel.fused_qk_rmsnorm_sumsq_cpu(q_shard, k_shard)
+            for q_shard, k_shard in zip(q_shards, k_shards)
+        ]
+        global_sum_sq = torch.stack(local_sum_sq).sum(dim=0)
+
+        shard_outputs = [
+            torch.ops.sgl_kernel.fused_qk_rmsnorm_apply_from_stats_cpu(
+                q_shard,
+                k_shard,
+                q_weight_shard,
+                k_weight_shard,
+                global_sum_sq,
+                tp_world_size,
+                eps,
+            )
+            for q_shard, k_shard, q_weight_shard, k_weight_shard in zip(
+                q_shards, k_shards, q_weight_shards, k_weight_shards
+            )
+        ]
+        q_out = torch.cat([output[0] for output in shard_outputs], dim=-1)
+        k_out = torch.cat([output[1] for output in shard_outputs], dim=-1)
+
+        ref_q_out = TestNorm()._forward_native(q, q_weight, eps)
+        ref_k_out = TestNorm()._forward_native(k, k_weight, eps)
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(q_out, ref_q_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(k_out, ref_k_out, atol=atol, rtol=rtol)
+
+        assert global_sum_sq.shape == (batch_size, 2)
+        assert global_sum_sq.dtype == torch.float32
+
+
 class TestLayerNorm:

    def _forward_native(
@@ -1,3 +1,4 @@
+import itertools
 import unittest

 import torch
@@ -207,6 +208,89 @@ class TestTopK(CustomTestCase):
            self._run_single_test(123, 256, 4, renormalize, torch.bfloat16)
            self._run_single_test(123, 160, 6, renormalize, torch.bfloat16)

+    def test_topk_softmax_mixed_input_dtypes(self):
+        torch.manual_seed(0)
+        hidden_states = torch.randn((17, 16), dtype=torch.bfloat16)
+        gating_output = torch.randn((17, 128), dtype=torch.float32)
+        correction_bias = torch.randn(128, dtype=torch.float32)
+
+        topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
+            hidden_states=hidden_states,
+            gating_output=gating_output,
+            topk=8,
+            renormalize=True,
+            correction_bias=correction_bias,
+        )
+
+        scores = torch.softmax(gating_output, dim=-1)
+        expected_ids = torch.topk(
+            scores + correction_bias.unsqueeze(0), k=8, dim=-1
+        ).indices
+        expected_weights = scores.gather(1, topk_ids.to(torch.int64))
+        expected_weights /= expected_weights.sum(dim=-1, keepdim=True)
+
+        self.assertEqual(
+            torch.sort(topk_ids.to(torch.int64), dim=-1).values.tolist(),
+            torch.sort(expected_ids, dim=-1).values.tolist(),
+        )
+        torch.testing.assert_close(topk_weights, expected_weights)
+
+    def test_topk_softmax_with_correction_bias(self):
+        """Bias must affect expert selection without becoming a routing weight."""
+        for num_tokens, num_experts, topk, with_bias, renormalize in itertools.product(
+            [1, 17, 128],
+            [16, 128, 384, 512],
+            [1, 2, 4, 8],
+            [False, True],
+            [False, True],
+        ):
+            torch.manual_seed(0)
+            hidden_states = torch.randn((num_tokens, 16), dtype=torch.bfloat16)
+            gating_output = torch.randn((num_tokens, num_experts), dtype=torch.bfloat16)
+            correction_bias = torch.randn(num_experts) if with_bias else None
+
+            topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
+                hidden_states=hidden_states,
+                gating_output=gating_output,
+                topk=topk,
+                renormalize=renormalize,
+                correction_bias=correction_bias,
+            )
+
+            scores = torch.softmax(gating_output.float(), dim=-1)
+            scores_for_choice = scores
+            if correction_bias is not None:
+                scores_for_choice = scores_for_choice + correction_bias.unsqueeze(0)
+            expected_choice_scores = torch.topk(
+                scores_for_choice, k=topk, dim=-1, sorted=True
+            ).values
+            selected_choice_scores = torch.sort(
+                scores_for_choice.gather(1, topk_ids.to(torch.int64)),
+                dim=-1,
+                descending=True,
+            ).values
+
+            expected_weights = scores.gather(1, topk_ids.to(torch.int64))
+            if renormalize:
+                expected_weights = expected_weights / expected_weights.sum(
+                    dim=-1, keepdim=True
+                )
+
+            self.assertEqual(topk_ids.dtype, torch.int32)
+            self.assertEqual(topk_weights.dtype, torch.float32)
+            self.assertTrue(torch.all((topk_ids >= 0) & (topk_ids < num_experts)))
+            sorted_ids = torch.sort(topk_ids, dim=-1).values
+            self.assertTrue(torch.all(sorted_ids[:, 1:] != sorted_ids[:, :-1]))
+            torch.testing.assert_close(
+                selected_choice_scores,
+                expected_choice_scores,
+                atol=1e-4,
+                rtol=1e-4,
+            )
+            torch.testing.assert_close(
+                topk_weights, expected_weights, atol=1e-4, rtol=1e-4
+            )
+

 class TestCustomTopK(CustomTestCase):
    def _run_single_test(
@@ -251,6 +335,79 @@ class TestCustomTopK(CustomTestCase):
                123, 32, 1, False, torch.bfloat16, native_custom_f, fused_custom_f
            )

+    def test_topk_sigmoid_with_correction_bias(self):
+        """Biased scores must select experts while returned weights stay unbiased."""
+        for num_tokens, num_experts, topk, with_bias, renormalize in itertools.product(
+            [1, 17, 128],
+            [16, 128, 256, 384, 512],
+            [1, 2, 4, 8],
+            [False, True],
+            [False, True],
+        ):
+            torch.manual_seed(0)
+            hidden_states = torch.randn((num_tokens, 16), dtype=torch.bfloat16)
+            gating_output = torch.randn((num_tokens, num_experts), dtype=torch.bfloat16)
+            correction_bias = torch.randn(num_experts) if with_bias else None
+
+            topk_weights, topk_ids = torch.ops.sgl_kernel.topk_sigmoid_cpu(
+                hidden_states=hidden_states,
+                gating_output=gating_output,
+                topk=topk,
+                renormalize=renormalize,
+                correction_bias=correction_bias,
+            )
+
+            scores = torch.sigmoid(gating_output.float())
+            scores_for_choice = scores
+            if correction_bias is not None:
+                scores_for_choice = scores_for_choice + correction_bias.unsqueeze(0)
+
+            expected_choice_scores = torch.topk(
+                scores_for_choice, k=topk, dim=-1
+            ).values
+            selected_choice_scores = torch.sort(
+                scores_for_choice.gather(1, topk_ids.to(torch.int64)),
+                dim=-1,
+                descending=True,
+            ).values
+
+            expected_weights = scores.gather(1, topk_ids.to(torch.int64))
+            if renormalize:
+                expected_weights /= expected_weights.sum(dim=-1, keepdim=True)
+
+            self.assertEqual(topk_ids.dtype, torch.int32)
+            self.assertEqual(topk_weights.dtype, torch.float32)
+            self.assertTrue(torch.equal(selected_choice_scores, expected_choice_scores))
+            torch.testing.assert_close(
+                topk_weights, expected_weights, atol=1e-4, rtol=1e-4
+            )
+
+    def test_topk_sigmoid_mixed_input_dtypes(self):
+        torch.manual_seed(0)
+        hidden_states = torch.randn((17, 16), dtype=torch.bfloat16)
+        gating_output = torch.randn((17, 256), dtype=torch.float32)
+
+        topk_weights, topk_ids = torch.ops.sgl_kernel.topk_sigmoid_cpu(
+            hidden_states=hidden_states,
+            gating_output=gating_output,
+            topk=8,
+            renormalize=True,
+            correction_bias=None,
+        )
+
+        scores = torch.sigmoid(gating_output)
+        expected_ids = torch.topk(scores, k=8, dim=-1).indices
+        expected_weights = scores.gather(1, topk_ids.to(torch.int64))
+        expected_weights /= expected_weights.sum(dim=-1, keepdim=True)
+
+        self.assertTrue(
+            torch.equal(
+                torch.sort(topk_ids.to(torch.int64), dim=-1).values,
+                torch.sort(expected_ids, dim=-1).values,
+            )
+        )
+        torch.testing.assert_close(topk_weights, expected_weights, atol=1e-5, rtol=1e-5)
+

 if __name__ == "__main__":
    unittest.main()