[Fix] MoE: require TopK layer_id to ensure routed expert captures

benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py

@@ -38,6 +38,7 @@ def fused_moe_triton_api(
 ):
     topk_op = TopK(
         top_k=topk,
+        layer_id=0,
         renormalize=False,
         use_grouped_topk=False,
         output_format=TopKOutputFormat.TRITON_KERNEL,

python/sglang/srt/layers/moe/topk.py

@@ -536,6 +536,8 @@ class TopK(BaseFusedOp):
     """
     Parameters:
     --top_k: The all number of top experts selected per token, including the fused shared expert(s).
+    --layer_id: global index of the owning layer. Required: routed-experts capture and
+      per-layer EPLB state are keyed by it.
     --num_fused_shared_experts: num of shared experts, can be activate both in TP or EP mode.
     --routed_scaling_factor: the scaling factor for routed experts in topk_weights.
     --fused_shared_experts_scaling_factor: scaling factor applied to the fused shared experts'
@@ -547,7 +549,7 @@ def __init__(
         self,
         top_k: int,
         *,
-        layer_id: Optional[int] = None,
+        layer_id: int,
         use_grouped_topk: bool = False,
         topk_group: Optional[int] = None,
         num_expert_group: Optional[int] = None,

python/sglang/srt/models/afmoe.py

@@ -152,6 +152,7 @@ def _custom_routing_function(
     def __init__(
         self,
         config: PretrainedConfig,
+        layer_id: int,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
     ):
@@ -234,6 +235,7 @@ def __init__(
         renormalize = self.route_norm if self.score_func == "sigmoid" else False
         self.topk = TopK(
             top_k=self.top_k,
+            layer_id=layer_id,
             renormalize=renormalize,
             use_grouped_topk=self.use_grouped_topk,
             num_expert_group=self.n_group if self.use_grouped_topk else None,
@@ -455,6 +457,7 @@ def __init__(
         if use_moe:
             self.mlp = AfmoeMoE(
                 config=config,
+                layer_id=layer_id,
                 quant_config=quant_config,
                 prefix=add_prefix("mlp", prefix),
             )

python/sglang/srt/models/bailing_moe.py

@@ -280,6 +280,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=self.top_k,
+            layer_id=self.layer_id,
             renormalize=self.norm_topk_prob,
             use_grouped_topk=self.use_grouped_topk,
             num_expert_group=self.num_expert_group,

python/sglang/srt/models/bailing_moe_linear.py

@@ -298,6 +298,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=self.top_k,
+            layer_id=self.layer_id,
             use_grouped_topk=self.use_grouped_topk,
             renormalize=self.norm_expert_prob,
             num_expert_group=self.num_expert_group,

python/sglang/srt/models/dbrx.py

@@ -97,6 +97,7 @@ class DbrxExperts(nn.Module):
     def __init__(
         self,
         config: DbrxConfig,
+        layer_id: int,
         quant_config: Optional[QuantizationConfig] = None,
         params_dtype: Optional[torch.dtype] = None,
         prefix: str = "",
@@ -115,6 +116,7 @@ def __init__(
         self.router = DbrxRouter(config, self.params_dtype)
         self.topk = TopK(
             self.top_k,
+            layer_id=layer_id,
             renormalize=True,
         )
         self.moe_runner_config = MoeRunnerConfig(inplace=True)
@@ -336,7 +338,7 @@ def __init__(
             quant_config=quant_config,
             prefix=add_prefix("norm_attn_norm", prefix),
         )
-        self.ffn = DbrxExperts(config, quant_config=quant_config)
+        self.ffn = DbrxExperts(config, layer_id, quant_config=quant_config)
 
     def forward(
         self,

python/sglang/srt/models/deepseek.py

@@ -109,6 +109,7 @@ class DeepseekMoE(nn.Module):
     def __init__(
         self,
         config: PretrainedConfig,
+        layer_id: int,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
     ):
@@ -125,6 +126,7 @@ def __init__(
             )
         self.topk = TopK(
             top_k=self.top_k,
+            layer_id=layer_id,
             renormalize=config.norm_topk_prob,
         )
         self.experts = nn.ModuleList(
@@ -339,6 +341,7 @@ def __init__(
         ):
             self.mlp = DeepseekMoE(
                 config=config,
+                layer_id=layer_id,
                 quant_config=quant_config,
                 prefix=add_prefix("mlp", prefix),
             )

python/sglang/srt/models/ernie45_moe_vl.py

@@ -206,6 +206,7 @@ def __init__(
 
             self.text_experts_topk = TopK(
                 top_k=config.moe_k,
+                layer_id=self.layer_id,
                 renormalize=True,
                 use_grouped_topk=False,
                 correction_bias=self.e_score_correction_bias[0],
@@ -237,6 +238,7 @@ def __init__(
 
             self.vision_experts_topk = TopK(
                 top_k=config.moe_k,
+                layer_id=self.layer_id,
                 renormalize=True,
                 use_grouped_topk=False,
                 correction_bias=self.e_score_correction_bias[1],

python/sglang/srt/models/exaone_moe.py

@@ -172,6 +172,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=config.num_experts_per_tok,
+            layer_id=self.layer_id,
             renormalize=config.norm_topk_prob,
             use_grouped_topk=True,
             num_expert_group=config.n_group,

python/sglang/srt/models/granitemoe.py

@@ -134,6 +134,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=top_k,
+            layer_id=layer_id,
             renormalize=True,
         )
 

python/sglang/srt/models/hunyuan_v3.py

@@ -120,6 +120,7 @@ def __init__(
         )
         self.topk = TopK(
             top_k=config.num_experts_per_tok,
+            layer_id=layer_id,
             use_grouped_topk=True,
             num_expert_group=1,
             topk_group=1,

python/sglang/srt/models/kimi_k3.py

@@ -494,6 +494,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=config.num_experts_per_token,
+            layer_id=self.layer_idx,
             renormalize=moe_renormalize,
             use_grouped_topk=True,
             num_expert_group=config.num_expert_group,

python/sglang/srt/models/kimi_linear.py

@@ -120,6 +120,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=config.num_experts_per_token,
+            layer_id=self.layer_idx,
             renormalize=moe_renormalize,
             use_grouped_topk=True,
             num_expert_group=config.num_expert_group,

python/sglang/srt/models/llada2.py

@@ -257,6 +257,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=self.top_k,
+            layer_id=self.layer_id,
             renormalize=self.norm_topk_prob,
             use_grouped_topk=self.use_grouped_topk,
             num_expert_group=self.num_expert_group,

python/sglang/srt/models/llama4.py

@@ -107,6 +107,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=self.top_k,
+            layer_id=layer_id,
             renormalize=False,
             custom_routing_function=Llama4MoE.custom_routing_function,
         )

python/sglang/srt/models/mimo_v2.py

@@ -426,6 +426,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=config.num_experts_per_tok,
+            layer_id=self.layer_id,
             renormalize=config.norm_topk_prob,
             use_grouped_topk=True,
             num_expert_group=config.n_group,

python/sglang/srt/models/minimax_m2.py

@@ -544,6 +544,7 @@ def __init__(
         )
         self.topk = TopK(
             top_k=config.num_experts_per_tok,
+            layer_id=layer_id,
             renormalize=True,
             scoring_func=config.scoring_func,
             correction_bias=self.e_score_correction_bias,

python/sglang/srt/models/mixtral.py

@@ -90,6 +90,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=top_k,
+            layer_id=layer_id,
             renormalize=True,
         )
 

python/sglang/srt/models/nemotron_h.py

@@ -223,6 +223,7 @@ def __init__(
         )
         self.topk = TopK(
             top_k=config.num_experts_per_tok,
+            layer_id=layer_idx,
             use_grouped_topk=True,
             topk_group=config.topk_group,
             num_expert_group=config.n_group,

python/sglang/srt/models/olmoe.py

@@ -82,6 +82,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=top_k,
+            layer_id=layer_id,
             renormalize=False,
         )
 

python/sglang/srt/models/phimoe.py

@@ -200,6 +200,7 @@ def __init__(
 
         self.topk = TopK(
             top_k=top_k,
+            layer_id=layer_id,
             renormalize=False,
             custom_routing_function=phimoe_routing_function,
         )

python/sglang/srt/models/xverse_moe.py

@@ -94,6 +94,7 @@ class XverseMoE(nn.Module):
     def __init__(
         self,
         config: PretrainedConfig,
+        layer_id: int,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
     ):
@@ -134,6 +135,7 @@ def __init__(
         )
         self.topk = TopK(
             top_k=self.top_k,
+            layer_id=layer_id,
             renormalize=getattr(self.config, "norm_topk_prob", False),
         )
 
@@ -305,6 +307,7 @@ def __init__(
         if config.num_experts is not None:
             self.mlp = XverseMoE(
                 config=config,
+                layer_id=layer_id,
                 quant_config=quant_config,
                 prefix=add_prefix("mlp", prefix),
             )

test/registered/moe/test_triton_fused_moe.py

@@ -106,6 +106,7 @@ def _test_case(self, m, n, k, e, topk, dtype):
 
         topk_op = TopK(
             top_k=topk,
+            layer_id=0,
             renormalize=False,
             use_grouped_topk=False,
         )