diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 911b32402..120f42e95 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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: diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index b6bb780eb..a7800d4b8 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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( diff --git a/test/registered/unit/eplb/test_deepep_waterfill_eplb.py b/test/registered/unit/eplb/test_deepep_waterfill_eplb.py new file mode 100644 index 000000000..c193a4e0e --- /dev/null +++ b/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()