diff --git a/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py b/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py index 6934650adb92..1dabf2f3582c 100644 --- a/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py +++ b/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, diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 3175d230dacc..faee873da700 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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, diff --git a/python/sglang/srt/models/afmoe.py b/python/sglang/srt/models/afmoe.py index ed5521f77174..e21b3e47ce72 100644 --- a/python/sglang/srt/models/afmoe.py +++ b/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), ) diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 4c7c94f612bf..d5436e5b7636 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/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, diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index c29338447319..e3c4c5210cc4 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/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, diff --git a/python/sglang/srt/models/dbrx.py b/python/sglang/srt/models/dbrx.py index d0425e603f8f..18037900beae 100644 --- a/python/sglang/srt/models/dbrx.py +++ b/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, diff --git a/python/sglang/srt/models/deepseek.py b/python/sglang/srt/models/deepseek.py index 6bb1996b60e4..a28f2a0a4db6 100644 --- a/python/sglang/srt/models/deepseek.py +++ b/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), ) diff --git a/python/sglang/srt/models/ernie45_moe_vl.py b/python/sglang/srt/models/ernie45_moe_vl.py index f15071207bc0..cdadc9314f45 100644 --- a/python/sglang/srt/models/ernie45_moe_vl.py +++ b/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], diff --git a/python/sglang/srt/models/exaone_moe.py b/python/sglang/srt/models/exaone_moe.py index 211a12de392c..01a488b7e84e 100755 --- a/python/sglang/srt/models/exaone_moe.py +++ b/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, diff --git a/python/sglang/srt/models/granitemoe.py b/python/sglang/srt/models/granitemoe.py index 5c6d4d91c797..0b5090a951f5 100644 --- a/python/sglang/srt/models/granitemoe.py +++ b/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, ) diff --git a/python/sglang/srt/models/hunyuan_v3.py b/python/sglang/srt/models/hunyuan_v3.py index f95840dd9976..7312e386258c 100644 --- a/python/sglang/srt/models/hunyuan_v3.py +++ b/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, diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index a8410480e401..72e780bc3fbc 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/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, diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index fd4537c4a0d3..730cc0e99a00 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/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, diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 4f4774c354be..ee26f5488eb5 100644 --- a/python/sglang/srt/models/llada2.py +++ b/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, diff --git a/python/sglang/srt/models/llama4.py b/python/sglang/srt/models/llama4.py index af4f259e1c03..fef08c87584b 100644 --- a/python/sglang/srt/models/llama4.py +++ b/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, ) diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index e1f27f560d2d..da7dfa9d4b27 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/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, diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index b0c968ba1360..4d9b12b2ae9f 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/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, diff --git a/python/sglang/srt/models/mixtral.py b/python/sglang/srt/models/mixtral.py index 4b2b5c90807d..5eb781d76233 100644 --- a/python/sglang/srt/models/mixtral.py +++ b/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, ) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index d20b290fb44b..878f0db2b967 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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, diff --git a/python/sglang/srt/models/olmoe.py b/python/sglang/srt/models/olmoe.py index c66156b33b46..ec0c2ece48d4 100644 --- a/python/sglang/srt/models/olmoe.py +++ b/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, ) diff --git a/python/sglang/srt/models/phimoe.py b/python/sglang/srt/models/phimoe.py index f9f2ab06fef4..e7e5f4342685 100644 --- a/python/sglang/srt/models/phimoe.py +++ b/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, ) diff --git a/python/sglang/srt/models/xverse_moe.py b/python/sglang/srt/models/xverse_moe.py index 898cf3abc42d..6dcb4e132078 100644 --- a/python/sglang/srt/models/xverse_moe.py +++ b/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), ) diff --git a/test/registered/moe/test_triton_fused_moe.py b/test/registered/moe/test_triton_fused_moe.py index 3591507b091d..9dcb63eedb7b 100644 --- a/test/registered/moe/test_triton_fused_moe.py +++ b/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, )