[1/N] MoE Refactor: refactor select_experts (#7966)

python/sglang/srt/models/olmoe.py

@@ -32,6 +32,7 @@ from sglang.srt.layers.linear import (
 )
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+from sglang.srt.layers.moe.topk import TopK
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.layers.rotary_embedding import get_rope
@@ -76,13 +77,16 @@ class OlmoeMoE(nn.Module):
             prefix=add_prefix("gate", prefix),
         )
 
+        self.topk = TopK(
+            top_k=top_k,
+            renormalize=False,
+        )
+
         self.experts = FusedMoE(
             num_experts=num_experts,
-            top_k=top_k,
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
             reduce_results=True,
-            renormalize=False,
             quant_config=quant_config,
             tp_size=tp_size,
             prefix=add_prefix("experts", prefix),
@@ -94,9 +98,8 @@ class OlmoeMoE(nn.Module):
         hidden_states = hidden_states.view(-1, self.hidden_size)
         # router_logits: (num_tokens, n_experts)
         router_logits, _ = self.gate(hidden_states)
-        final_hidden_states = self.experts(
-            hidden_states=hidden_states, router_logits=router_logits
-        )
+        topk_output = self.topk(hidden_states, router_logits)
+        final_hidden_states = self.experts(hidden_states, topk_output)
         return final_hidden_states.view(orig_shape)