aclnnNsaSelectedAttentionInfer

产品支持情况

[object Object]undefined

功能说明

• 接口功能：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$的矩阵。

函数原型

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

aclnnNsaSelectedAttentionInferGetWorkspaceSize

• 参数说明

  [object Object]

• 返回值

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

  第一段接口完成入参校验，出现以下场景时报错：

  [object Object]

aclnnNsaSelectedAttentionInfer

• 参数说明

  [object Object]

• 返回值

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

约束说明

• 确定性计算：
  • aclnnNsaSelectedAttentionInfer默认确定性实现。
• 支持B轴小于等于3072。
• 仅支持paged attention。
• 多token推理场景下，仅支持query的S轴最大等于4，并且此时要求每个batch单独的actualQSeqLen <= actualSelKvSeqLen。

调用示例

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