CANN/cann-bench：AddRmsNormDynamicQuant算子API描述

AddRmsNormDynamicQuant 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介Add、RMSNorm 和动态量化的融合。主要应用场景大语言模型推理中残差连接 归一化 量化的融合加速Transformer 模型中 RMSNorm 前的残差加法与后处理量化一体化INT8/INT4 低精度推理的动态量化预处理算子特征难度等级L3FusedComposite多输入多输出融合 Add、RMSNorm 和动态量化三个操作输入 x1、x2 为 ND 格式张量gamma 为缩放参数2. 算子定义数学公式$$ y, xOut, scaleOut \text{quantize}(\text{rmsnorm}(x_1 x_2) \times \gamma) $$具体步骤Add 操作$xOut x_1 x_2$RMSNorm$\text{normalized} \frac{xOut}{\sqrt{\text{mean}(xOut^2) \epsilon}} \times \gamma$动态量化根据 dst_type 将归一化结果量化为 INT8 或 INT4同时输出量化 scale3. 接口规范算子原型cann_bench.add_rms_norm_dynamic_quant(Tensor x1, Tensor x2, Tensor gamma, float epsilon, int dst_type) - (Tensor y, Tensor xOut, Tensor scaleOut)输入参数说明参数类型默认值描述x1Tensor必选第 1 个输入张量x2Tensor必选第 2 个输入张量gammaTensor必选缩放参数epsilonfloat1e-6epsilon 值dst_typeint0目标数据类型 (0:DT_INT8, 1:DT_INT4)输出参数Shapedtype描述y与输入 x1 相同int8 / int4量化后的输出张量xOut与输入 x1 相同float16 / bfloat16Add 结果x1 x2scaleOut标量float32量化使用的 scale 值数据类型输入 (x1, x2, gamma) dtype输出 y dtype输出 xOut dtype输出 scaleOut dtypefloat16int8 / int4float16float32bfloat16int8 / int4bfloat16float32注意INT4 量化dst_type1的输出值范围为 [-8, 7]由于 PyTorch 不支持 int4 dtype实际存储为 int8 dtype。规则与约束x1 和 x2 的 shape 和 dtype 必须一致gamma 的 dtype 须与 x1、x2 一致x1 为 ND 格式dst_type 取值0 表示 DT_INT81 表示 DT_INT4epsilon 用于 RMSNorm 的数值稳定性默认 1e-6scaleOut 为 float32 类型标量4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差MERE采样点中相对误差平均值$$ \text{MERE} \text{avg}(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$最大相对误差MARE采样点中相对误差最大值$$ \text{MARE} \max(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$通过标准数据类型FLOAT16BFLOAT16FLOAT32HiFLOAT32FLOAT8 E4M3FLOAT8 E5M2通过阈值(Threshold)2^-102^-72^-132^-112^-32^-2当平均相对误差 MERE Threshold最大相对误差 MARE 10 * Threshold 时判定为通过。5. 标准 Golden 代码import torch AddRmsNormDynamicQuant 算子 Torch Golden 参考实现 Add、RMSNorm 和动态量化的融合 公式y, xOut, scaleOut quantize(rmsnorm(x1 x2) * gamma) def add_rms_norm_dynamic_quant( x1: torch.Tensor, x2: torch.Tensor, gamma: torch.Tensor, epsilon: float 1e-6, dst_type: int 0 ) - tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Add、RMSNorm 和动态量化的融合 公式y, xOut, scaleOut quantize(rmsnorm(x1 x2) * gamma) Args: x1: 第 1 个输入张量 x2: 第 2 个输入张量 gamma: 缩放参数 epsilon: epsilon 值 dst_type: 目标数据类型 (0:DT_INT8, 1:DT_INT4) Returns: y: 量化后的输出张量 xOut: Add 结果x1 x2 scaleOut: 量化使用的 scale 值 # Add 操作 xOut x1 x2 # RMSNorm variance xOut.pow(2).mean(-1, keepdimTrue) rms torch.sqrt(variance epsilon) normalized xOut / rms y_norm normalized * gamma # 动态量化 # 将 y_norm 转换为 float32 以保证 scale 计算精度和 dtype 正确 y_norm_f32 y_norm.float() if dst_type 0: # INT8 scale (127.0 / y_norm_f32.abs().max()).to(torch.float32) y torch.clamp((y_norm_f32 * scale.item()).round(), -128, 127).to(torch.int8) else: # INT4 (存储为 int8值范围 [-8, 7]) scale (7.0 / y_norm_f32.abs().max()).to(torch.float32) y torch.clamp((y_norm_f32 * scale.item()).round(), -8, 7).to(torch.int8) return y, xOut, scale6. 额外信息算子调用示例import torch import cann_bench x1 torch.randn(2, 4096, dtypetorch.float16, devicenpu) x2 torch.randn(2, 4096, dtypetorch.float16, devicenpu) gamma torch.ones(4096, dtypetorch.float16, devicenpu) y, xOut, scaleOut cann_bench.add_rms_norm_dynamic_quant(x1, x2, gamma, 1e-6, 0) # INT8 量化 y, xOut, scaleOut cann_bench.add_rms_norm_dynamic_quant(x1, x2, gamma, 1e-6, 1) # INT4 量化【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考