[RL] support bf16 moe tp permute, group_gemm, unpermute#7189
[RL] support bf16 moe tp permute, group_gemm, unpermute#7189ckl117 wants to merge 3 commits intoPaddlePaddle:developfrom
Conversation
|
Thanks for your contribution! |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #7189 +/- ##
==========================================
Coverage ? 73.44%
==========================================
Files ? 383
Lines ? 53640
Branches ? 8411
==========================================
Hits ? 39398
Misses ? 11566
Partials ? 2676
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
9f1ab8c to
8eb2f1e
Compare
fastdeploy-bot
left a comment
There was a problem hiding this comment.
🤖 AI Code Review |
2026-04-10 18:09 CST
📋 Review 摘要
PR 概述:支持 bf16 MoE 的 tp permute、group_gemm、unpermute 功能,使用 paddlefleet_ops.deep_gemm.m_grouped_bf16_gemm_nn_contiguous 替代原有的 compute_ffn 路径。
变更范围:model_executor/layers/moe/fused_moe_cutlass_backend.py
影响面 Tag:[RL] [OP]
📝 PR 规范检查
PR 标题包含 [RL] Tag,符合规范。
但 PR 描述中的 Motivation 和 Modifications 部分未填写,建议补充:
标题建议(已符合规范,可保持不变):
[RL] support bf16 moe tp permute, group_gemm, unpermute
描述模板(建议补充):
## Motivation
优化 bf16 MoE 在 TP 模式下的性能,通过使用 `paddlefleet_ops.deep_gemm.m_grouped_bf16_gemm_nn_contiguous` 和 `fused_swiglu_scale` 算子替代原有的 `moe_expert_ffn` 路径。
## Modifications
1. 新增 `m_grouped_bf16_gemm_nn_contiguous` 函数封装 `paddlefleet_ops.deep_gemm.m_grouped_bf16_gemm_nn_contiguous`
2. 在 `apply_ep_prefill` 和 `apply_tp` 中,当 `FD_USE_PHI_MOE_PERMUTE=True` 且 `moe_quant_type=="w16a16"` 时:
- 使用 `moe_permute(return_expert_indices=True)` 获取 `expert_idx_per_token`
- 使用 `m_grouped_bf16_gemm_nn_contiguous` + `fused_swiglu_scale` 替代 `compute_ffn`
- 使用 `moe_unpermute(using_weighted_combine=False)`
3. 在 `apply_tp` 中新增对 `layer.with_bias` 的支持
4. 移除不再使用的 `count_tokens_per_expert_func` 导入问题
| 级别 | 文件 | 概述 |
|---|---|---|
| 🟡 建议 | fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py:181 |
apply_ep_prefill 中缺少对 layer.with_bias 的处理 |
总体评价
代码整体思路清晰,使用 paddlefleet_ops.deep_gemm 的新算子替代原有实现以优化性能。新参数 return_expert_indices 和 using_weighted_combine=False 在其他后端文件中已有使用,符合设计模式。
建议补充 PR 描述,并修复 apply_ep_prefill 中 bias 处理的不一致性问题。
| out = paddlefleet_ops.fused_swiglu_scale(out, dst_weights) | ||
| ffn_out = m_grouped_bf16_gemm_nn_contiguous( | ||
| out, getattr(layer, self.added_weight_attrs[1]), expert_idx_per_token | ||
| ) |
There was a problem hiding this comment.
🟡 建议 apply_ep_prefill 中缺少对 layer.with_bias 的处理
在 apply_ep_prefill 的新代码路径(FD_USE_PHI_MOE_PERMUTE and self.moe_quant_type == "w16a16")中,没有处理 layer.with_bias 的情况,与 apply_tp 中的处理不一致。
原始 compute_ffn 方法在第 115-117 行有 bias 处理逻辑,apply_tp 在第 370-372 行也有相同的处理,建议在 apply_ep_prefill 中添加相同的 bias 处理代码以保证一致性。
Motivation
Modifications
Usage or Command
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.