[FEAT][ROCM] enable fused shared expert for Rocm (#12201)

Co-authored-by: ZLkanyo009 <4071250045@qq.com>
Co-authored-by: HAI <hixiao@gmail.com>

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -69,6 +69,7 @@ if get_moe_runner_backend().is_flashinfer_trtllm():
 _is_hip = is_hip()
 _is_cpu_amx_available = cpu_has_amx_support()
 _is_cpu = is_cpu()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
 logger = logging.getLogger(__name__)
 
@@ -166,8 +167,12 @@ class FusedMoE(torch.nn.Module):
         self.moe_ep_rank = get_moe_expert_parallel_rank()
         self.moe_tp_size = get_moe_tensor_parallel_world_size()
         self.moe_tp_rank = get_moe_tensor_parallel_rank()
-        assert num_experts % self.moe_ep_size == 0
+        assert (num_experts - num_fused_shared_experts) % self.moe_ep_size == 0
-        self.num_local_experts = num_experts // self.moe_ep_size
+        self.num_local_experts = (
+            num_experts - num_fused_shared_experts
+        ) // self.moe_ep_size + num_fused_shared_experts
+
+        self.expert_mask_gpu = None
 
         assert intermediate_size % self.moe_tp_size == 0
         self.intermediate_size_per_partition = intermediate_size // self.moe_tp_size
@@ -469,10 +474,18 @@ class FusedMoE(torch.nn.Module):
             expert_data.copy_(loaded_weight)
 
     def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int:
-        start_idx = self.moe_ep_rank * self.num_local_experts
+        num_global_routed_experts = self.num_experts - self.num_fused_shared_experts
-        end_idx = (self.moe_ep_rank + 1) * self.num_local_experts
+        num_local_routed_experts = (
+            self.num_local_experts - self.num_fused_shared_experts
+        )
+        start_idx = self.moe_ep_rank * num_local_routed_experts
+        end_idx = (self.moe_ep_rank + 1) * num_local_routed_experts
         if start_idx <= expert_id < end_idx:
             return expert_id - start_idx
+        elif (
+            self.num_fused_shared_experts > 0 and expert_id >= num_global_routed_experts
+        ):
+            return expert_id - num_global_routed_experts + num_local_routed_experts
         else:
             return -1
 
@@ -546,7 +559,7 @@ class FusedMoE(torch.nn.Module):
         # WARN: This makes the `expert_id` mean "local" and "global" in different cases
         if not getattr(param, "_sglang_require_global_experts", False):
             expert_id = self._map_global_expert_id_to_local_expert_id(expert_id)
-            if expert_id == -1:
+            if expert_id < 0 or expert_id >= self.num_local_experts:
                 return
 
         if isinstance(
@@ -844,6 +857,15 @@ class FusedMoE(torch.nn.Module):
         dispatch_output = self.dispatcher.dispatch(
             hidden_states=hidden_states, topk_output=topk_output
         )
+        if _use_aiter and self.dispatcher.local_expert_mapping is not None:
+            self.expert_mask_gpu = (
+                (
+                    (self.dispatcher.local_expert_mapping >= 0)
+                    & (self.dispatcher.local_expert_mapping < self.num_local_experts)
+                )
+                .to(torch.int32)
+                .to(device="cuda")
+            )
 
         combine_input = self.run_moe_core(
             dispatch_output=dispatch_output,

python/sglang/srt/layers/moe/token_dispatcher/standard.py

@@ -18,6 +18,10 @@ from sglang.srt.layers.moe.token_dispatcher.base import (
 )
 from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker
 from sglang.srt.layers.moe.utils import get_moe_runner_backend
+from sglang.srt.utils import get_bool_env_var, is_hip
+
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.topk import TopKOutput
@@ -58,7 +62,10 @@ class StandardDispatcher(BaseDispatcher):
             get_moe_runner_backend().is_flashinfer_cutlass()
         )
         self.num_experts = moe_runner_config.num_experts
-        self.num_local_experts = moe_runner_config.num_local_experts
+        self.num_local_shared_experts = moe_runner_config.num_fused_shared_experts
+        self.num_local_routed_experts = (
+            moe_runner_config.num_local_experts - self.num_local_shared_experts
+        )
         self.moe_ep_rank = get_moe_expert_parallel_rank()
         self.local_expert_mapping = None
 
@@ -77,13 +84,24 @@ class StandardDispatcher(BaseDispatcher):
                 )
                 self.local_expert_mapping[
                     self.moe_ep_rank
-                    * self.num_local_experts : (self.moe_ep_rank + 1)
+                    * self.num_local_routed_experts : (self.moe_ep_rank + 1)
-                    * self.num_local_experts
+                    * self.num_local_routed_experts
                 ] = torch.arange(
-                    0, self.num_local_experts, dtype=torch.int32, device="cuda"
+                    0, self.num_local_routed_experts, dtype=torch.int32, device="cuda"
                 )
 
-        if self.local_expert_mapping is not None:
+                if self.num_local_shared_experts > 0:
+                    self.local_expert_mapping[-self.num_local_shared_experts :] = (
+                        torch.arange(
+                            self.num_local_routed_experts,
+                            self.num_local_routed_experts
+                            + self.num_local_shared_experts,
+                            dtype=torch.int32,
+                            device="cpu",
+                        )
+                    )
+
+        if self.local_expert_mapping is not None and not _use_aiter:
             if TopKOutputChecker.format_is_standard(topk_output):
                 topk_output = topk_output._replace(
                     topk_ids=self.local_expert_mapping[topk_output.topk_ids]

python/sglang/srt/layers/moe/topk.py

@@ -100,6 +100,7 @@ class TopKConfig:
     torch_native: bool = False
     routed_scaling_factor: Optional[float] = None
     apply_routed_scaling_factor_on_output: bool = False
+    fused_shared_experts_scaling_factor: Optional[float] = None
     output_format: Optional[TopKOutputFormat] = None
 
 
@@ -190,6 +191,13 @@ class BypassedTopKOutput(NamedTuple):
 
 
 class TopK(CustomOp):
+    """
+    Parameters:
+    --top_k: The all number of top experts selected per token, including the fused shared expert(s).
+    --num_fused_shared_experts: num of shared experts, can be activate both in TP or EP mode.
+    --routed_scaling_factor: the scaling factor for routed experts in topk_weights.
+    --fused_shared_experts_scaling_factor: scaling factor for fused shared experts on AMD-platform.
+    """
 
     def __init__(
         self,
@@ -207,6 +215,7 @@ class TopK(CustomOp):
         routed_scaling_factor: Optional[float] = None,
         apply_routed_scaling_factor_on_output: Optional[bool] = False,
         output_format: Optional[TopKOutputFormat] = None,
+        fused_shared_experts_scaling_factor: Optional[float] = None,
     ):
         # NOTE: scoring_func is not used for now, but we keep it for future use
         # see https://github.com/sgl-project/sglang/pull/4505 for more details
@@ -226,6 +235,7 @@ class TopK(CustomOp):
             correction_bias=correction_bias,
             routed_scaling_factor=routed_scaling_factor,
             apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
+            fused_shared_experts_scaling_factor=fused_shared_experts_scaling_factor,
             output_format=output_format,
         )
 
@@ -554,7 +564,10 @@ def grouped_topk_gpu(
             dtype=topk_ids.dtype,
             device=topk_ids.device,
         )
-        topk_weights[:, -1] = topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
+        if routed_scaling_factor is not None:
+            topk_weights[:, -1] = (
+                topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
+            )
 
     if renormalize:
         topk_weights_sum = (
@@ -698,7 +711,10 @@ def biased_grouped_topk_impl(
             dtype=topk_ids.dtype,
             device=topk_ids.device,
         )
-        topk_weights[:, -1] = topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
+        if routed_scaling_factor is not None:
+            topk_weights[:, -1] = (
+                topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
+            )
 
     if renormalize:
         topk_weights_sum = (
@@ -753,9 +769,6 @@ def biased_grouped_topk_gpu(
     expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
     apply_routed_scaling_factor_on_output: Optional[bool] = False,
 ):
-    assert (
-        routed_scaling_factor is not None
-    ), "routed_scaling_factor is required for biased_grouped_topk"
     # TODO: moe_fused_gate kernel is not supported for num_fused_shared_experts > 0 now.
     if (
         _is_cuda
@@ -770,7 +783,7 @@ def biased_grouped_topk_gpu(
             topk_group,
             topk,
             num_fused_shared_experts,
-            routed_scaling_factor,
+            routed_scaling_factor if routed_scaling_factor is not None else 1.0,
             apply_routed_scaling_factor_on_output,
         )
         # TODO merge into kernel
@@ -798,7 +811,7 @@ def biased_grouped_topk_gpu(
             num_expert_group,
             topk_group,
             renormalize,
-            routed_scaling_factor,
+            routed_scaling_factor if routed_scaling_factor is not None else 1.0,
         )
         return topk_weights, topk_ids
     else:
@@ -897,6 +910,9 @@ def select_experts(
     apply_routed_scaling_factor_on_output = (
         topk_config.apply_routed_scaling_factor_on_output
     )
+    fused_shared_experts_scaling_factor = (
+        topk_config.fused_shared_experts_scaling_factor
+    )
 
     router_logits, correction_bias = (
         expert_location_dispatch.transform_select_experts_inputs(
@@ -907,6 +923,8 @@ def select_experts(
     )
 
     # DeepSeek V2/V3/R1 series models use grouped_top_k
+    # remove num_fused_shared_experts from grouped_topk/biased_grouped_topk
+    num_routed_topk = top_k - num_fused_shared_experts
     if use_grouped_topk:
         assert topk_group is not None
         assert num_expert_group is not None
@@ -914,7 +932,7 @@ def select_experts(
             topk_weights, topk_ids = grouped_topk(
                 hidden_states=hidden_states,
                 gating_output=router_logits,
-                topk=top_k,
+                topk=num_routed_topk if _use_aiter else top_k,
                 renormalize=renormalize,
                 num_expert_group=num_expert_group,
                 topk_group=topk_group,
@@ -929,7 +947,7 @@ def select_experts(
                 hidden_states=hidden_states,
                 gating_output=router_logits,
                 correction_bias=correction_bias,
-                topk=top_k,
+                topk=num_routed_topk if _use_aiter else top_k,
                 renormalize=renormalize,
                 num_expert_group=num_expert_group,
                 topk_group=topk_group,
@@ -948,7 +966,7 @@ def select_experts(
         topk_weights, topk_ids = fused_topk_native(
             hidden_states=hidden_states,
             gating_output=router_logits,
-            topk=top_k,
+            topk=num_routed_topk if _use_aiter else top_k,
             renormalize=renormalize,
             correction_bias=correction_bias,
         )
@@ -958,7 +976,7 @@ def select_experts(
         topk_weights, topk_ids = fused_topk(
             hidden_states=hidden_states,
             gating_output=router_logits,
-            topk=top_k,
+            topk=num_routed_topk if _use_aiter else top_k,
             renormalize=renormalize,
             num_token_non_padded=num_token_non_padded,
             expert_location_dispatch_info=expert_location_dispatch_info,
@@ -972,10 +990,45 @@ def select_experts(
         topk_weights, topk_ids = custom_routing_function(
             hidden_states=hidden_states,
             gating_output=router_logits,
-            topk=top_k,
+            topk=num_routed_topk if _use_aiter else top_k,
             renormalize=renormalize,
         )
 
+    # TODO: fused ops of shared experts in topk function itself when num_fused_shared_experts > 0.
+    if num_fused_shared_experts > 0 and _use_aiter:
+        M, N = router_logits.shape
+        scale_factor = (
+            1.0
+            if fused_shared_experts_scaling_factor is None
+            else fused_shared_experts_scaling_factor
+        )
+
+        topk_ids = torch.cat(
+            [
+                topk_ids,
+                torch.arange(
+                    N,
+                    N + num_fused_shared_experts,
+                    dtype=topk_ids.dtype,
+                    device=topk_ids.device,
+                ).expand(M, -1),
+            ],
+            dim=1,
+        )
+
+        topk_weights = torch.cat(
+            [
+                topk_weights,
+                torch.full(
+                    (topk_weights.size(0), num_fused_shared_experts),
+                    scale_factor,
+                    dtype=topk_weights.dtype,
+                    device=topk_weights.device,
+                ),
+            ],
+            dim=1,
+        )
+
     get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
 
     return StandardTopKOutput(topk_weights, topk_ids, router_logits)

python/sglang/srt/layers/quantization/fp8.py

@@ -1301,7 +1301,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                         if activation == "silu"
                         else ActivationType.Gelu
                     ),
-                    expert_mask=None,
+                    expert_mask=layer.expert_mask_gpu,
                 )
             else:
                 return fused_moe(
@@ -1318,6 +1318,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                         if activation == "silu"
                         else ActivationType.Gelu
                     ),
+                    expert_mask=layer.expert_mask_gpu,
                 )
         return None
 

python/sglang/srt/layers/quantization/unquant.py

@@ -289,6 +289,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
                         if moe_runner_config.activation == "silu"
                         else ActivationType.Gelu
                     ),
+                    expert_mask=layer.expert_mask_gpu,
                 )
                 return StandardCombineInput(hidden_states=output)
             else:

python/sglang/srt/models/deepseek_v2.py

@@ -36,6 +36,7 @@ from sglang.srt.configs.model_config import (
     is_deepseek_nsa,
 )
 from sglang.srt.distributed import (
+    divide,
     get_moe_expert_parallel_world_size,
     get_pp_group,
     get_tensor_model_parallel_world_size,
@@ -593,6 +594,7 @@ class DeepseekV2MoE(nn.Module):
     ):
         super().__init__()
         self.tp_size = get_tensor_model_parallel_world_size()
+        self.moe_ep_size = get_moe_expert_parallel_world_size()
         self.routed_scaling_factor = config.routed_scaling_factor
         self.n_shared_experts = config.n_shared_experts
         self.num_fused_shared_experts = (
@@ -624,6 +626,14 @@ class DeepseekV2MoE(nn.Module):
             is_nextn=is_nextn,
         )
 
+        # scaling factor for fused shared experts on AMD-platform.
+        fused_shared_experts_scaling_factor = None
+        if self.moe_ep_size > 1 and self.num_fused_shared_experts > 0:
+            # if enable_ep_moe tp_szie == ep_size, every gpu get shared experts gemm output
+            # so we scale with 1 / self.moe_ep_size in ep mode which will make it equalation as in tp mode
+            # with fused_shared_experts
+            fused_shared_experts_scaling_factor = 1.0 / float(self.moe_ep_size)
+
         self.experts = get_moe_impl_class(quant_config)(
             num_experts=config.n_routed_experts
             + self.num_fused_shared_experts
@@ -649,6 +659,7 @@ class DeepseekV2MoE(nn.Module):
             quant_config=quant_config,
             routed_scaling_factor=self.routed_scaling_factor,
             apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk,
+            fused_shared_experts_scaling_factor=fused_shared_experts_scaling_factor,
             # Some Fp4 MoE backends require the output format to be bypassed but the MTP layers are unquantized
             # and requires the output format to be standard. We use quant_config to determine the output format.
             output_format=TopKOutputFormat.STANDARD if quant_config is None else None,
@@ -2978,9 +2989,14 @@ class DeepseekV2Model(nn.Module):
             allocate_size = 0
             for i in range(len(self.layers)):
                 if isinstance(self.layers[i].mlp, DeepseekV2MoE):
-                    allocate_size = self.layers[
+                    tp_size = get_tensor_model_parallel_world_size()
-                        i
+                    intermediate_size = (
-                    ].mlp.shared_experts.gate_up_proj.output_size_per_partition
+                        config.moe_intermediate_size * config.n_shared_experts
+                    )
+                    share_expert_output_size_per_partition = divide(
+                        intermediate_size * 2, tp_size
+                    )
+                    allocate_size = share_expert_output_size_per_partition
                     break
 
             self.gemm_output_zero_allocator_size = (
@@ -3158,15 +3174,24 @@ class DeepseekV2ForCausalLM(nn.Module):
         # Only Deepseek V3/R1 can use shared experts fusion optimization now.
         disable_reason = None
         if (
-            not _is_cuda
+            self.config.architectures[0] != architecture
-            or torch.cuda.get_device_capability("cuda") < (8, 0)
-            or self.config.architectures[0] != architecture
             or self.config.n_routed_experts != 256
             or self.config.n_shared_experts != 1
         ):
-            disable_reason = "Only Deepseek V3/R1 on NV-platform with capability >= 80 can use shared experts fusion optimization."
+            disable_reason = "Config not support fused shared expert(s)."
-        elif get_moe_expert_parallel_world_size() > 1:
+        elif (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and (
-            disable_reason = "Deepseek V3/R1 can not use shared experts fusion optimization under expert parallelism."
+            not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
+        ):
+            disable_reason = (
+                "Only Deepseek V3/R1 on NV-platform with capability >= 80 "
+                "or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization."
+            )
+        elif get_moe_expert_parallel_world_size() > 1 and (
+            not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
+        ):
+            disable_reason = "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism."
+        elif disable_reason is None and get_moe_a2a_backend().is_deepep():
+            disable_reason = "Deepseek V3/R1 can not use shared experts fusion optimization under deepep expert parallelism."
         elif self.quant_config and self.quant_config.get_name() == "w4afp8":
             disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts."