aclnnFusedFloydAttention

产品支持情况

[object Object]undefined

功能说明

• 接口功能：训练场景下，使用FloydAttention算法实现多维自注意力的计算。

• 计算公式：

  注意力的正向计算公式如下：

  $$weights = Softmax(attenMask + scale*(einsum(query, key1^T) + einsum(query, key2^T)))$$ $$attention\_out = einsum(weights, value1) + einsum(weights, value2)$$

函数原型

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

aclnnFusedFloydAttentionGetWorkspaceSize

• 参数说明

  [object Object]

• 返回值

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

  第一段接口完成入参校验，若出现以下错误码，则对应原因为：

  [object Object]

aclnnFusedFloydAttention

• 参数说明

  [object Object]

• 返回值

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

约束说明[object Object][object Object]

• 该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配

• 关于数据shape的约束，其中：

  • B：取值范围为1~2K。
  • H：取值范围为1~256。
  • N：取值范围为16~1M且N%16==0。
  • M：取值范围为128~1M且M%128==0。
  • K：取值范围为128~1M且K%128==0。
  • D：取值范围为16~128。

• query与key1的第0/2/4轴需相同。

• key1与value1 shape需相同。

• key2与value2 shape需相同。

• softmaxMax与softmaxSum shape需相同。

调用示例

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