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

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:
+        require_global_experts = getattr(param, "_sglang_require_global_experts", False)
-            # This is a shared expert.
+        shared_expert_id = (
-            physical_expert_ids = [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

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

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