fix group topk

Author: hiworldwzj

lightllm/common/fused_moe/grouped_topk.py

@@ -118,6 +118,7 @@ def grouped_topk_kernel(
     EXPERT_GROUP_NUM: tl.constexpr,  # tl.next_power_two_of(group_num)
     EXPERT_GROUP_SIZE: tl.constexpr,  # tl.next_power_two_of(group_expert_num)
     RENORMALIZE: tl.constexpr,
+    GROUP_SCORE_USED_TOPK_NUM: tl.constexpr,
 ):
     token_index = tl.program_id(axis=0)
     offs_n = tl.arange(0, EXPERT_BLOCK_SIZE)
@@ -148,7 +149,15 @@ def grouped_topk_kernel(
         other=-10000000.0,
     )  # [group, group_size]
 
-    group_value = tl.max(group_scores, axis=1)  # [group,]
+    group_value = tl.sum(
+        tl.where(
+            (offs_group < group_num)[:, None] & (offs_group_v < GROUP_SCORE_USED_TOPK_NUM)[None, :],
+            tl.sort(group_scores, dim=1, descending=True),
+            0.0,
+        ),
+        axis=1,
+    )
+
     sorted_group_value = tl.sort(group_value, descending=True)
     group_topk_value = tl.sum(tl.where(offs_group == group_topk_num - 1, sorted_group_value, 0.0))
     mask_group_scores = tl.where(
@@ -198,6 +207,7 @@ def triton_grouped_topk(
     num_expert_group: int = 0,
     topk_group: int = 0,
     scoring_func: str = "softmax",
+    group_score_used_topk_num=2,
 ):
 
     if correction_bias is not None:
@@ -239,6 +249,7 @@ def triton_grouped_topk(
         EXPERT_GROUP_NUM=triton.next_power_of_2(num_expert_group),
         EXPERT_GROUP_SIZE=triton.next_power_of_2(total_expert_num // num_expert_group),
         RENORMALIZE=renormalize,
+        GROUP_SCORE_USED_TOPK_NUM=group_score_used_topk_num,
         num_warps=1,
         num_stages=1,
     )

lightllm/common/fused_moe/topk_select.py

@@ -93,6 +93,7 @@ def grouped_topk(
     return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
 
 
+# biased_grouped_topk adapt from sgl-project/sglang/python/sglang/srt/layers/moe/topk.py
 def biased_grouped_topk(
     hidden_states: torch.Tensor,
     gating_output: torch.Tensor,
@@ -196,7 +197,12 @@ def select_experts(
                 scoring_func=scoring_func,
             )
         else:
-            topk_weights, topk_ids = biased_grouped_topk(
+            group_score_topk_num = 1
+            # for deepseek v3
+            if topk_group == 4 and num_expert_group == 8 and top_k == 8:
+                group_score_topk_num = 2
+
+            topk_weights, topk_ids = triton_grouped_topk(
                 hidden_states=hidden_states,
                 gating_output=router_logits,
                 correction_bias=correction_bias,
@@ -205,7 +211,9 @@ def select_experts(
                 num_expert_group=num_expert_group,
                 topk_group=topk_group,
                 scoring_func=scoring_func,
+                group_score_used_topk_num=group_score_topk_num,
             )
+
     elif custom_routing_function is None:
         topk_weights, topk_ids = fused_topk(
             hidden_states=hidden_states, gating_output=router_logits, topk=top_k, renormalize=renormalize

unit_tests/common/fused_moe/test_grouped_topk.py

@@ -2,8 +2,8 @@
 import time
 import pytest
 import numpy as np
-from lightllm.common.fused_moe.topk_select import grouped_topk
 from lightllm.common.fused_moe.grouped_topk import triton_grouped_topk
+from lightllm.common.fused_moe.topk_select import biased_grouped_topk as grouped_topk
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -21,7 +21,9 @@
     [
         (*a, b, c)
         for a in [(256, 4, 8, 8), (160, 3, 8, 6)]
-        for b in ["softmax", "sigmoid"]
+        for b in [
+            "sigmoid",
+        ]
         for c in [1, 8, 256, 1024, 2048, 4096, 8192]
     ],
 )