aclnnNsaCompressAttentionInfer

产品支持情况

[object Object]undefined

功能说明

• 接口功能：Native Sparse Attention推理过程中，Compress Attention的计算。
• 计算公式：

[object Object]$$P_{cmp}= Softmax(scale * query · key^T) \\ attentionOut = P_{cmp} · value\\ P_{slc}[j] = \sum\limits_{m=0}^{l'/d -1} \sum\limits_{n = 0}^{l/d -1} P_{cmp} [l'/d * j -m - n]\\ P_{slc'} = \sum\limits_{g=1}^{G} P_{slc} ^g,\quad \text{其中 } G = \text{GroupSize（分组大小），即：} G = \frac{\text{numHeads}}{\text{numKeyValueHeads}} \\ topkIndices = topk(P_{slc'})\\$$ [object Object]

函数原型

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

aclnnNsaCompressAttentionInferGetWorkspaceSize

• 参数说明

  [object Object]

• 返回值

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

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

  [object Object]

aclnnNsaCompressAttentionInfer

• 参数说明

  [object Object]

• 返回值

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

约束说明

• 确定性计算：
  • aclnnNsaCompressAttentionInfer默认确定性实现。

• 参数query仅支持TND、BSND输入。T是B和S合轴紧密排列的数据（每个batch的actualQSeqLenOptional）、B（batch）表示输入样本批量大小、S（qSeqlen）表示输入样本序列长度、N（numHeads）表示多头数、D（headDimsQK）表示隐藏层最小的单元尺寸。
• 压缩前的kvSeqlen的上限可以表示为：actualSelKvSeqLenCeil=(actualCmpKvSeqLenOptional-1)*compressBlockStride+compressBlockSize，需要满足actualSelKvSeqLenCeil/selectBlockSize<=4096，且需要满足selectBlockCount<=actualSelKvSeqLenCeil/selectBlockSize。如果actualSelKvSeqLenOptional不满足actualCmpKvSeqLenOptional=(actualSelKvSeqLenOptional-compressBlockSize)/compressBlockStride+1，或者actualCmpKvSeqLenOptional的长度和blockTableOptional的batch维度不同，则会默认进入单token推理场景。
• 多token推理场景下，actualQSeqLenOptional参数必须传入，actualQSeqLenOptional的长度必须和blockTableOptional的batch维度相等，仅支持query的S轴最大等于4，并且此时要求每个batch单独的actualQSeqLenOptional<=actualSelKvSeqLenOptional。如果actualQSeqLenOptional的长度和blockTableOptional的batch维度不同，或者actualQSeqLenOptional的值小于1或者大于4，则会默认进入单token推理场景。

调用示例

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