[Bugfix] Fix deepseekv3 grouped topk error

vllm-project · Feb 18, 2025 · 76896ef · 76896ef
1 parent e2603fe
commit 76896ef
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -946,8 +946,15 @@ def grouped_topk(hidden_states: torch.Tensor,
         scores = scores + e_score_correction_bias.unsqueeze(0)
 
     num_token = scores.shape[0]
-    group_scores = scores.view(num_token, num_expert_group,
-                               -1).max(dim=-1).values  # [n, n_group]
+    if e_score_correction_bias is not None:
+        group_scores = (
+            scores.view(num_token, num_expert_group, -1)
+            .topk(2, dim=-1)[0]
+            .sum(dim=-1)
+        )
+    else:
+        group_scores = scores.view(num_token, num_expert_group,
+                                   -1).max(dim=-1).values  # [n, n_group]
     group_idx = torch.topk(group_scores, k=topk_group, dim=-1,
                            sorted=False)[1]  # [n, top_k_group]
     group_mask = torch.zeros_like(group_scores)  # [n, n_group]