Support Waterfill with MegaMoE backend

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -53,7 +53,10 @@
     TopKOutput,
     TopKOutputChecker,
 )
-from sglang.srt.layers.moe.utils import RoutingMethodType, is_deepep_class_backend
+from sglang.srt.layers.moe.utils import (
+    RoutingMethodType,
+    use_rank_local_fused_shared_experts,
+)
 from sglang.srt.layers.quantization.base_config import (
     FusedMoEMethodBase,
     QuantizationConfig,
@@ -199,10 +202,10 @@ def __init__(
         self.moe_tp_size = get_moe_tensor_parallel_world_size()
         self.moe_tp_rank = get_moe_tensor_parallel_rank()
 
-        # DeepEP: each rank has its own shared expert slot, so total shared
-        # weight slots = num_fused_shared_experts * ep_size.
+        # Rank-local shared dispatch: each EP rank has its own shared expert
+        # slot, so total shared weight slots = num_fused_shared_experts * ep_size.
         # AMD/Standard: shared experts are global, slots = num_fused_shared_experts.
-        if num_fused_shared_experts > 0 and is_deepep_class_backend():
+        if num_fused_shared_experts > 0 and use_rank_local_fused_shared_experts():
             num_shared_slots = num_fused_shared_experts * self.moe_ep_size
         else:
             num_shared_slots = num_fused_shared_experts
@@ -764,7 +767,7 @@ def weight_loader(
         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():
+            if require_global_experts and use_rank_local_fused_shared_experts():
                 physical_expert_ids = [
                     rank * self.num_local_experts
                     + self._num_local_routed

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

@@ -97,7 +97,7 @@ def routing(
 )
 from sglang.srt.layers.dp_attention import is_allocation_symmetric
 from sglang.srt.layers.moe import get_moe_runner_backend
-from sglang.srt.layers.moe.utils import is_deepep_class_backend
+from sglang.srt.layers.moe.utils import use_rank_local_fused_shared_experts
 from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
 from sglang.srt.utils import (
@@ -520,7 +520,10 @@ def empty_topk_output(self, device: torch.device) -> TopKOutput:
         # FIXME: router_logits should be of size (0, num_experts)
         router_logits = torch.empty((0, topk), dtype=torch.float32, device=device)
         topk_output = StandardTopKOutput(topk_weights, topk_ids, router_logits)
-        if self.topk_config.num_fused_shared_experts > 0 and is_deepep_class_backend():
+        if (
+            self.topk_config.num_fused_shared_experts > 0
+            and use_rank_local_fused_shared_experts()
+        ):
             n = self.topk_config.num_fused_shared_experts
             topk_output = topk_output._replace(
                 topk_ids=topk_output.topk_ids.new_empty(
@@ -1394,7 +1397,7 @@ def _post_process_topk_ids(
         # EPLB dispatch must only remap the routed expert columns.
         # The shared expert column (value = n_routed_experts) would be out-of-bounds
         # for the logical-to-physical dispatch table.
-        if num_fused_shared_experts > 0 and is_deepep_class_backend():
+        if num_fused_shared_experts > 0 and use_rank_local_fused_shared_experts():
             shared_cols = topk_ids[:, -num_fused_shared_experts:]
             routed_cols = topk_ids[:, :-num_fused_shared_experts]
             routed_cols = _biased_grouped_topk_postprocess(
@@ -1429,7 +1432,7 @@ 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():
+    if num_fused_shared_experts > 0 and use_rank_local_fused_shared_experts():
         num_physical_routed_experts = (
             expert_location_dispatch_info.num_physical_experts
             if expert_location_dispatch_info is not None

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

@@ -355,10 +355,20 @@ def is_sbo_enabled() -> bool:
     return IS_SBO_ENABLED
 
 
-def is_deepep_class_backend() -> bool:
-    """Check if the MoE backend is DeepEP-family (DeepEP, Mooncake, or Mori)."""
+def use_rank_local_fused_shared_experts() -> bool:
+    """Whether fused shared experts need one physical slot per EP rank."""
     b = get_moe_a2a_backend()
-    return b.is_deepep() or b.is_mooncake() or b.is_mori()
+    if b.is_deepep() or b.is_mooncake() or b.is_mori():
+        return True
+    try:
+        server_args = get_global_server_args()
+    except ValueError:
+        return False
+    return (
+        b.is_megamoe()
+        and server_args is not None
+        and server_args.enable_deepep_waterfill
+    )
 
 
 def is_flashinfer_cutedsl_v1_path() -> bool:

python/sglang/srt/models/deepseek_v2.py

@@ -106,9 +106,9 @@
 from sglang.srt.layers.moe.utils import (
     RoutingMethodType,
     filter_moe_weight_param_global_expert,
-    is_deepep_class_backend,
     is_sbo_enabled,
     is_tbo_enabled,
+    use_rank_local_fused_shared_experts,
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.quantization.fp8 import Fp8Config
@@ -537,18 +537,20 @@ def __init__(
         # mlp.shared_experts → mlp.experts.256 when > 0.
         self.num_fused_shared_experts = 0 if _fusion_disabled else n_shared_experts
 
-        # DeepEP shared expert fusion: shared expert is fused into the same MoE kernel
-        # as a local expert at the home EP rank. Expert layout is expanded from 256
-        # routed to 256+EP_size (e.g. 272 for EP=16). TopK handles interleaving.
-        _is_deepep_fusion = (
-            is_deepep_class_backend() and self.num_fused_shared_experts > 0
+        # Rank-local shared expert fusion: shared expert is fused into the same
+        # MoE kernel as a local expert at each EP rank. Expert layout is expanded
+        # from 256 routed to 256+EP_size (e.g. 272 for EP=16). TopK handles
+        # interleaving.
+        _use_rank_local_shared_slots = (
+            use_rank_local_fused_shared_experts()
+            and self.num_fused_shared_experts > 0
         )
 
-        if _is_deepep_fusion:
+        if _use_rank_local_shared_slots:
             # 256 routed + EP_size shared slots = 272 experts total (for EP=16)
             num_experts_for_moe = config.n_routed_experts + self.moe_ep_size
             top_k_for_moe = config.num_experts_per_tok + 1  # 8 routed + 1 shared
-            # Interleaving for DeepEP dispatch is handled by TopK internally.
+            # Interleaving for rank-local shared dispatch is handled by TopK internally.
         else:
             num_experts_for_moe = (
                 config.n_routed_experts + self.num_fused_shared_experts
@@ -2301,9 +2303,11 @@ def __init__(
             for i in range(len(self.layers)):
                 if isinstance(self.layers[i].mlp, DeepseekV2MoE):
                     # tp_size = get_tensor_model_parallel_world_size()
-                    is_a2a_moe = is_deepep_class_backend()
+                    uses_rank_local_shared_slots = use_rank_local_fused_shared_experts()
                     tp_size = (
-                        1 if is_a2a_moe else get_tensor_model_parallel_world_size()
+                        1
+                        if uses_rank_local_shared_slots
+                        else get_tensor_model_parallel_world_size()
                     )
                     intermediate_size = (
                         config.moe_intermediate_size * config.n_shared_experts
@@ -2573,8 +2577,8 @@ def determine_num_fused_shared_experts(
             pass
         elif is_sbo_enabled() or is_tbo_enabled():
             disable_reason = "SBO/TBO enabled: incompatible with fusing shared expert into MoE kernel."
-        elif is_deepep_class_backend():
-            disable_reason = "DeepEP: fusion off by default (use --enforce-shared-experts-fusion to enable)."
+        elif use_rank_local_fused_shared_experts():
+            disable_reason = "Rank-local shared expert dispatch: fusion off by default (use --enforce-shared-experts-fusion to enable)."
         elif (
             self.config.architectures[0] != architecture
             # Allow-list of n_routed_experts values that have been validated

python/sglang/srt/server_args.py

@@ -3445,13 +3445,6 @@ def _validate_cutedsl_a2a_token_budget(self):
             )
 
     def _handle_a2a_moe(self):
-        if self.enable_deepep_waterfill and self.moe_a2a_backend != "deepep":
-            logger.warning(
-                "moe_a2a_backend is overridden to 'deepep' because DeepEP "
-                "Waterfill requires the DeepEP backend."
-            )
-            self.moe_a2a_backend = "deepep"
-
         if (
             envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get()
             and self.moe_a2a_backend != "megamoe"
@@ -3462,6 +3455,24 @@ def _handle_a2a_moe(self):
                 "auto-configuring --moe-a2a-backend megamoe."
             )
 
+        if self.enable_deepep_waterfill and self.moe_a2a_backend not in (
+            "deepep",
+            "megamoe",
+        ):
+            logger.warning(
+                "moe_a2a_backend is overridden to 'deepep' because DeepEP "
+                "Waterfill requires the DeepEP backend."
+            )
+            self.moe_a2a_backend = "deepep"
+
+        if self.enable_deepep_waterfill:
+            if self.disable_shared_experts_fusion:
+                logger.warning(
+                    "disable_shared_experts_fusion is overridden to False because DeepEP Waterfill requires shared expert fusion."
+                )
+                self.disable_shared_experts_fusion = False
+            self.enforce_shared_experts_fusion = True
+
         if self.moe_a2a_backend == "megamoe":
             self.ep_size = self.tp_size
             if not envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.is_set():
@@ -3480,12 +3491,6 @@ def _handle_a2a_moe(self):
                 f"DeepEP MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]."
             )
             if self.enable_deepep_waterfill:
-                if self.disable_shared_experts_fusion:
-                    logger.warning(
-                        "disable_shared_experts_fusion is overridden to False because DeepEP Waterfill requires shared expert fusion."
-                    )
-                    self.disable_shared_experts_fusion = False
-                self.enforce_shared_experts_fusion = True
                 logger.info(
                     "DeepEP Waterfill is enabled. Shared expert will be dispatched through DeepEP for load balancing."
                 )

test/registered/unit/server_args/test_server_args.py

@@ -548,6 +548,33 @@ def test_waterfill_overrides_moe_a2a_backend_to_deepep(self):
         self.assertEqual(server_args.moe_a2a_backend, "deepep")
         self.assertTrue(server_args.enforce_shared_experts_fusion)
 
+    def test_waterfill_keeps_explicit_megamoe_backend(self):
+        server_args = ServerArgs(
+            model_path="dummy",
+            moe_a2a_backend="megamoe",
+            enable_deepep_waterfill=True,
+            disable_shared_experts_fusion=True,
+        )
+        # dummy-model path short-circuits __post_init__; invoke the handler directly.
+        server_args._handle_a2a_moe()
+
+        self.assertEqual(server_args.moe_a2a_backend, "megamoe")
+        self.assertFalse(server_args.disable_shared_experts_fusion)
+        self.assertTrue(server_args.enforce_shared_experts_fusion)
+
+    def test_waterfill_uses_megamoe_env(self):
+        server_args = ServerArgs(
+            model_path="dummy",
+            moe_a2a_backend="none",
+            enable_deepep_waterfill=True,
+        )
+        # dummy-model path short-circuits __post_init__; invoke the handler directly.
+        with server_args_module.envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.override(True):
+            server_args._handle_a2a_moe()
+
+        self.assertEqual(server_args.moe_a2a_backend, "megamoe")
+        self.assertTrue(server_args.enforce_shared_experts_fusion)
+
     def test_waterfill_supports_deepep_low_latency_mode(self):
         server_args = ServerArgs(
             model_path="dummy",