CANN/ops-transformer FlashAttentionScore算子

FlashAttentionScore【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品√Atlas A3 推理系列产品×Atlas A2 训练系列产品√Atlas A2 推理系列产品×功能说明算子功能训练场景下使用FlashAttention算法实现self-attention自注意力的计算pseType1时需要先add再mul。pseType≠1时需要先mul再add。计算公式注意力的正向计算公式如下pseType1时公式如下 $$ attention_outDropout(Softmax(Mask(scale*(pse(queryd_scale_q)(keyd_scale_k)^T), atten_mask)), keep_prob)(value*d_scale_v) $$pseType≠1时公式如下 $$ attention_outDropout(Softmax(Mask(scale*((queryd_scale_q)(keyd_scale_k)^T) pse),atten_mask),keep_prob)(value*d_scale_v) $$参数说明参数名输入/输出/属性描述数据类型数据格式query输入公式中的输入query。BFLOAT16、FLOAT16、FLOAT、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8NDkey输入公式中的输入key。BFLOAT16、FLOAT16、FLOAT、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8NDvalue输入公式中的输入value。BFLOAT16、FLOAT16、FLOAT、FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8NDrealShiftOptional可选输入公式中的pse表示位置编码。BFLOAT16、FLOAT16、FLOATNDdropMaskOptional可选输入公式中的Dropout表示数据丢弃掩码。取值为1代表保留该数据为0代表丢弃该数据。UINT8NDattenMaskOptional可选输入公式中的atten_mask表示注意力掩码取值为1代表该位不参与计算不生效为0代表该位参与计算。BOOL、UINT8NDdScaleQOptional可选输入公式中的d_scale_qFP8场景下query的全量化参数。FLOATNDdScaleKOptional可选输入公式中的d_scale_kFP8场景下key的全量化参数。FLOATNDdScaleVOptional可选输入公式中的d_scale_vFP8场景下value的全量化参数。FLOATNDscaleValue可选属性公式中的scale表示缩放系数作为计算流中Muls的scalar值。默认值为1.0。DOUBLE-keepProb可选属性公式中的keep_prob表示数据需要保留的概率。默认值为1.0。DOUBLE-pseType可选属性控制add与mul的执行次序支持配置值为0、1、2、3。默认值为1。INT64-softmaxMaxOut输出Softmax计算的Max中间结果用于反向计算。FLOATNDsoftmaxSumOut输出Softmax计算的Sum中间结果用于反向计算。FLOATNDattentionOut输出公式中的attention_out。BFLOAT16、FLOAT16、FLOATNDAtlas A2 训练系列产品 不支持FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8三种数据类型。Atlas A3 训练系列产品 不支持FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8三种数据类型。约束说明输入query、key、value、realShiftOptional的数据类型必须一致。输入query、key、value的inputLayout必须一致。关于数据shape的约束以inputLayout的BSND、BNSD为例BSH、SBH下HN*D其中B取值范围为1~2M。当prefixOptional的时候B最大支持2K。N取值范围为1~256。S取值范围为1~1M。D取值范围为1~768。输入query、key、value类型为FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8时D取值范围为1~128。输入query、key、value类型为FLOAT8_E5M2、FLOAT8_E4M3FN、HIFLOAT8时, 不支持queryRopeOptional、keyRopeOptional、realShiftOptional、attenMaskOptional、dropMaskOptional、keepProb、pseType等相关可选参数。keepProb的取值范围为(0, 1]。部分场景下如果计算量过大可能会导致算子执行超时aicore error类型报错errorStr为timeout or trap error此时建议做轴切分处理注这里的计算量会受B、S、N、D等参数的影响值越大计算量越大。pseType为2或3的时候当前只支持Sq和Skv等长。调用说明调用方式调用样例说明aclnn调用test_aclnn_flash_attention_score非TND场景通过aclnnFlashAttentionScore接口方式调用FlashAttention算子。参考资源算子设计原理【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考