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 94a4b3991ad8..c798a06dcb59 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index e21e5f5102ff..e6ff2cb06d49 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index c6ac46ee8c9d..cb8dbf57c832 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/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: diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index b6bb780ebdc3..dd0ca4b81eea 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7f2b2f298604..6ca54a81b01a 100644 --- a/python/sglang/srt/server_args.py +++ b/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." ) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 6464d65e9576..7a99c491f6fd 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/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",