EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[Bugfix] Fix scores mask for moe topk (#3705)

Browse Source
This commit is contained in:
chenxiaobing
2025-02-21 02:17:23 +08:00
committed by GitHub
parent ddcf9fe3be
commit d5d80ab477
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/moe/topk.py
View File

@@ -141,7 +141,9 @@ def biased_grouped_topk(
.expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
.reshape(num_token, -1) .reshape(num_token, -1)
) # [n, e] ) # [n, e]
tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e] tmp_scores = scores_for_choice.masked_fill(
~score_mask.bool(), float("-inf")
) # [n, e]
_, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False) _, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)
topk_weights = scores.gather(1, topk_ids) topk_weights = scores.gather(1, topk_ids)