[MoE] Decouple Mega MoE from DeepEP backend - #25406
Conversation
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
|
/rerun-stage stage-c-test-dsv4-4-gpu-b200 |
github-actions
Bot
added
documentation
|
Stage granularity is too coarse — a stage usually doesn't map to one feature, so rerunning a stage re-pays the cost of unrelated tests. If you don't know which exact test files to rerun, you shouldn't be using Use one of these instead:
AMD CI: stage-level dispatch is still available via Actions UI → PR Test (AMD) / PR Test ROCm 7.2 (AMD) → Run workflow → pick a stage from the dropdown. |
![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 integrates 'megamoe' as a standard MoE all-to-all backend, replacing the previous environment-variable-driven configuration. Key changes include updating the backend enum, automating server argument adjustments for expert parallelism when megamoe is active, and updating documentation and tests. A critical bug was identified in the FP8 quantization logic where a required function is used without being imported, which will lead to a runtime error.
| 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(): |
There was a problem hiding this comment.
The function get_moe_a2a_backend() is used here but it does not appear to be imported in this file. This will lead to a NameError at runtime when processing weights for FP8 MoE models. You should add the import, either at the top of the file or locally within this method to maintain consistency with the other local imports in this block.
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
if get_moe_a2a_backend().is_megamoe():
|
/rerun-test test_deepseek_v4_flash_fp4_megamoe_b200.py |
|
🚀 |
|
/rerun-test test_deepseek_v4_flash_fp4_megamoe_b200.py |
|
🚀 |
|
/tag-and-rerun-ci |
1b31bff to
338a7e0
Compare
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.
- 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)
When set, auto-configure --moe-a2a-backend megamoe so existing scripts using the env var continue to work.
…, remove JIT_NORM, add MegaMoE toggle to cookbook
152d545 to
056432a
Compare
|
/rerun-test test_deepseek_v4_flash_fp4_b200.py |
|
/rerun-test test_deepseek_v4_flash_fp4_megamoe_b200.py |
|
🚀 |
|
🚀 |
|
/rerun-test test_deepep_small.py |
|
🚀 |
…AL branches `topk_fused_transform` is annotated `__cluster_dims__(1, kClusterSize=8, 1)` via FUSED_COMBINE_KERNEL, but its SMALL (`seq_len <= Small::kMax1PassLength`) and TRIVIAL (`seq_len <= K`) branches let 7 of every 8 cluster blocks `return` immediately without ever calling `cluster.sync()`. On SM90 and SM100 with TP=4/DP=4 (the configuration covered by CI) this is tolerated, but on B300 / SM100 at TP=8/DP=8 with DeepEP it surfaces as `CUDA_ERROR_ILLEGAL_ADDRESS` reported at downstream sync points (CUDA graph replay, dense FP8 GEMM, etc.) -- the actual failing kernel is async. Reproducer: deepseek-ai/DeepSeek-V4-Pro on 8x B300 SXM6, `--tp 8 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep` with `SGLANG_OPT_USE_TOPK_V2=1` (the default since sgl-project#25406). conc=32 ISL=8192 OSL=1024 triggers the fused 1-stage path (batch_size_per_rank=4 <= kNumClusters=15) and crashes within ~10 seconds of decode. Workaround: `SGLANG_OPT_USE_TOPK_V2=0` falls back to the v1 kernel. Fix: restructure both branches so rank-0 does the work and all 8 ranks meet at `cooperative_groups::this_cluster().sync()` before exit, honoring the implicit cluster contract. Validated end-to-end on a B300 box: SGLANG_OPT_USE_TOPK_V2=1 + the patch + the same workload completes cleanly at 958.5 tok/s/gpu, vs 963.7 with the TOPK_V2=0 fallback (well within run-to-run noise -- the 7 dummy ranks were already pinned to the cluster's SM allotment whether they early-returned or sat at the barrier).
2 participants
Summary
megamoeas an--moe-a2a-backendchoice so users can enable Mega MoE without thedeep_eplibrary or--moe-a2a-backend deepep.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOEenv var as backward-compatible shortcut.Re-land of #24884
CI States
Latest PR Test (Base): Run #25952622517⚠️ Not enabled — add
Latest PR Test (Extra):
run-ci-extralabel to opt in.