CANN/atvoss RmsNorm算子样例

RmsNorm算子样例

【免费下载链接】atvossATVOSS（Ascend C Templates for Vector Operator Subroutines）是一套基于Ascend C开发的Vector算子库，致力于为昇腾硬件上的Vector类融合算子提供极简、高效、高性能、高拓展的编程方式。项目地址: https://gitcode.com/cann/atvoss

概述

样例概述：本样例介绍了利用ATVOSS实现RmsNorm单算子并完成功能验证

• 算子功能：RmsNorm算子是大模型常用的归一化操作，相比LayerNorm算子，其去掉了减去均值的部分
• 调用方式：Kernel直调

样例支持的产品

• Ascend 950PR/Ascend 950DT

算子描述

• 算子数学计算公式： $$ \operatorname{RmsNorm}(x_i)=\frac{x_i}{\operatorname{Rms}(\mathbf{x})} g_i, \quad \text { where } \operatorname{Rms}(\mathbf{x})=\sqrt{\frac{1}{n} \sum_{i=1}^n x_i^2} $$

• 算子规格：

• 当前只支持二维输入
• 总的输入Shape(M, N)要满足：
  • M < 8160，N <= 7168
  • N需要32元素对齐
• Tile块的Shape(m, n)，要满足n = N，m * n <=7168
• 目前只支持float类型

目录结构

RmsNorm样例算子的编译和运行

• 编译 在代码仓根目录下执行：

bash scripts/build.sh -DSOC=ascend950 rms_norm

• 运行 在代码仓目录下执行：

output/bin/rms_norm --help // 查看帮助 output/bin/rms_norm --shape=16,32 // 运行样例

【免费下载链接】atvossATVOSS（Ascend C Templates for Vector Operator Subroutines）是一套基于Ascend C开发的Vector算子库，致力于为昇腾硬件上的Vector类融合算子提供极简、高效、高性能、高拓展的编程方式。项目地址: https://gitcode.com/cann/atvoss