[Optimization] Use triton qk_norm both in Prefill and Decode.#7213
[Optimization] Use triton qk_norm both in Prefill and Decode.#7213zhoutianzi666 merged 1 commit intoPaddlePaddle:developfrom
Conversation
|
Thanks for your contribution! |
|
“liuruian” 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. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #7213 +/- ##
==========================================
Coverage ? 73.55%
==========================================
Files ? 383
Lines ? 53513
Branches ? 8378
==========================================
Hits ? 39359
Misses ? 11408
Partials ? 2746
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:
|
fastdeploy-bot
left a comment
fastdeploy-bot
left a comment
There was a problem hiding this comment.
🤖 AI Code Review | 2026-04-08
📋 Review 摘要
PR 概述：移除 QKRMSNorm 中的 step_use_cudagraph 条件，使 Triton qk_rmsnorm_fused kernel 在 Prefill 和 Decode 两个阶段都使用，以提升性能。
变更范围：fastdeploy/model_executor/layers/normalization.py
影响面 Tag：[Optimization] [OP]
📝 PR 规范检查
标题：✅ 符合规范，包含 [Optimization] Tag
Motivation：✅ 说明了优化目的和性能提升预期
Modifications：✅ 简要说明使用 QKRMSNorm 替代 paddle 散 Op
Checklist 问题：
- ❌ "Provide accuracy results" 未勾选，未提供精度测试结果
- ❌ "Add unit tests" 未勾选
问题
| 级别 | 文件 | 概述 |
|---|---|---|
| 🟡 建议 | normalization.py:344 |
建议补充 Prefill 阶段的大 batch 精度验证 |
总体评价
代码变更逻辑清晰，qk_rmsnorm_fused 已有单元测试验证。但考虑到 Prefill 阶段 batch size 通常远大于测试场景（测试最大 8192，实际可能更大），建议补充大 batch 下的精度对比测试，确保 Triton kernel 在 Prefill 场景的数值正确性。
| proxy_rmsnorm=None, | ||
| ) -> paddle.Tensor: | ||
| if proxy_rmsnorm is None and self.qk_norm_fused and forward_meta.step_use_cudagraph: | ||
| if proxy_rmsnorm is None and self.qk_norm_fused: |
There was a problem hiding this comment.
🟡 建议 建议补充 Prefill 阶段的精度验证结果。
虽然 qk_rmsnorm_fused 已有单元测试验证，但该测试主要针对小 batch size（如 128）。Prefill 阶段 batch size 通常更大（可能达 4096+），建议补充大 batch 下的精度对比测试，确保 Triton kernel 在 Prefill 场景下的数值正确性。
K11OntheBoat
force-pushed
the
fused_rms_triton
branch
from
a3d1fcf to
fde4f8b
Compare
fastdeploy-bot
left a comment
fastdeploy-bot
left a comment
There was a problem hiding this comment.
🤖 AI Code Review |
2026-04-08 12:09 CST
📋 Review 摘要
PR 概述：移除 QKRMSNorm.forward() 中的 forward_meta.step_use_cudagraph 条件，使 prefill 阶段也能使用 qk_rmsnorm_fused triton kernel
变更范围：model_executor/layers/normalization.py、tests/e2e/test_Qwen3VL_serving.py
影响面 Tag：[Optimization] [OP]
📝 PR 规范检查
标题：✅ 符合规范，包含有效 Tag [Optimization]
描述：
- Motivation 和 Modifications 填写完整
- Checklist 未勾选（特别是"Add unit tests"、"Provide accuracy results"）
- 缺少具体的性能 benchmark 数据（如
tests/operators/test_qk_rmsnorm_fused.py中的 benchmark 结果）
建议标题：无需修改
建议描述补充：
## Accuracy Tests
- [x] 使用 `tests/operators/test_qk_rmsnorm_fused.py` 验证 triton kernel 与 paddle 实现的数值一致性
- [ ] 提供 prefill 阶段使用前后的端到端准确性对比数据
## Benchmark Results
参考 `tests/operators/test_qk_rmsnorm_fused.py::TestQKNorm.test_qk_norm_result` 的测试结果：
问题
| 级别 | 文件 | 概述 |
|---|---|---|
| 🟡 建议 | tests/e2e/test_Qwen3VL_serving.py:176 |
测试文件文本变更应与核心修改分离 |
| 🟡 建议 | fastdeploy/model_executor/layers/normalization.py:344 |
缺少准确性测试结果 |
| ❓ 疑问 | fastdeploy/model_executor/layers/normalization.py:344 |
历史限制原因未说明 |
总体评价
变更逻辑清晰，qk_rmsnorm_fused triton kernel 支持任意 token 数量，适用于 prefill 和 decode 阶段。但建议补充准确性测试数据和修改背景说明。
|
|
||
| # base result | ||
| content2 = "视频中手机支架的颜色是黑色的。" | ||
| content2 = "视频中手机支架的颜色是黑色。" |
There was a problem hiding this comment.
🟡 建议 测试文件变更应与核心修改分离到单独 PR
测试文件中将预期结果从"黑色的"改为"黑色"是一个独立的文本格式调整（去掉助词"的"），与本次 PR 的核心优化（移除 step_use_cudagraph 条件）无关。
建议将此文本格式调整移到单独的 PR，保持 PR 的单一职责原则。
| proxy_rmsnorm=None, | ||
| ) -> paddle.Tensor: | ||
| if proxy_rmsnorm is None and self.qk_norm_fused and forward_meta.step_use_cudagraph: | ||
| if proxy_rmsnorm is None and self.qk_norm_fused: |
There was a problem hiding this comment.
🟡 建议 缺少准确性测试结果
PR 声称有性能提升（"部分模型 单Kernel部分加速2～7倍"），但 Checklist 中未提供具体的准确性测试结果。
需要说明：
- prefill 阶段使用 triton kernel 与原来 paddle 实现的数值精度是否一致
- 是否有端到端测试对比修改前后的模型输出准确性
建议提供具体的 accuracy 测试数据，例如使用标准数据集（如 MMLU）的对比结果。
| proxy_rmsnorm=None, | ||
| ) -> paddle.Tensor: | ||
| if proxy_rmsnorm is None and self.qk_norm_fused and forward_meta.step_use_cudagraph: | ||
| if proxy_rmsnorm is None and self.qk_norm_fused: |
There was a problem hiding this comment.
❓ 疑问 历史限制原因未说明
根据 git 历史，提交 #6080（fix opt qknorm）添加了 forward_meta.step_use_cudagraph 条件，限制了只在 decode 阶段使用 qk_rmsnorm_fused。
PR 描述没有说明：
- 为什么 [BugFix] Fix qk_norm optimization #6080 要限制为 decode 阶段（是否有已知的 bug 或性能问题？）
- 现在移除限制是否安全，是否解决了 [BugFix] Fix qk_norm optimization #6080 当时遇到的问题
建议在 PR 描述中补充相关背景信息。
K11OntheBoat
force-pushed
the
fused_rms_triton
branch
from
fde4f8b to
c7b49c3
Compare
fastdeploy-bot
left a comment
fastdeploy-bot
left a comment
There was a problem hiding this comment.
🤖 AI Code Review |
2026-04-08 16:43 CST
📋 Review 摘要
PR 概述：移除 QKRMSNorm 融合算子的 step_use_cudagraph 条件限制，使其在 Prefill 和 Decode 阶段都能使用 Triton 优化算子
变更范围：model_executor/layers/normalization.py、tests/e2e/test_Qwen3VL_serving.py
影响面 Tag：[OP] [Optimization]
📝 PR 规范检查
PR 标题包含有效 [Optimization] 标签，Motivation 和 Modifications 已填写，符合规范。
问题
| 级别 | 文件 | 概述 |
|---|---|---|
| 🟡 建议 | tests/e2e/test_Qwen3VL_serving.py:176 |
测试期望值变更原因未说明 |
总体评价
核心变更是移除 QKRMSNorm.forward() 中对 forward_meta.step_use_cudagraph 的检查，使得 Triton 融合算子在 Prefill 阶段也能使用。从代码逻辑看：
- 正确性：
qk_rmsnorm_fused是独立的 Triton kernel，不依赖 CUDA Graph，移除此条件是安全的 - 兼容性：调用方
qwen3.py传递forward_meta，而glm4_moe.py不传递forward_meta，两者都能正常工作 - 性能：根据 PR 描述，该变更可带来 2-7 倍的 Kernel 加速，对 Prefill 阶段效果显著
主要建议是说明测试期望值变更的原因。
|
|
||
| # base result | ||
| content2 = "视频中手机支架的颜色是黑色的。" | ||
| content2 = "视频中手机支架的颜色是黑色。" |
There was a problem hiding this comment.
🟡 建议 请说明测试期望值变更的原因。
修改从 "黑色。" 到 "黑色"，请问是因为使用融合算子后模型输出发生了变化，还是仅仅是文本调整？如果模型输出确实有变化，建议确认这是否符合预期，并考虑增加对应的精度验证测试。
Co-authored-by: “liuruian” <liuruian@baidu.com>
|
✅ Cherry-pick successful! Created PR: #7305 |
Co-authored-by: “liuruian” <liuruian@baidu.com>
|
✅ Cherry-pick successful! Created PR: #7306 |
6 participants
Motivation
Prefill 阶段使用QKRMSNorm融合算子. 部分模型 单Kernel部分加速2～7倍. Prefill 空泡较大的模型单次Forward可加速2倍左右.
Modifications
使用QKRMSNorm 替代paddle 散Op
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.