NsaCompressWithCache

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	×
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×
Kirin X90 处理器系列产品	√
Kirin 9030 处理器系列产品	√

功能说明

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

  1. 检查act_seq_lens是否有满足$s\ge compressBlockSize$ 且 $(s-compressBlockSize)\%stride==0$的序列长度；
  2. 找到满足序列长度的batchIdx，根据block_table找到该batch的后compress_block_size个token压缩；
  3. 执行压缩算法；
  4. 根据slot_mapping写回到output_cache中。

• 计算公式：

$$\begin{gathered} compressIdx=(s-compressBlockSize)/stride \\ outputCacheRef[slotMapping[i]]=input[compressIdx\ast stride:compressIdx\ast stride+compressBlockSize]\ast weight[:] \end{gathered}$$

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
s	属性	当前batch的token长度。	INT64	-
compressBlockSize	属性	压缩滑窗大小。	INT64	-
stride	属性	两次压缩滑窗间隔大小。	INT64	-
weight	输入	k/v值的压缩weight。	BFLOAT16、FLOAT16	ND
input	输入	k/v值的cache。	BFLOAT16、FLOAT16	ND
slotMapping	输入	每个batch尾部压缩数据存储的位置的索引。	INT32	ND
outputCacheRef	输入/输出	输出的cache。	BFLOAT16、FLOAT16	ND

• Kirin X90/Kirin 9030 处理器系列产品: 不支持BFLOAT16。

约束说明

• input和weight满足broadcast关系，input的第三维大小与weight的第二维大小相等。
• compressBlockSize、stride必须是16的整数倍，且compressBlockSize>=stride，compressBlockSize<=64。
• actSeqLenType目前仅支持取值1。
• layoutOptional取值可以是BSH、SBH、BSND、BNSD、TND，但是不会生效。
• pageBlockSize只能是64或者128。
• headDim是16的整数倍，且headDim<=256。
• 不支持input/weight/outputCache为空输入。
• slotMapping的值无重复，否则会导致计算结果不稳定。
• blockTableOptional的值不超过blockNum，否则会发生越界。
• actSeqLenOptional的值不应该超过序列最大长度。
• headNum<=64，且headNum>50时headNum%2=0。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_nsa_compress_with_cache	通过aclnnNsaCompressWithCache接口方式调用NsaCompressWithCache算子。