From aead0ef5e55303ed956f287207c96d444ae655b8 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Thu, 13 Nov 2025 16:41:40 +0800 Subject: [PATCH] [FEAT][ROCM] enable fused shared expert for Rocm (#12201) Co-authored-by: ZLkanyo009 <4071250045@qq.com> Co-authored-by: HAI --- .../srt/layers/moe/fused_moe_triton/layer.py | 32 ++++++-- .../layers/moe/token_dispatcher/standard.py | 28 +++++-- python/sglang/srt/layers/moe/topk.py | 77 ++++++++++++++++--- python/sglang/srt/layers/quantization/fp8.py | 3 +- .../sglang/srt/layers/quantization/unquant.py | 1 + python/sglang/srt/models/deepseek_v2.py | 43 ++++++++--- 6 files changed, 152 insertions(+), 32 deletions(-) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 9a931c70b..51d1d39a0 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/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 - self.num_local_experts = num_experts // self.moe_ep_size + assert (num_experts - num_fused_shared_experts) % self.moe_ep_size == 0 + 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 - end_idx = (self.moe_ep_rank + 1) * self.num_local_experts + num_global_routed_experts = self.num_experts - self.num_fused_shared_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, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index d01ade923..26105a27c 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/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_experts + * self.num_local_routed_experts : (self.moe_ep_rank + 1) + * 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] diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 006552e37..e42488192 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 78dfbf711..8cab59729 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 26a83c560..6c6c728e2 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/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: diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index ed19fe474..3d5ea2c65 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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[ - i - ].mlp.shared_experts.gate_up_proj.output_size_per_partition + tp_size = get_tensor_model_parallel_world_size() + intermediate_size = ( + 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 - or torch.cuda.get_device_capability("cuda") < (8, 0) - or self.config.architectures[0] != architecture + 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." - elif get_moe_expert_parallel_world_size() > 1: - disable_reason = "Deepseek V3/R1 can not use shared experts fusion optimization under expert parallelism." + disable_reason = "Config not support fused shared expert(s)." + elif (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and ( + 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."