aclnnFlashAttentionScoreVX

产品支持情况

产品	是否支持
[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]	×

功能说明

算子功能：训练场景下，使用FlashAttention算法实现self-attention（自注意力）的计算。该接口query、key、value参数支持多个长度相等或者多个长度不相等的sequence
- 该接口相较于undefined接口，新增支持数据类型FLOAT8_E5M2、FLOAT8_E4M3FN，调整Dropout功能：
- 该接口相较于undefined接口，调整Dropout功能：
计算公式：

注意力的正向计算公式如下：
- psetype=1时，与undefined计算公式相同。
- psetype=其他取值时，公式如下：
  $attention\_out=Dropout(Softmax(Mask(scale*(query*key^T) + pse),atten\_mask),keep\_prob)*value$

函数原型

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

aclnnStatus aclnnFlashAttentionScoreVXGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *realShiftOptional, const aclTensor *dropMaskOptional, const aclTensor *paddingMaskOptional, const aclTensor *attenMaskOptional, const aclTensor *queryRopeOptional, const aclTensor *keyRopeOptional, const aclTensor *dScaleQOptional, const aclTensor *dScaleKOptional, const aclTensor *dScaleVOptional, const aclIntArray *prefixOptional, const aclIntArray *actualSeqQLenOptional, const aclIntArray *actualSeqKvLenOptional, const aclIntArray *qStartIdxOptional, const aclIntArray *kvStartIdxOptional, double scaleValue, double keepProb, int64_t preTokens, int64_t nextTokens, int64_t seed, int64_t offset, int64_t headNum, char *inputLayout, int64_t innerPrecise, int64_t sparseMode, int64_t outDtype, int64_t pseType,int64_t seed, int64_t offset, aclTensor *softmaxMaxOut, aclTensor *softmaxSumOut, aclTensor *softmaxOutOut, aclTensor *attentionOutOut, uint64_t *workspaceSize, aclOpExecutor **executor)
aclnnStatus aclnnFlashAttentionScoreVX(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)

aclnnFlashAttentionScoreVXGetWorkspaceSize

