From 37f030a0de6d98a2ca545b30a40ed401d08f4715 Mon Sep 17 00:00:00 2001 From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Date: Fri, 15 May 2026 02:01:44 +0800 Subject: [PATCH] [MoE] Decouple Mega MoE from DeepEP backend (#24884) --- .../autoregressive/deepseek-v4-deployment.jsx | 26 ++++------------- .../srt/layers/moe/fused_moe_triton/layer.py | 2 +- python/sglang/srt/layers/moe/mega_moe.py | 3 +- .../srt/layers/moe/moe_runner/deep_gemm.py | 1 - python/sglang/srt/layers/moe/utils.py | 4 +++ python/sglang/srt/layers/quantization/fp8.py | 2 +- python/sglang/srt/server_args.py | 29 ++++++++++++++++++- ...test_deepseek_v4_flash_fp4_megamoe_b200.py | 10 ++----- 8 files changed, 44 insertions(+), 33 deletions(-) diff --git a/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx b/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx index 79b79b863..b6efab022 100644 --- a/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx +++ b/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx @@ -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 \\ 

{
 `# 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`
         }
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 56eaf8a1e..97a5bfc3d 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -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() diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 93f4d9a15..e60d7602e 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -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 diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index da6f13fcd..61af5533f 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -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( diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index e05167da9..fbca714d4 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 99d86a56e..f1d5cc4a9 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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, ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 66aff8edf..4ff4826de 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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`") diff --git a/test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py b/test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py index 1c4df29c4..ca4ecf200 100644 --- a/test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py +++ b/test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py @@ -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",