From 0dfec3bd2274558758464ad26b56c27f61753c52 Mon Sep 17 00:00:00 2001 From: set Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 01/40] Fix weight shape mismatch assertion in ReplicatedLinear class to raise ValueError with detailed information --- python/sglang/srt/layers/linear.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 7af9eb004008..0c7b80760779 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -266,7 +266,11 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): param.dtype == loaded_weight.dtype ), "init para dtype and loaded weight dtype should be the same" - assert param.size() == loaded_weight.size() + if param.size() != loaded_weight.size(): + raise ValueError( + f"Weight shape mismatch: param {param.shape} vs loaded {loaded_weight.shape}. " + f"param.numel()={param.numel()}, loaded.numel()={loaded_weight.numel()}" + ) param.data.copy_(loaded_weight) def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: From e6e9e184eced1f48b22a04e98a3b86c9104fdf2e Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 02/40] Add shared expert fusion support in Qwen2 MoE block - Introduced logic to determine the number of fused shared experts based on configuration. - Updated top-k output handling to include shared expert weights when applicable. - Enhanced logging for shared expert fusion status during initialization and weight loading. - Refactored methods to accommodate the new shared expert functionality. --- python/sglang/srt/models/qwen2_moe.py | 87 +++++++++++++++++++++++++-- 1 file changed, 81 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 8f3475a24323..4cfd9af52e97 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -61,7 +61,7 @@ ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton import FusedMoE -from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.moe.topk import StandardTopKOutput, TopK, TopKOutputChecker from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, @@ -90,6 +90,7 @@ is_cpu, is_cuda, make_layers, + rank0_log, use_intel_amx_backend, ) from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -166,6 +167,7 @@ def __init__( super().__init__() self.tp_size = get_tensor_model_parallel_world_size() self.layer_id = layer_id + self.num_experts = config.num_experts self.alt_stream = alt_stream if self.tp_size > config.num_experts: raise ValueError( @@ -173,6 +175,15 @@ def __init__( f"the number of experts {config.num_experts}." ) + self.num_fused_shared_experts = self._determine_num_fused_shared_experts( + config + ) + if self.num_fused_shared_experts > 0: + rank0_log( + "Shared experts fusion enabled for Qwen2 MoE " + "(topk+1 experts per token)." + ) + self.topk = TopK( top_k=config.num_experts_per_tok, renormalize=config.norm_topk_prob, @@ -181,14 +192,16 @@ def __init__( self.experts = get_moe_impl_class(quant_config)( layer_id=self.layer_id, - top_k=config.num_experts_per_tok, + top_k=config.num_experts_per_tok + self.num_fused_shared_experts, num_experts=config.num_experts + + self.num_fused_shared_experts + get_global_server_args().ep_num_redundant_experts, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, quant_config=quant_config, prefix=add_prefix("experts", prefix), routing_method_type=RoutingMethodType.RenormalizeNaive, + num_fused_shared_experts=self.num_fused_shared_experts, ) self.gate = ReplicatedLinear( @@ -198,7 +211,10 @@ def __init__( quant_config=None, prefix=add_prefix("gate", prefix), ) - if config.shared_expert_intermediate_size > 0: + if ( + config.shared_expert_intermediate_size > 0 + and self.num_fused_shared_experts == 0 + ): self.shared_expert = Qwen2MoeMLP( hidden_size=config.hidden_size, intermediate_size=config.shared_expert_intermediate_size, @@ -234,6 +250,18 @@ def __init__( self.top_k = config.num_experts_per_tok self.is_nextn = is_nextn + def _determine_num_fused_shared_experts(self, config: PretrainedConfig) -> int: + """Determine if shared expert can be fused with router experts (topk+1). + Fusion requires shared_expert_intermediate_size == moe_intermediate_size. + """ + if get_global_server_args().disable_shared_experts_fusion: + return 0 + if getattr(config, "shared_expert_intermediate_size", 0) <= 0: + return 0 + if config.shared_expert_intermediate_size != config.moe_intermediate_size: + return 0 + return 1 + def get_moe_weights(self): return [ x.data @@ -244,6 +272,44 @@ def get_moe_weights(self): ) ] + def _get_shared_expert_weights(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Return sigmoid(shared_expert_gate) for fused shared expert weights.""" + if self.num_fused_shared_experts <= 0 or self.shared_expert_gate is None: + return None + shared_out = self.shared_expert_gate(hidden_states) + shared_logits = shared_out[0] if isinstance(shared_out, tuple) else shared_out + return F.sigmoid(shared_logits) + + def _append_shared_to_topk_output( + self, + topk_output: StandardTopKOutput, + hidden_states: torch.Tensor, + ) -> StandardTopKOutput: + """Append shared expert ids and weights to topk output before fused MoE.""" + if self.num_fused_shared_experts <= 0: + return topk_output + shared_weights = self._get_shared_expert_weights(hidden_states) + if shared_weights is None: + return topk_output + M = topk_output.topk_ids.shape[0] + shared_expert_id = self.num_experts # 512 for Qwen3.5 MoE + shared_ids = torch.full( + (M, self.num_fused_shared_experts), + shared_expert_id, + dtype=topk_output.topk_ids.dtype, + device=topk_output.topk_ids.device, + ) + shared_weights = shared_weights.expand( + M, self.num_fused_shared_experts + ).to(topk_output.topk_weights.dtype) + fused_topk_ids = torch.cat([topk_output.topk_ids, shared_ids], dim=-1) + fused_topk_weights = torch.cat([topk_output.topk_weights, shared_weights], dim=-1) + return StandardTopKOutput( + topk_weights=fused_topk_weights, + topk_ids=fused_topk_ids, + router_logits=topk_output.router_logits, + ) + def _forward_shared_experts(self, hidden_states: torch.Tensor): shared_output = None if self.shared_expert is not None: @@ -296,9 +362,14 @@ def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBat return final_hidden_states def _forward_router_experts(self, hidden_states: torch.Tensor): - # router_logits: (num_tokens, n_experts) - router_logits, _ = self.gate(hidden_states) - topk_output = self.topk(hidden_states, router_logits) + gate_logits, _ = self.gate(hidden_states) + topk_output = self.topk(hidden_states, gate_logits) + if self.num_fused_shared_experts > 0 and TopKOutputChecker.format_is_standard( + topk_output + ): + topk_output = self._append_shared_to_topk_output( + topk_output, hidden_states + ) return self.experts(hidden_states, topk_output) def forward_normal_dual_stream( @@ -850,6 +921,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: + raw_weight_name = name layer_id = get_layer_id(name) if ( layer_id is not None @@ -881,6 +953,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): if name not in params_dict: continue + rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) @@ -891,6 +964,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): if weight_name not in name: continue name = name.replace(weight_name, param_name) + rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = param.weight_loader weight_loader( @@ -909,6 +983,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): continue if name in params_dict.keys(): + rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader From 4eeddb063f27fe35ea9a4125f0d0947a0f5f3b95 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 03/40] Add support for dynamic expert count in Qwen3 MoE model - Implemented a method to retrieve the number of fused shared experts for weight loading. - Updated weight loading logic to accommodate the dynamic calculation of total experts. - Enhanced logging to track weight mapping during the loading process. - Refactored expert parameters mapping to include shared expert configurations. --- python/sglang/srt/models/qwen3_5.py | 109 +++++++++++++++++++++++++--- 1 file changed, 99 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 2f430c2b9b28..8c4ec8d3219c 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -91,6 +91,7 @@ is_gfx95_supported, is_npu, make_layers, + rank0_log, set_weight_attrs, ) from sglang.srt.utils.hf_transformers_utils import get_processor, get_rope_config @@ -1108,6 +1109,15 @@ def __init__( ) -> None: super().__init__(config=config, quant_config=quant_config, prefix=prefix) + def _get_num_fused_shared_experts(self): + if not ( + hasattr(self.model, "layers") + and len(self.model.layers) > 0 + and hasattr(self.model.layers[0].mlp, "num_fused_shared_experts") + ): + return 0 + return self.model.layers[0].mlp.num_fused_shared_experts + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ # (param_name, shard_name, shard_id) @@ -1123,13 +1133,17 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("in_proj_ba.", "in_proj_a.", 1), ] + num_experts_base = self.config.num_experts + num_fused_shared_experts = self._get_num_fused_shared_experts() + num_experts = num_experts_base + num_fused_shared_experts + # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) expert_params_mapping = FusedMoE.make_expert_params_mapping( ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, + num_experts=num_experts, ) # Skip loading extra parameters for GPTQ/modelopt models. @@ -1152,8 +1166,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("experts.w2_weight", "experts.down_proj", 0, "w2"), ] - num_experts = self.config.num_experts - def load_fused_expert_weights( name: str, params_dict: dict, @@ -1181,6 +1193,7 @@ def load_fused_expert_weights( params_dict = dict(self.named_parameters(remove_duplicate=False)) for name, loaded_weight in weights: + raw_weight_name = name if "rotary_emb.inv_freq" in name: continue if "mtp" in name: @@ -1225,6 +1238,7 @@ def load_fused_expert_weights( if name not in params_dict: continue + rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) @@ -1244,6 +1258,7 @@ def load_fused_expert_weights( if is_fused_expert: if "experts.gate_up_proj" in name: loaded_weight = loaded_weight.chunk(2, dim=-2) + rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped} (w1)") load_fused_expert_weights( name_mapped, params_dict, @@ -1251,6 +1266,7 @@ def load_fused_expert_weights( "w1", num_experts, ) + rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped} (w3)") load_fused_expert_weights( name_mapped, params_dict, @@ -1259,6 +1275,7 @@ def load_fused_expert_weights( num_experts, ) else: + rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped}") load_fused_expert_weights( name_mapped, params_dict, @@ -1273,6 +1290,7 @@ def load_fused_expert_weights( and name_mapped not in params_dict ): continue + rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped}") param = params_dict[name_mapped] # We should ask the weight loader to return success or # not here since otherwise we may skip experts with @@ -1297,6 +1315,7 @@ def load_fused_expert_weights( continue if name in params_dict.keys(): + rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader @@ -1466,6 +1485,15 @@ def __init__( self.deepstack_visual_indexes = self.visual.deepstack_visual_indexes + def _get_num_fused_shared_experts(self): + if not ( + hasattr(self.model, "layers") + and len(self.model.layers) > 0 + and hasattr(self.model.layers[0].mlp, "num_fused_shared_experts") + ): + return 0 + return self.model.layers[0].mlp.num_fused_shared_experts + def get_embed_and_head(self): embed = self.model.embed_tokens.weight if self.pp_group.is_first_rank else None head = self.lm_head.weight if self.pp_group.is_last_rank else None @@ -1496,13 +1524,17 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("in_proj_ba.", "in_proj_a.", 1), ] + text_config = getattr(self.config, "text_config", self.config) + num_experts_base = text_config.num_experts + num_fused_shared_experts = self._get_num_fused_shared_experts() + # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) expert_params_mapping = FusedMoE.make_expert_params_mapping( ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, + num_experts=num_experts_base + num_fused_shared_experts, ) # Skip loading extra parameters for GPTQ/modelopt models. @@ -1522,8 +1554,37 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), ("experts.w2_weight", "experts.down_proj", 0, "w2"), ] - - num_experts = self.config.num_experts + if num_fused_shared_experts > 0: + # Shared expert: checkpoint may use experts.512.gate_proj/up_proj/down_proj + # (separate) or experts.512.gate_up_proj (combined) + # param_name uses "experts.w13_" / "experts.w2_" so replace yields experts.w13_weight + fused_expert_params_mapping += [ + ( + "experts.w13_", + f"experts.{num_experts_base}.gate_proj.", + num_experts_base, + "w1", + ), + ( + "experts.w13_", + f"experts.{num_experts_base}.up_proj.", + num_experts_base, + "w3", + ), + ( + "experts.w13_", + f"experts.{num_experts_base}.gate_up_proj.", + num_experts_base, + "w1", + ), + ( + "experts.w2_", + f"experts.{num_experts_base}.down_proj.", + num_experts_base, + "w2", + ), + ] + num_experts = num_experts_base + num_fused_shared_experts def load_fused_expert_weights( name: str, @@ -1552,6 +1613,7 @@ def load_fused_expert_weights( params_dict = dict(self.named_parameters(remove_duplicate=False)) for name, loaded_weight in weights: + raw_weight_name = name if "rotary_emb.inv_freq" in name: continue if "mtp" in name: @@ -1580,6 +1642,15 @@ def load_fused_expert_weights( ): continue + # Remap shared expert to fused expert index when shared experts are fused + if num_fused_shared_experts > 0 and "mlp.shared_expert." in name: + # Only replace shared_expert submodule params (gate_up_proj, down_proj), + # not shared_expert_gate (which would incorrectly become experts.512_gate). + name = name.replace( + "mlp.shared_expert.", + f"mlp.experts.{num_experts_base}.", + ) + for param_name, weight_name, shard_id in stacked_params_mapping: if name.endswith("experts.gate_up_proj") or name.endswith( "experts.down_proj" @@ -1609,6 +1680,7 @@ def load_fused_expert_weights( if name not in params_dict: continue + rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) @@ -1635,23 +1707,39 @@ def load_fused_expert_weights( params_dict, loaded_weight[0], "w1", - num_experts, + num_experts_base, ) + rank0_log(f"load_weights: {raw_weight_name}[0] -> {name_mapped} (w1) expert_id: {expert_id}") load_fused_expert_weights( name_mapped, params_dict, loaded_weight[1], "w3", - num_experts, + num_experts_base, ) - else: + rank0_log(f"load_weights: {raw_weight_name}[1] -> {name_mapped} (w3) expert_id: {expert_id}") + elif "experts.down_proj" in name: load_fused_expert_weights( name_mapped, params_dict, loaded_weight, shard_id, - num_experts, + num_experts_base, ) + rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped} ({shard_id}) expert_id: {expert_id}") + else: + param = params_dict[name_mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + param = params_dict[name_mapped] + if f"{num_experts_base}.gate_up_proj" in name: + loaded_weight = loaded_weight.chunk(2, dim=-2) + weight_loader(param, loaded_weight[0], name_mapped, "w1", expert_id) + rank0_log(f"load_weights: {raw_weight_name}[0] -> {name_mapped} (w1) expert_id: {expert_id}") + weight_loader(param, loaded_weight[1], name_mapped, "w3", expert_id) + rank0_log(f"load_weights: {raw_weight_name}[1] -> {name_mapped} (w3) expert_id: {expert_id}") + else: + weight_loader(param, loaded_weight, name_mapped, shard_id, expert_id) + rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped} ({shard_id}) expert_id: {expert_id}") else: # Skip loading extra parameters for GPTQ models. if ( @@ -1659,6 +1747,7 @@ def load_fused_expert_weights( and name_mapped not in params_dict ): continue + rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped}") param = params_dict[name_mapped] # We should ask the weight loader to return success or # not here since otherwise we may skip experts with From e9609cf15797a7365e504667339593eaad286506 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 04/40] Remove redundant logging statements in Qwen3 MoE model weight loading process to streamline code and improve readability. --- python/sglang/srt/models/qwen3_5.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 8c4ec8d3219c..c92e3c0b4499 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -91,7 +91,6 @@ is_gfx95_supported, is_npu, make_layers, - rank0_log, set_weight_attrs, ) from sglang.srt.utils.hf_transformers_utils import get_processor, get_rope_config @@ -1238,7 +1237,6 @@ def load_fused_expert_weights( if name not in params_dict: continue - rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) @@ -1258,7 +1256,6 @@ def load_fused_expert_weights( if is_fused_expert: if "experts.gate_up_proj" in name: loaded_weight = loaded_weight.chunk(2, dim=-2) - rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped} (w1)") load_fused_expert_weights( name_mapped, params_dict, @@ -1266,7 +1263,6 @@ def load_fused_expert_weights( "w1", num_experts, ) - rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped} (w3)") load_fused_expert_weights( name_mapped, params_dict, @@ -1275,7 +1271,6 @@ def load_fused_expert_weights( num_experts, ) else: - rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped}") load_fused_expert_weights( name_mapped, params_dict, @@ -1290,7 +1285,6 @@ def load_fused_expert_weights( and name_mapped not in params_dict ): continue - rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped}") param = params_dict[name_mapped] # We should ask the weight loader to return success or # not here since otherwise we may skip experts with @@ -1315,7 +1309,6 @@ def load_fused_expert_weights( continue if name in params_dict.keys(): - rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader @@ -1680,7 +1673,6 @@ def load_fused_expert_weights( if name not in params_dict: continue - rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) @@ -1709,7 +1701,6 @@ def load_fused_expert_weights( "w1", num_experts_base, ) - rank0_log(f"load_weights: {raw_weight_name}[0] -> {name_mapped} (w1) expert_id: {expert_id}") load_fused_expert_weights( name_mapped, params_dict, @@ -1717,7 +1708,6 @@ def load_fused_expert_weights( "w3", num_experts_base, ) - rank0_log(f"load_weights: {raw_weight_name}[1] -> {name_mapped} (w3) expert_id: {expert_id}") elif "experts.down_proj" in name: load_fused_expert_weights( name_mapped, @@ -1726,7 +1716,6 @@ def load_fused_expert_weights( shard_id, num_experts_base, ) - rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped} ({shard_id}) expert_id: {expert_id}") else: param = params_dict[name_mapped] weight_loader = getattr(param, "weight_loader", default_weight_loader) @@ -1734,12 +1723,9 @@ def load_fused_expert_weights( if f"{num_experts_base}.gate_up_proj" in name: loaded_weight = loaded_weight.chunk(2, dim=-2) weight_loader(param, loaded_weight[0], name_mapped, "w1", expert_id) - rank0_log(f"load_weights: {raw_weight_name}[0] -> {name_mapped} (w1) expert_id: {expert_id}") weight_loader(param, loaded_weight[1], name_mapped, "w3", expert_id) - rank0_log(f"load_weights: {raw_weight_name}[1] -> {name_mapped} (w3) expert_id: {expert_id}") else: weight_loader(param, loaded_weight, name_mapped, shard_id, expert_id) - rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped} ({shard_id}) expert_id: {expert_id}") else: # Skip loading extra parameters for GPTQ models. if ( @@ -1747,7 +1733,6 @@ def load_fused_expert_weights( and name_mapped not in params_dict ): continue - rank0_log(f"load_weights: {raw_weight_name} -> {name_mapped}") param = params_dict[name_mapped] # We should ask the weight loader to return success or # not here since otherwise we may skip experts with From 5d3334805f923b0e496cec4c0dab1a6b672c72ee Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 05/40] Add support for Aiter in Qwen2 MoE model - Introduced a new boolean environment variable to control Aiter usage. - Updated Qwen2MoeSparseMoeBlock to include support for shared experts. - Added checks to ensure compatibility with Aiter and prevent failures on unsupported platforms. --- python/sglang/srt/models/qwen2_moe.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 4cfd9af52e97..d7e9ad74dad2 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -89,6 +89,8 @@ cpu_has_amx_support, is_cpu, is_cuda, + get_bool_env_var, + is_hip, make_layers, rank0_log, use_intel_amx_backend, @@ -100,6 +102,8 @@ _is_cuda = is_cuda() _is_cpu = is_cpu() _is_cpu_amx_available = cpu_has_amx_support() +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip class Qwen2MoeMLP(nn.Module): @@ -163,7 +167,9 @@ def __init__( alt_stream: Optional[torch.cuda.Stream] = None, prefix: str = "", is_nextn: bool = False, + support_shared_experts: bool = False, ): + # By default, fuse shared experts to routed experts are not supported except Qwen3.5 MoE for now. super().__init__() self.tp_size = get_tensor_model_parallel_world_size() self.layer_id = layer_id @@ -260,6 +266,9 @@ def _determine_num_fused_shared_experts(self, config: PretrainedConfig) -> int: return 0 if config.shared_expert_intermediate_size != config.moe_intermediate_size: return 0 + if not _use_aiter: + # Only Aiter is tested for now. This is to avoid failure on other platform. + return 0 return 1 def get_moe_weights(self): From 9f247c4360c5774c684373c2dbf7631f1833b38b Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 06/40] Add support for shared expert fusion in Qwen3 decoder layers - Updated Qwen3_5LinearDecoderLayer and Qwen3_5AttentionDecoderLayer to include support for shared expert fusion. - Enhanced the initialization parameters to accommodate the new functionality. --- python/sglang/srt/models/qwen3_5.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index c92e3c0b4499..4be175a3ef7a 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -528,6 +528,7 @@ def __init__( alt_stream=alt_stream, prefix=add_prefix("mlp", prefix.replace(".linear_attn", "")), is_nextn=is_nextn, + support_shared_expert_fusion=True, ) is_layer_sparse = True is_previous_layer_sparse = True @@ -732,6 +733,7 @@ def __init__( alt_stream=alt_stream, prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), is_nextn=is_nextn, + support_shared_expert_fusion=True, ) is_layer_sparse = True is_previous_layer_sparse = True From 045e50ea0a378200d6d6cfd7c81ab46141df6e1b Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 07/40] Refactor support for shared expert fusion in Qwen2 MoE block - Renamed parameter for clarity from `support_shared_experts` to `support_shared_expert_fusion`. - Updated logging messages to generalize shared expert fusion status. - Added a check in `_determine_num_fused_shared_experts` to respect the new parameter for controlling fusion behavior. --- python/sglang/srt/models/qwen2_moe.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index d7e9ad74dad2..34cbcf198fd9 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -167,10 +167,11 @@ def __init__( alt_stream: Optional[torch.cuda.Stream] = None, prefix: str = "", is_nextn: bool = False, - support_shared_experts: bool = False, + support_shared_expert_fusion: bool = False, ): # By default, fuse shared experts to routed experts are not supported except Qwen3.5 MoE for now. super().__init__() + self.support_shared_expert_fusion = support_shared_expert_fusion self.tp_size = get_tensor_model_parallel_world_size() self.layer_id = layer_id self.num_experts = config.num_experts @@ -186,8 +187,7 @@ def __init__( ) if self.num_fused_shared_experts > 0: rank0_log( - "Shared experts fusion enabled for Qwen2 MoE " - "(topk+1 experts per token)." + "Shared experts fusion enabled (topk+1 experts per token)." ) self.topk = TopK( @@ -258,8 +258,11 @@ def __init__( def _determine_num_fused_shared_experts(self, config: PretrainedConfig) -> int: """Determine if shared expert can be fused with router experts (topk+1). - Fusion requires shared_expert_intermediate_size == moe_intermediate_size. + Fusion requires shared_expert_intermediate_size == moe_intermediate_size, + support_shared_expert_fusion=True (e.g. Qwen3.5 MoE), and Aiter on HIP. """ + if not self.support_shared_expert_fusion: + return 0 if get_global_server_args().disable_shared_experts_fusion: return 0 if getattr(config, "shared_expert_intermediate_size", 0) <= 0: From a5e909e0144e5392c4bd9ef85209414066fb796e Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 08/40] Refactor shared expert fusion checks in Qwen2 MoE block - Consolidated multiple conditions into a single if statement for clarity and efficiency. - Added a check for `quant_config` to enhance compatibility with quantization settings. - Streamlined the logic to determine the number of fused shared experts based on updated parameters. --- python/sglang/srt/models/qwen2_moe.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 34cbcf198fd9..a329ae40921e 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -261,16 +261,13 @@ def _determine_num_fused_shared_experts(self, config: PretrainedConfig) -> int: Fusion requires shared_expert_intermediate_size == moe_intermediate_size, support_shared_expert_fusion=True (e.g. Qwen3.5 MoE), and Aiter on HIP. """ - if not self.support_shared_expert_fusion: - return 0 - if get_global_server_args().disable_shared_experts_fusion: - return 0 - if getattr(config, "shared_expert_intermediate_size", 0) <= 0: - return 0 - if config.shared_expert_intermediate_size != config.moe_intermediate_size: - return 0 - if not _use_aiter: - # Only Aiter is tested for now. This is to avoid failure on other platform. + if (not self.support_shared_expert_fusion + or get_global_server_args().disable_shared_experts_fusion is True + or getattr(config, "shared_expert_intermediate_size", 0) <= 0 + or config.shared_expert_intermediate_size != config.moe_intermediate_size + or not _use_aiter + or getattr(config, "quant_config", None) is not None + ): return 0 return 1 From 17419e41a3e868b29b09257b525fe1b13691a220 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 09/40] Enhance shared expert fusion logic in Qwen2 MoE block - Updated `_determine_num_fused_shared_experts` method to accept `quant_config` as an additional parameter for improved flexibility. - Adjusted conditions to check for `quant_config` directly, streamlining the logic for determining the number of fused shared experts. - Ensured compatibility with quantization settings in the expert fusion process. --- python/sglang/srt/models/qwen2_moe.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index a329ae40921e..b12987bb5ddf 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -183,7 +183,8 @@ def __init__( ) self.num_fused_shared_experts = self._determine_num_fused_shared_experts( - config + config, + quant_config, ) if self.num_fused_shared_experts > 0: rank0_log( @@ -256,7 +257,7 @@ def __init__( self.top_k = config.num_experts_per_tok self.is_nextn = is_nextn - def _determine_num_fused_shared_experts(self, config: PretrainedConfig) -> int: + def _determine_num_fused_shared_experts(self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig]) -> int: """Determine if shared expert can be fused with router experts (topk+1). Fusion requires shared_expert_intermediate_size == moe_intermediate_size, support_shared_expert_fusion=True (e.g. Qwen3.5 MoE), and Aiter on HIP. @@ -266,7 +267,7 @@ def _determine_num_fused_shared_experts(self, config: PretrainedConfig) -> int: or getattr(config, "shared_expert_intermediate_size", 0) <= 0 or config.shared_expert_intermediate_size != config.moe_intermediate_size or not _use_aiter - or getattr(config, "quant_config", None) is not None + or quant_config is not None ): return 0 return 1 From 02b36a5f51c548849746f482d1bfa9abf20e918a Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 10/40] Remove unused `quant_config` check in Qwen2 MoE block logic - Eliminated the condition checking for `quant_config` in the `_determine_num_fused_shared_experts` method to simplify the logic. - This change streamlines the decision-making process for determining the number of fused shared experts, enhancing code clarity. --- python/sglang/srt/models/qwen2_moe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index b12987bb5ddf..f91b545a57d1 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -267,7 +267,6 @@ def _determine_num_fused_shared_experts(self, config: PretrainedConfig, quant_co or getattr(config, "shared_expert_intermediate_size", 0) <= 0 or config.shared_expert_intermediate_size != config.moe_intermediate_size or not _use_aiter - or quant_config is not None ): return 0 return 1 From c026e29faa7fecee18973932b90911fee4fa6e16 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 11/40] Remove redundant logging statements in weight loading process of Qwen2 MoE model - Eliminated unnecessary logging of raw weight names during the weight loading process to enhance code clarity and reduce verbosity. - This change streamlines the weight loading logic while maintaining functionality. --- python/sglang/srt/models/qwen2_moe.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index f91b545a57d1..28493969a164 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -930,7 +930,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: - raw_weight_name = name layer_id = get_layer_id(name) if ( layer_id is not None @@ -962,7 +961,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): if name not in params_dict: continue - rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) @@ -973,7 +971,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): if weight_name not in name: continue name = name.replace(weight_name, param_name) - rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = param.weight_loader weight_loader( @@ -992,7 +989,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): continue if name in params_dict.keys(): - rank0_log(f"load_weights: {raw_weight_name} -> {name}") param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader From e474ed8bd91b67f935b3155df7b55dd844b52c37 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 12/40] disable fp8/fp4 --- python/sglang/srt/models/qwen2_moe.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 28493969a164..9c05323c64a0 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -267,6 +267,7 @@ def _determine_num_fused_shared_experts(self, config: PretrainedConfig, quant_co or getattr(config, "shared_expert_intermediate_size", 0) <= 0 or config.shared_expert_intermediate_size != config.moe_intermediate_size or not _use_aiter + or quant_config is not None ): return 0 return 1 From 44d6f8d1a572330d84f4c721095d7a26cfbf2741 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 13/40] revert changes in linear --- python/sglang/srt/layers/linear.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 0c7b80760779..7af9eb004008 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -266,11 +266,7 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): param.dtype == loaded_weight.dtype ), "init para dtype and loaded weight dtype should be the same" - if param.size() != loaded_weight.size(): - raise ValueError( - f"Weight shape mismatch: param {param.shape} vs loaded {loaded_weight.shape}. " - f"param.numel()={param.numel()}, loaded.numel()={loaded_weight.numel()}" - ) + assert param.size() == loaded_weight.size() param.data.copy_(loaded_weight) def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: From f61e947f439e83f791e17097e8bc9865b73a7b31 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 14/40] Remove unused logging statement in Qwen2 MoE block - Eliminated the `rank0_log` statement related to shared experts fusion to reduce verbosity and enhance code clarity. - This change simplifies the logging process while maintaining the functionality of the Qwen2 MoE model. --- python/sglang/srt/models/qwen2_moe.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 9c05323c64a0..454fc203f077 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -92,7 +92,6 @@ get_bool_env_var, is_hip, make_layers, - rank0_log, use_intel_amx_backend, ) from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -186,10 +185,6 @@ def __init__( config, quant_config, ) - if self.num_fused_shared_experts > 0: - rank0_log( - "Shared experts fusion enabled (topk+1 experts per token)." - ) self.topk = TopK( top_k=config.num_experts_per_tok, From f394d45307df7581923b20ed2979fe9a6c14a798 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 15/40] fix lint error --- python/sglang/srt/models/qwen2_moe.py | 23 ++++++++++++---------- python/sglang/srt/models/qwen3_5.py | 28 +++++++++++++++++++++++---- 2 files changed, 37 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 454fc203f077..55e17b6e396e 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -87,9 +87,9 @@ from sglang.srt.utils import ( add_prefix, cpu_has_amx_support, + get_bool_env_var, is_cpu, is_cuda, - get_bool_env_var, is_hip, make_layers, use_intel_amx_backend, @@ -252,12 +252,15 @@ def __init__( self.top_k = config.num_experts_per_tok self.is_nextn = is_nextn - def _determine_num_fused_shared_experts(self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig]) -> int: + def _determine_num_fused_shared_experts( + self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] + ) -> int: """Determine if shared expert can be fused with router experts (topk+1). Fusion requires shared_expert_intermediate_size == moe_intermediate_size, support_shared_expert_fusion=True (e.g. Qwen3.5 MoE), and Aiter on HIP. """ - if (not self.support_shared_expert_fusion + if ( + not self.support_shared_expert_fusion or get_global_server_args().disable_shared_experts_fusion is True or getattr(config, "shared_expert_intermediate_size", 0) <= 0 or config.shared_expert_intermediate_size != config.moe_intermediate_size @@ -304,11 +307,13 @@ def _append_shared_to_topk_output( dtype=topk_output.topk_ids.dtype, device=topk_output.topk_ids.device, ) - shared_weights = shared_weights.expand( - M, self.num_fused_shared_experts - ).to(topk_output.topk_weights.dtype) + shared_weights = shared_weights.expand(M, self.num_fused_shared_experts).to( + topk_output.topk_weights.dtype + ) fused_topk_ids = torch.cat([topk_output.topk_ids, shared_ids], dim=-1) - fused_topk_weights = torch.cat([topk_output.topk_weights, shared_weights], dim=-1) + fused_topk_weights = torch.cat( + [topk_output.topk_weights, shared_weights], dim=-1 + ) return StandardTopKOutput( topk_weights=fused_topk_weights, topk_ids=fused_topk_ids, @@ -372,9 +377,7 @@ def _forward_router_experts(self, hidden_states: torch.Tensor): if self.num_fused_shared_experts > 0 and TopKOutputChecker.format_is_standard( topk_output ): - topk_output = self._append_shared_to_topk_output( - topk_output, hidden_states - ) + topk_output = self._append_shared_to_topk_output(topk_output, hidden_states) return self.experts(hidden_states, topk_output) def forward_normal_dual_stream( diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 4be175a3ef7a..f323031c9800 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1720,14 +1720,34 @@ def load_fused_expert_weights( ) else: param = params_dict[name_mapped] - weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) param = params_dict[name_mapped] if f"{num_experts_base}.gate_up_proj" in name: loaded_weight = loaded_weight.chunk(2, dim=-2) - weight_loader(param, loaded_weight[0], name_mapped, "w1", expert_id) - weight_loader(param, loaded_weight[1], name_mapped, "w3", expert_id) + weight_loader( + param, + loaded_weight[0], + name_mapped, + "w1", + expert_id, + ) + weight_loader( + param, + loaded_weight[1], + name_mapped, + "w3", + expert_id, + ) else: - weight_loader(param, loaded_weight, name_mapped, shard_id, expert_id) + weight_loader( + param, + loaded_weight, + name_mapped, + shard_id, + expert_id, + ) else: # Skip loading extra parameters for GPTQ models. if ( From c5072c4bf9fd94256e0a6854cefd1722457a2e87 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 16/40] Update router logic in Qwen2 MoE block - Renamed `gate_logits` to `router_logits` for clarity in the `_forward_router_experts` method. - This change improves code readability while maintaining the functionality of the expert routing process. --- python/sglang/srt/models/qwen2_moe.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 55e17b6e396e..0b96dd16921d 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -372,8 +372,9 @@ def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBat return final_hidden_states def _forward_router_experts(self, hidden_states: torch.Tensor): - gate_logits, _ = self.gate(hidden_states) - topk_output = self.topk(hidden_states, gate_logits) + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) if self.num_fused_shared_experts > 0 and TopKOutputChecker.format_is_standard( topk_output ): From 0221fbae0a206202a9044dcd9e4ba610b637b55b Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 17/40] Add check for deep expert backend in Qwen2 MoE block - Introduced a condition in the `_determine_num_fused_shared_experts` method to check if the backend is a deep expert using `get_moe_a2a_backend().is_deepep()`. - This enhancement improves the logic for determining the number of fused shared experts based on the backend type, ensuring better compatibility with deep expert configurations. --- python/sglang/srt/models/qwen2_moe.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 0b96dd16921d..32d932362b17 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -266,6 +266,7 @@ def _determine_num_fused_shared_experts( or config.shared_expert_intermediate_size != config.moe_intermediate_size or not _use_aiter or quant_config is not None + or get_moe_a2a_backend().is_deepep() ): return 0 return 1 From e73b78a5d180e233d168df6f1d6e2275d5d0c0c6 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 18/40] Refactor expert routing logic in Qwen2 MoE block - Updated the `_forward_router_experts` method to improve the handling of router logits, enhancing clarity and maintainability. - This change ensures better alignment with the overall architecture of the Qwen2 MoE model while preserving existing functionality. --- python/sglang/srt/models/qwen3_5.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index f323031c9800..b257e9eeb345 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1194,7 +1194,6 @@ def load_fused_expert_weights( params_dict = dict(self.named_parameters(remove_duplicate=False)) for name, loaded_weight in weights: - raw_weight_name = name if "rotary_emb.inv_freq" in name: continue if "mtp" in name: From 2d0d51ed4569191097370e884ff3cacd4a9665a1 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 19/40] Remove unused variable in weight loading process of Qwen3.5 MoE model - Eliminated the `raw_weight_name` variable from the weight loading loop to enhance code clarity and reduce unnecessary complexity. - This change streamlines the weight loading logic while maintaining existing functionality. --- python/sglang/srt/models/qwen3_5.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index b257e9eeb345..3f5c8a32e5a3 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1607,7 +1607,6 @@ def load_fused_expert_weights( params_dict = dict(self.named_parameters(remove_duplicate=False)) for name, loaded_weight in weights: - raw_weight_name = name if "rotary_emb.inv_freq" in name: continue if "mtp" in name: From fbb68ad0691df75a2d7e109f1acc645a6d908f4d Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 20/40] revert changes in Qwen3_5MoeForCausalLM --- python/sglang/srt/models/qwen3_5.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 3f5c8a32e5a3..b4a5847e9f40 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1110,14 +1110,6 @@ def __init__( ) -> None: super().__init__(config=config, quant_config=quant_config, prefix=prefix) - def _get_num_fused_shared_experts(self): - if not ( - hasattr(self.model, "layers") - and len(self.model.layers) > 0 - and hasattr(self.model.layers[0].mlp, "num_fused_shared_experts") - ): - return 0 - return self.model.layers[0].mlp.num_fused_shared_experts def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ @@ -1134,9 +1126,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("in_proj_ba.", "in_proj_a.", 1), ] - num_experts_base = self.config.num_experts - num_fused_shared_experts = self._get_num_fused_shared_experts() - num_experts = num_experts_base + num_fused_shared_experts # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) @@ -1144,7 +1133,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=num_experts, + num_experts=self.config.num_experts, ) # Skip loading extra parameters for GPTQ/modelopt models. From 76f111feea200dbf750e83b2fd46bf98b365f2cc Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 21/40] reformat --- python/sglang/srt/models/qwen3_5.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index b4a5847e9f40..c264cb29f50f 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1110,7 +1110,6 @@ def __init__( ) -> None: super().__init__(config=config, quant_config=quant_config, prefix=prefix) - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ # (param_name, shard_name, shard_id) @@ -1126,7 +1125,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("in_proj_ba.", "in_proj_a.", 1), ] - # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) expert_params_mapping = FusedMoE.make_expert_params_mapping( From d8267ae00ad966cad99f3a87cbc46c8a6f96b183 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 22/40] Add num_experts variable initialization in Qwen3.5 MoE model - Introduced the `num_experts` variable to store the number of experts from the model configuration, enhancing clarity for future weight loading processes. - This change prepares the code for potential future modifications related to expert handling while maintaining existing functionality. --- python/sglang/srt/models/qwen3_5.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index c264cb29f50f..047f4ac59fbd 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1153,7 +1153,9 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), ("experts.w2_weight", "experts.down_proj", 0, "w2"), ] - + + num_experts = self.config.num_experts + def load_fused_expert_weights( name: str, params_dict: dict, From 97b9615aa0d1894c5c6b8947e9cd4540d5bedf4c Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 23/40] fix lint --- python/sglang/srt/models/qwen3_5.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 047f4ac59fbd..5f7596a84aa5 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1153,9 +1153,9 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), ("experts.w2_weight", "experts.down_proj", 0, "w2"), ] - + num_experts = self.config.num_experts - + def load_fused_expert_weights( name: str, params_dict: dict, From 4193219ce2fc86e91f54935386ae3da37646889e Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 24/40] [AMD] Enable fused_moe refactor --- python/sglang/srt/models/qwen2_moe.py | 53 ++++++++++++------------ python/sglang/srt/models/qwen3_5.py | 59 +++++++++++++++++++-------- 2 files changed, 68 insertions(+), 44 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 32d932362b17..cfadf9ee0c41 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -168,23 +168,22 @@ def __init__( is_nextn: bool = False, support_shared_expert_fusion: bool = False, ): - # By default, fuse shared experts to routed experts are not supported except Qwen3.5 MoE for now. super().__init__() - self.support_shared_expert_fusion = support_shared_expert_fusion self.tp_size = get_tensor_model_parallel_world_size() self.layer_id = layer_id - self.num_experts = config.num_experts self.alt_stream = alt_stream if self.tp_size > config.num_experts: raise ValueError( f"Tensor parallel size {self.tp_size} is greater than " f"the number of experts {config.num_experts}." ) - - self.num_fused_shared_experts = self._determine_num_fused_shared_experts( - config, - quant_config, - ) + self.num_experts = config.num_experts + self.num_fused_shared_experts = 0 + self.enable_fused_moe = True if (support_shared_expert_fusion and _use_aiter) else False + if self.enable_fused_moe: + self.num_fused_shared_experts = self._determine_num_fused_shared_experts( + config, quant_config + ) self.topk = TopK( top_k=config.num_experts_per_tok, @@ -194,10 +193,18 @@ def __init__( self.experts = get_moe_impl_class(quant_config)( layer_id=self.layer_id, - top_k=config.num_experts_per_tok + self.num_fused_shared_experts, - num_experts=config.num_experts - + self.num_fused_shared_experts - + get_global_server_args().ep_num_redundant_experts, + top_k=( + config.num_experts_per_tok + if not self.enable_fused_moe + else config.num_experts_per_tok + self.num_fused_shared_experts + ), + num_experts=( + config.num_experts + get_global_server_args().ep_num_redundant_experts + if not self.enable_fused_moe + else config.num_experts + + get_global_server_args().ep_num_redundant_experts + + self.num_fused_shared_experts + ), hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, quant_config=quant_config, @@ -213,10 +220,7 @@ def __init__( quant_config=None, prefix=add_prefix("gate", prefix), ) - if ( - config.shared_expert_intermediate_size > 0 - and self.num_fused_shared_experts == 0 - ): + if config.shared_expert_intermediate_size > 0: self.shared_expert = Qwen2MoeMLP( hidden_size=config.hidden_size, intermediate_size=config.shared_expert_intermediate_size, @@ -256,19 +260,18 @@ def _determine_num_fused_shared_experts( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] ) -> int: """Determine if shared expert can be fused with router experts (topk+1). - Fusion requires shared_expert_intermediate_size == moe_intermediate_size, - support_shared_expert_fusion=True (e.g. Qwen3.5 MoE), and Aiter on HIP. + Fusion requires matching intermediate sizes, a supported backend, no + quantization, and no DeepEP. """ if ( - not self.support_shared_expert_fusion - or get_global_server_args().disable_shared_experts_fusion is True + get_global_server_args().disable_shared_experts_fusion is True or getattr(config, "shared_expert_intermediate_size", 0) <= 0 or config.shared_expert_intermediate_size != config.moe_intermediate_size - or not _use_aiter - or quant_config is not None or get_moe_a2a_backend().is_deepep() ): + # Falled back to non-fused version, return 0 shared_experts. return 0 + # Return 1 shared_expert for qwen model. return 1 def get_moe_weights(self): @@ -301,7 +304,7 @@ def _append_shared_to_topk_output( if shared_weights is None: return topk_output M = topk_output.topk_ids.shape[0] - shared_expert_id = self.num_experts # 512 for Qwen3.5 MoE + shared_expert_id = self.num_experts shared_ids = torch.full( (M, self.num_fused_shared_experts), shared_expert_id, @@ -376,9 +379,7 @@ def _forward_router_experts(self, hidden_states: torch.Tensor): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) - if self.num_fused_shared_experts > 0 and TopKOutputChecker.format_is_standard( - topk_output - ): + if self.enable_fused_moe and TopKOutputChecker.format_is_standard(topk_output): topk_output = self._append_shared_to_topk_output(topk_output, hidden_states) return self.experts(hidden_states, topk_output) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 5f7596a84aa5..4385c1c75f60 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -86,9 +86,11 @@ LazyValue, add_prefix, cpu_has_amx_support, + get_bool_env_var, is_cpu, is_cuda, is_gfx95_supported, + is_hip, is_npu, make_layers, set_weight_attrs, @@ -100,6 +102,8 @@ _is_npu = is_npu() _is_cpu = is_cpu() _is_gfx95 = is_gfx95_supported() +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_amx_available = cpu_has_amx_support() @@ -1467,6 +1471,10 @@ def __init__( self.is_mrope_enabled = "mrope_section" in rope_config self.deepstack_visual_indexes = self.visual.deepstack_visual_indexes + self.enable_fused_moe = True if _use_aiter else False + self.num_fused_shared_experts = ( + 0 if not enable_fused_moe else self._get_num_fused_shared_experts() + ) def _get_num_fused_shared_experts(self): if not ( @@ -1509,7 +1517,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): text_config = getattr(self.config, "text_config", self.config) num_experts_base = text_config.num_experts - num_fused_shared_experts = self._get_num_fused_shared_experts() # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) @@ -1517,7 +1524,11 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=num_experts_base + num_fused_shared_experts, + num_experts=( + self.config.num_experts + if not self.enable_fused_moe + else num_experts_base + self.num_fused_shared_experts + ), ) # Skip loading extra parameters for GPTQ/modelopt models. @@ -1537,10 +1548,9 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), ("experts.w2_weight", "experts.down_proj", 0, "w2"), ] - if num_fused_shared_experts > 0: - # Shared expert: checkpoint may use experts.512.gate_proj/up_proj/down_proj - # (separate) or experts.512.gate_up_proj (combined) - # param_name uses "experts.w13_" / "experts.w2_" so replace yields experts.w13_weight + + num_experts = self.config.num_experts + if self.enable_fused_moe: fused_expert_params_mapping += [ ( "experts.w13_", @@ -1567,7 +1577,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): "w2", ), ] - num_experts = num_experts_base + num_fused_shared_experts + num_experts = num_experts_base + self.num_fused_shared_experts def load_fused_expert_weights( name: str, @@ -1624,14 +1634,12 @@ def load_fused_expert_weights( ): continue - # Remap shared expert to fused expert index when shared experts are fused - if num_fused_shared_experts > 0 and "mlp.shared_expert." in name: - # Only replace shared_expert submodule params (gate_up_proj, down_proj), - # not shared_expert_gate (which would incorrectly become experts.512_gate). - name = name.replace( - "mlp.shared_expert.", - f"mlp.experts.{num_experts_base}.", - ) + if self.enable_fused_moe: + if "mlp.shared_expert." in name: + name = name.replace( + "mlp.shared_expert.", + f"mlp.experts.{num_experts_base}.", + ) for param_name, weight_name, shard_id in stacked_params_mapping: if name.endswith("experts.gate_up_proj") or name.endswith( @@ -1688,14 +1696,24 @@ def load_fused_expert_weights( params_dict, loaded_weight[0], "w1", - num_experts_base, + ( + num_experts + if not self.enable_fused_moe + else num_experts_base + + self.num_fused_shared_experts + ), ) load_fused_expert_weights( name_mapped, params_dict, loaded_weight[1], "w3", - num_experts_base, + ( + num_experts + if not self.enable_fused_moe + else num_experts_base + + self.num_fused_shared_experts + ), ) elif "experts.down_proj" in name: load_fused_expert_weights( @@ -1703,7 +1721,12 @@ def load_fused_expert_weights( params_dict, loaded_weight, shard_id, - num_experts_base, + ( + num_experts + if not self.enable_fused_moe + else num_experts_base + + self.num_fused_shared_experts + ), ) else: param = params_dict[name_mapped] From 84fb0c325a04c3fd2bbcc428bd4b78705bfb051e Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 25/40] [Fix] Fix typo issue --- python/sglang/srt/models/qwen3_5.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 4385c1c75f60..dfb1d7dd4faa 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1473,7 +1473,7 @@ def __init__( self.deepstack_visual_indexes = self.visual.deepstack_visual_indexes self.enable_fused_moe = True if _use_aiter else False self.num_fused_shared_experts = ( - 0 if not enable_fused_moe else self._get_num_fused_shared_experts() + 0 if not self.enable_fused_moe else self._get_num_fused_shared_experts() ) def _get_num_fused_shared_experts(self): From 2c8ae8b7383d61d56a5d13bac238f642866adbcd Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 26/40] Add back condition to get rid of re-compute self.shared_expert in fused_moe case --- python/sglang/srt/models/qwen2_moe.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index cfadf9ee0c41..ecf99afe7929 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -220,7 +220,12 @@ def __init__( quant_config=None, prefix=add_prefix("gate", prefix), ) - if config.shared_expert_intermediate_size > 0: + # When enable_fused_moe is True, the shared expert is fused into the MoE + # kernel as an extra expert (appended via _append_shared_to_topk_output), + # so we must NOT create a separate shared_expert MLP here — otherwise + # _forward_shared_experts() would compute it a second time and its + # output would be double-counted in forward(). + if (config.shared_expert_intermediate_size > 0 and not self.enable_fused_moe): self.shared_expert = Qwen2MoeMLP( hidden_size=config.hidden_size, intermediate_size=config.shared_expert_intermediate_size, From 38d3b1c6e99805fa00ab22560069f26a9c992077 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 27/40] Remove unused expert weight mappings in Qwen3.5 MoE model - Eliminated redundant weight mappings for `gate_proj` and `up_proj` in the fused expert parameters, streamlining the weight loading process. - This change enhances code clarity and reduces complexity while maintaining existing functionality. --- python/sglang/srt/models/qwen3_5.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index dfb1d7dd4faa..d64d40295a77 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1552,18 +1552,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): num_experts = self.config.num_experts if self.enable_fused_moe: fused_expert_params_mapping += [ - ( - "experts.w13_", - f"experts.{num_experts_base}.gate_proj.", - num_experts_base, - "w1", - ), - ( - "experts.w13_", - f"experts.{num_experts_base}.up_proj.", - num_experts_base, - "w3", - ), ( "experts.w13_", f"experts.{num_experts_base}.gate_up_proj.", From 735fecdbe1c7c730d8ef02ec30e1eaaf0f44af14 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 28/40] Refactor expert handling in Qwen3.5 MoE model - Consolidated the initialization of the `num_experts` variable to improve clarity and consistency in weight loading processes. - Updated references to `num_experts` throughout the code to ensure accurate mapping of shared experts when fused, enhancing the overall functionality of the model. - Added comments to clarify the logic for loading fused expert weights, improving code maintainability. --- python/sglang/srt/models/qwen3_5.py | 45 +++++++++++++++++++---------- 1 file changed, 29 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index d64d40295a77..6a10addeb904 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1515,8 +1515,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("in_proj_ba.", "in_proj_a.", 1), ] - text_config = getattr(self.config, "text_config", self.config) - num_experts_base = text_config.num_experts + num_experts = self.config.num_experts # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) @@ -1525,9 +1524,9 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", num_experts=( - self.config.num_experts + num_experts if not self.enable_fused_moe - else num_experts_base + self.num_fused_shared_experts + else num_experts + self.num_fused_shared_experts # map shared experts to routed experts ), ) @@ -1549,23 +1548,27 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("experts.w2_weight", "experts.down_proj", 0, "w2"), ] - num_experts = self.config.num_experts if self.enable_fused_moe: + """ + When shared experts are fused, we need to map the shared experts to routed experts. + + mlp.share_expert.gate_up_proj.weight --> experts.512.gate_up_proj.weight -> experts.w13_weight, expert_id = 512 + mlp.share_expert.down_proj.weight --> experts.512.down_proj.weight -> experts.w2_weight, expert_id = 512 + """ fused_expert_params_mapping += [ ( "experts.w13_", - f"experts.{num_experts_base}.gate_up_proj.", - num_experts_base, + f"experts.{num_experts}.gate_up_proj.", + num_experts, "w1", ), ( "experts.w2_", - f"experts.{num_experts_base}.down_proj.", - num_experts_base, + f"experts.{num_experts}.down_proj.", + num_experts, "w2", ), ] - num_experts = num_experts_base + self.num_fused_shared_experts def load_fused_expert_weights( name: str, @@ -1624,9 +1627,10 @@ def load_fused_expert_weights( if self.enable_fused_moe: if "mlp.shared_expert." in name: + # Firstly map mlp.shared_expert.xx_proj to mlp.experts.512.xx_proj name = name.replace( "mlp.shared_expert.", - f"mlp.experts.{num_experts_base}.", + f"mlp.experts.{num_experts}.", ) for param_name, weight_name, shard_id in stacked_params_mapping: @@ -1677,7 +1681,10 @@ def load_fused_expert_weights( is_expert_weight = True name_mapped = name.replace(weight_name, param_name) if is_fused_expert: + # is_fused_expert is True, the checkpoint contains gate_up_proj and down_proj for each expert if "experts.gate_up_proj" in name: + # experts.gate_up_proj contains all 512 routed experts, excluding shared experts + # split into w1 and w3 loaded_weight = loaded_weight.chunk(2, dim=-2) load_fused_expert_weights( name_mapped, @@ -1687,7 +1694,7 @@ def load_fused_expert_weights( ( num_experts if not self.enable_fused_moe - else num_experts_base + else num_experts + self.num_fused_shared_experts ), ) @@ -1699,11 +1706,12 @@ def load_fused_expert_weights( ( num_experts if not self.enable_fused_moe - else num_experts_base + else num_experts + self.num_fused_shared_experts ), ) elif "experts.down_proj" in name: + # experts.down_proj contains all 512 routed experts, excluding shared experts load_fused_expert_weights( name_mapped, params_dict, @@ -1712,18 +1720,21 @@ def load_fused_expert_weights( ( num_experts if not self.enable_fused_moe - else num_experts_base + else num_experts + self.num_fused_shared_experts ), ) - else: + elif self.enable_fused_moe: + # shared experts should be loaded to experts.w13_weight and experts.w2_weight param = params_dict[name_mapped] weight_loader = getattr( param, "weight_loader", default_weight_loader ) param = params_dict[name_mapped] - if f"{num_experts_base}.gate_up_proj" in name: + if f"{num_experts}.gate_up_proj" in name: + # split into w1 and w3 loaded_weight = loaded_weight.chunk(2, dim=-2) + # load to experts.w13_weight, shard_id = w1, expert_id = 512 weight_loader( param, loaded_weight[0], @@ -1731,6 +1742,7 @@ def load_fused_expert_weights( "w1", expert_id, ) + # load to experts.w13_weight, shard_id = w3, expert_id = 512 weight_loader( param, loaded_weight[1], @@ -1739,6 +1751,7 @@ def load_fused_expert_weights( expert_id, ) else: + # load to experts.w2_weight, shard_id = w2, expert_id = 512 weight_loader( param, loaded_weight, From 35a3a29364381f77a02005934e3d910d34c34cd9 Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 29/40] Fix lint test --- python/sglang/srt/models/qwen2_moe.py | 6 ++++-- python/sglang/srt/models/qwen3_5.py | 14 ++++++-------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index ecf99afe7929..b3845930a68c 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -179,7 +179,9 @@ def __init__( ) self.num_experts = config.num_experts self.num_fused_shared_experts = 0 - self.enable_fused_moe = True if (support_shared_expert_fusion and _use_aiter) else False + self.enable_fused_moe = ( + True if (support_shared_expert_fusion and _use_aiter) else False + ) if self.enable_fused_moe: self.num_fused_shared_experts = self._determine_num_fused_shared_experts( config, quant_config @@ -225,7 +227,7 @@ def __init__( # so we must NOT create a separate shared_expert MLP here — otherwise # _forward_shared_experts() would compute it a second time and its # output would be double-counted in forward(). - if (config.shared_expert_intermediate_size > 0 and not self.enable_fused_moe): + if config.shared_expert_intermediate_size > 0 and not self.enable_fused_moe: self.shared_expert = Qwen2MoeMLP( hidden_size=config.hidden_size, intermediate_size=config.shared_expert_intermediate_size, diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 6a10addeb904..fa691d21d131 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1525,8 +1525,9 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ckpt_up_proj_name="up_proj", num_experts=( num_experts - if not self.enable_fused_moe - else num_experts + self.num_fused_shared_experts # map shared experts to routed experts + if not self.enable_fused_moe + else num_experts + + self.num_fused_shared_experts # map shared experts to routed experts ), ) @@ -1694,8 +1695,7 @@ def load_fused_expert_weights( ( num_experts if not self.enable_fused_moe - else num_experts - + self.num_fused_shared_experts + else num_experts + self.num_fused_shared_experts ), ) load_fused_expert_weights( @@ -1706,8 +1706,7 @@ def load_fused_expert_weights( ( num_experts if not self.enable_fused_moe - else num_experts - + self.num_fused_shared_experts + else num_experts + self.num_fused_shared_experts ), ) elif "experts.down_proj" in name: @@ -1720,8 +1719,7 @@ def load_fused_expert_weights( ( num_experts if not self.enable_fused_moe - else num_experts - + self.num_fused_shared_experts + else num_experts + self.num_fused_shared_experts ), ) elif self.enable_fused_moe: From 37666f28529a771a6abc1fcce1973b060d1d7d33 Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 30/40] Fix lint test --- python/sglang/srt/models/qwen2_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index b3845930a68c..b02597480fff 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -276,7 +276,7 @@ def _determine_num_fused_shared_experts( or config.shared_expert_intermediate_size != config.moe_intermediate_size or get_moe_a2a_backend().is_deepep() ): - # Falled back to non-fused version, return 0 shared_experts. + # Fall back to non-fused version, return 0 shared_experts. return 0 # Return 1 shared_expert for qwen model. return 1 From 1217df6954db30795cb3f9c14a427b8e097b998b Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 31/40] Refactor weight loading logic in Qwen3.5 MoE model - Simplified the weight loading process by removing conditional checks for `num_experts` related to fused MoE, ensuring a more straightforward implementation. - Enhanced code clarity and maintainability by streamlining the parameters passed during weight loading. --- python/sglang/srt/models/qwen3_5.py | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index fa691d21d131..3ec9c240e178 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1692,22 +1692,14 @@ def load_fused_expert_weights( params_dict, loaded_weight[0], "w1", - ( - num_experts - if not self.enable_fused_moe - else num_experts + self.num_fused_shared_experts - ), + num_experts, ) load_fused_expert_weights( name_mapped, params_dict, loaded_weight[1], "w3", - ( - num_experts - if not self.enable_fused_moe - else num_experts + self.num_fused_shared_experts - ), + num_experts, ) elif "experts.down_proj" in name: # experts.down_proj contains all 512 routed experts, excluding shared experts @@ -1716,11 +1708,7 @@ def load_fused_expert_weights( params_dict, loaded_weight, shard_id, - ( - num_experts - if not self.enable_fused_moe - else num_experts + self.num_fused_shared_experts - ), + num_experts ) elif self.enable_fused_moe: # shared experts should be loaded to experts.w13_weight and experts.w2_weight From 9a4673a4b4516eaa974aa89458a365d34c4f4403 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:51:20 -0500 Subject: [PATCH 32/40] Add shared expert fusion capability in Qwen2 MoE model - Introduced a new function `can_fuse_shared_expert` to determine if shared experts can be fused based on configuration and server arguments. - Updated the initialization of `enable_shared_expert_fusion` and `num_fused_shared_experts` to reflect the new fusion logic. - Refactored related code sections to ensure correct handling of shared experts during weight loading and processing, improving overall model functionality and maintainability. --- python/sglang/srt/models/qwen2_moe.py | 89 +++++++++++++-------------- python/sglang/srt/models/qwen3_5.py | 31 +++++++--- 2 files changed, 65 insertions(+), 55 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index b02597480fff..b6f796460709 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -105,6 +105,23 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +def can_fuse_shared_expert( + config: PretrainedConfig, +) -> bool: + """Whether the shared expert may be fused as an extra MoE expert (Qwen3.5 + Aiter). + + Caller must still gate on ``support_shared_expert_fusion`` and ``_use_aiter``. + """ + if ( + get_global_server_args().disable_shared_experts_fusion is True + or getattr(config, "shared_expert_intermediate_size", 0) <= 0 + or config.shared_expert_intermediate_size != config.moe_intermediate_size + or get_moe_a2a_backend().is_deepep() + ): + return False + return True + + class Qwen2MoeMLP(nn.Module): def __init__( self, @@ -178,14 +195,16 @@ def __init__( f"the number of experts {config.num_experts}." ) self.num_experts = config.num_experts - self.num_fused_shared_experts = 0 - self.enable_fused_moe = ( - True if (support_shared_expert_fusion and _use_aiter) else False + self.num_shared_experts = 1 + self.enable_shared_expert_fusion = ( + support_shared_expert_fusion + and _use_aiter + and can_fuse_shared_expert(config) + ) + + self.num_fused_shared_experts = ( + self.num_shared_experts if self.enable_shared_expert_fusion else 0 ) - if self.enable_fused_moe: - self.num_fused_shared_experts = self._determine_num_fused_shared_experts( - config, quant_config - ) self.topk = TopK( top_k=config.num_experts_per_tok, @@ -193,20 +212,13 @@ def __init__( layer_id=layer_id, ) + ep_redundant = get_global_server_args().ep_num_redundant_experts self.experts = get_moe_impl_class(quant_config)( layer_id=self.layer_id, - top_k=( - config.num_experts_per_tok - if not self.enable_fused_moe - else config.num_experts_per_tok + self.num_fused_shared_experts - ), - num_experts=( - config.num_experts + get_global_server_args().ep_num_redundant_experts - if not self.enable_fused_moe - else config.num_experts - + get_global_server_args().ep_num_redundant_experts - + self.num_fused_shared_experts - ), + top_k=config.num_experts_per_tok + self.num_fused_shared_experts, + num_experts=config.num_experts + + ep_redundant + + self.num_fused_shared_experts, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, quant_config=quant_config, @@ -222,12 +234,13 @@ def __init__( quant_config=None, prefix=add_prefix("gate", prefix), ) - # When enable_fused_moe is True, the shared expert is fused into the MoE - # kernel as an extra expert (appended via _append_shared_to_topk_output), - # so we must NOT create a separate shared_expert MLP here — otherwise - # _forward_shared_experts() would compute it a second time and its - # output would be double-counted in forward(). - if config.shared_expert_intermediate_size > 0 and not self.enable_fused_moe: + # When enable_shared_expert_fusion, the shared expert runs inside the MoE kernel + # (via _append_shared_to_topk_output); a separate shared_expert MLP would + # double-count. If fusion is off (num_fused_shared_experts == 0), keep shared_expert. + if ( + config.shared_expert_intermediate_size > 0 + and not self.enable_shared_expert_fusion + ): self.shared_expert = Qwen2MoeMLP( hidden_size=config.hidden_size, intermediate_size=config.shared_expert_intermediate_size, @@ -263,24 +276,6 @@ def __init__( self.top_k = config.num_experts_per_tok self.is_nextn = is_nextn - def _determine_num_fused_shared_experts( - self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] - ) -> int: - """Determine if shared expert can be fused with router experts (topk+1). - Fusion requires matching intermediate sizes, a supported backend, no - quantization, and no DeepEP. - """ - if ( - get_global_server_args().disable_shared_experts_fusion is True - or getattr(config, "shared_expert_intermediate_size", 0) <= 0 - or config.shared_expert_intermediate_size != config.moe_intermediate_size - or get_moe_a2a_backend().is_deepep() - ): - # Fall back to non-fused version, return 0 shared_experts. - return 0 - # Return 1 shared_expert for qwen model. - return 1 - def get_moe_weights(self): return [ x.data @@ -293,7 +288,7 @@ def get_moe_weights(self): def _get_shared_expert_weights(self, hidden_states: torch.Tensor) -> torch.Tensor: """Return sigmoid(shared_expert_gate) for fused shared expert weights.""" - if self.num_fused_shared_experts <= 0 or self.shared_expert_gate is None: + if not self.enable_shared_expert_fusion or self.shared_expert_gate is None: return None shared_out = self.shared_expert_gate(hidden_states) shared_logits = shared_out[0] if isinstance(shared_out, tuple) else shared_out @@ -305,7 +300,7 @@ def _append_shared_to_topk_output( hidden_states: torch.Tensor, ) -> StandardTopKOutput: """Append shared expert ids and weights to topk output before fused MoE.""" - if self.num_fused_shared_experts <= 0: + if not self.enable_shared_expert_fusion: return topk_output shared_weights = self._get_shared_expert_weights(hidden_states) if shared_weights is None: @@ -386,7 +381,9 @@ def _forward_router_experts(self, hidden_states: torch.Tensor): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) - if self.enable_fused_moe and TopKOutputChecker.format_is_standard(topk_output): + if self.enable_shared_expert_fusion and TopKOutputChecker.format_is_standard( + topk_output + ): topk_output = self._append_shared_to_topk_output(topk_output, hidden_states) return self.experts(hidden_states, topk_output) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 3ec9c240e178..56dab45de032 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1471,11 +1471,12 @@ def __init__( self.is_mrope_enabled = "mrope_section" in rope_config self.deepstack_visual_indexes = self.visual.deepstack_visual_indexes - self.enable_fused_moe = True if _use_aiter else False self.num_fused_shared_experts = ( - 0 if not self.enable_fused_moe else self._get_num_fused_shared_experts() + self._get_num_fused_shared_experts() if _use_aiter else 0 ) + self.enable_shared_expert_fusion = self.num_fused_shared_experts > 0 + def _get_num_fused_shared_experts(self): if not ( hasattr(self.model, "layers") @@ -1525,9 +1526,8 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ckpt_up_proj_name="up_proj", num_experts=( num_experts - if not self.enable_fused_moe - else num_experts - + self.num_fused_shared_experts # map shared experts to routed experts + if not self.enable_shared_expert_fusion + else num_experts + self.num_fused_shared_experts ), ) @@ -1549,7 +1549,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("experts.w2_weight", "experts.down_proj", 0, "w2"), ] - if self.enable_fused_moe: + if self.enable_shared_expert_fusion: """ When shared experts are fused, we need to map the shared experts to routed experts. @@ -1569,6 +1569,19 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): num_experts, "w2", ), + ## shared experts may contain gate_proj and up_proj instead of gate_up_proj + ( + "experts.w13_", + f"experts.{num_experts}.gate_proj.", + num_experts, + "w1", + ), + ( + "experts.w13_", + f"experts.{num_experts}.up_proj.", + num_experts, + "w3", + ), ] def load_fused_expert_weights( @@ -1626,7 +1639,7 @@ def load_fused_expert_weights( ): continue - if self.enable_fused_moe: + if self.enable_shared_expert_fusion: if "mlp.shared_expert." in name: # Firstly map mlp.shared_expert.xx_proj to mlp.experts.512.xx_proj name = name.replace( @@ -1708,9 +1721,9 @@ def load_fused_expert_weights( params_dict, loaded_weight, shard_id, - num_experts + num_experts, ) - elif self.enable_fused_moe: + elif self.enable_shared_expert_fusion: # shared experts should be loaded to experts.w13_weight and experts.w2_weight param = params_dict[name_mapped] weight_loader = getattr( From d3d8eba5f07a929efb4b43c588d1195f0c176f7b Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 00:57:47 -0500 Subject: [PATCH 33/40] Update comments in Qwen3.5 MoE model weight loading for clarity - Enhanced comments to specify loading behavior for `down_proj`, `gate_proj`, and `up_proj` in the weight loading process. - Improved code documentation to aid understanding of expert weight handling in the model. --- python/sglang/srt/models/qwen3_5.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 56dab45de032..c101f5a5e2ed 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1750,7 +1750,8 @@ def load_fused_expert_weights( expert_id, ) else: - # load to experts.w2_weight, shard_id = w2, expert_id = 512 + # load down_proj to experts.w2_weight, shard_id = w2, expert_id = 512 + # Or load gate_proj and up_proj to experts.w13_weight, shard_id = w1/w3, expert_id = 512 weight_loader( param, loaded_weight, From 25a7e043a6e484355d28c20918be9d1a0031d8d4 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 01:48:48 -0500 Subject: [PATCH 34/40] Refactor shared expert initialization in Qwen2 MoE model - Updated the logic for determining the number of shared experts based on configuration settings, allowing for more flexible expert handling. - Defaulted `enable_shared_expert_fusion` to False and adjusted its initialization to depend on the `_use_aiter` flag, improving clarity and maintainability of the code. - Enhanced comments to clarify the conditions under which shared expert fusion is enabled. --- python/sglang/srt/models/qwen2_moe.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index b6f796460709..8501525d39dc 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -195,15 +195,24 @@ def __init__( f"the number of experts {config.num_experts}." ) self.num_experts = config.num_experts - self.num_shared_experts = 1 - self.enable_shared_expert_fusion = ( - support_shared_expert_fusion - and _use_aiter - and can_fuse_shared_expert(config) - ) + num_shared_experts = 0 + if config.get("n_shared_experts", 0) > 0: + # config defines the number of shared experts + num_shared_experts = config.n_shared_experts + elif config.get("shared_expert_intermediate_size", 0) > 0: + # n_shared_experts is not defined, but shared_expert_intermediate_size is defined, so we use 1 as the number of shared experts + num_shared_experts = 1 + self.num_shared_experts = num_shared_experts + self.enable_shared_expert_fusion = False # default to False + if _use_aiter: + # enable shared expert fusion when use aiter + self.enable_shared_expert_fusion = ( + support_shared_expert_fusion + and can_fuse_shared_expert(config) + ) self.num_fused_shared_experts = ( - self.num_shared_experts if self.enable_shared_expert_fusion else 0 + num_shared_experts if self.enable_shared_expert_fusion else 0 ) self.topk = TopK( From db8887693702e29fbf24660e0238dd04630dde75 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 01:49:37 -0500 Subject: [PATCH 35/40] Refactor shared expert initialization logic in Qwen2 MoE model - Adjusted the initialization of `num_shared_experts` to ensure it defaults to 0 when no configuration is provided, enhancing clarity and robustness. - Improved the handling of shared expert configuration settings, allowing for more flexible expert management in the model. --- python/sglang/srt/models/qwen2_moe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 8501525d39dc..9ae9ca54d086 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -195,13 +195,14 @@ def __init__( f"the number of experts {config.num_experts}." ) self.num_experts = config.num_experts - num_shared_experts = 0 if config.get("n_shared_experts", 0) > 0: # config defines the number of shared experts num_shared_experts = config.n_shared_experts elif config.get("shared_expert_intermediate_size", 0) > 0: # n_shared_experts is not defined, but shared_expert_intermediate_size is defined, so we use 1 as the number of shared experts num_shared_experts = 1 + else: + num_shared_experts = 0 self.num_shared_experts = num_shared_experts self.enable_shared_expert_fusion = False # default to False if _use_aiter: From 564af0181760ba002407d04d45da83fec3c27e36 Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 01:52:06 -0500 Subject: [PATCH 36/40] Refactor shared expert initialization in Qwen2 MoE model - Cleaned up the initialization logic for `num_shared_experts` and `enable_shared_expert_fusion`, improving code clarity and maintainability. - Enhanced comments to clarify the conditions for shared expert configuration, ensuring better understanding of the model's behavior. --- python/sglang/srt/models/qwen2_moe.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 9ae9ca54d086..5156b70e277b 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -197,19 +197,18 @@ def __init__( self.num_experts = config.num_experts if config.get("n_shared_experts", 0) > 0: # config defines the number of shared experts - num_shared_experts = config.n_shared_experts + num_shared_experts = config.n_shared_experts elif config.get("shared_expert_intermediate_size", 0) > 0: - # n_shared_experts is not defined, but shared_expert_intermediate_size is defined, so we use 1 as the number of shared experts + # n_shared_experts is not defined, but shared_expert_intermediate_size is defined, so we use 1 as the number of shared experts num_shared_experts = 1 else: num_shared_experts = 0 self.num_shared_experts = num_shared_experts - self.enable_shared_expert_fusion = False # default to False + self.enable_shared_expert_fusion = False # default to False if _use_aiter: # enable shared expert fusion when use aiter self.enable_shared_expert_fusion = ( - support_shared_expert_fusion - and can_fuse_shared_expert(config) + support_shared_expert_fusion and can_fuse_shared_expert(config) ) self.num_fused_shared_experts = ( From ab548117ac4f8ccfba85f65f3be17e2d5c8ac97b Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 01:53:21 -0500 Subject: [PATCH 37/40] Refactor shared expert configuration checks in Qwen2 MoE model - Updated the initialization logic for `num_shared_experts` to use `hasattr` for better attribute checking, enhancing robustness and clarity. - Improved conditions for determining shared expert settings, ensuring more flexible configuration handling in the model. --- python/sglang/srt/models/qwen2_moe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 5156b70e277b..78b10244fc3c 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -195,10 +195,10 @@ def __init__( f"the number of experts {config.num_experts}." ) self.num_experts = config.num_experts - if config.get("n_shared_experts", 0) > 0: + if hasattr(config, "n_shared_experts"): # config defines the number of shared experts num_shared_experts = config.n_shared_experts - elif config.get("shared_expert_intermediate_size", 0) > 0: + elif hasattr(config, "shared_expert_intermediate_size") and config.shared_expert_intermediate_size > 0: # n_shared_experts is not defined, but shared_expert_intermediate_size is defined, so we use 1 as the number of shared experts num_shared_experts = 1 else: From e0e720bda9b7b8f00a9ee5abc5441c78bf2e024b Mon Sep 17 00:00:00 2001 From: "Chen, Todd" Date: Thu, 2 Apr 2026 01:56:20 -0500 Subject: [PATCH 38/40] Refactor expert initialization in Qwen2 MoE model - Updated the logic for calculating the total number of experts by directly calling `get_global_server_args().ep_num_redundant_experts`, improving code clarity and maintainability. - Enhanced the initialization of the `experts` attribute to streamline the configuration process for expert management in the model. --- python/sglang/srt/models/qwen2_moe.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 78b10244fc3c..fb11b54c2be7 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -221,12 +221,11 @@ def __init__( layer_id=layer_id, ) - ep_redundant = get_global_server_args().ep_num_redundant_experts self.experts = get_moe_impl_class(quant_config)( layer_id=self.layer_id, top_k=config.num_experts_per_tok + self.num_fused_shared_experts, num_experts=config.num_experts - + ep_redundant + + get_global_server_args().ep_num_redundant_experts + self.num_fused_shared_experts, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, From 81769558095c83a1e554529faca0542b3b01804e Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Thu, 2 Apr 2026 07:25:18 +0000 Subject: [PATCH 39/40] Refactor --- python/sglang/srt/models/qwen2_moe.py | 40 +++++++++++++++------------ python/sglang/srt/models/qwen3_5.py | 6 ++-- 2 files changed, 26 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index fb11b54c2be7..fa4985ccabc2 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -195,15 +195,17 @@ def __init__( f"the number of experts {config.num_experts}." ) self.num_experts = config.num_experts + self.num_shared_experts = 0 if hasattr(config, "n_shared_experts"): # config defines the number of shared experts - num_shared_experts = config.n_shared_experts - elif hasattr(config, "shared_expert_intermediate_size") and config.shared_expert_intermediate_size > 0: + self.num_shared_experts = config.n_shared_experts + elif ( + hasattr(config, "shared_expert_intermediate_size") + and config.shared_expert_intermediate_size > 0 + ): # n_shared_experts is not defined, but shared_expert_intermediate_size is defined, so we use 1 as the number of shared experts - num_shared_experts = 1 - else: - num_shared_experts = 0 - self.num_shared_experts = num_shared_experts + self.num_shared_experts = 1 + self.enable_shared_expert_fusion = False # default to False if _use_aiter: # enable shared expert fusion when use aiter @@ -211,10 +213,6 @@ def __init__( support_shared_expert_fusion and can_fuse_shared_expert(config) ) - self.num_fused_shared_experts = ( - num_shared_experts if self.enable_shared_expert_fusion else 0 - ) - self.topk = TopK( top_k=config.num_experts_per_tok, renormalize=config.norm_topk_prob, @@ -223,16 +221,24 @@ def __init__( self.experts = get_moe_impl_class(quant_config)( layer_id=self.layer_id, - top_k=config.num_experts_per_tok + self.num_fused_shared_experts, - num_experts=config.num_experts - + get_global_server_args().ep_num_redundant_experts - + self.num_fused_shared_experts, + top_k=( + config.num_experts_per_tok + if not self.enable_shared_expert_fusion + else config.num_experts_per_tok + self.num_shared_experts + ), + num_experts=( + config.num_experts + get_global_server_args().ep_num_redundant_experts + if not self.enable_shared_expert_fusion + else config.num_experts + + get_global_server_args().ep_num_redundant_experts + + self.num_shared_experts + ), hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, quant_config=quant_config, prefix=add_prefix("experts", prefix), routing_method_type=RoutingMethodType.RenormalizeNaive, - num_fused_shared_experts=self.num_fused_shared_experts, + num_fused_shared_experts=self.num_shared_experts, ) self.gate = ReplicatedLinear( @@ -316,12 +322,12 @@ def _append_shared_to_topk_output( M = topk_output.topk_ids.shape[0] shared_expert_id = self.num_experts shared_ids = torch.full( - (M, self.num_fused_shared_experts), + (M, self.num_shared_experts), shared_expert_id, dtype=topk_output.topk_ids.dtype, device=topk_output.topk_ids.device, ) - shared_weights = shared_weights.expand(M, self.num_fused_shared_experts).to( + shared_weights = shared_weights.expand(M, self.num_shared_experts).to( topk_output.topk_weights.dtype ) fused_topk_ids = torch.cat([topk_output.topk_ids, shared_ids], dim=-1) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index c101f5a5e2ed..7a7918454e63 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1471,9 +1471,9 @@ def __init__( self.is_mrope_enabled = "mrope_section" in rope_config self.deepstack_visual_indexes = self.visual.deepstack_visual_indexes - self.num_fused_shared_experts = ( - self._get_num_fused_shared_experts() if _use_aiter else 0 - ) + self.num_fused_shared_experts = 0 + if _use_aiter: + self.num_fused_shared_experts = self._get_num_fused_shared_experts() self.enable_shared_expert_fusion = self.num_fused_shared_experts > 0 From c2dafdeeafffabd40fa5e2906eb7957bb905305d Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Thu, 2 Apr 2026 08:33:44 +0000 Subject: [PATCH 40/40] Accuracy fix for fp8 after refactor --- python/sglang/srt/models/qwen2_moe.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index fa4985ccabc2..f172e8ceb917 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -196,6 +196,7 @@ def __init__( ) self.num_experts = config.num_experts self.num_shared_experts = 0 + self.num_fused_shared_experts = 0 if hasattr(config, "n_shared_experts"): # config defines the number of shared experts self.num_shared_experts = config.n_shared_experts @@ -212,6 +213,8 @@ def __init__( self.enable_shared_expert_fusion = ( support_shared_expert_fusion and can_fuse_shared_expert(config) ) + if self.enable_shared_expert_fusion: + self.num_fused_shared_experts = self.num_shared_experts self.topk = TopK( top_k=config.num_experts_per_tok, @@ -224,21 +227,21 @@ def __init__( top_k=( config.num_experts_per_tok if not self.enable_shared_expert_fusion - else config.num_experts_per_tok + self.num_shared_experts + else config.num_experts_per_tok + self.num_fused_shared_experts ), num_experts=( config.num_experts + get_global_server_args().ep_num_redundant_experts if not self.enable_shared_expert_fusion else config.num_experts + get_global_server_args().ep_num_redundant_experts - + self.num_shared_experts + + self.num_fused_shared_experts ), hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, quant_config=quant_config, prefix=add_prefix("experts", prefix), routing_method_type=RoutingMethodType.RenormalizeNaive, - num_fused_shared_experts=self.num_shared_experts, + num_fused_shared_experts=self.num_fused_shared_experts, ) self.gate = ReplicatedLinear( @@ -322,12 +325,12 @@ def _append_shared_to_topk_output( M = topk_output.topk_ids.shape[0] shared_expert_id = self.num_experts shared_ids = torch.full( - (M, self.num_shared_experts), + (M, self.num_fused_shared_experts), shared_expert_id, dtype=topk_output.topk_ids.dtype, device=topk_output.topk_ids.device, ) - shared_weights = shared_weights.expand(M, self.num_shared_experts).to( + shared_weights = shared_weights.expand(M, self.num_fused_shared_experts).to( topk_output.topk_weights.dtype ) fused_topk_ids = torch.cat([topk_output.topk_ids, shared_ids], dim=-1)