[AMD] Load fused shared experts for Qwen4-Exp and Qwen3.5 MTP (#38878)

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

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
             )