aclnnSwinAttentionScoreQuant

产品支持情况

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

功能说明

• 算子功能：完成swin-transformer场景的Attention计算，相较于SwinAttentionScore算子，支持int8量化功能
• 计算公式如下：

$$out= Softmax(QK^T + bias1 + bias2)V$$

函数原型

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

• aclnnStatus aclnnSwinAttentionScoreQuantGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *scaleQuant, const aclTensor *scaleDequant1, const aclTensor *scaleDequant2, const aclTensor *biasQuantOptional, const aclTensor *biasDequant1Optional, const aclTensor *biasDequant2Optional, const aclTensor *paddingMask1Optional, const aclTensor *paddingMask2Optional, bool queryTranspose, bool keyTranspose, bool valueTranspose, int64_t softmaxAxes, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnSwinAttentionScoreQuant(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnSwinAttentionScoreQuantGetWorkspaceSize

• 参数说明：

  • query(aclTensor*,计算输入)：表示输入样本的查询张量，Device侧的aclTensor，公式中的Q，维度支持四维，输入维度[N,C,S,H]需要和key、value保持一致，其中N代表batch size，C为通道深度，S为序列长度，H为headNum，S<=1024，H=32/64，NC维度支持任意值，数据类型支持INT8，不支持undefined，undefined支持ND。
  • key(aclTensor*,计算输入): 表示输入样本的每个位置的特征张量，Device侧的aclTensor，公式中的K，维度支持四维，输入维度[N,C,S,H]需要和query、value保持一致，S<=1024，H=32/64，NC维度支持任意值，数据类型支持INT8，不支持undefined，undefined支持ND。
  • value(aclTensor*,计算输入): 表示计算注意力后的每个位置值张量，Device侧的aclTensor，公式中的V，维度支持四维，输入维度[N,C,S,H]需要和query、key保持一致，S<=1024，H=32/64，NC维度支持任意值，数据类型支持INT8，不支持undefined，undefined支持ND。
  • scaleQuant(aclTensor*,计算输入): 表示对注意力softmax归一化后进行量化的尺度缩放张量，Device侧的aclTensor，维度支持二维，输入维度[1,S]，S<=1024，数据类型支持FLOAT16，不支持undefined，undefined支持ND。
  • scaleDequant1(aclTensor*,计算输入)：表示计算注意力时进行反量化的尺度缩放张量，Device侧的aclTensor，维度支持二维，输入维度[1,S]，S<=1024，数据类型支持UINT64，不支持undefined，undefined支持ND。
  • scaleDequant2(aclTensor*,计算输入): 表示计算注意力后的输出进行反量化的尺度缩放张量，Device侧的aclTensor，维度支持二维，输入维度[1,H]，H=32/64，数据类型支持UINT64，不支持undefined，undefined支持ND。
  • biasQuantOptional(aclTensor*,计算输入): 表示对注意力softmax归一化后进行量化的偏移张量，Device侧的aclTensor，维度支持二维，输入维度[1,S]，S<=1024，数据类型支持FLOAT16，不支持undefined，undefined支持ND。
  • biasDequant1Optional(aclTensor*,计算输入): 表示计算注意力时进行反量化的偏移张量，Device侧的aclTensor，维度支持二维，输入维度[1,S]，S<=1024，数据类型支持INT32，不支持undefined，undefined支持ND。
  • biasDequant2Optional(aclTensor*,计算输入)：表示计算注意力后的输出进行反量化的偏移张量 Device侧的aclTensor，维度支持二维，输入维度[1,H]，H=32/64，数据类型支持INT32，不支持undefined，undefined支持ND。
  • paddingMask1Optional(aclTensor*,计算输入): Device侧的aclTensor，公式中的bias1，支持输入nullptr，或者四维的tensor，维度[1,C,S,S]，S<=1024，C维度支持任意值，数据类型支持FLOAT16，不支持undefined，undefined支持ND。
  • paddingMask2Optional(aclTensor*,计算输入): 预留参数，公式中的bias2，当前仅支持输入nullptr。
  • queryTranspose(bool，计算输入): Host侧的bool，表示query是否转置，当前仅支持不转置false。
  • keyTranspose(bool，计算输入): Host侧的bool，表示key是否转置，当前仅支持不转置false。
  • valueTranspose(bool，计算输入): Host侧的bool，表示value是否转置，当前仅支持不转置false。
  • softmaxAxes(int，计算输入): Host侧的int，用于指定softmax计算的维度，当前仅支持取-1（即tensor的最后一维）。
  • out(aclTensor*, 计算输出)：Device侧的aclTensor，维度支持四维，输出维度[N,C,S,H]，S<=1024，H=32/64，数据类型支持FLOAT16，不支持undefined，undefined支持ND。
  • workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。
  • executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。

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

  [object Object]

aclnnSwinAttentionScoreQuant

• 参数说明

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

• 返回值

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

约束说明

• QKV输入维度是[N,C,S,H]的情况下，S<=1024，H=32/64，NC维度支持任意值
• 不支持维度是[N,C,S,H]的QKV转置后输入
• 只支持非对称量化
• 不支持加bias2的功能
• 只支持对QK^T + bias1 + bias2的最后一维进行softmax操作

调用示例

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