NsaCompressAttentionInfer

产品支持情况

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

功能说明

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

$$\begin{gathered} P_{cmp}=Softmax(scale\ast query\cdot ke{y}^T) \\ attentionOut=P_{cmp}\cdot value \\ {P}_{slc}[j]=\sum _{m=0}^{l^{\prime }/d-1}\sum _{n=0}^{l/d-1}{P}_{cmp}[l^{\prime }/d\ast j-m-n] \\ {P}_{sl{c}^{\prime }}=\sum _{g=1}^G{P}_{slc}^g,\text{其中 }G=\text{GroupSize（分组大小），即：}G=\frac{\text{numHeads}}{\text{numKeyValueHeads}} \\ topkIndices=topk(P_{sl{c}^{\prime }}) \end{gathered}$$

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
query	输入	公式中的输入query。	FLOAT16、BFLOAT16	ND
key	输入	公式中的输入key。	FLOAT16、BFLOAT16	ND
value	输入	公式中的输入value。	FLOAT16、BFLOAT16	ND
scale	输入	公式中的输入scale，代表attention计算的缩放系数。	DOUBLE	-
l'	属性	公式中的输入l'，代表select阶段的block大小。	INT64	-
l	属性	公式中的输入l，代表compress阶段的block大小。	INT64	-
d	属性	公式中的输入d，代表两次压缩间的滑窗间隔大小。	INT64	-
attentionOut	输出	公式中的attentionOut，attention计算的结果。	FLOAT16、BFLOAT16	ND
topkIndices	输出	公式中的topkIndices，重要性得分最高的几个block的索引。	INT32	-

约束说明

• 参数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。
  • 多token推理场景下，仅支持query的S轴最大等于4，并且此时要求每个batch单独的actualQSeqLen<=actualSelKvSeqLen。
  • 仅支持paged attention。
  • 仅支持selectBlockSize取值为16的整数倍，最大支持到128。
  • selectBlockCount上限满足selectBlockCount*selectBlockSize<=MaxKvSeqlen，MaxKvSeqlen=Max(actualSelKvSeqLenOptional)。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_nsa_compress_attention_infer	通过aclnnNsaCompressAttentionInfer接口方式调用NsaCompressAttentionInfer算子。