[perf] use optimized topk_softmax + renormalize (lightop)

vllm/model_executor/layers/fused_moe/router/fused_topk_router.py

@@ -9,6 +9,8 @@ from vllm._aiter_ops import rocm_aiter_ops
 from vllm.distributed.eplb.eplb_state import EplbLayerState
 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
 from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
+import vllm.envs as envs
+from lightop import op as op
 
 
 def vllm_topk_softmax(
@@ -18,6 +20,15 @@ def vllm_topk_softmax(
     gating_output: torch.Tensor,
     renormalize: bool = False,
 ) -> tuple[torch.Tensor, ...]:
+    if envs.VLLM_USE_TOPK_RENORM and renormalize is True:
+        op.topk_softmax(
+            topk_weights,
+            topk_indices,
+            token_expert_indices,
+            gating_output,
+            renormalize,
+        )
+    else:
         ops.topk_softmax(
             topk_weights,
             topk_indices,