diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 31d066a58..db1f79fd6 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -293,8 +293,9 @@ class Qwen3_5ForCausalLMMTP(nn.Module): "_input_scale", ) - # fused experts: experts.w13_weight / experts.w2_weight - is_fused_expert = False + # Fused checkpoint tensors: experts.gate_up_proj / experts.down_proj. + # The checkpoint interleaves these with separate shared-expert tensors, + # so picking one mapping must not affect the next weight. fused_expert_params_mapping = [ ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), ("experts.w2_weight", "experts.down_proj", 0, "w2"), @@ -370,13 +371,17 @@ class Qwen3_5ForCausalLMMTP(nn.Module): f"mlp.experts.{num_experts}.", ) + is_fused_expert = ( + "experts.gate_up_proj" in name or "experts.down_proj" in name + ) + current_expert_params_mapping = ( + fused_expert_params_mapping + if is_fused_expert + else expert_params_mapping + ) + # 1) Process stacked parameters (q_proj/k_proj/v_proj & gate_proj/up_proj) for param_name, weight_name, shard_id in stacked_params_mapping: - # Check if this is a fused expert weight - if "experts.gate_up_proj" in name or "experts.down_proj" in name: - is_fused_expert = True - expert_params_mapping = fused_expert_params_mapping - # Skip non-matching weights if weight_name not in name: continue @@ -406,7 +411,7 @@ class Qwen3_5ForCausalLMMTP(nn.Module): # 2) Process MoE expert weights (including fused experts) is_expert_weight = False - for mapping in expert_params_mapping: + for mapping in current_expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue diff --git a/python/sglang/srt/models/qwen4_exp.py b/python/sglang/srt/models/qwen4_exp.py index 21f120af1..224d2b393 100644 --- a/python/sglang/srt/models/qwen4_exp.py +++ b/python/sglang/srt/models/qwen4_exp.py @@ -69,7 +69,9 @@ from sglang.srt.models.qwen4_exp_ple_table import ( make_ple_file_rss_trimmer, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import logger +from sglang.srt.utils import get_bool_env_var, is_hip, logger + +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() # Decode/verify-sized batches only: at prefill sizes both chains are compute # bound and serializing them on one stream is faster than contending. @@ -1836,12 +1838,21 @@ class Qwen4ExpForConditionalGeneration(Qwen3VLForConditionalGeneration): ] num_experts = getattr(self.config, "num_experts", None) + # A fused shared expert lives in routed slot `num_experts`, so the + # mapping has to cover one more expert than the config declares. + num_fused_shared_experts = 0 + if _use_aiter: + for module in self.modules(): + fused = getattr(module, "num_fused_shared_experts", 0) + if fused: + num_fused_shared_experts = fused + break 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=num_experts, + num_experts=num_experts + num_fused_shared_experts, ) if num_experts is not None else [] @@ -2037,6 +2048,17 @@ class Qwen4ExpForConditionalGeneration(Qwen3VLForConditionalGeneration): ): continue + if ( + _use_aiter + and num_fused_shared_experts > 0 + and "mlp.shared_expert." in name + ): + # Map mlp.shared_expert.xx_proj to mlp.experts.{num_experts}.xx_proj + name = name.replace( + "mlp.shared_expert.", + f"mlp.experts.{num_experts}.", + ) + is_fused_expert = ( "experts.gate_up_proj" in name or "experts.down_proj" in name )