feat: add GammaAddRmsNorm custom operator.#10
Open
Fengfengst123 wants to merge 1 commit into
Open
Conversation
LMX-xin
reviewed
Jul 15, 2026
| @@ -0,0 +1,163 @@ | |||
| { | |||
| "op_type": "GammaAddRmsNorm", | |||
Collaborator
There was a problem hiding this comment.
其他算子都是在def函数中直接设置兼容的config就行
Collaborator
|
这个是cann中的算子，还是你们自己写的，是cann版本的算子的话，当前cann包中没有吗 |
LMX-xin
changed the title
Fengfengst123
force-pushed
the
codex/gamma-add-rmsnorm-upstream-draft
branch
3 times, most recently
from
6cf02eb to
eefcb2e
Compare
Contributor
|
需要补充算子的测试代码 |
Fengfengst123
force-pushed
the
codex/gamma-add-rmsnorm-upstream-draft
branch
2 times, most recently
from
ce18890 to
1f03b79
Compare
Contributor
Author
这个是ops-nn的AddRMSNorm, 本次现在qwen3.5的路径中，有一处gamma +1 的Add操作 + 加 AddRmsNorm，本次加了一个参数，做了融合，然后省去了这一部分Add的时间 |
Contributor
Author
好的，已添加 |
Sinle4Cat
reviewed
Jul 16, 2026
| namespace OpTiling { | ||
| static const gert::Shape g_vec_1_shape = {1}; | ||
|
|
||
| static bool IsRegbaseSocVersion(platform_ascendc::SocVersion version) |
Contributor
Author
There was a problem hiding this comment.
本地用的cann包版本太低了，AI自动做了下兼容，已经修复去掉了
Contributor
There was a problem hiding this comment.
不涉及 310p 9030 x90 不要在这里体现
| this->AICore().AddConfig("ascend910_93"); | ||
|
|
||
| OpAICoreConfig config310P; | ||
| config310P.Input("x1") |
yingxudeng
reviewed
Jul 16, 2026
| uint32_t curElementByte = dtypeByteIterator->second; | ||
| numColAlign = CeilDiv(numCol * curElementByte, ubBlockSize) * ubBlockSize / curElementByte; | ||
|
|
||
| // 计算 二分累加 分界点 |
Collaborator
There was a problem hiding this comment.
Please keep all comments in English.
Fengfengst123
force-pushed
the
codex/gamma-add-rmsnorm-upstream-draft
branch
from
1f03b79 to
8d587d9
Compare
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
4 participants
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.
功能概述
GammaAddRmsNorm自定义算子及 ACLNN 接口aclnnGammaAddRmsNorm，将x = x1 + x2与 RMSNorm 融合执行。addGammaOffset，支持在 kernel 内执行 Gemma 风格的gamma + 1，避免额外的 Add 算子下发。ascend910b)、Ascend 910C (ascend910_93) 和 Ascend 950 (ascend950)的自定义算子构建列表。
xllm-ops 中独立构建，不依赖外部 ops-nn 源码树。
性能测试
在本机 Ascend 910B3 上使用 ATK
performance_device测量 BF16 devicetask latency。候选链路为
GammaAddRmsNorm(addGammaOffset=true)，原链路为Add(gamma, 1) + AddRmsNorm。两条链路使用相同的 shape、dtype、输入范围和随机种子；每个 case 执行 100 次 warmup、20 次测量，并由 ATK 重复 10 轮。
模型相关 shape（归一化维度
D >= 1024）结果如下：Qwen3.5 实际 decode 小 shape 的本机 msprof device-task 结果如下：
精度验证
GammaAddRmsNorm精度通过，与原链路
Add(gamma, 1) + AddRmsNorm一致。Summary
GammaAddRmsNormcustom operator and theaclnnGammaAddRmsNormACLNN API, fusingx = x1 + x2with RMSNorm.addGammaOffsetattribute to apply the Gemma-stylegamma + 1inside the kernel and eliminate a separate Add launch.ascend910b), Ascend 910C(
ascend910_93), and Ascend 950 (ascend950) custom-op build lists.operator builds independently in xllm-ops without an external ops-nn source
tree.
Performance
BF16 device-task latency was measured locally on Ascend 910B3 with ATK
performance_device. The candidate isGammaAddRmsNorm(addGammaOffset=true)and the previous chain isAdd(gamma, 1) + AddRmsNorm. Both paths use identical shapes, dtypes, inputranges, and random seeds. Each case uses 100 warmups, 20 measured iterations,
and 10 ATK performance repeats.
Results for model-relevant shapes with normalized dimension
D >= 1024:Qwen3.5 decode small-shape device-task results from the local msprof runs:
Accuracy validation
GammaAddRmsNormaccuracy matches the previousAdd(gamma, 1) + AddRmsNormchain.