aclnnNsaSelectedAttention

产品支持情况

[object Object]undefined

功能说明

• 接口功能：训练场景下，实现NativeSparseAttention算法中selected-attention（选择注意力）的计算。

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

  $$selected\_key = Gather(key, topk\_indices[i]),0<=i<selected\_block\_count \\ selected\_value = Gather(value, topk\_indices[i]),0<=i<selected\_block\_count$$ $$attention\_out = Softmax(Mask(scale * (query @ selected\_key^T), atten\_mask)) @ selected\_value$$

函数原型

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

aclnnNsaSelectedAttentionGetWorkspaceSize

• 参数说明：

  [object Object]

• 返回值：

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

  第一段接口完成入参校验，出现以下场景时报错：

  [object Object]

aclnnNsaSelectedAttention

• 参数说明：

  [object Object]

• 返回值：

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

约束说明

• 确定性计算：

  • aclnnNsaSelectedAttention默认确定性实现。

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

• 输入query、key、value的batchsize必须相等，即要求传入的actualSeqQLenOptional和actualSeqKvLenOptional具有相同的长度。

• 输入query、key、value的D：Head-Dim必须满足（D_q == D_k && D_k >= D_v）。

• 输入query、key、value的数据类型必须一致。

• 输入query、key、value的inputLayout必须一致。

• selectBlockCount：支持[1~128]。 总计选择的大小[object Object] < 128*64(8K)

• Layout为TND时，每个Batch的S2都要大于总计选择的大小[object Object]

• 支持输入query的N和key/value的N不相等，但必须成比例关系，即N_q / N_kv必须是非0整数，称为G（group），且需满足G <= 32。

• 当attenMaskOptional输入为nullptr时，sparseMode参数不生效，固定为全计算。

• 关于数据shape的约束，以inputLayout的TND举例（注：T等于各batch S的长度累加和。当各batch的S相等时，T=B*S）。其中：

  • B（Batchsize）：取值范围为1~1024。
  • N（Head-Num）：取值范围为1~128。
  • G（Group）：取值范围为1~32。
  • S（Seq-Length）：取值范围为1~128K。同时需要满足S_kv >= selectedBlockSize * selectedBlockCount，且S_kv长度为selectedBlockSize的整数倍。
  • D（Head-Dim）：D_qk=192，D_v=128。

调用示例

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