ATB自注意力C++示例

加速库SelfAttentionOperation C Demo【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost介绍该目录下为加速库SelfAttentionOperation C调用示例。使用说明首先source 对应的CANN和nnal包的安装路径source [cann安装路径]/set_env.sh 默认source /usr/local/Ascend/ascend-toolkit/set_env.shsource [nnal安装路径]/set_env.sh默认source /usr/local/Ascend/nnal/atb/set_env.sh①. 如果使用加速库源码编译source [加速库源码路径]/output/atb/set_env.sh例如 source ./ascend-transformer-boost/output/atb/set_env.sh运行demobash build.sh注意使用cxx_abi0默认时设置D_GLIBCXX_USE_CXX11_ABI为0即g -D_GLIBCXX_USE_CXX11_ABI0 -I ...使用cxx_abi1时更改D_GLIBCXX_USE_CXX11_ABI为1即g -D_GLIBCXX_USE_CXX11_ABI1 -I ...生成的二进制文件***_demo可以额外传入一个int参数作为deviceId默认为0如./self_attention_encoder_demo 0提供的build脚本仅用于编译和运行self_attention_encoder_demo.cpp如需编译其他demo需要替换“self_attention_encoder_demo”为对应的cpp文件名额外说明示例中生成的数据不代表实际场景如需数据生成参考请查看根目录下的python用例目录 tests/apitest/opstest/python/operations/self_attention/产品支持情况SelfAttention在Atlas A2/A3系列仅支持部分场景且Encoder场景在Atlas 推理系列产品上调用与Atlas A2/A3有所区别场景说明提供demo分别对应不同产品/场景具体场景区别可以参见官网编译运行时需要对应更改build脚本self_attention_encoder_demo.cpp场景FA Encoder基础场景分开传入keyCacheKvalue和CacheV默认编译脚本可编译运行该demo仅支持在Atlas A2/A3上运行demo中使用全量的上三角mask演示参数设置成员名称取值headNum32kvHeadNum32qkScale1/sqrt(128)calcTypeENCODERkernelTypeKERNELTYPE_DEFAULTmaskTypeMASK_TYPE_NORM数据规格tensor名字数据类型数据格式维度信息cpu/npuqueryfloat16nd[160,2048]npukeyfloat16nd[160,2048]npuvaluefloat16nd[160,2048]npucacheKfloat16nz[1, 10, 1024, 2048]npucacheVfloat16nz[1, 10, 1024, 2048]npuattentionMaskfloat16nd[10, 1024, 1024]nputokenOffsetint32nd[10]cpuseqLenint32nd[10]cpulayerIdint32nd[1]npuOutputoutputfloat16nd[160, 2048]npuqkv第一维度为总词元长度即sum(seqlen)qkv第二维度headNumheadSize合轴实际为headNum(32) $\times$ headSize(128)self_attention_encoder_inference_demo.cpp场景fa encoder基础场景在Atlas推理系列上的实现分开传入keyCacheKvalue和CacheV该demo仅支持在Atlas推理系列上运行参数设置成员名称取值headNum16kvHeadNum16qkScale1/sqrt(128)calcTypeENCODERkernelTypeKERNELTYPE_DEFAULTmaskTypeMASK_TYPE_UNDEFINED数据规格tensor名字数据类型数据格式维度信息cpu/npuqueryfloat16nd[16, 256]npukeyfloat16nd[16, 256]npuvaluefloat16nd[16, 256]npucacheKfloat16nz[1, 1, 16, 256, 16]npucacheVfloat16nz[1, 1, 16, 256, 16]nputokenOffsetint32nd[1]cpuseqLenint32nd[1]cpulayerIdint32nd[1]npuOutputoutputfloat16nd[16, 256]npuself_attention_pa_encoder_demo.cpp场景FA使用PA Encoder的场景使用FA输入只需传入keyvalue传入不同的headNumkvHeadNum且headNum可以被kvHeadNum整除时开启GQAGrouped Query Attention。该demo仅支持在Atlas A2/A3系列上运行demo中使用全量的上三角mask演示参数设置成员名称取值headNum32kvHeadNum16qkScale1/sqrt(128)calcTypePA_ENCODERkernelTypeKERNELTYPE_DEFAULTmaskTypeMASK_TYPE_NORM数据规格tensor名字数据类型数据格式维度信息cpu/npuqueryfloat16nd[3072, 32, 128]npukeyfloat16nd[3072, 16, 128]npuvaluefloat16nd[3072, 16, 128]npuattentionMaskfloat16nd[4, 1024, 1024]npuseqLenint32nd[4]cpuOutputoutputfloat16nd[3072, 32, 128]npuself_attention_pa_encoder_qwen_demo.cpp场景FA使用PA Encoder的场景使用FA输入只需传入keyvalue传入不同的headNumkvHeadNum且headNum可以被kvHeadNum整除时开启GQAGrouped Query Attention。该demo仅支持在Atlas A2/A3系列上运行demo中为了应对长序列的情况使用压缩上三角mask演示参数设置成员名称取值headNum5kvHeadNum1qkScale1/sqrt(128)isTriuMask1calcTypePA_ENCODERkernelTypeKERNELTYPE_DEFAULTmaskTypeMASK_TYPE_NORM_COMPRESS数据规格tensor名字数据类型数据格式维度信息cpu/npuqueryfloat16nd[1024, 5, 128]npukeyfloat16nd[1024, 1, 128]npuvaluefloat16nd[1024, 1, 128]npuattentionMaskfloat16nd[128, 128]npuseqLenint32nd[1]cpuOutputoutputfloat16nd[1024, 5, 128]npuself_attention_prefix_encoder_demo.cpp场景FA使用Prefix Encoder的场景传入PA的依据blockTables存放的keyvalue此场景支持q和kv不等长但是要求 $$\forall i \lt len(seqLen)kvSeqLen[i] - seqLen[i] 0 \ (mod \ 128) $$该demo仅支持在Atlas A2/A3系列上运行demo中使用Alibi上三角mask叠加bias slopes演示参数设置成员名称取值headNum32kvHeadNum8qkScale1/sqrt(128)isTriuMask1calcTypePREFIX_ENCODERkernelTypeKERNELTYPE_HIGH_PRECISIONmaskTypeMASK_TYPE_ALIBI_COMPRESS数据规格tensor名字数据类型数据格式维度信息cpu/npuqueryfloat16nd[96, 32, 128]npukeyfloat16nd[480, 128, 8, 128]npuvaluefloat16nd[480, 128, 8, 128]npublockTablesfloat16nd[4, 4]npumaskfloat16nd[32, 96, 128]npuseqLenint32nd[4]cpukvSeqLenint32nd[4]cpuslopesfloat32nd[128]npuOutputoutputfloat16nd[96, 32, 128]npu【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考