[Feature] Support MegaMoE (#7943)#8038
Conversation
* [Feature] Support MegaMoE * update code * fix code style * fix code style * fix test * fix test * fix test * fix test * fix typo * fix code style * fix test * fix xpu test * fix test * fix test * fix code style * fix typo * fix typo
|
root seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## release/online/20260415 #8038 +/- ##
==========================================================
Coverage ? 72.02%
==========================================================
Files ? 389
Lines ? 54428
Branches ? 8526
==========================================================
Hits ? 39201
Misses ? 12519
Partials ? 2708
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
PaddlePaddle-bot
left a comment
PaddlePaddle-bot
left a comment
There was a problem hiding this comment.
🤖 Paddle-CI-Agent | pr_review |
2026-06-11 15:01:22
📋 Review 摘要
PR 概述：支持 MegaMoE / wfp4afp8 路径，包含新增 pre-dispatch CUDA op、DeepGEMM MegaMoE backend、量化配置和 CLI/worker 透传。
变更范围：custom_ops、MoE/EP backend、quantization、engine/config、相关测试。
影响面 Tag：OP Quantization FDConfig Engine
问题
| 级别 | 文件 | 概述 |
|---|---|---|
| 🔴 Bug | custom_ops/gpu_ops/mega_moe_pre_dispatch.cu:43 |
warp reduce 使用全掩码，合法小 hidden/非整 warp 输入会触发未定义行为 |
历史 Findings 修复情况
| Finding | 问题 | 状态 |
|---|---|---|
| F1 | MegaMoE 默认 full quant config 与 torch 格式量化解析冲突 | |
| F2 | apply_mega_moe() 未处理 topk_ids_hookfunc / shared_experts |
📝 PR 规范检查
标题当前为 [Feature] Support MegaMoE (#7943)，目标分支是 release/online/20260415，按仓库 release 分支规则需要使用 [Cherry-Pick] 前缀；PR 描述保留了模板但 Motivation / Modifications / Usage or Command / Accuracy Tests 均为空。
标题建议（可直接复制）：
[Cherry-Pick][Feature] Support MegaMoE (#7943)
PR 描述建议（点击展开，可直接复制）
## Motivation
Support MegaMoE on Blackwell/SM100 for MoE models on the release/online/20260415 branch, using wfp4afp8 for MoE and block_wise_fp8 for dense Linear layers. This PR is a cherry-pick from #7943.
## Modifications
- Added `--enable-mega-moe` configuration and worker propagation.
- Added the `mega_moe_pre_dispatch` CUDA custom op and operator tests.
- Added `wfp4afp8` quantization config and DeepGEMM MegaMoE backend integration.
- Updated DeepEP/MoE top-k selection calls and safetensors weight iteration ordering.
- Replaced torch-proxy compatibility calls with `paddle.enable_compat`.
## Usage or Command
Use `--enable-mega-moe --enable-expert-parallel` on supported SM100 devices. Full end-to-end launch command: N/A.
## Accuracy Tests
N/A. Accuracy results were not provided in the PR description.
## Checklist
- [x] Add at least a tag in the PR title.
- Tag list: [`[FDConfig]`,`[APIServer]`,`[Engine]`, `[Scheduler]`, `[PD Disaggregation]`, `[Executor]`, `[Graph Optimization]`, `[Speculative Decoding]`, `[RL]`, `[Models]`, `[Quantization]`, `[Loader]`, `[OP]`, `[KVCache]`, `[DataProcessor]`, `[BugFix]`, `[Docs]`, `[CI]`, `[Optimization]`, `[Feature]`, `[Benchmark]`, `[Others]`, `[XPU]`, `[HPU]`, `[GCU]`, `[DCU]`, `[Iluvatar]`, `[Metax]`]
- You can add new tags based on the PR content, but the semantics must be clear.
- [ ] Format your code, run `pre-commit` before commit.
- [x] Add unit tests. Please write the reason in this PR if no unit tests.
- [ ] Provide accuracy results.
- [ ] If the current PR is submitting to the `release` branch, make sure the PR has been submitted to the `develop` branch, then cherry-pick it to the `release` branch with the `[Cherry-Pick]` PR tag.总体评价
核心 MegaMoE 路径仍有一个新增 CUDA op 边界正确性问题；另外两条历史 finding 在当前代码中仍未修复，建议合入前一并处理。
| "kNumThreads must be a power of 2"); | ||
| #pragma unroll | ||
| for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) { | ||
| value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, mask, WARP_SIZE)); |
There was a problem hiding this comment.
🔴 Bug WarpReduceMax 固定使用全 warp 掩码，但 host 侧 shape 校验允许非整 warp 的合法输入。
num_threads_i64 = hidden_i64 / 8 只检查了 <= 1024 和 >= top_k，因此 group_size=32, hidden=128 会以 16 个线程启动一个 block。此时 __shfl_xor_sync(0xffffffffu, ..., WARP_SIZE) 把未参与的 lane 也放进同步掩码，结果是未定义的，可能导致 scale 错误或 kernel hang；现有单测只覆盖了 hidden_size=7168 这种整 warp 情况。
建议修复方式：在 WarpReduceMax 内使用 active mask 并按 kNumThreads 分组，例如 const unsigned m = __activemask(); __shfl_xor_sync(m, value, offset, kNumThreads)；或者在 host 侧显式拒绝 num_threads_i64 % WARP_SIZE != 0，并补充 group_size=32, hidden=128 的算子测试。
9f2f09d
into
PaddlePaddle:release/online/20260415
|
Thanks for your contribution! |
6 participants
Motivation
Support MegaMoE, cherry-pick from 7943
Modifications
Support MegaMoE
Usage or Command
refer to 7943
Accuracy Tests
Checklist
[FDConfig],[APIServer],[Engine],[Scheduler],[PD Disaggregation],[Executor],[Graph Optimization],[Speculative Decoding],[RL],[Models],[Quantization],[Loader],[OP],[KVCache],[DataProcessor],[BugFix],[Docs],[CI],[Optimization],[Feature],[Benchmark],[Others],[XPU],[HPU],[GCU],[DCU],[Iluvatar],[Metax]]pre-commitbefore commit.releasebranch, make sure the PR has been submitted to thedevelopbranch, then cherry-pick it to thereleasebranch with the[Cherry-Pick]PR tag.