aclnnAttentionUpdate

支持的产品型号

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]	√
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

• 算子功能：将各SP域PA算子的输出的中间结果lse，localOut两个局部变量结果更新成全局结果。
• 计算公式：输入$lse_i$和$O_i$、输出$O$。 $$lse_{sum} = \text{sum}lse_i$$ $$lse = \sum_i \text{exp}(lse_i - lse_{sum})$$ $$lse_m = lse_{sum} + \text{log}(lse)$$ $$O = \sum_i O_i \cdot \text{exp}(lse_i - lse_m)$$

函数原型

每个算子分为，必须先调用“aclnnAttentionUpdateGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnAttentionUpdate”接口执行计算。

• aclnnStatus aclnnAttentionUpdateGetWorkspaceSize(const aclTensorList* lse, const aclTensorList* localOut, int64_t updateType, aclTensor* out, aclTensor* lseOut, uint64_t* workspaceSize, aclOpExecutor** executor)

• aclnnStatus aclnnAttentionUpdate(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnAttentionUpdateGetWorkspaceSize

• 参数说明：

  • lse（aclTensorList*,计算输入）：各SP域的局部lse，公式中的$lse_i$，Device侧的tensorList，tensorList长度为sp，每个Tensor shape为(batch * seqLen * headNum)。数据类型支持FLOAT32，支持为ND。
  • localOut（aclTensorList*,计算输入）：各SP域的局部attentionOut，公式中的$O_i$，Device侧的tensorList，tensorList长度为sp，每个Tensor shape为(batch * seqLen * headNum, headDim)。数据类型支持FLOAT32，支持为ND。
  • updateType（int64_t，计算输入）：指定需要执行的操作，当前只支持0(DECODE_UPDATE)。
  • out（aclTensor*，计算输出）：输出的tensor，公式中的$O$，shape为(batch * seqLen * headNum, headDim)。数据类型支持FLOAT32，支持为ND。
  • lseOut（aclTensor*，计算输出）：预留参数，暂未使用，当前只支持nullptr。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

  aclnnStatus：返回状态码，具体参见。

  [object Object]

aclnnAttentionUpdate

• 参数说明：

  • workspace(void *，入参)：在Device侧申请的workspace内存地址。
  • workspaceSize(uint64_t，入参)：在Device侧申请的workspace大小，由第一段接口aclnnAttentionUpdateGetWorkspaceSize获取。
  • executor(aclOpExecutor *，入参)：op执行器，包含了算子计算流程。
  • stream(aclrtStream，入参)：指定执行任务的Stream。

• 返回值：

  aclnnStatus：返回状态码，具体参见。

约束说明

• 序列并行的并行度sp取值范围[1, 16]
• headDim取值范围[8, 512]且是8的倍数
• 不支持
• 不支持空Tensor

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考。