[Feature] Support MegaMoE - #7943
Merged
Merged
Conversation
Wanglongzhi2001
had a problem deploying
to
Metax_ci
May 27, 2026 11:30 — with 
GitHub Actions
Failure
GitHub Actions
Failure
|
Thanks for your contribution! |
Contributor
There was a problem hiding this comment.
Pull request overview
该 PR 旨在为 FastDeploy 的 MoE 路径引入 MegaMoE 支持：新增一个用于 pre-dispatch/量化准备的 CUDA 自定义算子，并在 DeepGEMM MoE backend 中接入对应的 buffer 与执行流程，同时补充算子级别的单测。
Changes:
- 新增
mega_moe_pre_dispatchCUDA 自定义算子，并在 custom ops 构建脚本中按架构条件编译进包。 - 在
DeepGemmFusedMoeMethod中新增 MegaMoE 相关 buffer 初始化、权重量化布局处理与apply_mage_moe执行路径。 - 增加
tests/operators/test_mega_moe_pre_dispatch.py用于校验 pre-dispatch 输出正确性。
另外：PR 标题格式符合要求，但当前 PR 描述的 Motivation/Modifications/Usage/Accuracy Tests 均未补全，建议补充（尤其是启用方式如环境变量、适用硬件/算子依赖与精度结果）。
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
tests/operators/test_mega_moe_pre_dispatch.py |
新增 MegaMoE pre-dispatch 单测（当前存在导入与分布式初始化可用性问题） |
fastdeploy/model_executor/layers/quantization/fp8_utils.py |
增加权重交织与 scale layout 变换辅助函数以适配 MegaMoE 权重格式 |
fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py |
DeepGEMM MoE backend 接入 MegaMoE：buffer、量化流程与执行入口 |
fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py |
在基类 apply 中根据环境变量切换到 MegaMoE 路径 |
fastdeploy/envs.py |
新增环境变量 FD_ENABLE_MAGE_MOE 开关 |
custom_ops/setup_ops.py |
将 mega_moe_pre_dispatch.cu 加入 SM100/103 构建源文件 |
custom_ops/gpu_ops/mega_moe_pre_dispatch.cu |
新增 MegaMoE pre-dispatch CUDA Kernel 与静态算子注册 |
| return ffn_out | ||
|
|
||
|
|
||
| @singleton |
CI报告基于以下代码生成（30分钟更新一次）: 1 Required任务 : 1/10 通过
2 失败详情🔴 xpu_8cards_case_test / run_xpu_8cards_cases — PR问题（置信度: 高）错误类型: PR问题 | 置信度: 高
关键日志:
修复建议:
关联变更: 🔴 Run Base Tests / base_tests — 环境问题（置信度: 高）错误类型: 环境问题 | 置信度: 高
关键日志:
修复建议:
关联变更: 未发现 PR 直接修改 🔴 Run FastDeploy LogProb Tests / run_tests_logprob — 环境问题（置信度: 高）错误类型: 环境问题 | 置信度: 高
关键日志:
修复建议:
关联变更: 未发现 LogProb baseline 或输出比对相关失败证据 🔴 Run Four Cards Tests / run_4_cards_tests — 环境问题（置信度: 高）错误类型: 环境问题 | 置信度: 高
关键日志:
修复建议:
关联变更: 未发现这些 e2e 测试文件被 PR 修改 🔴 Run Stable Tests / stable_tests — 环境问题（置信度: 高）错误类型: 环境问题 | 置信度: 高
关键日志:
修复建议:
关联变更: 未发现 stable case 脚本被 PR 修改 🔴 Approval — 需要 Approval（置信度: 高）该 Job 需要人工 Approval，完成审批后 CI 才会继续执行。 🔴 Extracted partial CE model tasks to run in CI. / run_ce_cases — 未知（置信度: 低）本轮 required 失败 job 中 cache miss 数量超过 5 个，按 Skill 预算限制未对该 Job 拉取深度日志；当前仅保留快速状态，错误摘要为“日志分析省略（超出失败 Job 量限制）”。 |
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 @@
## develop #7943 +/- ##
==========================================
Coverage ? 67.48%
==========================================
Files ? 471
Lines ? 66356
Branches ? 10215
==========================================
Hits ? 44782
Misses ? 18722
Partials ? 2852
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:
|
Wanglongzhi2001
had a problem deploying
to
Metax_ci
June 8, 2026 08:35 — with
GitHub Actions
Failure
GitHub Actions
Failure
Comment on lines
+36
to
+45
| __device__ __forceinline__ float WarpReduceMax(float value) { | ||
| static_assert(kNumThreads >= 1 && kNumThreads <= WARP_SIZE, | ||
| "kNumThreads must be in [1, 32]"); | ||
| static_assert((kNumThreads & (kNumThreads - 1)) == 0, | ||
| "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)); | ||
| } | ||
| return value; |
Comment on lines
+41
to
+44
| #pragma unroll | ||
| for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) { | ||
| value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, mask, WARP_SIZE)); | ||
| } |
Comment on lines
+54
to
+56
| def get_quant_method(self, layer) -> Optional[QuantMethodBase]: | ||
| logger.info("Currently only support DeepGEMMMegaMoE for wfp4afp8") | ||
| if isinstance(layer, FusedMoE): |
Wanglongzhi2001
had a problem deploying
to
Metax_ci
June 8, 2026 10:11 — with
GitHub Actions
Failure
GitHub Actions
Failure
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 17 changed files in this pull request and generated 6 comments.
Comments suppressed due to low confidence (1)
custom_ops/setup_ops.py:526
- 这里的注释格式是
pass # ...，会触发 pre-commit 的 flake8/pep8（#前应至少两个空格）。建议改为pass # ...。
# Add SM100 specific sources if any, e.g., for new hardware intrinsics
# sources += ["gpu_ops/cutlass_kernels/w8a8/c4x_sm100.cu"] # Example
pass # No SM100 specific sources identified yet beyond what CUTLASS handles
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
Wanglongzhi2001
had a problem deploying
to
Metax_ci
June 9, 2026 06:51 — with
GitHub Actions
Failure
GitHub Actions
Failure
Wanglongzhi2001
had a problem deploying
to
Metax_ci
June 9, 2026 09:02 — with
GitHub Actions
Failure
GitHub Actions
Failure
Wanglongzhi2001
had a problem deploying
to
Metax_ci
June 9, 2026 11:24 — with
GitHub Actions
Failure
GitHub Actions
Failure
Comment on lines
+76
to
+79
| @classmethod | ||
| def setUpClass(cls): | ||
| paddle.seed(2025) | ||
|
|
Comment on lines
+67
to
+71
| x_view = x_padded.reshape([m, padded_n // gran_k, gran_k]) | ||
| x_amax = x_view.abs().astype("float32").amax(axis=2).reshape([m, padded_n // gran_k]).clip(min=1e-4) | ||
| sf = x_amax / 448.0 | ||
| sf = ceil_to_ue8m0(sf) if use_ue8m0 else sf | ||
| x_fp8 = (x_view * (1.0 / sf.unsqueeze(2))).astype(paddle.float8_e4m3fn).reshape([m, padded_n])[:, :n] |
Comment on lines
+54
to
+56
| def get_quant_method(self, layer) -> Optional[QuantMethodBase]: | ||
| logger.info("Currently only support DeepGEMMMegaMoE for wfp4afp8") | ||
| if isinstance(layer, FusedMoE): |
Wanglongzhi2001
had a problem deploying
to
Metax_ci
June 9, 2026 13:14 — with
GitHub Actions
Failure
GitHub Actions
Failure
lizhenyun01
reviewed
Jun 10, 2026
| ) | ||
| parser.add_argument( | ||
| "--enable_mega_moe", | ||
| "--enable-mega-moe", |
| from ..moe import FusedMoE | ||
| from .quant_base import QuantConfigBase, QuantMethodBase | ||
|
|
||
| QUANT_SCALING_FACTOR = 6 |
Collaborator
There was a problem hiding this comment.
可以看看Copilot的review是否有必要
Comment on lines
+76
to
+103
| def _get_mega_moe_quantization_config(): | ||
| return { | ||
| "quantization": "mix_quant", | ||
| "kv_cache_quant_type": "block_wise_fp8", | ||
| "dense_quant_type": "block_wise_fp8", | ||
| "moe_quant_type": "wfp4afp8", | ||
| "is_quantized": False, | ||
| } | ||
|
|
||
|
|
||
| def parse_quant_config(args, model_config, is_ernie, is_v1_loader): | ||
| if args.quantization is not None and isinstance(args.quantization, str): | ||
| args.quantization = parse_quantization(args.quantization) | ||
|
|
||
| enable_mega_moe = getattr(args, "enable_mega_moe", False) | ||
| if enable_mega_moe: | ||
| mega_moe_quantization_config = _get_mega_moe_quantization_config() | ||
|
|
||
| if args.quantization is None and model_config.quantization_config is None: | ||
| args.quantization = mega_moe_quantization_config | ||
| if args.quantization is not None and not _is_mega_moe_quantization_config(args.quantization): | ||
| raise ValueError("--enable-mega-moe requires moe_quant_type=wfp4afp8.") | ||
| if model_config.quantization_config is not None and not _is_mega_moe_quantization_config( | ||
| model_config.quantization_config | ||
| ): | ||
| raise ValueError( | ||
| "--enable-mega-moe conflicts with model quantization_config. It requires moe_quant_type=wfp4afp8." | ||
| ) |
Wanglongzhi2001
had a problem deploying
to
Metax_ci
June 10, 2026 06:45 — with
GitHub Actions
Failure
GitHub Actions
Failure
Comment on lines
+42
to
+44
| for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) { | ||
| value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, mask, WARP_SIZE)); | ||
| } |
Comment on lines
+226
to
+227
| PD_CHECK(num_tokens_i64 <= num_max_tokens_per_rank, | ||
| "num_tokens must not exceed padded_max"); |
Comment on lines
+21
to
+25
| from ..moe import FusedMoE | ||
| from .quant_base import QuantConfigBase, QuantMethodBase | ||
|
|
||
| QUANT_SCALING_FACTOR = 6 | ||
|
|
Wanglongzhi2001
had a problem deploying
to
Metax_ci
June 10, 2026 07:17 — with
GitHub Actions
Failure
GitHub Actions
Failure
PaddlePaddle-bot
suggested changes
Jun 10, 2026
PaddlePaddle-bot
left a comment
PaddlePaddle-bot
left a comment
There was a problem hiding this comment.
🤖 Paddle-CI-Agent | pr_review |
2026-06-10 15:19:34
📋 Review 摘要
PR 概述：新增 Blackwell MegaMoE wfp4afp8 路径、pre-dispatch CUDA op、量化配置和 worker 参数传递。
变更范围：custom op、MoE DeepGEMM backend、quantization、Engine/worker 配置、算子测试。
影响面 Tag：[OP] [Quantization] [FDConfig] [Engine]
问题
建议拆分方案：
- PR 1: MegaMoE pre-dispatch 算子与测试 —
custom_ops/gpu_ops/mega_moe_pre_dispatch.cu,custom_ops/setup_ops.py,tests/operators/test_mega_moe_pre_dispatch.py - PR 2: MegaMoE 配置与进程参数传递 —
fastdeploy/config.py,fastdeploy/engine/args_utils.py,fastdeploy/engine/engine.py,fastdeploy/worker/worker_process.py,tests/engine/test_engine.py - PR 3: wfp4afp8 quant 与 DeepGEMM MegaMoE backend —
fastdeploy/model_executor/layers/moe/ep.py,fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py,fastdeploy/model_executor/layers/quantization/*.py,tests/model_executor/test_ep.py
| 级别 | 文件 | 概述 |
|---|---|---|
| 🔴 Bug | fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py:1359 |
MegaMoE 前向忽略 shared experts / latent projections / routing hook 等 MoE 语义，相关模型输出会不正确 |
| 🟡 建议 | PR整体 | 变更量较大，建议按算子、配置、backend 三块拆分 |
历史 Findings 修复情况
| Finding | 问题 | 状态 |
|---|---|---|
| F1 | --enable-mega-moe 默认量化配置写入时序 |
|
| F2 | wfp4afp8 未限制 EP MegaMoE 路径 |
📝 PR 规范检查
标题已有官方 [Feature] Tag；描述仍缺少 Accuracy Tests 实际内容，Checklist 仍未按现状勾选。
总体评价
本轮按风险优先审查了 MegaMoE 前向、量化配置时序、custom op 注册与测试路径，未全量覆盖其余机械改动。当前存在会导致带 shared experts 或 latent projection 的 MoE 模型输出错误的阻塞问题，需要修复后再合入。
lizhenyun01
approved these changes
Jun 10, 2026
5 tasks
Jiang-Jia-Jun
pushed a commit
that referenced
this pull request
Jun 11, 2026
* [Feature] Support MegaMoE (#7943) * [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 * delete log * fix code style * fix paddle compat * fix code style --------- Co-authored-by: root <root@tjzj-inf-sci-k8s-bzz2-00005.tjzj.baidu.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
支持 MegaMoE (wfp4afp8)
Modifications
支持 MegaMoE (wfp4afp8)
Usage or Command
拉取下面两个 wheel 包进行编译，或者联系本人获取编译好的 wheel 包
在服务启动参数中加入下面的参数（两种方法皆可）
或者
然后模型 config.json 中加入：
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.