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
This commit is contained in:
LuminolT
2026-07-06 10:26:43 +08:00
parent 6f96a1c269
commit beb2162028
4 changed files with 48 additions and 1 deletions

View File

@@ -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)

View File

@@ -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()

View File

@@ -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:

View File

@@ -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`")