From 490fa9fa44c305056a57e0d5238d21b8739cc102 Mon Sep 17 00:00:00 2001
From: Nicolas Castet <26874160+nvcastet@users.noreply.github.com>
Date: Tue, 7 Apr 2026 03:38:38 -0500
Subject: [PATCH] [Perf] Restore torch.compile fusion for topk postprocessing
 (#21771)

---
 python/sglang/srt/layers/moe/topk.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py
index 892dcebae..bb6691814 100644
--- a/python/sglang/srt/layers/moe/topk.py
+++ b/python/sglang/srt/layers/moe/topk.py
@@ -956,8 +956,9 @@ def _post_process_topk_ids(
         topk_ids=topk_ids,
     )
     if _is_cuda:
-        topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
-        _mask_topk_ids_padded_region(topk_ids, num_token_non_padded)
+        topk_ids = _biased_grouped_topk_postprocess(
+            topk_ids, expert_location_dispatch_info, num_token_non_padded
+        )
 
     if num_fused_shared_experts > 0 and _use_aiter:
         M, N = router_logits.shape