chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: ZLkanyo009 <4071250045@qq.com>
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
Ling Zhang
2025-11-13 16:41:40 +08:00
committed by GitHub
parent 6664083522
commit aead0ef5e5
6 changed files with 152 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -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,

28
python/sglang/srt/layers/moe/token_dispatcher/standard.py
View File

@@ -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]

77
python/sglang/srt/layers/moe/topk.py
View File

@@ -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)

3
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -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

1
python/sglang/srt/layers/quantization/unquant.py
View File

@@ -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:

43
python/sglang/srt/models/deepseek_v2.py
View File

@@ -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."