aclnnNsaCompressAttentionInfer

产品支持情况

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

功能说明

• 算子功能：Native Sparse Attention推理过程中，Compress Attention的计算。

• 计算公式：

  self-attention（自注意力）利用输入样本自身的关系构建了一种注意力模型。其原理是假设有一个长度为$n$的输入样本序列$x$，$x$的每个元素都是一个$d$维向量，可以将每个$d$维向量看作一个token embedding，将这样一条序列经过3个权重矩阵变换得到3个维度为$n*d$的矩阵。

  Compress Attention的计算由三阶段的attention计算、importance score与topK三个过程融合而成，首先，$query$和$key^T$的乘积进行softmax得到注意力分数$P_{cmp}$

  $$ P_{cmp}= Softmax(scale * query · key^T)

  $$ 一方面，注意力分数$P_{cmp}$与$value$相乘得到自注意力的结果 $$ attentionOut = P_{cmp} · value $$

  另一方面，注意力分数$P_{cmp}$被用于计算selection block的重要性分数$P_{slc}[j]$ $$ P_{slc}[j] = \sum\limits_{m=0}^{l'/d -1} \sum\limits_{n = 0}^{l/d -1} P_{cmp} [l'/d * j -m - n]

  $$ 接着，重要性分数在group内进行累加，得到共享的重要性分数$P_{slc'}$ $$

  P_{slc'} = \sum\limits_{h=1}^{H} P_{slc} ^h

  $$ 最后，选出重要性分数最高的K个selection block，其下标数组$topkIndices$为topK的输出 $$

  topkIndices = topk(P_{slc'}) $$

函数原型

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

• aclnnStatus aclnnNsaCompressAttentionInferGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *attentionMaskOptional, const aclTensor *blockTableOptional, const aclIntArray *actualQSeqLenOptional, const aclIntArray *actualCmpKvSeqLenOptional, const aclIntArray *actualSelKvSeqLenOptional, const aclTensor *topKMaskOptional, int64_t numHeads, int64_t numKeyValueHeads, int64_t selectBlockSize, int64_t selectBlockCount, int64_t compressBlockSize, int64_t compressBlockStride, double scaleValue, char *layoutOptional, int64_t pageBlockSize, int64_t sparseMode, const aclTensor *output, const aclTensor *topKOutput, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnNsaCompressAttentionInfer(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnNsaCompressAttentionInferGetWorkspaceSize

• 参数说明：

  • query（aclTensor *，计算输入）：Device侧的aclTensor，attention结构的Query输入，undefined支持ND，数据类型支持FLOAT16、BFLOAT16，支持输入的维度是3维和4维，不支持undefined。

  • key（aclTensor *，计算输入）：Device侧的aclTensor，attention结构的Key输入，undefined支持ND，数据类型支持FLOAT16、BFLOAT16，支持输入的维度是3维，不支持undefined。

  • value（aclTensor *，计算输入）：Device侧的aclTensor，attention结构的Value输入，undefined支持ND，数据类型支持FLOAT16、BFLOAT16，支持输入的维度是3维，不支持undefined。

  • attentionMaskOptional（aclTensor *，计算输入）：Device侧的aclTensor，可选参数，表示attention掩码矩阵，undefined支持ND，如不使用该功能时可传入nullptr。预留参数，暂未使用。

  • blockTableOptional（aclTensor *，计算输入）：Device侧的aclTensor，可选参数，表示paged attention中KV存储使用的block映射表，数据类型支持INT32，undefined支持ND，不支持undefined，当前算子只支持paged attention，该参数必须传入。

  • actualQSeqLenOptional（aclIntArray *，计算输入）：Host侧的aclIntArray，可选参数，表示query的S轴实际长度，undefined支持ND，数据类型支持INT64，如不使用该功能时可传入nullptr。预留参数，暂未使用。

  • actualCmpKvSeqLenOptional（aclIntArray *，计算输入）：Host侧的aclIntArray，可选参数，表示经过压缩后的key和value的S轴实际长度，也即该算子处理的key和value的S轴实际长度，undefined支持ND，数据类型支持INT64，由于该算子当前只支持paged attention，因此该参数必须传入。

  • actualSelKvSeqLenOptional（aclIntArray *，计算输入）：Host侧的aclIntArray，可选参数，表示压缩前的key和value的S轴实际长度，undefined支持ND，数据类型支持INT64，如不使用该功能时可传入nullptr。预留参数，暂未使用。

  • topKMaskOptional（aclTensor *，计算输入）：Device侧的aclTensor，可选参数，表示topK计算中的掩码矩阵，undefined支持ND，如不使用该功能时可传入nullptr。预留参数，暂未使用。

  • numHeads（int64_t，计算输入 ）：Host侧的int64_t，代表head个数，数据类型支持INT64。

  • numKeyValueHeads（int64_t，计算输入 ）：Host侧的int64_t，代表kvHead个数，数据类型支持INT64。

  • selectBlockSize（int64_t，计算输入 ）：Host侧的int64_t，代表select阶段的block大小，在计算importance score时使用，数据类型支持INT64。

  • selectBlockCount（int64_t，计算输入 ）：Host侧的int64_t，代表topK阶段需要保留的block数量，数据类型支持INT64。

  • compressBlockSize（int64_t，计算输入 ）：Host侧的int64_t，代表压缩时的滑窗大小，数据类型支持INT64。

  • compressBlockStride（int64_t，计算输入 ）：Host侧的int64_t，代表两次压缩间的滑窗间隔大小，数据类型支持INT64。

  • scaleValue（double，计算输入）：Host侧的double，公式中d开根号的倒数，代表缩放系数，作为计算流中Muls的scalar值，数据类型支持DOUBLE。

  • layoutOptional（char *，计算输入）：Host侧的字符指针，用于标识输入query、key、value的数据排布格式，当前支持取值“TND”和“BSND”。

    说明： query的数据排布格式当前支持（B, N1, D1）和（B, 1, N1, D1）, B（Batch）表示Batch数，N1（Head-Num）表示numHeads，D1表示headSizeQK。 key的数据排布格式当前（paged attention）支持（blocknum, blocksize, N2 * D1），N2（Head-Num）表示numKeyValueHeads, D1表示headSizeQK。 value的数据排布格式当前（paged attention）支持（blocknum, blocksize, N2 * D2），N2（Head-Num）表示numKeyValueHeads, D2表示headSizeVO。 blockTableOptional的数据排布格式当前支持（B, maxBlockPerQuery），B（Batch）表示Batch数，maxBlockPerQuery表示最长的keyValueSeqLen需要占用的block数量。

  • pageBlockSize（int64_t，计算输入 ）：Host侧的int64_t，代表blockTable中一个block的大小，数据类型支持INT64。

  • sparseMode（int64_t，计算输入）：Host侧的int64_t，表示sparse的模式，控制有attentionMask输入时的稀疏计算。数据类型支持INT64，预留参数，暂未使用。

  • output（aclTensor *，计算输出）：Device侧的aclTensor，attention的输出，undefined支持ND。数据类型支持FLOAT16、BFLOAT16。

  • topKOutput（aclTensor *，计算输出）：Device侧的aclTensor，topK的输出，undefined支持ND，数据类型支持INT32。

  • workspaceSize（uint64_t *，出参）：返回用户需要在Device侧申请的workspace大小。

  • executor（aclOpExecutor **，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnNsaCompressAttentionInfer

• 参数说明：

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

• 返回值：

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

约束说明

• 参数query中的N和numHeads值相等，key、value的N和numKeyValueHeads值相等，并且numHeads是numKeyValueHeads的倍数关系。
• 参数query中的D和key的D(H/numKeyValueHeads)值相等，value的D(H/numKeyValueHeads)和output的D值相等。
• 参数query中的B是[1, 10000]区间内的整数。
• 参数query中的B与blockTable中的B与actualCmpKvSeqLenOptional数组的长度相等。
• 参数key中的numBlocks和参数value中的numBlocks值相等。
• 参数key中的blockSize、参数value中的blockSize和pageBlockSize值相等。
• query，key，value输入，功能使用限制如下：
  • 支持query的N轴必须是key/value的N轴（H/D）的整数倍。
  • 支持query的N轴与key/value的N轴（H/D）的比值（即GQA中的group大小）小于等于128，且128是group的整数倍。
  • 支持query与Key的D轴小于等于192。
  • 支持value的D轴小于等于128。
  • 支持query与Key的D轴大于等于value的D轴。
  • 支持key与value的blockSize小于等于128，且是16的整数倍。
  • 仅支持query，key，value输入的数据类型完全相同，为FLOAT16或BFLOAT16。
  • 仅支持query的S轴等于1。
  • 仅支持paged attention。
  • 仅支持key/value的S轴小于等于8192。
  • 仅支持compressBlockSize取值16、32、48、64、80、96、112、118。
  • 仅支持compressBlockStride取值16、32、48、64。
  • 仅支持selectBlockSize取值16、32、48、64、80、96、112、118。
  • 仅支持compressBlockSize大于等于compressBlockStride , selectBlockSize大于等于compressBlockSize , selectBlockSize是compressBlockStride的整数倍。
  • 压缩前的kvSeqlen的上限可以表示为：NoCmpKvSeqlenCeil =（cmpKvSeqlen - 1）* compressBlockStride + compressBlockSize，需要满足NoCmpKvSeqlenCeil / selectBlockSize <= 4096，且需要满足selectBlockCount <= NoCmpKvSeqlenCeil / selectBlockSize。

调用示例

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