[MoE] Decouple Mega MoE from DeepEP backend - #24884
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces auto-configuration for the expert parallel size when Mega MoE is enabled and no specific all-to-all backend is selected. The reviewer suggested refining the activation condition to ensure tp_size > 1, which avoids redundant operations on single-GPU setups, and recommended using a warning log level for consistency with other backend handlers.
| if ( | ||
| envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get() | ||
| and self.moe_a2a_backend == "none" | ||
| and self.ep_size == 1 | ||
| ): | ||
| self.ep_size = self.tp_size | ||
| 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}]." | ||
| ) |
There was a problem hiding this comment.
Consider adding self.tp_size > 1 to the condition and using logger.warning for consistency with other backends in this function (e.g., DeepEP, Mooncake, Mori). Adding the tp_size check avoids redundant logging and no-op assignments on single-GPU setups where expert parallelism is not applicable, while ensuring the adjustment is visible to users in multi-GPU configurations.
| if ( | |
| envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get() | |
| and self.moe_a2a_backend == "none" | |
| and self.ep_size == 1 | |
| ): | |
| self.ep_size = self.tp_size | |
| 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 ( | |
| envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get() | |
| and self.moe_a2a_backend == "none" | |
| and self.ep_size == 1 | |
| and self.tp_size > 1 | |
| ): | |
| self.ep_size = self.tp_size | |
| logger.warning( | |
| f"Mega MoE is enabled. The expert parallel size is adjusted " | |
| f"to be the same as the tensor parallel size[{self.tp_size}]." | |
| ) |
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
/rerun-stage stage-c-test-dsv4-4-gpu-b200 |
|
🚀 Triggered |
Auto-configure EP when SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE is set, so users no longer need to pass --moe-a2a-backend=deepep or install the deep_ep library to use Mega MoE.
When SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE is enabled together with ep_size > 1, both the Mega MoE path (decode) and the normal/DeepEP path (prefill fallback) share the same MoE weights. Auto-set SGLANG_OPT_FIX_MEGA_MOE_MEMORY=True so they share one weight copy instead of duplicating, saving significant GPU memory.
Add megamoe as a first-class a2a backend. It uses symmetric memory for fused all-to-all + GEMM (decode), and falls back to StandardDispatcher (all-reduce EP) for large prefills. User experience: python -m sglang.launch_server ... --moe-a2a-backend megamoe
Replace env var reads with get_moe_a2a_backend().is_megamoe() checks. The --moe-a2a-backend megamoe flag is now the sole way to enable Mega MoE.
944404c to
1554138
Compare
- Use --moe-a2a-backend megamoe instead of deepep + env var - Remove dead env vars (SGLANG_OPT_USE_FAST_MASK_EP, SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN, SGLANG_OPT_FIX_NEXTN_MEGA_MOE) - Remove redundant env vars (SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2, SGLANG_OPT_FIX_MEGA_MOE_MEMORY)
github-actions
Bot
added
documentation
|
/rerun-stage stage-c-test-dsv4-4-gpu-b200 |
|
🚀 Triggered |
When set, auto-configure --moe-a2a-backend megamoe so existing scripts using the env var continue to work.
|
The dpskv4 B200 test has passed. After that, only one backward-compatibility commit for an environment variable was added, so we can merge now. |
|
/rerun-test test_deepseek_v4_flash_fp4_megamoe_b200.py |
This reverts commit 9f1293a.
|
Dispatched on this branch: |
3 participants
Summary
ep_size = tp_size) whenSGLANG_OPT_USE_DEEPGEMM_MEGA_MOE=1is set, so users no longer need--moe-a2a-backend=deepepor thedeep_eplibrary to use Mega MoE.deep_gemm.fp8_fp4_mega_moe()for its own all-to-all communication and never touches the DeepEP dispatcher, so the dependency was unnecessary.