minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support Waterfill with MegaMoE backend (#27350)

Browse Source
This commit is contained in:
xutizhou
2026-07-13 03:56:46 -07:00
committed by GitHub
parent 874fc07d9b
commit eb31b5310c
14 changed files with 87 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
+6 -6
View File
@@ -1986,12 +1986,12 @@ def _moe_runner_fusion_disable(view: Any) -> dict:
def _a2a_fusion_adjustments(view: Any) -> dict:
"""A2A-backend-driven shared-experts fusion adjustments, declared at the
legacy write slots in _handle_a2a_moe: DeepEP Waterfill requires the
legacy write slots in _handle_a2a_moe: Waterfill requires the
fusion enabled; FlashInfer A2A requires it disabled."""
if view.moe_a2a_backend == "deepep" and view.enable_deepep_waterfill:
if view.moe_a2a_backend in ("deepep", "megamoe") and view.enable_waterfill:
if view.disable_shared_experts_fusion:
logger.warning(
"disable_shared_experts_fusion is overridden to False because DeepEP Waterfill requires shared expert fusion."
"disable_shared_experts_fusion is overridden to False because Waterfill requires shared expert fusion."
)
return {"disable_shared_experts_fusion": False}
return {}
@@ -2027,10 +2027,10 @@ _A2A_EP_SPANNING_BACKENDS = frozenset(
def _a2a_backend_overrides(view: Any) -> dict:
moe_a2a_backend = view.moe_a2a_backend
if view.enable_deepep_waterfill and moe_a2a_backend != "deepep":
if view.enable_waterfill and moe_a2a_backend not in ("deepep", "megamoe"):
logger.warning(
"moe_a2a_backend is overridden to 'deepep' because DeepEP "
"Waterfill requires the DeepEP backend."
"moe_a2a_backend is overridden to 'deepep' because Waterfill "
"requires the DeepEP or MegaMOE backend."
)
moe_a2a_backend = "deepep"
if envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get() and moe_a2a_backend != "megamoe":
python/sglang/srt/environ.py
+2 -2
View File
@@ -645,8 +645,8 @@ class Envs:
SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32)
SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False)
# Force dynamic DeepEP Waterfill with runtime EP all-reduce instead of the
# default static local-batch path.
# Force dynamic Waterfill with runtime EP all-reduce instead of the default
# static local-batch path.
SGLANG_DISABLE_STATIC_WATERFILL = EnvBool(False)
# NIXL-EP
python/sglang/srt/layers/moe/hash_topk.py
+11 -11
View File
@@ -46,12 +46,12 @@ class HashTopK(nn.Module):
self.layer_id = layer_id
from sglang.srt.runtime_context import get_server_args
self.enable_deepep_waterfill = (
num_fused_shared_experts > 0 and get_server_args().enable_deepep_waterfill
self.enable_waterfill = (
num_fused_shared_experts > 0 and get_server_args().enable_waterfill
)
self.deepep_waterfill_balancer = None
self.waterfill_balancer = None
if self.enable_deepep_waterfill:
if self.enable_waterfill:
# Waterfill appends the shared expert after EPLB maps routed IDs.
topk -= num_fused_shared_experts
num_fused_shared_experts = 0
@@ -119,18 +119,18 @@ class HashTopK(nn.Module):
(0, topk_output.topk_weights.shape[-1] + n)
),
)
return self._apply_deepep_waterfill(topk_output, num_tokens=0)
return self._apply_waterfill(topk_output, num_tokens=0)
def _apply_deepep_waterfill(
def _apply_waterfill(
self, topk_output: StandardTopKOutput, num_tokens: int
) -> StandardTopKOutput:
if self.enable_deepep_waterfill and self.deepep_waterfill_balancer is None:
if self.enable_waterfill and self.waterfill_balancer is None:
raise RuntimeError(
"DeepEP waterfill HashTopK must be prepared by ModelRunner before forward."
"Waterfill HashTopK must be prepared by ModelRunner before forward."
)
if self.deepep_waterfill_balancer is None:
if self.waterfill_balancer is None:
return topk_output
return self.deepep_waterfill_balancer.expand_topk(topk_output, num_tokens)
return self.waterfill_balancer.expand_topk(topk_output, num_tokens)
def _forward_torch(
self, router_logits: torch.Tensor, input_ids: torch.Tensor
@@ -267,7 +267,7 @@ class HashTopK(nn.Module):
topk_output = StandardTopKOutput(
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits
)
topk_output = self._apply_deepep_waterfill(topk_output, hidden_states.shape[0])
topk_output = self._apply_waterfill(topk_output, hidden_states.shape[0])
if is_hip():
_zero_topk_weights_padded_region(
topk_output.topk_weights, num_token_non_padded
python/sglang/srt/layers/moe/topk.py
+13 -15
View File
@@ -397,12 +397,12 @@ class TopK(MultiPlatformOp):
self.layer_id = layer_id
from sglang.srt.runtime_context import get_server_args
self.enable_deepep_waterfill = (
num_fused_shared_experts > 0 and get_server_args().enable_deepep_waterfill
self.enable_waterfill = (
num_fused_shared_experts > 0 and get_server_args().enable_waterfill
)
self.deepep_waterfill_balancer = None
if self.enable_deepep_waterfill:
self.waterfill_balancer = None
if self.enable_waterfill:
# TODO(ch-wan): Refactor shared-expert fusion and routed TopK fusion.
top_k -= num_fused_shared_experts
num_fused_shared_experts = 0
@@ -428,17 +428,15 @@ class TopK(MultiPlatformOp):
allow_routed_experts_capture=allow_routed_experts_capture,
)
def _apply_deepep_waterfill(
self, topk_output: TopKOutput, num_tokens: int
) -> TopKOutput:
if self.enable_deepep_waterfill and self.deepep_waterfill_balancer is None:
def _apply_waterfill(self, topk_output: TopKOutput, num_tokens: int) -> TopKOutput:
if self.enable_waterfill and self.waterfill_balancer is None:
raise RuntimeError(
"DeepEP waterfill TopK must be prepared by ModelRunner before forward."
"Waterfill TopK must be prepared by ModelRunner before forward."
)
if self.deepep_waterfill_balancer is None:
if self.waterfill_balancer is None:
return topk_output
assert TopKOutputChecker.format_is_standard(topk_output)
return self.deepep_waterfill_balancer.expand_topk(topk_output, num_tokens)
return self.waterfill_balancer.expand_topk(topk_output, num_tokens)
def forward_native(
self,
@@ -457,7 +455,7 @@ class TopK(MultiPlatformOp):
num_token_non_padded=num_token_non_padded,
expert_location_dispatch_info=expert_location_dispatch_info,
)
return self._apply_deepep_waterfill(topk_output, hidden_states.shape[0])
return self._apply_waterfill(topk_output, hidden_states.shape[0])
def forward_cuda(
self,
@@ -521,7 +519,7 @@ class TopK(MultiPlatformOp):
num_token_non_padded=num_token_non_padded,
expert_location_dispatch_info=expert_location_dispatch_info,
)
return self._apply_deepep_waterfill(topk_output, hidden_states.shape[0])
return self._apply_waterfill(topk_output, hidden_states.shape[0])
def forward_cpu(
self,
@@ -539,7 +537,7 @@ class TopK(MultiPlatformOp):
num_token_non_padded=num_token_non_padded,
expert_location_dispatch_info=expert_location_dispatch_info,
)
return self._apply_deepep_waterfill(topk_output, hidden_states.shape[0])
return self._apply_waterfill(topk_output, hidden_states.shape[0])
def forward_npu(
self,
@@ -604,7 +602,7 @@ class TopK(MultiPlatformOp):
(0, topk_output.topk_weights.shape[-1] + n)
),
)
return self._apply_deepep_waterfill(topk_output, 0)
return self._apply_waterfill(topk_output, 0)
def forward_xpu(
self,
python/sglang/srt/layers/moe/deepep_waterfill.py → python/sglang/srt/layers/moe/waterfill.py
+6 -6
View File
@@ -11,7 +11,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""DeepEP Waterfill: shared expert as 9th routed expert, dispatched to least-loaded rank."""
"""Waterfill: shared expert as 9th routed expert, dispatched to least-loaded rank."""
from typing import NamedTuple, Optional, Tuple
@@ -29,7 +29,7 @@ _LOCAL_PREF_DENOM = 10
class WaterfillDispatchPlan(NamedTuple):
"""Inputs needed by the fused DeepEP Waterfill expansion path."""
"""Inputs needed by the fused Waterfill expansion path."""
# Effective rank load consumed by the fused kernel.
rank_load: Tensor
@@ -275,10 +275,10 @@ def materialize_waterfill_dispatch_fused(
allow_all_ranks: bool = False,
target_total: int = 0,
) -> Tuple[Tensor, Tensor]:
"""Run fused Waterfill rank selection and DeepEP TopK expansion.
"""Run fused Waterfill rank selection and TopK expansion.
The Triton kernel intentionally selects each token's shared-expert rank and
writes the expanded DeepEP TopK layout in one pass.
writes the expanded TopK layout in one pass.
"""
num_tokens = topk_ids.shape[0]
topk = topk_ids.shape[1]
@@ -358,7 +358,7 @@ def expand_topk_with_shared_expert(
return expanded_topk_ids, expanded_topk_weights
class DeepEPWaterfillBalancer:
class WaterfillBalancer:
"""Waterfill load balancer: shared expert fused as real routed expert (topk 8→9)."""
MIN_BATCH_FOR_BALANCE = 64
@@ -484,7 +484,7 @@ class DeepEPWaterfillBalancer:
return self._build_static_dispatch_plan(local_routed_counts)
global_routed_counts, local_tokens_per_rank = (
DeepEPWaterfillBalancer._all_reduce_dynamic_rank_load(
WaterfillBalancer._all_reduce_dynamic_rank_load(
local_routed_counts, num_tokens
)
)
python/sglang/srt/model_executor/model_runner.py
+6 -11
View File
@@ -1631,10 +1631,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
for module in self.model.modules():
if not isinstance(module, (TopK, HashTopK)):
continue
if (
not module.enable_deepep_waterfill
or module.deepep_waterfill_balancer is not None
):
if not module.enable_waterfill or module.waterfill_balancer is not None:
continue
if num_routed_experts is None:
num_routed_experts = getattr(
@@ -1642,14 +1639,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
if num_routed_experts is None:
raise ValueError(
"DeepEP waterfill requires model config n_routed_experts."
"Waterfill requires model config n_routed_experts."
)
if balancer_cls is None:
from sglang.srt.layers.moe.deepep_waterfill import (
DeepEPWaterfillBalancer,
)
from sglang.srt.layers.moe.waterfill import WaterfillBalancer
balancer_cls = DeepEPWaterfillBalancer
balancer_cls = WaterfillBalancer
# 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.
@@ -1660,7 +1655,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
routed_scaling_factor = module.topk_config.routed_scaling_factor
else:
routed_scaling_factor = module.routed_scaling_factor
module.deepep_waterfill_balancer = balancer_cls(
module.waterfill_balancer = balancer_cls(
num_routed_experts=num_physical_routed_experts,
world_size=self.moe_ep_size,
rank=self.moe_ep_rank,
@@ -1672,7 +1667,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
num_prepared += 1
if num_prepared:
log_info_on_rank0(
logger, f"Prepared {num_prepared} DeepEP waterfill TopK modules."
logger, f"Prepared {num_prepared} Waterfill TopK modules."
)
def _init_lplb_solvers(self):
python/sglang/srt/server_args.py
+7 -8
View File
@@ -1977,9 +1977,9 @@ class ServerArgs:
Optional[str],
"The InfiniBand devices for Mooncake Backend transfer, accepts multiple comma-separated devices (e.g., --mooncake-ib-device mlx5_0,mlx5_1). Default is None, which triggers automatic device detection when Mooncake Backend is enabled.",
] = None
enable_deepep_waterfill: A[
enable_waterfill: A[
bool,
"Enable DeepEP Waterfill: dispatch the shared expert as the 9th routed expert to the least-loaded EP rank. Automatically sets --moe-a2a-backend deepep, implicitly enables shared-expert fusion, and supports --deepep-mode auto, normal, or low_latency. Use auto or low_latency for production decode so CUDA graph remains enabled. Supported on DeepSeek-V3/R1 with EP >= 2.",
"Enable Waterfill: dispatch the fused shared expert as an extra routed expert slot to the least-loaded EP rank. Supports DeepEP and MegaMOE MoE A2A backends, implicitly enables shared-expert fusion, and supports --deepep-mode auto, normal, or low_latency when used with DeepEP. Use auto or low_latency for production DeepEP decode so CUDA graph remains enabled. Supported on DeepSeek-V3/R1 with EP >= 2.",
] = False
elastic_ep_rejoin: A[
bool,
@@ -1992,7 +1992,7 @@ class ServerArgs:
disable_shared_experts_fusion: A[
bool,
Arg(
help="Disable the built-in shared experts fusion optimization for DeepSeek V3/R1. Note: DeepEP Waterfill (--enable-deepep-waterfill) still routes shared expert through DeepEP as an extra MoE slot, so shared expert is not separated from the MoE path when Waterfill is enabled.",
help="Disable the built-in shared experts fusion optimization for DeepSeek V3/R1. Note: Waterfill (--enable-waterfill) routes the shared expert as an extra MoE slot, so the shared expert is not separated from the MoE path when Waterfill is enabled.",
resolvable=True,
),
] = False
@@ -5492,6 +5492,10 @@ class ServerArgs:
run_post_process_pass(self, _a2a_fusion_adjustments)
a2a_backend = resolved_view(self).moe_a2a_backend
if self.enable_waterfill:
self.enforce_shared_experts_fusion = True
logger.info(f"Waterfill is enabled with moe_a2a_backend='{a2a_backend}'.")
if a2a_backend == "megamoe":
if not envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.is_set():
envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.set(True)
@@ -5508,11 +5512,6 @@ class ServerArgs:
logger.warning(
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:
self.enforce_shared_experts_fusion = True
logger.info(
"DeepEP Waterfill is enabled. Shared expert will be dispatched through DeepEP for load balancing."
)
if a2a_backend == "mooncake":
logger.warning(