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 d8fc286fb..3c7c141b4 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -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 diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 93740f329..732ed7ff7 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -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, ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 075c57bdd..55850e5ea 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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,