diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index c3ae9af25..a806a4052 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -524,7 +524,7 @@ def biased_grouped_topk_gpu( topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device) topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) aiter_biased_grouped_topk( - gating_output, + gating_output.to(dtype=torch.float32), correction_bias, topk_weights, topk_ids,