加速库RingMLA C Demo【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost介绍该目录下为加速库RingMLA 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如./ring_mla_demo 0额外说明示例中生成的数据不代表实际场景如需数据生成参考请查看根目录下的python用例目录 tests/apitest/opstest/python/operations/ring_mla/产品支持情况RingMLA仅Atlas A2/A3系列场景说明RingMLA基础场景对于querykey分别传入带与不带rope转置的矩阵传入固定shape512x512的上三角mask默认编译脚本可编译运行该demo仅支持在Atlas A2/A3系列上运行Demo分段参数/tensor规格设置因为RingMLA和其他的算子较为不同第一次运算时不带前次生成的prevOut和prevLse但是从第二次起使用时需要带上这里分两段描述第一轮参数设置成员名称取值calcTypeCALC_TYPE_FISRT_RINGheadNum16kvHeadNum8qkScale1/sqrt(192)kernelTypeKERNELTYPE_DEFAULTmaskTypeMASK_TYPE_TRIU注意qkScale设置值为RingMLA做rope转置前querykey合一的headSize即128(nope) 64(rope) 192数据规格tensor名字数据类型数据格式维度信息cpu/npuqueryNopebf16nd[1228, 16, 128]npuqueryRopebf16nd[1228, 16, 64]npukeyNopebf16nd[828, 8, 128]npukeyRopebf16nd[828, 8, 64]npuvaluebf16nd[828, 8, 128]npumaskbf16nd[512, 512]npuseqLenint32nd[2, 3]cpuOutputoutputbf16nd[1228, 16, 128]npusoftmaxLsefloatnd[16, 1228]npuq第一维度为总词元长度对应sum(seqlen[0])kv第一维度对应sum(seqlen[1])第二轮 第二轮中会使用第一轮新生成的output和softmaxLse进行计算。参数设置成员名称取值calcTypeCALC_TYPE_DEFAULTheadNum16kvHeadNum8qkScale1/sqrt(192)kernelTypeKERNELTYPE_DEFAULTmaskTypeMASK_TYPE_TRIU需要额外更改param里的calcType为CALC_TYPE_DEFAULT其他保持一致数据规格tensor名字数据类型数据格式维度信息cpu/npuqueryNopebf16nd[1228, 16, 128]npuqueryRopebf16nd[1228, 16, 64]npukeyNopebf16nd[828, 8, 128]npukeyRopebf16nd[828, 8, 64]npuvaluebf16nd[828, 8, 128]npumaskbf16nd[512, 512]npuseqLenbf16nd[2, 3]cpuprevOutbf16nd[1228, 16, 128]npuprevLsefloatnd[16, 1228]npuOutputoutputbf16nd[1228, 16, 128]npusoftmaxLsefloatnd[16, 1228]npu第二轮使用第一轮新生成的output和softmaxLse来作为prevOut和prevLse。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考