From b2c6db0cc42992687f9ed2c80dce7058c3009df6 Mon Sep 17 00:00:00 2001 From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Date: Sat, 16 May 2026 15:18:43 +0800 Subject: [PATCH] [MoE] Decouple Mega MoE from DeepEP backend (#25406) --- python/sglang/srt/environ.py | 5 ++-- .../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/models/deepseek_v2.py | 1 + python/sglang/srt/server_args.py | 29 ++++++++++++++++++- test/manual/dsv4/test_b200_flash.py | 1 - test/manual/dsv4/test_b200_pro.py | 1 - test/manual/dsv4/test_b300_flash.py | 1 - test/manual/dsv4/test_b300_pro.py | 1 - test/manual/dsv4/test_dsv4_flash_mtp_tp8.py | 2 -- test/manual/dsv4/test_dsv4_pro_mtp.py | 2 -- test/manual/dsv4/test_gb300_flash.py | 1 - test/manual/dsv4/test_gb300_pro.py | 1 - test/manual/dsv4/test_h200_fp8_flash.py | 1 - ...test_deepseek_v4_flash_fp4_megamoe_b200.py | 12 ++------ 18 files changed, 41 insertions(+), 29 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5380d528b..b4bf8d308 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -584,7 +584,7 @@ class Envs: SGLANG_OPT_USE_TILELANG_MHC_PRE = EnvBool(True) SGLANG_OPT_USE_TILELANG_MHC_POST = EnvBool(True) SGLANG_OPT_USE_TILELANG_INDEXER = EnvBool(False) - SGLANG_OPT_USE_JIT_INDEXER_METADATA = EnvBool(False) + SGLANG_OPT_USE_JIT_INDEXER_METADATA = EnvBool(True) SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False) SGLANG_OPT_USE_COMPRESSOR_V2 = EnvBool(True) SGLANG_FP8_PAGED_MQA_LOGITS_TORCH = EnvBool(False) @@ -617,13 +617,12 @@ class Envs: # TopK SGLANG_OPT_USE_FUSED_HASH_TOPK = EnvBool(True) SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK = EnvBool(True) - SGLANG_OPT_USE_TOPK_V2 = EnvBool(False) + SGLANG_OPT_USE_TOPK_V2 = EnvBool(True) # GEMM / kernel fusion SGLANG_OPT_FP8_WO_A_GEMM = EnvBool(True) SGLANG_OPT_BF16_FP32_GEMM_ALGO = EnvStr("cublas") SGLANG_OPT_USE_JIT_EP_ACTIVATION = EnvBool(True) - SGLANG_OPT_USE_JIT_NORM = EnvBool(False) SGLANG_OPT_FUSE_WQA_WKV = EnvBool(True) SGLANG_OPT_SWIGLU_CLAMP_FUSION = EnvBool(True) 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 f3b188ef8..943ac4f17 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 1c13f7be9..9574bd2da 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 105c7b7b0..7c1843ea2 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/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 8b7055f8c..5c834a595 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -601,6 +601,7 @@ class DeepseekV2MoE(nn.Module): or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_flashinfer() + or get_moe_a2a_backend().is_megamoe() or should_use_flashinfer_cutlass_moe_fp4_allgather() or envs.SGLANG_SHARED_EXPERT_TP1.get() ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 379da02a6..f7d7cf67b 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 = [ @@ -610,7 +611,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" flashinfer_mxfp4_moe_precision: Literal["default", "bf16"] = "default" @@ -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/manual/dsv4/test_b200_flash.py b/test/manual/dsv4/test_b200_flash.py index 3c8281161..05d738bcf 100644 --- a/test/manual/dsv4/test_b200_flash.py +++ b/test/manual/dsv4/test_b200_flash.py @@ -100,7 +100,6 @@ class TestB200FlashCP(DSV4FlashAime25TestBase): DEEPEP_LARGE_SMS_CONFIG, ] EXTRA_ENV = { - "SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024", } diff --git a/test/manual/dsv4/test_b200_pro.py b/test/manual/dsv4/test_b200_pro.py index 67eaadd21..eba9b6546 100644 --- a/test/manual/dsv4/test_b200_pro.py +++ b/test/manual/dsv4/test_b200_pro.py @@ -116,7 +116,6 @@ class TestB200ProCP(DSV4ProAime25TestBase): DEEPEP_LARGE_SMS_CONFIG, ] EXTRA_ENV = { - "SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256", } diff --git a/test/manual/dsv4/test_b300_flash.py b/test/manual/dsv4/test_b300_flash.py index 261bac476..4e800526d 100644 --- a/test/manual/dsv4/test_b300_flash.py +++ b/test/manual/dsv4/test_b300_flash.py @@ -102,7 +102,6 @@ class TestB300FlashCP(DSV4FlashAime25TestBase): DEEPEP_LARGE_SMS_CONFIG, ] EXTRA_ENV = { - "SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024", } diff --git a/test/manual/dsv4/test_b300_pro.py b/test/manual/dsv4/test_b300_pro.py index 1ed254e52..701a42a39 100644 --- a/test/manual/dsv4/test_b300_pro.py +++ b/test/manual/dsv4/test_b300_pro.py @@ -118,7 +118,6 @@ class TestB300ProCP(DSV4ProAime25TestBase): DEEPEP_LARGE_SMS_CONFIG, ] EXTRA_ENV = { - "SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256", } diff --git a/test/manual/dsv4/test_dsv4_flash_mtp_tp8.py b/test/manual/dsv4/test_dsv4_flash_mtp_tp8.py index 4913f0bb9..87c63f9c3 100644 --- a/test/manual/dsv4/test_dsv4_flash_mtp_tp8.py +++ b/test/manual/dsv4/test_dsv4_flash_mtp_tp8.py @@ -24,9 +24,7 @@ DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8" DSV4_FLASH_BASE_ENV = { "SGLANG_ENABLE_SPEC_V2": "1", - "SGLANG_OPT_USE_TOPK_V2": "1", "SGLANG_DSV4_FP4_EXPERTS": "0", - "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "0", } DSV4_FLASH_SERVER_ARGS = [ diff --git a/test/manual/dsv4/test_dsv4_pro_mtp.py b/test/manual/dsv4/test_dsv4_pro_mtp.py index d989f22db..7e5cf62ae 100644 --- a/test/manual/dsv4/test_dsv4_pro_mtp.py +++ b/test/manual/dsv4/test_dsv4_pro_mtp.py @@ -41,9 +41,7 @@ HONGLOUMENG_PATH = os.environ.get( DSV4_PRO_BASE_ENV = { "SGLANG_ENABLE_SPEC_V2": "1", - "SGLANG_OPT_USE_TOPK_V2": "1", "SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2": "1", - "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "0", } DSV4_PRO_SERVER_ARGS = [ diff --git a/test/manual/dsv4/test_gb300_flash.py b/test/manual/dsv4/test_gb300_flash.py index 4e7a7ec5e..ef997f50e 100644 --- a/test/manual/dsv4/test_gb300_flash.py +++ b/test/manual/dsv4/test_gb300_flash.py @@ -100,7 +100,6 @@ class TestGB300FlashCP(DSV4FlashAime25TestBase): DEEPEP_LARGE_SMS_CONFIG, ] EXTRA_ENV = { - "SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024", } diff --git a/test/manual/dsv4/test_gb300_pro.py b/test/manual/dsv4/test_gb300_pro.py index 2e186591b..1c53129b2 100644 --- a/test/manual/dsv4/test_gb300_pro.py +++ b/test/manual/dsv4/test_gb300_pro.py @@ -118,7 +118,6 @@ class TestGB300ProCP(DSV4ProAime25TestBase): DEEPEP_LARGE_SMS_CONFIG, ] EXTRA_ENV = { - "SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256", } diff --git a/test/manual/dsv4/test_h200_fp8_flash.py b/test/manual/dsv4/test_h200_fp8_flash.py index 2aacca9ae..fe53ded83 100644 --- a/test/manual/dsv4/test_h200_fp8_flash.py +++ b/test/manual/dsv4/test_h200_fp8_flash.py @@ -112,7 +112,6 @@ class TestH200Fp8FlashCP(DSV4FlashAime25TestBase): ] EXTRA_ENV = { **H200_FP8_ENV, - "SGLANG_OPT_USE_JIT_INDEXER_METADATA": "1", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024", } 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 6cf20c8b3..baab29a2a 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,20 +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", "SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND": "1", } @@ -81,7 +73,7 @@ class TestDSV4FlashFP4B200W4A8MegaMoE(ServerSanityMixin, CustomTestCase): "4", "--enable-dp-attention", "--moe-a2a-backend", - "deepep", + "megamoe", "--speculative-algorithm", "EAGLE", "--speculative-num-steps", @@ -122,7 +114,7 @@ class TestDSV4FlashFP4B200W4A4MegaMoE(ServerSanityMixin, CustomTestCase): "4", "--enable-dp-attention", "--moe-a2a-backend", - "deepep", + "megamoe", "--speculative-algorithm", "EAGLE", "--speculative-num-steps",