参数说明：
- query（aclTensor*，计算输入）：Device侧的aclTensor，数据类型支持FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT16、BFLOAT16、FLOAT32。数据类型与key/value的数据类型一致，undefined支持ND；综合约束请见undefined。
- key（aclTensor*，计算输入）：Device侧的aclTensor，数据类型支持FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT16、BFLOAT16、FLOAT32，数据类型与query/value的数据类型一致，undefined支持ND；综合约束请见undefined。
- value（aclTensor*，计算输入）：Device侧的aclTensor，数据类型支持FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT16、BFLOAT16、FLOAT32，数据类型与query/key的数据类型一致，undefined支持ND；综合约束请见undefined。
- realShiftOptional（aclTensor*，计算输入）：Device侧的aclTensor，公式中的输入pse。数据类型支持FLOAT16、BFLOAT16、FLOAT32，数据类型与query的数据类型一致。undefined支持ND，输入shape类型需为[B,N,S,S]、[B,N,1,S]、[1,N,S,S]。如果Sq大于1024的每个batch的Sq与Skv等长且是sparseMode为0、2、3的下三角掩码场景，可使能alibi位置编码压缩，此时只需要输入原始PSE最后1024行，实现内存优化，即alibi_compress = ori_pse[:, :, -1024:, :]，参数每个batch不相同时，输入BNHSkv(H=1024)，每个batch相同时，输入1NHSkv(H=1024)。后续章节如无特殊说明，S表示query或key、value的sequence length，Sq表示query的sequence length，Skv表示key、value的sequence length，SS表示Sq*Skv。如果psetype为2或3的时候，数据类型需为FLOAT32, 对应shape支持范围是[B,N]或[N]。如不使用该参数可传入nullptr。
- dropMaskOptional（aclTensor*，计算输入）：Device侧的aclTensor，公式中的dropout，数据类型支持UINT8（标识8个1bit BOOL），undefined支持ND；综合约束请见undefined。如不使用该参数可传入nullptr。
- paddingMaskOptional（aclTensor*，计算输入）：预留参数，暂未使用。
- attenMaskOptional（aclTensor*，计算输入）：Device侧的aclTensor，公式中的atten_mask，数据类型支持BOOL、UINT8，undefined支持ND，输入shape类型需为[B,N,S,S]、[B,1,S,S]、[1,1,S,S]、[S,S]，varlen场景只支持SS格式，SS分别是maxSq和maxSkv；综合约束请见undefined。如不使用该参数可传入nullptr。
- queryRopeOptional （aclTensor*，计算输入）：Device侧的aclTensor，数据类型支持FLOAT16、BFLOAT16。数据类型与query的数据类型一致，undefined支持ND；综合约束请见undefined。
- keyRopeOptional （aclTensor*，计算输入）：Device侧的aclTensor，数据类型支持FLOAT16、BFLOAT16。数据类型与key的数据类型一致，undefined支持ND；综合约束请见undefined。
- dScaleQOptional （aclTensor*，计算输入）：Device侧的aclTensor ，量化参数，图片中第一个MUL的结果按128分块乘以该系数，可选参数，数据类型支持FLOAT32， undefined支持ND；输入shape为[B, N1, S1/128, 1]。只在量化场景下生效。
- dScaleKOptional （aclTensor*，计算输入）：Device侧的aclTensor ，量化参数，图片中第一个MUL的结果按128分块乘以该系数，可选参数，数据类型支持FLOAT32， undefined支持ND；输入shape为[B, N2, S2/128, 1]。只在量化场景下生效。
- dScaleVOptional （aclTensor*，计算输入）：Device侧的aclTensor ，量化参数，图片中第二个MUL的结果按128分块乘以该系数，可选参数，数据类型支持FLOAT32， undefined支持ND；输入shape为[B, N2, S2/128, 1]。只在量化场景下生效。
- prefixOptional（aclIntArray*，计算输入）：Host侧的aclIntArray，数据类型支持INT64，代表prefix稀疏计算场景每个Batch的N值，undefined支持ND；综合约束请见undefined。如不使用该参数可传入nullptr。
- actualSeqQLenOptional（aclIntArray*，计算输入）：Host侧的aclIntArray，数据类型支持INT64，undefined支持ND，描述了每个Batch对应的query S大小。只在sequence长度不相等的场景下需要输入，其余场景下输入nullptr。
- actualSeqKvLenOptional（aclIntArray*，计算输入）：Host侧的aclIntArray，数据类型支持INT64，undefined支持ND，描述了每个Batch对应的key/value S大小。只在sequence长度不相等的场景下需要输入，其余场景下输入nullptr。
- qStartIdxOptional（aclIntArray*，计算输入）：Host侧的aclIntArray，数据类型支持INT64，代表外切场景，当前分块的query的sequence在全局中的起始索引，undefined支持ND；综合约束请见undefined。如不使用该参数可传入nullptr。
- kvStartIdxOptional（aclIntArray*，计算输入）：Host侧的aclIntArray，数据类型支持INT64，代表外切场景，当前分块的key和value的sequence在全局中的起始索引，undefined支持ND；综合约束请见undefined。如不使用该参数可传入nullptr。
- scaleValue（double，计算输入）：Host侧的double，公式中的scale，代表缩放系数，作为计算流中Muls的scalar值，数据类型支持DOUBLE，一般设置为D^-0.5。D的含义见下文说明。
- keepProb（double，计算输入）：Host侧的double，数据类型支持DOUBLE，代表dropMaskOptional中1的比例。用户不特意指定时建议传入1.0。
- preTokens（int64_t，计算输入）：Host侧的int64_t，数据类型支持INT64，用于稀疏计算，表示slides window的左边界；综合约束请见undefined。用户不特意指定时建议传入2147483647。
- nextTokens（int64_t，计算输入）：Host侧的int64_t，数据类型支持INT64，用于稀疏计算，表示slides window的右边界；综合约束请见undefined。用户不特意指定时建议传入2147483647。
- headNum（int64_t，计算输入）：Host侧的int64_t，数据类型支持INT64，代表单卡的head个数；综合约束请见undefined。
- inputLayout（string*，计算输入）：Host侧的string，数据类型支持String，代表输入query、key、value的数据排布格式，支持BSH、SBH、BSND、BNSD、TND。
  
  说明： query、key、value数据排布格式支持从多种维度解读，其中T是B和S合轴紧密排列的数据（每个batch的SeqLenQ和SeqLenKV）、B（Batch）表示输入样本批量大小、S（Seq-Length）表示输入样本序列长度、H（Head-Size）表示隐藏层的大小、N（Head-Num）表示多头数、D（Head-Dim）表示隐藏层最小的单元尺寸，且满足D=H/N。
- innerPrecise（int64_t，计算输入）：Host侧的int64_t，数据类型支持INT64，用于提升精度，默认配置为0即可。
  
  说明： 当前0、1为保留配置值，当计算过程中存在整行mask进而导致精度有损失时，可以尝试将该参数配置为2以提升精度，但是该配置会导致性能下降。在非TND数据排布下，如果算子可判断出存在无效行场景，会自动使能无效行计算，例如sparseMode为3，Sq > Skv场景。
