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

Support EPLB in FusedMoE (#8448)

Browse Source
This commit is contained in:
Cheng Wan
2025-07-29 16:02:41 -07:00
committed by GitHub
parent 1992ef9ba7
commit 9effeb5bdd
15 changed files with 107 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -325,6 +325,7 @@ class DeepseekV2MoE(nn.Module):
num_experts=config.n_routed_experts
+ self.num_fused_shared_experts
+ global_server_args_dict["ep_num_redundant_experts"],
num_fused_shared_experts=self.num_fused_shared_experts,
top_k=config.num_experts_per_tok + self.num_fused_shared_experts,
hidden_size=config.hidden_size,
intermediate_size=config.moe_intermediate_size,
@@ -2112,6 +2113,7 @@ class DeepseekV2ForCausalLM(nn.Module):
if disable_reason is not None:
global_server_args_dict["disable_shared_experts_fusion"] = True
self.num_fused_shared_experts = 0
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",

4
python/sglang/srt/models/glm4_moe.py
View File

@@ -434,6 +434,7 @@ class Glm4MoeSparseMoeBlock(DeepseekV2MoE):
num_experts=config.n_routed_experts
+ self.num_fused_shared_experts
+ global_server_args_dict["ep_num_redundant_experts"],
num_fused_shared_experts=self.num_fused_shared_experts,
top_k=config.num_experts_per_tok + self.num_fused_shared_experts,
hidden_size=config.hidden_size,
intermediate_size=config.moe_intermediate_size,
@@ -740,10 +741,11 @@ class Glm4MoeForCausalLM(DeepseekV2ForCausalLM):
global_server_args_dict["enable_deepep_moe"]
or global_server_args_dict["enable_ep_moe"]
):
disable_reason = "Deepseek GLM-4.5 can not use shared experts fusion optimization when in deepep_moe or ep_moe mode."
disable_reason = "Deepseek and GLM-4.5 can not use shared experts fusion optimization when in deepep_moe or ep_moe mode."
if disable_reason is not None:
global_server_args_dict["disable_shared_experts_fusion"] = True
self.num_fused_shared_experts = 0
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",

3
python/sglang/srt/models/granitemoe.py
View File

@@ -43,6 +43,7 @@ class GraniteMoeMoE(nn.Module):
top_k: int,
hidden_size: int,
intermediate_size: int,
layer_id: int,
params_dtype: Optional[torch.dtype] = None,
quant_config: Optional[QuantizationConfig] = None,
tp_size: Optional[int] = None,
@@ -71,6 +72,7 @@ class GraniteMoeMoE(nn.Module):
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
layer_id=layer_id,
params_dtype=params_dtype,
reduce_results=True,
quant_config=quant_config,
@@ -203,6 +205,7 @@ class GraniteMoeDecoderLayer(nn.Module):
top_k=config.num_experts_per_tok,
hidden_size=config.hidden_size,
intermediate_size=config.intermediate_size,
layer_id=layer_id,
quant_config=quant_config,
prefix=f"{prefix}.block_sparse_moe",
)

3
python/sglang/srt/models/grok.py
View File

@@ -78,6 +78,7 @@ class Grok1MoE(nn.Module):
def __init__(
self,
config: PretrainedConfig,
layer_id: int,
num_experts: int,
top_k: int,
hidden_size: int,
@@ -128,6 +129,7 @@ class Grok1MoE(nn.Module):
self.experts = MoEImpl(
num_experts=num_experts,
top_k=top_k,
layer_id=layer_id,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
params_dtype=params_dtype,
@@ -331,6 +333,7 @@ class Grok1DecoderLayer(nn.Module):
)
self.block_sparse_moe = Grok1MoE(
config=config,
layer_id=layer_id,
num_experts=config.num_local_experts,
top_k=config.num_experts_per_tok,
hidden_size=config.hidden_size,

1
python/sglang/srt/models/hunyuan.py
View File

@@ -163,6 +163,7 @@ class HunYuanSparseMoeBlock(nn.Module):
hidden_size=config.hidden_size,
intermediate_size=intermediate_size,
reduce_results=False,
layer_id=layer_id,
quant_config=quant_config,
)

3
python/sglang/srt/models/llama4.py
View File

@@ -87,6 +87,7 @@ class Llama4MoE(nn.Module):
def __init__(
self,
config: Llama4TextConfig,
layer_id: int,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
@@ -114,6 +115,7 @@ class Llama4MoE(nn.Module):
num_experts=config.num_local_experts,
hidden_size=config.hidden_size,
intermediate_size=intermediate_size_moe,
layer_id=layer_id,
reduce_results=False,
quant_config=quant_config,
apply_router_weight_on_input=True,
@@ -373,6 +375,7 @@ class Llama4DecoderLayer(nn.Module):
if is_moe_layer:
self.feed_forward = Llama4MoE(
config=config,
layer_id=layer_id,
quant_config=quant_config,
prefix=add_prefix("feed_forward", prefix),
)

3
python/sglang/srt/models/mixtral.py
View File

@@ -69,6 +69,7 @@ class MixtralMoE(nn.Module):
top_k: int,
hidden_size: int,
intermediate_size: int,
layer_id: int,
params_dtype: Optional[torch.dtype] = None,
quant_config: Optional[QuantizationConfig] = None,
tp_size: Optional[int] = None,
@@ -97,6 +98,7 @@ class MixtralMoE(nn.Module):
self.experts = MoEImpl(
num_experts=num_experts,
top_k=top_k,
layer_id=layer_id,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
params_dtype=params_dtype,
@@ -226,6 +228,7 @@ class MixtralDecoderLayer(nn.Module):
top_k=config.num_experts_per_tok,
hidden_size=config.hidden_size,
intermediate_size=config.intermediate_size,
layer_id=layer_id,
quant_config=quant_config,
prefix=add_prefix("block_sparse_moe", prefix),
)

3
python/sglang/srt/models/olmoe.py
View File

@@ -63,6 +63,7 @@ class OlmoeMoE(nn.Module):
params_dtype: Optional[torch.dtype] = None,
quant_config: Optional[QuantizationConfig] = None,
tp_size: Optional[int] = None,
layer_id: int = 0,
prefix: str = "",
):
super().__init__()
@@ -89,6 +90,7 @@ class OlmoeMoE(nn.Module):
reduce_results=True,
quant_config=quant_config,
tp_size=tp_size,
layer_id=layer_id,
prefix=add_prefix("experts", prefix),
)
@@ -224,6 +226,7 @@ class OlmoeDecoderLayer(nn.Module):
top_k=config.num_experts_per_tok,
hidden_size=config.hidden_size,
intermediate_size=config.intermediate_size,
layer_id=layer_id,
quant_config=quant_config,
prefix=add_prefix("mlp", prefix),
)

1
python/sglang/srt/models/phimoe.py
View File

@@ -210,6 +210,7 @@ class PhiMoE(nn.Module):
self.experts = FusedMoE(
num_experts=num_experts,
top_k=top_k,
layer_id=layer_id,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
reduce_results=True,