Remove type conversion and fix id map in topk (#7759)

8b1942c6 · Ke Bao · GitHub · 489934be · 8b1942c6
Unverified Commit 8b1942c6 authored Jul 04, 2025 by Ke Bao Committed by GitHub Jul 03, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

python/sglang/srt/layers/moe/topk.py python/sglang/srt/layers/moe/topk.py +2 -1

No files found.
--- a/python/sglang/srt/layers/moe/topk.py
+++ b/python/sglang/srt/layers/moe/topk.py
@@ -112,10 +112,11 @@ def fused_topk(
    topk_softmax(
        topk_weights,
        topk_ids,
-        gating_output.float(),
+        gating_output,
        renormalize,
    )

+    topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
    _mask_topk_ids_padded_region(topk_ids, num_token_non_padded)
    return topk_weights, topk_ids