aclnnNsaSelectedAttentionInfer

产品支持情况

产品	是否支持
[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推理过程中，Selected Attention的计算。

• 计算公式：

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

  Selected Attention的计算由topK索引取数与attention计算融合而成，外加paged attention取kvCache。首先，通过$topkIndices$索引从$key$中取出$key_topk$，从$value$中取出$value\_topk$，计算self_attention公式如下：

  $$ Attention(query,key,value)=Softmax(\frac{query · key_topk^T}{\sqrt{d}})value_topk $$ 其中$query$和$key_topk^T$的乘积代表输入$x$的注意力，为避免该值变得过大，通常除以$d$的开根号进行缩放，并对每行进行softmax归一化，与$value\_topk$相乘后得到一个$n*d$的矩阵。

函数原型

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

• aclnnStatus aclnnNsaSelectedAttentionInferGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *topkIndices, const aclTensor *attenMaskOptional, const aclTensor *blockTableOptional, const aclIntArray *actualQSeqLenOptional, const aclIntArray *actualKvSeqLenOptional, char *layoutOptional, int64_t numHeads, int64_t numKeyValueHeads, int64_t selectBlockSize, int64_t selectBlockCount, int64_t pageBlockSize, double scaleValue, int64_t sparseMode, aclTensor *output, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnNsaSelectedAttentionInfer(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)

aclnnNsaSelectedAttentionInferGetWorkspaceSize

• 参数说明：

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

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

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

  • topkIndices （aclTensor *，计算输入）：Device侧的aclTensor， NSA里的topK索引，undefined支持ND，数据类型支持INT32，支持输入的维度是3维，不支持undefined。

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

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

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

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

  • inputLayoutOptional（char *，计算输入）：Host侧的字符指针，用于标识输入query、key、value的数据排布格式，当前支持BSH/BSND/TND，当不传入该参数时，默认为“BSND”，分别对应query、key、value 3/4维。

    说明： query的数据排布格式中，B即Batch，S即Seq-Length，N（Head-Num）表示多头数、D（Head-Dim）表示隐藏层最小的单元尺寸，且满足D=H/N。key和value的数据排布格式当前（paged attention）支持（blocknum, blocksize, H），（blocknum, blocksize, N, D），H（Head-Size）表示隐藏层的大小，H = N * D。

  • 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。

  • pageBlockSize（int64_t，计算输入 ）：Host侧的int64_t，代表paged attention的block大小，在kv cache取数时使用，数据类型支持INT64。

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

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

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

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

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

• 返回值：

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

  [object Object]

aclnnNsaSelectedAttentionInfer

• 参数说明：

  • workspace（void*，入参）：在Device侧申请的workspace内存地址。
  • workspaceSize（uint64_t，入参）：在Device侧申请的workspace大小，由第一段接口aclnnNsaSelectedAttentionInferGetWorkspaceSize获取。
  • 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，key，value输入，功能使用限制如下：
  • 支持B轴小于等于3072；
  • 支持key/value的N轴小于等于256；
  • 支持query的N轴与key/value的N轴（H/D）的比值（即GQA中的group大小）小于等于16；
  • 支持query与Key的D轴等于192；
  • 支持value的D轴等于128；
  • 支持Key与Value的blockSize等于64或128；
  • 仅支持query的S轴等于1。
  • 仅支持paged attention。
  • 仅支持selectBlockSize取值为16的整数倍，最大支持到128。
  • selectBlockCount上限满足selectBlockCount * selectBlockSize <= MaxKvSeqlen，MaxKvSeqlen = Max(actualSelKvSeqLenOptional)。

调用示例

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