From d5d80ab477de9c6f2dc4dc6c766de475749ba68f Mon Sep 17 00:00:00 2001 From: chenxiaobing <22113491+Chen-XiaoBing@users.noreply.github.com> Date: Fri, 21 Feb 2025 02:17:23 +0800 Subject: [PATCH] [Bugfix] Fix scores mask for moe topk (#3705) --- python/sglang/srt/layers/moe/topk.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index b0b57d68d..91ca00c6e 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -141,7 +141,9 @@ def biased_grouped_topk( .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) .reshape(num_token, -1) ) # [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_weights = scores.gather(1, topk_ids)