diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 130f9aad4..5e7dd7961 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -39,10 +39,13 @@ from sglang.srt.runtime_context import ( get_parallel, get_spec, ) -from sglang.srt.utils import add_prefix, is_npu +from sglang.srt.utils import add_prefix, get_bool_env_var, is_hip, is_npu logger = logging.getLogger(__name__) +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + def _mtp_quant_config(quant_config): """The quantization the MTP module itself is built with. @@ -244,12 +247,20 @@ class Qwen3_5ForCausalLMMTP(nn.Module): # Params for MoE experts (non-fused/fused) num_experts = getattr(self.config, "num_experts", None) + # A fused shared expert lives in routed slot `num_experts`. + 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 if num_experts is not None: 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, ) else: expert_params_mapping = [] @@ -317,6 +328,17 @@ class Qwen3_5ForCausalLMMTP(nn.Module): if ".self_attn." in name: name = name.replace(".self_attn", "") + 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}.", + ) + # 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