AttentionUpdate

English | 简体中文

Product Support

Product	Supported
Ascend 950PR&950DT products	√
Atlas A3 products	√
Atlas A2 products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Function Description

• Operator function: Updates the intermediate results lse and localOut output by the PA operators in each SP domain into the global results.
• Computation formula: Inputs $ls{e}_i$ and $O_i$, and output $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)$$

Parameter Description

Parameter name	Input/Output/Attribute	Description	Data type	Data format
lsei	Input	Local lse of each SP domain.	FLOAT32	ND
Oi	Input	Local attentionout of each SP domain.	FLOAT32, FLOAT16, BFLOAT16	ND
lsem	Output	Updated global lse.	FLOAT32	ND
O	Output	Updated global attentionout.	FLOAT32, FLOAT16, BFLOAT16	ND

Constraint Description

• For Atlas A2 products and Atlas A3 products: O_i and O of the FLOAT32, FLOAT16, and BFLOAT16 types are supported.
• For Ascend 950PR&950DT products: O_i and O of the FLOAT32, FLOAT16, and BFLOAT16 types are supported, and O_i and O must have the same data type.
• The value range of the sequence parallelism degree sp:
  • Atlas A2 products and Atlas A3 products: [1, 128].
  • Ascend 950PR&950DT products: [1, 16].
• headDim ranges from [8, 512] and must be a multiple of 8.
• localout supports non-contiguous tensors.
• Empty input with T=0 is supported. In this case, the output is empty and no computation is performed.

Invocation Description

Invocation method	Sample code	Description
aclnn interface	test_aclnn_attention_update	Invokes the AttentionUpdate operator through the aclnnAttentionUpdate interface.