A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
[Bugfix] Fix deepseekv3 grouped topk error (#13474)
Signed-off-by: Chen-XiaoBing <chenxb002@whu.edu.cn>
chenxiaobing committed
Feb 20, 2025 at 22:47 UTC
ed6e9075d31e32c8548b480a47d1ffb77da1f54c
1 file changed
+8
-5
vllm/model_executor/layers/fused_moe/fused_moe.py
+8
-5
@@ -939,15 +939,17 @@ def grouped_topk(hidden_states: torch.Tensor,
939
else:
940
raise ValueError(f"Unsupported scoring function: {scoring_func}")
941
942
+ num_token = scores.shape[0]
943
if e_score_correction_bias is not None:
944
# Store original scores before applying correction bias. We use biased
945
# scores for expert selection but original scores for routing weights
946
original_scores = scores
947
scores = scores + e_score_correction_bias.unsqueeze(0)
947
-
948
- num_token = scores.shape[0]
949
- group_scores = scores.view(num_token, num_expert_group,
950
- -1).max(dim=-1).values # [n, n_group]
948
+ group_scores = (scores.view(num_token, num_expert_group,
949
+ -1).topk(2, dim=-1)[0].sum(dim=-1))
950
+ else:
951
+ group_scores = scores.view(num_token, num_expert_group,
952
+ -1).max(dim=-1).values # [n, n_group]
953
group_idx = torch.topk(group_scores, k=topk_group, dim=-1,
954
sorted=False)[1] # [n, top_k_group]
955
group_mask = torch.zeros_like(group_scores) # [n, n_group]
@@ -955,7 +957,8 @@ def grouped_topk(hidden_states: torch.Tensor,
957
score_mask = group_mask.unsqueeze(-1).expand(
958
num_token, num_expert_group,
959
scores.shape[-1] // num_expert_group).reshape(num_token, -1) # [n, e]
958
- tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e]
960
+ tmp_scores = scores.masked_fill(~score_mask.bool(),
961
+ float("-inf")) # [n, e]
962
963
if e_score_correction_bias is not None:
964
topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1]