aclnnNsaCompressAttention

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	√
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品[object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

算子功能：NSA中compress attention以及select topk索引计算。论文：undefined
计算公式：压缩block大小： $l$ ，select block大小： $l'$ ，压缩stride大小： $d$

P_{cmp} = Softmax(query*key^T) \\

attentionOut = Softmax(atten\_mask(scale*query*key^T, atten\_mask)))*value

P_{slc}[j] = \sum_{m=0}^{l'/d-1}\sum_{n=0}^{l/d-1}P_{cmp} [l'/d*j-m-n],

P_{slc'} = \sum_{h=1}^{H}P_{slc}^{h}

P_{slc'} = topk\_mask(P_{slc'})

topkIndices = topk(P_{slc'})

NsaCompressAttention输入query、key、value的数据排布格式支持从多种维度排布解读，可通过inputLayout传入，当前仅支持TND。

B：表示输入样本批量大小（Batch）
T：B和S合轴紧密排列的长度
S：表示输入样本序列长度（Seq-Length）
H：表示隐藏层的大小（Head-Size）
N：表示多头数（Head-Num）
D：表示隐藏层最小的单元尺寸，需满足D=H/N（Head-Dim）

函数原型

aclnnStatus aclnnNsaCompressAttentionGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *attenMaskOptional, const aclTensor *topkMaskOptional, const aclIntArray *actualSeqQLenOptional, const aclIntArray *actualCmpSeqKvLenOptional, const aclIntArray *actualSelSeqKvLenOptional, double scaleValue, int64_t headNum, char *inputLayout, int64_t sparseMode, int64_t compressBlockSize, int64_t compressStride, int64_t selectBlockSize, int64_t selectBlockCount, const aclTensor *softmaxMaxOut, const aclTensor *softmaxSumOut, const aclTensor *attentionOutOut, const aclTensor *topkIndicesOut, uint64_t *workspaceSize, aclOpExecutor **executor);
aclnnStatus aclnnNsaCompressAttention(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream);

aclnnNsaCompressAttentionGetWorkspaceSize

参数说明：
- query（aclTensor *，计算输入）：Device侧的aclTensor，公式中的query，数据类型支持FLOAT16、BFLOAT16。undefined支持ND；综合约束请见undefined。
- key（aclTensor *，计算输入）：Device侧的aclTensor，公式中的key，数据类型支持FLOAT16、BFLOAT16。undefined支持ND；综合约束请见undefined。
- value（aclTensor *，计算输入）：Device侧的aclTensor，公式中的value，数据类型支持FLOAT16、BFLOAT16。undefined支持ND；综合约束请见undefined。
- attenMaskOptional（aclTensor *，计算输入）：Device侧的aclTensor，公式中的atten_mask，数据类型支持BOOL，undefined支持ND，输入shape需为[S,S]，TND场景只支持SS格式，SS分别是max(Sq)和max(CmqSkv)；综合约束请见undefined。
- actualSeqQLenOptional（aclIntArray*，计算输入）：Host侧的aclIntArray，数据类型支持INT64，undefined支持ND，描述了每个Batch对应的query S大小(Sq)；综合约束请见undefined。
- actualCmpSeqKvLenOptional（aclIntArray*，计算输入）：Host侧的aclIntArray，数据类型支持INT64，undefined支持ND，描述了compress attention的每个Batch对应的key/value S大小(CmpSkv)；综合约束请见undefined。
- actualSelSeqKvLenOptional（aclIntArray*，计算输入）：Host侧的aclIntArray，数据类型支持INT64，undefined支持ND，描述了经importance score计算压缩后的每个Batch对应的key/value S大小(SelSkv)，对应公式中的P_slc'第0维；综合约束请见undefined。
- topkMaskOptional（aclTensor *，计算输入）：Device侧的aclTensor，公式中的topk_mask，数据类型支持BOOL，undefined支持ND，输入shape类型需为[S,S]，TND场景只支持SS格式，SS分别是max(Sq)和max(SelSkv)；综合约束请见undefined。如不使用该参数可传入nullptr。
- scaleValue（double，计算输入）：Host侧的double。公式中的scale，代表缩放系数，数据类型支持DOUBLE，一般设置为D^-0.5。
- headNum（int64_t，计算输入）：Host侧的int64_t，数据类型支持INT64，代表query的head个数。
- inputLayout（string*，计算输入）：Host侧的string，数据类型支持String，代表输入query、key、value的数据排布格式，当前支持TND。
- sparseMode（int64_t，计算输入）：Host侧的int64_t。数据类型支持INT64。当前仅支持0和1；sparse不同模式的详细说明请参见undefined。
- compressBlockSize（int64_t，计算输入）：Host侧的int64_t，压缩滑窗大小，对应公式中的l。
- compressStride（int64_t，计算输入）：Host侧的int64_t，两次压缩滑窗间隔大小，对应公式中的d。
- selectBlockSize（int64_t，计算输入）：Host侧的int64_t，选择块大小，对应公式中的l'。
- selectBlockCount（int64_t，计算输入）：Host侧的int64_t，选择块个数，对应公式中topK选择个数。
- softmaxMaxOut（aclTensor*，计算输出）：Device侧的aclTensor，Softmax计算的Max中间结果，用于反向计算。数据类型支持FLOAT，输出的shape类型为[T,N,8]。undefined支持ND。
- softmaxSumOut（aclTensor*，计算输出）：Device侧的aclTensor，Softmax计算的Max中间结果，用于反向计算。数据类型支持FLOAT，输出的shape类型为[T,N,8]。undefined支持ND。
- attentionOut（aclTensor*，计算输出）：Device侧的aclTensor，公式中的attentionOut，数据类型支持FLOAT16、BFLOAT16，数据类型和shape类型与query保持一致，undefined支持ND。
- topkIndicesOut（aclTensor*，计算输出）：Device侧的aclTensor，公式中的topkIndices，数据类型支持INT32，输出的shape类型为[T,N2,selectBlockCount]。
- workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
- executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。
返回值：

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

[object Object]

aclnnNsaCompressAttention

参数说明：
- workspace（void*，入参）：在Device侧申请的workspace内存地址。
- workspaceSize（uint64_t，入参）：在Device侧申请的workspace大小，由第一段接口aclnnNsaCompressAttentionGetWorkspaceSize获取。
- executor（aclOpExecutor*，入参）：op执行器，包含了算子计算流程。
- stream（aclrtStream，入参）：指定执行任务的Stream。
返回值：

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

约束说明

该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。
compressBlockSize、compressStride、selectBlockSize必须是16的整数倍，并且满足：compressBlockSize>=compressStride && selectBlockSize>=compressBlockSize && selectBlockSize%compressStride==0
compressBlockSize：16对齐，支持到128
compressStride：16对齐，支持到64
selectBlockSize：16对齐，支持到128
selectBlockCount：支持[1~32] && selectBlockCount <= min(SelSkv)
actualSeqQLenOptional, actualCmpSeqKvLenOptional, actualSelSeqKvLenOptional需要是前缀和模式；且TND格式下必须传入。
由于UB限制，CmpSkv需要满足以下约束：CmpSkv <= 14000
SelSkv = CeilDiv(CmpSkv, selectBlockSize // compressStride)
layoutOptional目前仅支持TND。
输入query、key、value的数据类型必须一致。
输入query、key、value的batchSize必须相等。
输入query、key、value的headDim必须满足：qD == kD && kD >= vD
输入query、key、value的inputLayout必须一致。
输入query的headNum为N1，输入key和value的headNum为N2，则N1 >= N2 && N1 % N2 == 0
设G = N1 / N2，G需要满足以下约束：G < 128 && 128 % G == 0
attenMask和topkMask的使用需符合论文描述

调用示例

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

[object Object]