CANNOpsTransformer注意力更新算子
AttentionUpdate
【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | √ |
| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ |
| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ |
| Atlas 200I/500 A2 推理产品 | × |
| Atlas 推理系列产品 | × |
| Atlas 训练系列产品 | × |
功能说明
- 算子功能：将各SP域PA算子的输出的中间结果lse，localOut两个局部变量结果更新成全局结果。
- 计算公式：输入$lse_i$和$O_i$、输出$O$。
$$ lse_{max} = \text{max}lse_i $$
$$ lse = \sum_i \text{exp}(lse_i - lse_{max}) $$
$$ lse_m = lse_{max} + \text{log}(lse) $$
$$ O = \sum_i O_i \cdot \text{exp}(lse_i - lse_m) $$
参数说明
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|---|---|---|---|---|
| lsei | 输入 | 各SP域的局部lse。 | FLOAT32 | ND |
| Oi | 输入 | 各SP域的局部attentionout。 | FLOAT32，FLOAT16，BFLOAT16 | ND |
| lsem | 输出 | 更新后的全局lse。 | FLOAT32 | ND |
| O | 输入 | 更新后的全局attentionout。 | FLOAT32，FLOAT16，BFLOAT16 | ND |
约束说明
- Atlas A2 训练系列产品/Atlas A2 推理系列产品 、 Atlas A3 训练系列产品/Atlas A3 推理系列产品 ：支持FLOAT32、FLOAT16、BFLOAT16的Oi和O。
- Ascend 950PR/Ascend 950DT ：支持FLOAT32、FLOAT16、BFLOAT16的Oi和O，且Oi和O数据类型相同。
- 序列并行的并行度sp取值范围[1, 16]。
- headDim取值范围[8, 512]且是8的倍数。
- 不支持非连续的Tensor。
- 支持空Tensor。
调用说明
| 调用方式 | 样例代码 | 说明 |
|---|---|---|
| aclnn接口 | test_aclnn_attention_update | 通过aclnnAttentionUpdate接口方式调用AttentionUpdate算子。 |
【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考