- sparseMode（int64_t，计算输入）：Host侧的int64_t。数据类型支持INT64。用户不特意指定时建议传入0，在sequence长度相等的场景下，支持配置值为0、1、2、3、4、5、6；在sequence长度不相等的场景下，支持配置值为0、1、2、3、4、6、7、8。当整网的attenMaskOptional都相同且shape小于2048*2048时，建议使用defaultMask模式，来减少内存使用量。sparse不同模式的详细说明请参见undefined。
- outDtypeOptional （int64_t，计算输入）：Host侧的int64_t。数据类型支持INT64。只在量化场景下生效， 0表示attenout输出是fp16格式，1表示输出是bf16格式。
- pseType（int64_t，计算输入）：Host侧的int64_t。数据类型支持INT64。用户不特意指定时可传入1，跟undefined实现一致，支持配置值为0、1、2、3。
  
  pseType 含义备注
  
  0 外部传入pse 先mul再add -
  
  1 外部传入pse 先add再mul 与undefined实现一致
  
  2 内部生成pse 先mul再add -
  
  3 内部生成pse 先mul再add再sqrt -
- seed（int64_t，计算输入）：Host侧的int64_t。数据类型支持INT64。
- offset（int64_t，计算输入）：Host侧的int64_t。数据类型支持INT64。keepProb小于1.0时，默认根据seed和offset生成DropoutMask。
- softmaxMaxOut（aclTensor*，计算输出）：Device侧的aclTensor，Softmax计算的Max中间结果，用于反向计算。数据类型支持FLOAT，输出的shape类型为[B,N,Sq,8]，undefined支持ND。
- softmaxSumOut（aclTensor*，计算输出）：Device侧的aclTensor，Softmax计算的Sum中间结果，用于反向计算。数据类型支持FLOAT，输出的shape类型为[B,N,Sq,8]，undefined支持ND。
- softmaxOutOut（aclTensor*，计算输出）：预留参数，暂未使用。
- attentionOutOut（aclTensor*，计算输出）：Device侧的aclTensor，数据类型支持FLOAT16、BFLOAT16、FLOAT32，数据类型和shape类型与query保持一致，undefined支持ND。
- workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
- executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。
返回值：

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

[object Object]

pseType	含义	备注
0	外部传入pse 先mul再add	-
1	外部传入pse 先add再mul	与undefined实现一致
2	内部生成pse 先mul再add	-
3	内部生成pse 先mul再add再sqrt	-

aclnnFlashAttentionScoreVX

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

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

约束说明

该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。
输入query、key、value的B：batchsize必须相等。
输入query、key、value的D：Head-Dim必须满足(qD == kD && kD >= vD)。
输入query、key、value的input_layout必须一致。
输入queryRopeOptional与query的输入shape仅在D维度不同，其他shape参数应该相同
输入keyRopeOptional与key的输入shape仅在D维度不同，其他shape参数应该相同
关于数据shape的约束，以inputLayout的TND、BSND、BNSD为例（BSH、SBH下H=N*D），其中：
- T(B*S)：取值范围为1~1M。
- B：取值范围为1~2M。带prefixOptional的时候B最大支持2K。
- N：取值范围为1~256。
- S：取值范围为1~1M。
- D：取值范围为1~512。
keepProb的取值范围为(0, 1]。
部分场景下，如果计算量过大可能会导致算子执行超时(aicore error类型报错，errorStr为：timeout or trap error)，此时建议做轴切分处理，注：这里的计算量会受B、S、N、D等参数的影响，值越大计算量越大。
band场景，preTokens和nextTokens之间必须要有交集。
prefixOptional稀疏计算场景，场景包括sequence长度相等的场景下sparseMode=5、sparseMode=6；sequence长度不相等的场景下sparseMode=6。这两种场景下，当Sq > Skv时，prefix的N值取值范围[0, Skv]；当Sq <= Skv时，prefix的N值取值范围[Skv-Sq, Skv]。当sparseModeOptional=5、prefix的N > Skv时，执行全计算。
realShiftOptional Sq大于1024时如果配置BNHS、1NHS，需要Sq和Skv等长。
sparse_mode=7时，不支持可选输入realShiftOptional。
sparse_mode=8时，当每个sequence的q、kv等长时支持可选输入realShiftOptional，针对全局做pse生成。支持q方向进行外切，需要外切前每个sequence的q、kv等长，外切后传入的actualSeqQLenOptional[0] - actualSeqKvLenOptional[0] + qStartIdxOptional - kvStartIdxOptional == 0（本功能属实验性功能）。
actualSeqQLenOptional输入支持某个Batch上的S长度为0，此时不支持可选输入realShiftOptional。
attenMaskOptional输入不支持补pad，即attenMaskOptional中不能存在某一行全1的场景。
支持actualSeqQLenOptional中某个Batch上的S长度为0；如果存在S为0的情况，不支持pse输入，假设真实的S长度为[2,2,0,2,2]，则传入的actualSeqQLenOptional为[2,4,4,6,8]。
pseType为2或3的时候，当前只支持Sq和Skv等长。

调用示例

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

[object Object]