[Optimization] Use triton qk_norm both in Prefill and Decode.

[K11OntheBoat]

Motivation

Prefill 阶段使用QKRMSNorm融合算子. 部分模型 单Kernel部分加速2～7倍. Prefill 空泡较大的模型单次Forward可加速2倍左右.

Modifications

使用QKRMSNorm 替代paddle 散Op

Usage or Command

Accuracy Tests

Checklist

• 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.
• 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.

[paddle-bot]

Thanks for your contribution!

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept 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-commenter]

Codecov Report

❌ Patch coverage is 0% with 1 line in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@ae2f9f4). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
fastdeploy/model_executor/layers/normalization.py	0.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             develop    #7213   +/-   ##
==========================================
  Coverage           ?   73.21%           
==========================================
  Files              ?      376           
  Lines              ?    52938           
  Branches           ?     8264           
==========================================
  Hits               ?    38761           
  Misses             ?    11461           
  Partials           ?     2716           

Flag	Coverage Δ
GPU	73.21% <0.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[fastdeploy-bot]

🤖 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 场景的数值正确性。

[fastdeploy-bot]

🟡 建议 建议补充 Prefill 阶段的精度验证结果。

虽然 qk_rmsnorm_fused 已有单元测试验证，但该测试主要针对小 batch size（如 128）。Prefill 阶段 batch size 通常更大（可能达 4096+），建议补充大 batch 下的精度对比测试，确保 Triton kernel 在 Prefill 场景下的数值正确性。