Support Waterfill with dynamic EPLB (#27150)

python/sglang/srt/layers/moe/topk.py

@@ -1464,7 +1464,7 @@ def _post_process_topk_ids(
     layer_id: int,
     num_token_non_padded: Optional[torch.Tensor] = None,
     expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
-) -> torch.Tensor:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     num_fused_shared_experts = topk_config.num_fused_shared_experts
     fused_shared_experts_scaling_factor = (
         topk_config.fused_shared_experts_scaling_factor
@@ -1474,6 +1474,7 @@ def _post_process_topk_ids(
             layer_id=layer_id,
             topk_indices=topk_ids,
         )
+    recorder_topk_ids = None
     if _is_cuda:
         # When shared experts are fused (appended as extra columns in topk_ids),
         # EPLB dispatch must only remap the routed expert columns.
@@ -1486,11 +1487,18 @@ def _post_process_topk_ids(
                 routed_cols, expert_location_dispatch_info, num_token_non_padded
             )
             topk_ids = torch.cat([routed_cols, shared_cols], dim=-1)
+            # ExpertDistributionRecorder tracks EPLB physical routed experts.
+            # DeepEP dispatch later inserts per-rank shared slots into topk_ids,
+            # so keep the routed physical ids separately for statistics.
+            recorder_topk_ids = routed_cols
         else:
             topk_ids = _biased_grouped_topk_postprocess(
                 topk_ids, expert_location_dispatch_info, num_token_non_padded
             )
 
+    if recorder_topk_ids is None:
+        recorder_topk_ids = topk_ids
+
     if num_fused_shared_experts > 0 and _use_aiter:
         M, N = router_logits.shape
         scale_factor = (
@@ -1528,7 +1536,7 @@ def _post_process_topk_ids(
             topk_config,
         )
 
-    return topk_ids, topk_weights
+    return topk_ids, topk_weights, recorder_topk_ids
 
 
 def select_experts(
@@ -1746,7 +1754,7 @@ def select_experts(
         if k > 0:
             topk_weights = torch.full_like(topk_weights, 1.0 / k)
 
-    topk_ids, topk_weights = _post_process_topk_ids(
+    topk_ids, topk_weights, recorder_topk_ids = _post_process_topk_ids(
         topk_ids=topk_ids,
         topk_weights=topk_weights,
         topk_config=topk_config,
@@ -1756,7 +1764,9 @@ def select_experts(
         expert_location_dispatch_info=expert_location_dispatch_info,
     )
 
-    get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
+    get_global_expert_distribution_recorder().on_select_experts(
+        topk_ids=recorder_topk_ids
+    )
 
     # ===== TO BE REFACTORED ====
     if packed_topk is not None:

python/sglang/srt/models/deepseek_v2.py

@@ -796,8 +796,14 @@ class DeepseekV2MoE(nn.Module):
         self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()
 
     def get_moe_weights(self):
+        # EPLB only rebalances physical routed experts. Fused shared expert
+        # slots live after each rank's routed slots and must stay stable.
+        num_local_experts_for_eplb = (
+            self.experts.num_local_experts - self.num_fused_shared_experts
+        )
+
         return [
-            x.data
+            x.data[:num_local_experts_for_eplb]
             for name, x in self.experts.named_parameters()
             if name not in ["correction_bias"]
             and filter_moe_weight_param_global_expert(

test/registered/unit/eplb/test_deepep_waterfill_eplb.py

@@ -0,0 +1,138 @@
+"""Unit tests for DeepEP Waterfill and EPLB updater compatibility."""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=7, suite="base-a-test-cpu")
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+from torch import nn
+
+from sglang.srt.layers.moe import topk as topk_module
+from sglang.srt.layers.moe.topk import TopKConfig
+from sglang.srt.models.deepseek_v2 import DeepseekV2MoE
+from sglang.test.test_utils import CustomTestCase
+
+
+class _FakeExpertParam(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.num_local_experts = 5
+        self.weight = nn.Parameter(
+            torch.arange(10, dtype=torch.float32).reshape(self.num_local_experts, 2)
+        )
+        self.correction_bias = nn.Parameter(torch.ones(self.num_local_experts))
+        self.global_scale = nn.Parameter(torch.ones(self.num_local_experts))
+        self.global_scale._sglang_require_global_experts = True
+
+
+class TestDeepEPWaterfillEPLB(CustomTestCase):
+    def test_deepseek_moe_get_moe_weights_excludes_fused_shared_slot(self):
+        experts = _FakeExpertParam()
+        moe = SimpleNamespace(num_fused_shared_experts=1, experts=experts)
+        shared_before = experts.weight.data[-1].clone()
+
+        weights = DeepseekV2MoE.get_moe_weights(moe)
+
+        self.assertEqual(len(weights), 1)
+        self.assertEqual(
+            weights[0].shape,
+            (experts.num_local_experts - moe.num_fused_shared_experts, 2),
+        )
+
+        weights[0][-1].zero_()
+        self.assertTrue(torch.equal(experts.weight.data[-2], torch.zeros(2)))
+        self.assertTrue(torch.equal(experts.weight.data[-1], shared_before))
+
+    def test_deepseek_moe_get_moe_weights_keeps_full_shape_without_fusion(self):
+        experts = _FakeExpertParam()
+        moe = SimpleNamespace(num_fused_shared_experts=0, experts=experts)
+        weights = DeepseekV2MoE.get_moe_weights(moe)
+
+        self.assertEqual(len(weights), 1)
+        self.assertEqual(weights[0].shape, (experts.num_local_experts, 2))
+
+    def test_topk_recorder_ids_exclude_deepep_fused_shared_slots(self):
+        topk_ids = torch.tensor([[0, 33, 263, 256]], dtype=torch.int32)
+        topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
+        topk_config = TopKConfig(
+            top_k=4,
+            num_fused_shared_experts=1,
+            routed_scaling_factor=1.0,
+        )
+        dispatch_info = SimpleNamespace(num_physical_experts=264)
+
+        def fake_eplb_postprocess(
+            ids, expert_location_dispatch_info, num_token_non_padded
+        ):
+            return ids
+
+        with (
+            patch.object(topk_module, "_is_cuda", True),
+            patch.object(topk_module, "_use_aiter", False),
+            patch.object(topk_module, "is_deepep_class_backend", return_value=True),
+            patch.object(
+                topk_module, "get_moe_expert_parallel_world_size", return_value=8
+            ),
+            patch.object(topk_module, "get_moe_expert_parallel_rank", return_value=7),
+            patch.object(
+                topk_module,
+                "_biased_grouped_topk_postprocess",
+                side_effect=fake_eplb_postprocess,
+            ),
+        ):
+            processed_ids, _, recorder_ids = topk_module._post_process_topk_ids(
+                topk_ids=topk_ids.clone(),
+                topk_weights=topk_weights.clone(),
+                topk_config=topk_config,
+                router_logits=torch.empty((1, 256)),
+                layer_id=0,
+                expert_location_dispatch_info=dispatch_info,
+            )
+
+        self.assertTrue(torch.equal(processed_ids, torch.tensor([[0, 34, 270, 271]])))
+        self.assertTrue(torch.equal(recorder_ids, torch.tensor([[0, 33, 263]])))
+
+    def test_topk_recorder_ids_match_dispatch_ids_for_non_deepep_fusion(self):
+        topk_ids = torch.tensor([[0, 33, 263, 256]], dtype=torch.int32)
+        topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
+        topk_config = TopKConfig(
+            top_k=4,
+            num_fused_shared_experts=1,
+            routed_scaling_factor=1.0,
+        )
+        dispatch_info = SimpleNamespace(num_physical_experts=264)
+
+        def fake_eplb_postprocess(
+            ids, expert_location_dispatch_info, num_token_non_padded
+        ):
+            return ids + 1
+
+        with (
+            patch.object(topk_module, "_is_cuda", True),
+            patch.object(topk_module, "_use_aiter", False),
+            patch.object(topk_module, "is_deepep_class_backend", return_value=False),
+            patch.object(
+                topk_module,
+                "_biased_grouped_topk_postprocess",
+                side_effect=fake_eplb_postprocess,
+            ),
+        ):
+            processed_ids, _, recorder_ids = topk_module._post_process_topk_ids(
+                topk_ids=topk_ids.clone(),
+                topk_weights=topk_weights.clone(),
+                topk_config=topk_config,
+                router_logits=torch.empty((1, 256)),
+                layer_id=0,
+                expert_location_dispatch_info=dispatch_info,
+            )
+
+        self.assertTrue(torch.equal(processed_ids, torch.tensor([[1, 34, 264, 257]])))
+        self.assertTrue(torch.equal(recorder_ids, processed_ids))
+
+
+if __name__ == "__main__":
+    unittest.main()