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

[MoE] Decouple Mega MoE from DeepEP backend (#24884)

Browse Source
This commit is contained in:
Yuhao Yang
2026-05-15 02:01:44 +08:00
committed by GitHub
parent 3fc60e5748
commit 37f030a0de
8 changed files with 44 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx
+6 -20
View File
@@ -429,7 +429,6 @@ export const DeepSeekV4Deployment = () => {
"SGLANG_OPT_USE_JIT_NORM=1",
"SGLANG_OPT_USE_JIT_INDEXER_METADATA=1",
"SGLANG_OPT_USE_TOPK_V2=1",
"SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=1",
);
}
} else if (recipe === "balanced") {
@@ -445,14 +444,6 @@ export const DeepSeekV4Deployment = () => {
"SGLANG_OPT_USE_JIT_NORM=1",
"SGLANG_OPT_USE_JIT_INDEXER_METADATA=1",
"SGLANG_OPT_USE_TOPK_V2=1",
"SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=1",
"SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN=1",
"SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE=0",
"SGLANG_OPT_USE_FAST_MASK_EP=1",
"SGLANG_OPT_FIX_MEGA_MOE_MEMORY=1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=4096",
"SGLANG_OPT_FIX_NEXTN_MEGA_MOE=1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=0",
);
} else {
recipeEnv.push(isBig
@@ -475,15 +466,9 @@ export const DeepSeekV4Deployment = () => {
"SGLANG_OPT_USE_JIT_NORM=1",
"SGLANG_OPT_USE_JIT_INDEXER_METADATA=1",
"SGLANG_OPT_USE_TOPK_V2=1",
"SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=1",
"SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN=1",
"SGLANG_OPT_USE_FAST_MASK_EP=1",
"SGLANG_OPT_FIX_MEGA_MOE_MEMORY=1",
"SGLANG_OPT_FIX_NEXTN_MEGA_MOE=1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=0",
"NVSHMEM_DISABLE_IB=1",
"SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW=1",
"SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE=1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=8320",
);
} else {
@@ -610,7 +595,11 @@ export const DeepSeekV4Deployment = () => {
flags.push(` --dp ${tp}`);
flags.push(" --enable-dp-attention");
if (multinode) flags.push(...multiNodeFlags(nnodes));
flags.push(" --moe-a2a-backend deepep");
if (isBig && hardware === "b200") {
flags.push(" --moe-a2a-backend megamoe");
} else {
flags.push(" --moe-a2a-backend deepep");
}
if (hardware === "h200" && isBig) {
flags.push(" --mem-fraction-static 0.88");
} else if (isBig && hardware === "gb300") {
@@ -993,12 +982,9 @@ python3 -m sglang_router.launch_router \\
</p>
<pre style={{ margin: "8px 0 0 0", padding: "8px 12px", background: isDark ? "#111827" : "#f5f5f5", borderRadius: "4px", fontSize: "12px", lineHeight: "1.5", overflowX: "auto" }}>{
`# Add this flag to the sglang serve command:
--moe-a2a-backend deepep
--moe-a2a-backend megamoe
# And set these env vars:
SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE=1
SGLANG_OPT_FIX_MEGA_MOE_MEMORY=1
SGLANG_OPT_FIX_NEXTN_MEGA_MOE=1
SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=8320
SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=0`
}</pre>
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+1 -1
View File
@@ -82,7 +82,7 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
a2a_backend = get_moe_a2a_backend()
if a2a_backend.is_none():
if a2a_backend.is_none() or a2a_backend.is_megamoe():
return StandardDispatcher(moe_runner_config)
elif (
a2a_backend.is_deepep()
python/sglang/srt/layers/moe/mega_moe.py
+2 -1
View File
@@ -25,6 +25,7 @@ from sglang.jit_kernel.deepseek_v4 import mega_moe_pre_dispatch
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers.dp_attention import get_dp_global_num_tokens
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
if TYPE_CHECKING:
@@ -94,7 +95,7 @@ def _get_mega_moe_symm_buffer(
def should_use_mega_moe(moe: "DeepseekV2MoE", hidden_states: torch.Tensor) -> bool:
if not envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get():
if not get_moe_a2a_backend().is_megamoe():
return False
if not getattr(moe.experts, "_mega_moe_weights_built", False):
return False
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
-1
View File
@@ -131,7 +131,6 @@ class DeepGemmRunnerCore(MoeRunnerCore):
if envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.get():
assert envs.SGLANG_OPT_SWIGLU_CLAMP_FUSION.get()
assert envs.SGLANG_OPT_USE_JIT_EP_ACTIVATION.get()
assert envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get()
self.use_swizzle = True
def run(
python/sglang/srt/layers/moe/utils.py
+4
View File
@@ -29,6 +29,7 @@ class MoeA2ABackend(Enum):
MORI = "mori"
ASCEND_FUSEEP = "ascend_fuseep"
FLASHINFER = "flashinfer"
MEGAMOE = "megamoe"
CUSTOMIZED = "customized"
@classmethod
@@ -61,6 +62,9 @@ class MoeA2ABackend(Enum):
def is_mori(self):
return self == MoeA2ABackend.MORI
def is_megamoe(self):
return self == MoeA2ABackend.MEGAMOE
def is_customized(self):
return self == MoeA2ABackend.CUSTOMIZED
python/sglang/srt/layers/quantization/fp8.py
+1 -1
View File
@@ -1193,7 +1193,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
layer.w13_weight.data = layer.w13_weight.data.view(torch.int8)
layer.w2_weight.data = layer.w2_weight.data.view(torch.int8)
if envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get():
if get_moe_a2a_backend().is_megamoe():
from sglang.srt.layers.moe.mega_moe import (
build_mega_moe_experts_weights,
)
python/sglang/srt/server_args.py
+28 -1
View File
@@ -209,6 +209,7 @@ MOE_A2A_BACKEND_CHOICES = [
"mori",
"ascend_fuseep",
"flashinfer",
"megamoe",
]
FP8_GEMM_RUNNER_BACKEND_CHOICES = [
@@ -609,7 +610,14 @@ class ServerArgs:
# Expert parallelism
ep_size: int = 1
moe_a2a_backend: Literal[
"none", "deepep", "mooncake", "nixl", "mori", "ascend_fuseep", "flashinfer"
"none",
"deepep",
"mooncake",
"nixl",
"mori",
"ascend_fuseep",
"flashinfer",
"megamoe",
] = "none"
moe_runner_backend: str = "auto"
record_nolora_graph: bool = True
@@ -3184,6 +3192,25 @@ class ServerArgs:
)
self.moe_a2a_backend = "deepep"
if (
envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get()
and self.moe_a2a_backend != "megamoe"
):
self.moe_a2a_backend = "megamoe"
logger.info(
"SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE is set, "
"auto-configuring --moe-a2a-backend megamoe."
)
if self.moe_a2a_backend == "megamoe":
self.ep_size = self.tp_size
if not envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.is_set():
envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.set(True)
logger.info(
f"Mega MoE is enabled. The expert parallel size is adjusted "
f"to be the same as the tensor parallel size[{self.tp_size}]."
)
if self.moe_a2a_backend == "deepep":
if self.deepep_mode == "normal":
logger.warning("Cuda graph is disabled because deepep_mode=`normal`")
test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py
+2 -8
View File
@@ -28,18 +28,12 @@ SERVER_LAUNCH_TIMEOUT = 3600
_W4A8_MEGAMOE_ENV = {
"SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE": "1",
"SGLANG_OPT_FIX_MEGA_MOE_MEMORY": "1",
"SGLANG_OPT_FIX_NEXTN_MEGA_MOE": "1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "4096",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "0",
}
_W4A4_MEGAMOE_ENV = {
"SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE": "1",
"SGLANG_OPT_FIX_MEGA_MOE_MEMORY": "1",
"SGLANG_OPT_FIX_NEXTN_MEGA_MOE": "1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "4096",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "0",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS": "1",
@@ -81,7 +75,7 @@ class TestDSV4FlashFP4B200W4A8MegaMoE(ServerSanityMixin, CustomTestCase):
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"megamoe",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
@@ -122,7 +116,7 @@ class TestDSV4FlashFP4B200W4A4MegaMoE(ServerSanityMixin, CustomTestCase):
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"megamoe",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",