From beb2162028021db6b14e92edd72715b7c0b9f7d7 Mon Sep 17 00:00:00 2001 From: LuminolT Date: Mon, 6 Jul 2026 10:26:43 +0800 Subject: [PATCH] feat(moe): add megamoe a2a backend plumbing Add the MegaMoE A2A backend name to server argument choices and MoE backend enums, wire the startup handling that forces ep_size to tp_size, and add the env flags used by the later MegaMoE forward and weight layout work. Treat MegaMoE as a dispatcher bypass so it does not enter DeepEP dispatcher logic.\n\nConstraint: Applies on top of origin/rc/v0.1 for GLM 5.2 MegaMoE adaptation.\nFeature-flag: --moe-a2a-backend megamoe and SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.\nConflict-hotspots: python/sglang/srt/server_args.py; python/sglang/srt/layers/moe/utils.py; python/sglang/srt/layers/moe/fused_moe_triton/layer.py\nScope-risk: medium\nTested: git diff --check\nTested: python3 -c ast.parse on modified Python files --- python/sglang/srt/environ.py | 7 +++++++ .../srt/layers/moe/fused_moe_triton/layer.py | 2 +- python/sglang/srt/layers/moe/utils.py | 20 +++++++++++++++++++ python/sglang/srt/server_args.py | 20 +++++++++++++++++++ 4 files changed, 48 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6d81dcf34..2142bcf03 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -506,6 +506,13 @@ class Envs: SGLANG_DG_USE_NVRTC = EnvBool(False) SGLANG_USE_DEEPGEMM_BMM = EnvBool(False) + # DeepGEMM MegaMoE + SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE = EnvBool(False) + SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK = EnvInt(1024) + SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS = EnvBool(False) + SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND = EnvBool(False) + SGLANG_OPT_FIX_MEGA_MOE_MEMORY = EnvBool(False) + # DeepSeek MHA Optimization SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD = EnvInt(8192) 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 f04479dbb..ec27422e8 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -94,7 +94,7 @@ logger = logging.getLogger(__name__) 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/utils.py b/python/sglang/srt/layers/moe/utils.py index b5600069f..065efef52 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -26,6 +26,7 @@ class MoeA2ABackend(Enum): MORI = "mori" ASCEND_FUSEEP = "ascend_fuseep" FLASHINFER = "flashinfer" + MEGAMOE = "megamoe" @classmethod def _missing_(cls, value): @@ -57,6 +58,9 @@ class MoeA2ABackend(Enum): def is_mori(self): return self == MoeA2ABackend.MORI + def is_megamoe(self): + return self == MoeA2ABackend.MEGAMOE + class MoeRunnerBackend(Enum): @@ -244,6 +248,22 @@ def get_deepep_config() -> str: return DEEPEP_CONFIG +def is_deepep_class_backend() -> bool: + """Check if the MoE backend uses DeepEP-family dispatcher semantics.""" + b = get_moe_a2a_backend() + return b.is_deepep() or b.is_mooncake() or b.is_mori() + + +def uses_per_rank_fused_shared_slots() -> bool: + """Check whether fused shared experts use per-rank physical slots.""" + return is_deepep_class_backend() or get_moe_a2a_backend().is_megamoe() + + +def has_per_rank_fused_shared_slots(num_fused_shared_experts: int) -> bool: + """Check whether this layer has fused shared experts in per-rank slots.""" + return num_fused_shared_experts > 0 and uses_per_rank_fused_shared_slots() + + def is_tbo_enabled() -> bool: global IS_TBO_ENABLED if IS_TBO_ENABLED is None: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6598a4b45..7b3c8747b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -218,6 +218,7 @@ MOE_A2A_BACKEND_CHOICES = [ "mori", "ascend_fuseep", "flashinfer", + "megamoe", ] FP8_GEMM_RUNNER_BACKEND_CHOICES = [ @@ -3030,6 +3031,25 @@ class ServerArgs: ) def _handle_a2a_moe(self): + 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.warning( + f"MegaMoE 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`")