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

Fix EPLB redundant experts with shared expert fusion and Waterfill (#25367)

Browse Source
This commit is contained in:
xutizhou
2026-05-20 22:58:08 -07:00
committed by GitHub
parent 847cbada9c
commit e8608bdcb5
3 changed files with 36 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+18 -6
View File
@@ -638,13 +638,25 @@ class FusedMoE(torch.nn.Module):
)
return
if self._has_fused_shared and expert_id >= self._num_global_routed:
# This is a shared expert.
physical_expert_ids = [expert_id]
require_global_experts = getattr(param, "_sglang_require_global_experts", False)
shared_expert_id = (
expert_id - global_expert_location_metadata.num_logical_experts
if self._has_fused_shared and expert_id is not None
else -1
)
if 0 <= shared_expert_id < self.num_fused_shared_experts:
# Checkpoint shared experts start after logical routed experts, while
# local fused MoE weights store them after physical routed experts.
if require_global_experts and is_deepep_class_backend():
physical_expert_ids = [
rank * self.num_local_experts
+ self._num_local_routed
+ shared_expert_id
for rank in range(self.moe_ep_size)
]
else:
physical_expert_ids = [self._num_global_routed + shared_expert_id]
else:
require_global_experts = getattr(
param, "_sglang_require_global_experts", False
)
physical_expert_ids = (
global_expert_location_metadata.logical_to_all_physical(
self.layer_id, expert_id, require_global_experts
python/sglang/srt/layers/moe/topk.py
+11 -3
View File
@@ -1242,7 +1242,7 @@ def _remap_topk_for_deepep(
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
num_fused_shared_experts: int,
n_routed_experts: int,
num_physical_routed_experts: int,
topk_config: TopKConfig,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Remap TopK output to DeepEP interleaved expert layout.
@@ -1260,7 +1260,10 @@ def _remap_topk_for_deepep(
ep_size = get_moe_expert_parallel_world_size()
ep_rank = get_moe_expert_parallel_rank()
num_local_routed = n_routed_experts // ep_size
# Static EPLB may add redundant physical experts. At this point routed
# topk_ids have already been remapped from logical to physical ids, so the
# DeepEP interleaved layout must use the physical routed count.
num_local_routed = num_physical_routed_experts // ep_size
num_local_experts = num_local_routed + num_fused_shared_experts
# Remap routed IDs: insert gaps for shared expert slots (single fused op)
@@ -1341,11 +1344,16 @@ def _post_process_topk_ids(
# DeepEP: remap to interleaved expert layout where each rank's shared
# expert has a unique ID for dispatch routing.
if num_fused_shared_experts > 0 and is_deepep_class_backend():
num_physical_routed_experts = (
expert_location_dispatch_info.num_physical_experts
if expert_location_dispatch_info is not None
else router_logits.shape[1]
)
topk_ids, topk_weights = _remap_topk_for_deepep(
topk_ids,
topk_weights,
num_fused_shared_experts,
router_logits.shape[1],
num_physical_routed_experts,
topk_config,
)
python/sglang/srt/model_executor/model_runner.py
+7 -1
View File
@@ -1447,8 +1447,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
balancer_cls = DeepEPWaterfillBalancer
# Static EPLB remaps TopK ids to physical expert ids before Waterfill.
# Redundant experts therefore need to be included in the per-rank
# expert count used for Waterfill's shared-expert slot remapping.
num_physical_routed_experts = (
num_routed_experts + self.server_args.ep_num_redundant_experts
)
module.deepep_waterfill_balancer = balancer_cls(
num_routed_experts=num_routed_experts,
num_routed_experts=num_physical_routed_experts,
world_size=self.moe_ep_size,
rank=self.moe_ep_rank,
layer_id=module.layer_id,