Support Waterfill with MegaMoE backend (#27350)
This commit is contained in:
@@ -198,7 +198,7 @@ sglang serve \
|
||||
|
||||
This BF16 setting applies only to the compressed attention state pools and reduces the GPU memory footprint of each compressed-state slot. It does not change model weight precision or the main KV cache dtype. With automatic pool sizing and no explicit capacity cap, the same memory budget holds more slots, and the startup log shows larger `c4_state` and `c128_state` pool sizes. Keep the default `float32` setting for the most conservative behavior.
|
||||
|
||||
**EPLB + DeepEP Waterfill (Experimental)**
|
||||
**EPLB + Waterfill (Experimental)**
|
||||
|
||||
For recorded/static EPLB reproduction, first record an expert-distribution file by following
|
||||
[Capture expert selection distribution in MoE models](../../../docs/basic_usage/native_api.mdx#capture-expert-selection-distribution-in-moe-models).
|
||||
@@ -211,7 +211,7 @@ For non-PD reproduction, use:
|
||||
--moe-a2a-backend deepep \
|
||||
--deepep-mode auto \
|
||||
--init-expert-location /path/to/expert_distribution_recorder_*.pt \
|
||||
--enable-deepep-waterfill
|
||||
--enable-waterfill
|
||||
```
|
||||
|
||||
For PD-Disagg reproduction, use `normal` mode on the prefill server and
|
||||
@@ -223,21 +223,21 @@ flag to both commands:
|
||||
--moe-a2a-backend deepep \
|
||||
--deepep-mode normal \
|
||||
--init-expert-location /path/to/expert_distribution_recorder_*.pt \
|
||||
--enable-deepep-waterfill
|
||||
--enable-waterfill
|
||||
|
||||
# decode
|
||||
--moe-a2a-backend deepep \
|
||||
--deepep-mode low_latency \
|
||||
--init-expert-location /path/to/expert_distribution_recorder_*.pt \
|
||||
--enable-deepep-waterfill
|
||||
--enable-waterfill
|
||||
```
|
||||
|
||||
You can also add `--ep-num-redundant-experts` and `--eplb-algorithm` to customize
|
||||
EPLB placement.
|
||||
|
||||
MegaMoE is not supported with this DeepEP Waterfill recipe yet. Waterfill routes
|
||||
the shared expert through DeepEP for load balancing, so `--enable-deepep-waterfill`
|
||||
requires `--moe-a2a-backend deepep`.
|
||||
Waterfill also supports MegaMOE. Use `--moe-a2a-backend megamoe
|
||||
--enable-waterfill` to keep the MegaMOE backend while applying Waterfill to the
|
||||
fused shared expert slot.
|
||||
|
||||
**FP4 Indexer (Experimental)**
|
||||
|
||||
|
||||
@@ -1869,8 +1869,8 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: str</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-deepep-waterfill`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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. By default, Waterfill uses the static local-batch path; set `SGLANG_DISABLE_STATIC_WATERFILL=1` to force dynamic Waterfill with runtime EP all-reduce.</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-waterfill`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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. By default, Waterfill uses the static local-batch path; set `SGLANG_DISABLE_STATIC_WATERFILL=1` to force dynamic Waterfill with runtime EP all-reduce.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
|
||||
</tr>
|
||||
|
||||
@@ -1733,7 +1733,7 @@ If the value is int8, you must also set the environment variable:DEEP_NORMAL_MOD
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-deepep-waterfill`</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-waterfill`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`False`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>bool flag (set to enable)</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
|
||||
|
||||
@@ -403,7 +403,7 @@ SGLang supports various environment variables that can be used to configure its
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DISABLE_STATIC_WATERFILL</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Force dynamic DeepEP waterfill with runtime EP all-reduce instead of the default static local-batch path.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Force dynamic Waterfill with runtime EP all-reduce instead of the default static local-batch path.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
+6
-6
@@ -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
|
||||
)
|
||||
)
|
||||
@@ -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):
|
||||
|
||||
@@ -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(
|
||||
|
||||
+2
-2
@@ -1,4 +1,4 @@
|
||||
"""Unit tests for DeepEP Waterfill and EPLB updater compatibility."""
|
||||
"""Unit tests for Waterfill and EPLB updater compatibility."""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
@@ -30,7 +30,7 @@ class _FakeExpertParam(nn.Module):
|
||||
self.global_scale._sglang_require_global_experts = True
|
||||
|
||||
|
||||
class TestDeepEPWaterfillEPLB(CustomTestCase):
|
||||
class TestWaterfillEPLB(CustomTestCase):
|
||||
def test_deepseek_moe_get_moe_weights_excludes_fused_shared_slot(self):
|
||||
experts = _FakeExpertParam()
|
||||
moe = SimpleNamespace(num_fused_shared_experts=1, experts=experts)
|
||||
@@ -973,12 +973,12 @@ class TestAdaptiveSpecArgs(CustomTestCase):
|
||||
self.assertEqual(args.speculative_num_draft_tokens, 4)
|
||||
|
||||
|
||||
class TestDeepEPWaterfillArgs(CustomTestCase):
|
||||
class TestWaterfillArgs(CustomTestCase):
|
||||
def test_waterfill_enforces_shared_experts_fusion(self):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
moe_a2a_backend="deepep",
|
||||
enable_deepep_waterfill=True,
|
||||
enable_waterfill=True,
|
||||
disable_shared_experts_fusion=True,
|
||||
)
|
||||
# dummy-model path short-circuits __post_init__; invoke the handler directly.
|
||||
@@ -995,7 +995,7 @@ class TestDeepEPWaterfillArgs(CustomTestCase):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
moe_a2a_backend="none",
|
||||
enable_deepep_waterfill=True,
|
||||
enable_waterfill=True,
|
||||
)
|
||||
# dummy-model path short-circuits __post_init__; invoke the handler directly.
|
||||
server_args._handle_a2a_moe()
|
||||
@@ -1006,11 +1006,27 @@ class TestDeepEPWaterfillArgs(CustomTestCase):
|
||||
self.assertEqual(resolved_view(server_args).moe_a2a_backend, "deepep")
|
||||
self.assertTrue(server_args.enforce_shared_experts_fusion)
|
||||
|
||||
def test_waterfill_keeps_megamoe_backend(self):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
moe_a2a_backend="megamoe",
|
||||
enable_waterfill=True,
|
||||
disable_shared_experts_fusion=True,
|
||||
)
|
||||
# dummy-model path short-circuits __post_init__; invoke the handler directly.
|
||||
server_args._handle_a2a_moe()
|
||||
|
||||
from sglang.srt.arg_groups.overrides import resolved_view
|
||||
|
||||
self.assertEqual(resolved_view(server_args).moe_a2a_backend, "megamoe")
|
||||
self.assertFalse(resolved_view(server_args).disable_shared_experts_fusion)
|
||||
self.assertTrue(server_args.enforce_shared_experts_fusion)
|
||||
|
||||
def test_waterfill_supports_deepep_low_latency_mode(self):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
moe_a2a_backend="deepep",
|
||||
enable_deepep_waterfill=True,
|
||||
enable_waterfill=True,
|
||||
deepep_mode="low_latency",
|
||||
)
|
||||
# dummy-model path short-circuits __post_init__; invoke the handler directly.
|
||||
@@ -1452,7 +1468,7 @@ class TestTwoBatchOverlapBackend(CustomTestCase):
|
||||
SGLANG_ENABLE_DP_TBO env: enabling DP TBO now needs no extra flag.
|
||||
|
||||
dummy-model short-circuits __post_init__, so the guard handler is invoked
|
||||
directly (same pattern as TestDeepEPWaterfillArgs)."""
|
||||
directly (same pattern as TestWaterfillArgs)."""
|
||||
|
||||
def _args(self, **overrides):
|
||||
args = ServerArgs(model_path="dummy")
|
||||
|
||||
@@ -1824,9 +1824,7 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
self.assertEqual(
|
||||
_a2a_backend_overrides(
|
||||
ResolvedView(
|
||||
SimpleNamespace(
|
||||
enable_deepep_waterfill=True, moe_a2a_backend="none"
|
||||
)
|
||||
SimpleNamespace(enable_waterfill=True, moe_a2a_backend="none")
|
||||
)
|
||||
),
|
||||
{"moe_a2a_backend": "deepep"},
|
||||
@@ -1836,9 +1834,7 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
self.assertEqual(
|
||||
_a2a_backend_overrides(
|
||||
ResolvedView(
|
||||
SimpleNamespace(
|
||||
enable_deepep_waterfill=True, moe_a2a_backend="none"
|
||||
)
|
||||
SimpleNamespace(enable_waterfill=True, moe_a2a_backend="none")
|
||||
)
|
||||
),
|
||||
{"moe_a2a_backend": "megamoe"},
|
||||
|
||||
Reference in New Issue
Block a user