AttentionUpdate

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：将各SP域PA算子的输出的中间结果lse，localOut两个局部变量结果更新成全局结果。
• 计算公式：输入$ls{e}_i$和$O_i$、输出$O$。

$$ls{e}_{max}=\text{max}ls{e}_i$$

$$lse=\sum _i\text{exp}(ls{e}_i-ls{e}_{max})$$

$$ls{e}_m=ls{e}_{max}+\text{log}(lse)$$

$$O=\sum _i{O}_i\cdot \text{exp}(ls{e}_i-ls{e}_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算子。