revise get_moe_scores (#3164)

test/operators/test_noaux_tc.py

@@ -0,0 +1,76 @@
+import unittest
+
+import paddle
+
+from fastdeploy.model_executor.ops.gpu import noaux_tc
+
+
+class TestMoeRouting(unittest.TestCase):
+    def setUp(self):
+        self.num_tokens = 10
+        self.num_experts = 64
+        self.gating_output = paddle.rand([self.num_tokens, self.num_experts])
+        self.e_score_correction_bias = paddle.rand([self.num_experts])
+        self.n_group = 8
+        self.topk_group = 4
+        self.top_k = 8
+        self.routed_scaling_factor = 1.5
+
+    def node_limit_routing(self, gate_probs):
+        """将所有专家分组, 只在topk_group个group内选择专家"""
+        assert len(gate_probs.shape) == 2
+        seq_length, n_experts = gate_probs.shape
+
+        group_scores = gate_probs.reshape([seq_length, 8, -1]).topk(2, axis=-1)[0].sum(axis=-1)
+        group_idx = paddle.topk(group_scores, k=4, axis=-1, sorted=True)[1]
+        group_mask = paddle.zeros_like(group_scores).put_along_axis(
+            group_idx, paddle.ones([], dtype="float32"), axis=-1
+        )
+        score_mask = group_mask.unsqueeze(-1).expand([seq_length, 8, n_experts // 8]).reshape([seq_length, -1])
+        gate_probs = gate_probs.masked_fill(~score_mask.astype(paddle.bool), float("-inf"))
+        return gate_probs
+
+    def ref_moe_routing(self):
+        scores = paddle.nn.functional.sigmoid(self.gating_output)
+        prob_for_choice = scores + self.e_score_correction_bias.unsqueeze(0)
+        prob_for_choice = self.node_limit_routing(prob_for_choice)
+        top_logits, topk_idx_ref = paddle.topk(prob_for_choice, self.top_k, axis=1)
+
+        token_num, top_k = topk_idx_ref.shape
+        _, num_expert = prob_for_choice.shape
+        topk_idx_expanded = paddle.unsqueeze(topk_idx_ref, axis=-1)
+        indices = paddle.concat(
+            [
+                paddle.arange(token_num, dtype="int64").unsqueeze(1).tile([1, top_k]).unsqueeze(-1),
+                topk_idx_expanded,
+            ],
+            axis=-1,
+        )
+        selected_gate_probs = paddle.gather_nd(scores, indices)
+
+        selected_gate_probs_sum = paddle.sum(selected_gate_probs, axis=1, keepdim=True)
+        topk_weights_ref = selected_gate_probs / selected_gate_probs_sum
+        topk_weights_ref = topk_weights_ref * self.routed_scaling_factor
+        return topk_weights_ref, topk_idx_ref
+
+    def test_moe_select(self):
+        scores = paddle.nn.functional.sigmoid(self.gating_output)
+        scores_with_bias = scores + self.e_score_correction_bias.unsqueeze(0)
+
+        scores, topk_values, topk_idx = noaux_tc(
+            scores,
+            scores_with_bias,
+            self.n_group,
+            self.topk_group,
+            self.top_k,
+            self.routed_scaling_factor,
+        )
+
+        ref_topk_values, ref_topk_idx = self.ref_moe_routing()
+
+        paddle.allclose(topk_values, ref_topk_values)
+        paddle.allclose(topk_idx.cast(int), ref_topk_idx.cast(int))
+
+
+if __name__ == "__main__":
+    unittest.main()