aclnnNsaSelectedAttention

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	×
[object Object]Atlas A2 训练系列产品[object Object]	√
[object Object]Atlas 800I A2 推理产品[object Object]	×
A200I A2 Box 异构组件	×
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品[object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

• 算子功能：训练场景下，实现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$$

函数原型

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

• aclnnStatus aclnnNsaSelectedAttentionGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *topkIndices, const aclTensor *attenMaskOptional,const aclIntArray *actualSeqQLenOptional, const aclIntArray *actualSeqKvLenOptional, double scaleValue, int64_t headNum, char *inputLayout, int64_t sparseMode, int64_t selectedBlockSize, int64_t selectedBlockCount, const aclTensor *softmaxMaxOut, const aclTensor *softmaxSumOut, const aclTensor *attentionOut, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnNsaSelectedAttention(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)

aclnnNsaSelectedAttentionGetWorkspaceSize

• 参数说明：

  • query（aclTensor *，计算输入）：Device侧的aclTensor，公式中的query，数据类型支持BFLOAT16、FLOAT16，数据类型与key/value的数据类型一致，undefined支持ND；综合约束请见undefined。
  • key（aclTensor *，计算输入）：Device侧的aclTensor，公式中的key，数据类型支持BFLOAT16、FLOAT16，数据类型与query/value的数据类型一致，undefined支持ND；综合约束请见undefined。
  • value（aclTensor *，计算输入）：Device侧的aclTensor，公式中的value，数据类型支持BFLOAT16、FLOAT16，数据类型与query/key的数据类型一致，undefined支持ND；综合约束请见undefined。
  • topkIndices（aclTensor *，计算输入）：Device侧的aclTensor，公式中的topk_indices，shape需为[T_q, N_kv, selected_block_count], 表示所选数据的索引，数据类型支持INT32，undefined支持ND，综合约束请见undefined。
  • attenMaskOptional（aclTensor *，计算输入）：Device侧的aclTensor，公式中的atten_mask，数据类型支持BOOL、UINT8取值为true/1代表该位不参与计算（不生效），为false/0代表该位参与计算，undefined支持ND，输入shape类型需为[S_q, S_kv]；综合约束请见undefined。
  • actualSeqQLenOptional（aclIntArray *，计算输入）：Host侧的aclIntArray，数据类型支持INT64，undefined支持ND，长度等于batchsize。该数组表示query每个Batch S的累加和长度，假设输入真实的S长度分别为[2,2,3,2]，则传入的actualSeqQLenOptional为[2,4,7,9]。在TND排布时需要输入，其余场景下输入nullptr。
  • actualSeqKvLenOptional（aclIntArray *，计算输入）：Host侧的aclIntArray，数据类型支持INT64，undefined支持ND，长度等于batchsize。该数组表示key/value每个Batch S的累加和长度，假设输入真实的S长度分别为[1024,1024,1024,1024]，则传入的actualSeqKvLenOptional为[1024,2048,3072,4096]。在TND排布时需要输入，其余场景下输入nullptr。
  • scaleValue（double，计算输入）：Host侧的double，公式中的scale，代表缩放系数，数据类型支持DOUBLE，一般设置为D^-0.5，其中D为输入query的head维度。
  • headNum（int64_t，计算输入）：Host侧的int64_t，代表head个数，即输入query的N轴长度，数据类型支持INT64；综合约束请见undefined。
  • inputLayout（string *，计算输入）：Host侧的string，数据类型支持String，代表输入query、key、value的数据排布格式，当前仅支持TND，其中T表示各batch S的长度累加和，N表示Head-Num，D表示Head-Dim。
  • selectedBlockSize（int64_t，计算输入）：Host侧的int64_t，表示select的每个block长度。
  • selectedBlockCount（int64_t，计算输入）：Host侧的int64_t，公式中的selected_block_count，表示select block的数量。
  • sparseMode（int64_t，计算输入）：Host侧的int64_t，表示sparse的模式。数据类型支持INT32。目前支持sparseMode=0或者2。sparse不同模式的详细说明请参见undefined。
  • softmaxMaxOut（aclTensor *，计算输出）：Device侧的aclTensor，Softmax计算的Max中间结果，用于反向计算。数据类型支持FLOAT，输出的shape类型为[T_q, N_q, 8]，undefined支持ND。
  • softmaxSumOut（aclTensor *，计算输出）：Device侧的aclTensor，Softmax计算的Sum中间结果，用于反向计算。数据类型支持FLOAT，输出的shape类型为[T_q, N_q, 8]，undefined支持ND。
  • attentionOut（aclTensor *，计算输出）：Device侧的aclTensor，计算公式的最终输出。数据类型支持BFLOAT16、FLOAT16，输出数据类型与query保持一致, shape类型为[T_q, N_q, D_v]，undefined支持ND。
  • workspaceSize（uint64_t *，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor **，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnNsaSelectedAttention

• 参数说明：

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

• 返回值：

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

约束说明

• 该接口与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的input_layout必须一致。

• sparseMode目前支持0和2。

• selectedBlockSize支持<=128且满足16的整数倍。

• selectedBlockCount支持<=32。

• inputLayout目前仅支持TND。

• 支持输入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。

调用示例

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