aclnnNsaSelectedAttentionGrad

产品支持情况

产品	是否支持
[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]	×

功能说明

• 算子功能：根据topkIndices对key和value选取大小为selectedBlockSize的数据重排，接着进行训练场景下计算注意力的反向输出。

• 计算公式：

  根据传入的topkIndice对keyIn和value选取数量为selectedBlockCount个大小为selectedBlockSize的数据重排，公式如下：

  $$selectedKey = Gather(key, topkIndices[i]),0<=i<selectedBlockCount \\ selectedValue = Gather(value, topkIndices[i]),0<=i<selectedBlockCount$$

  接着，进行注意力机制的反向计算，计算公式为：

  $$V=P^TdY$$ $$Q=\frac{((dS)*K)}{\sqrt{d}}$$ $$K=\frac{((dS)^T*Q)}{\sqrt{d}}$$

函数原型

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

• aclnnStatus aclnnNsaSelectedAttentionGradGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *attentionOut, const aclTensor *attentionOutGrad, const aclTensor *softmaxMax, const aclTensor *softmaxSum, const aclTensor *topkIndices, const aclIntArray *actualSeqQLenOptional,const aclIntArray *actualSeqKvLenOptional, const aclTensor *attenMaskOptional, double scaleValue, int64_t selectedBlockSize, int64_t selectedBlockCount, int64_t headNum, char *inputLayout, int64_t sparseMode, aclTensor *dqOut, aclTensor *dkOut, aclTensor *dvOut, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnNsaSelectedAttentionGrad(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnNsaSelectedAttentionGradGetWorkspaceSize

• 参数说明：

  • query（aclTensor *，计算输入）：Device侧的aclTensor，公式中的输入Q，数据类型支持BFLOAT16和FLOAT16，shape支持3维，undefined支持ND，综合约束请见undefined。
  • key（aclTensor *，计算输入）：Device侧的aclTensor，公式中的输入K，数据类型支持BFLOAT16和FLOAT16，数据类型与query一致，shape支持3维，undefined支持ND，综合约束请见undefined。
  • value（aclTensor *，计算输入）：Device侧的aclTensor，公式中的输入V，数据类型支持BFLOAT16和FLOAT16，数据类型与query一致，shape支持3维，undefined支持ND，综合约束请见undefined。
  • attentionOut（aclTensor *，计算输入）：Device侧的aclTensor，注意力正向计算的最终输出，数据类型支持BFLOAT16和FLOAT16，数据类型与query一致，shape支持3维，undefined支持ND。
  • attentionOutGrad（aclTensor *，计算输入）：Device侧的aclTensor，公式中的输入dY，数据类型支持BFLOAT16和FLOAT16，数据类型与query一致，shape支持3维，undefined支持ND，综合约束请见undefined。
  • softmaxMax（aclTensor *，计算输入）：Device侧的aclTensor，注意力正向计算的中间输出，数据类型支持FLOAT，shape支持3维，undefined支持ND。综合约束请见undefined。
  • softmaxSum（aclTensor *，计算输入）：Device侧的aclTensor，注意力正向计算的中间输出，数据类型支持FLOAT，shape支持3维，undefined支持ND；综合约束请见undefined。
  • topkIndices（aclTensor *，计算输入）：Device侧的aclTensor，公式中的所选择KV的索引，数据类型支持INT32，shape支持3维，undefined支持ND，综合约束请见undefined。
  • actualSeqQLenOptional（aclIntArray *，计算输入）：数据类型支持INT64，描述了每个Batch对应的query S大小；综合约束请见undefined。
  • actualSeqKvLenOptional（aclIntArray *，计算输入）：数据类型支持INT64，描述了每个Batch对应的key/value S大小。
  • attenMaskOptional（aclTensor *，计算输入）：Device侧的aclTensor，0代表保留，1代表掩掉的mask矩阵，预留参数，暂不支持。
  • scaleValue（double，计算输入）：公式中d开根号的倒数，代表缩放系数，作为计算流中Muls的scalar值，数据类型支持DOUBLE。一般设置为$D^{-0.5}$。
  • selectedBlockSize（int64_t，计算输入）：所选择的块的大小，数据类型支持INT64，目前支持16~128，需对齐16。
  • selectedBlockCount（int64_t，计算输入）：所选择的块的数量，数据类型支持INT64，目前支持1~32。
  • headNum（int64_t，计算输入）：Host侧的int64_t，代表输入query的N轴长度，数据类型支持INT64；综合约束请见undefined。
  • inputLayout（string *，计算输入）：Host侧的string，代表输入query、keyIn、value的数据排布格式，支持TND。
    • query、keyIn、value数据排布格式支持从多种维度解读，其中B（Batch）表示输入样本批量大小、S（Seq-Length）表示输入样本序列长度、H（Head-Size）表示隐藏层的大小、N（Head-Num）表示多头数、G（Group）表示组数，仅对Q生效，D（Head-Dim）表示隐藏层最小的单元尺寸，且满足D=H/N。
  • sparseMode（int64_t，计算输入）：Host侧的int，表示sparse的模式。数据类型支持INT32。目前仅支持sparseMode=0或2。sparse不同模式的详细说明请参见undefined。
  • dqOut（aclTensor *，计算输出）：Device侧的aclTensor，公式中的dQ，表示query的梯度，计算输出，数据类型支持BFLOAT16和FLOAT16，数据类型与query一致，shape支持3维，shape必须与query的shape一致,undefined支持ND。
  • dkOut（aclTensor *，计算输出）：Device侧的aclTensor，公式中的dK，表示key的梯度，计算输出，数据类型支持BFLOAT16和FLOAT16，数据类型与query一致，shape支持3维，shape必须与key的shape一致，undefined支持ND。
  • dvOut（aclTensor *，计算输出）：Device侧的aclTensor，公式中的dV，表示value的梯度，计算输出，数据类型支持BFLOAT16和FLOAT16，数据类型与query一致，shape支持3维，shape必须与value的shape一致，undefined支持ND。
  • workspaceSize（uint64_t *，出参）：返回用户需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor **，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnNsaSelectedAttentionGrad

• 参数说明：

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

• 返回值：

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

约束说明

• 该接口与pytorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。
• 输入query、key、value、attentionOut、attentionOutGrad的B（batchsize）必须相等。
• 输入key、value的N（numHead）必须一致。
• 输入query、attentionOut、attentionOutGrad的N（numHead）必须一致。
• 输入value、attentionOut、attentionOutGrad的D（HeadDim）必须一致。
• 输入query、key、value、attentionOut、attentionOutGrad的inputLayout必须一致。
• 关于数据shape的约束，以inputLayout的TND举例。其中：
  • T1：取值范围为1~2M。T1表示query所有batch下S的和。
  • T2：取值范围为1~2M。T2表示key、value所有batch下S的和。
  • B：取值范围为1~2M。
  • N1：取值范围为1~128。表示query的headNum。N1必须为N2的整数倍。
  • N2：取值范围为1~128。表示key、value的headNum。
  • G：取值范围为1~32。G = N1 / N2
  • S：取值范围为1~128K。对于key、value的S 必须大于等于selectedBlockSize * selectedBlockCount, 且必须为selectedBlockSize的整数倍。
  • D：取值范围为192或128，支持K和V的D（HeadDim）不相等。
• 关于softmaxMax与softmaxSum参数shape的约束：[T1, N1, 8]。
• 关于topkIndices参数shape的约束：[T1, N2, selectedBlockCount]。

参数解释请参见算子执行接口。

调用示例

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