aclnnNsaCompressWithCache

产品支持情况

[object Object]undefined

功能说明

• 接口功能：用于Native-Sparse-Attention推理阶段的KV压缩，每次推理每个batch会产生一个新的token，每当某个batch的token数量凑满一个compress_block时，该算子会将该batch的后compress_block个token压缩成一个compress_token。
• 计算公式：

$$compressIdx=(s-compressBlockSize)/stride\\ outputCacheRef[slotMapping[i]] = input[compressIdx*stride : compressIdx*stride+compressBlockSize]*weight[:]$$

函数原型

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

aclnnNsaCompressWithCacheGetWorkspaceSize

• 参数说明

  [object Object]

• 返回值

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

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

  [object Object]

aclnnNsaCompressAttentionInfer

• 参数说明

  [object Object]

• 返回值

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

约束说明

• 确定性计算：
  • aclnnNsaCompressWithCache默认确定性实现。
• outputCache的N和D和input一致，而且要满足result_len>(blockNum*pageBlockSize-compressBlockSize)/compressStride。
• page attention场景下input的shape支持[blockNum,pageBlockSize,N,D]，其余场景下input的shape支持[T,N,D]。

调用示例

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