From fb272d27dbe1600fdcb32a5166948b1395bbd1a0 Mon Sep 17 00:00:00 2001
From: YC Yen-Ching Tseng <yctseng@amd.com>
Date: Sat, 25 Apr 2026 13:18:08 +0800
Subject: [PATCH] [AMD] Optimize MiniMax-M2.5 - use aiter biased_grouped_topk
 for sigmoid scoring in MoE routing (#23611)

---
 python/sglang/srt/layers/moe/topk.py | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py
index b7e4a9158..02fac83a0 100644
--- a/python/sglang/srt/layers/moe/topk.py
+++ b/python/sglang/srt/layers/moe/topk.py
@@ -538,13 +538,24 @@ def fused_topk(
                 renormalize,
             )
     elif scoring_func == "sigmoid":
-        topk_sigmoid(
-            topk_weights,
-            topk_ids,
-            gating_output,
-            renormalize,
-            correction_bias,
-        )
+        if _use_aiter and correction_bias is not None:
+            aiter_biased_grouped_topk(
+                gating_output,
+                correction_bias.to(dtype=gating_output.dtype),
+                topk_weights,
+                topk_ids,
+                num_expert_group=1,
+                topk_group=1,
+                need_renorm=renormalize,
+            )
+        else:
+            topk_sigmoid(
+                topk_weights,
+                topk_ids,
+                gating_output,
+                renormalize,
+                correction_bias,
+            )
     else:
         raise ValueError(f"Invalid scoring function: {scoring_func}")