[Perf] Restore torch.compile fusion for topk postprocessing (#21771)

2026-04-07 01:38:38 -07:00
parent 3d3a32c0b9
commit 490fa9fa44
1 changed files with 3 additions and 2 deletions
@@ -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