Ascend Transformer Boost【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost [2025/09] Ascend Transformer Boost项目首次上线。一、什么是ATBATB介绍Ascend Transformer Boost加速库下文简称为ATB加速库是一款高效、可靠的加速库基于华为Ascend AI处理器专门为Transformer模型的训练和推理而设计。具体的工作原理可以参考ATB加速原理。软件架构ATB的架构图如上图所示其接口功能主要分成三部分提供经过优化的融合算子Operation用户可以根据需求使用对应的算子完成计算功能。提供图算子机制用户根据模型设计对应的图算子使用加速库提供的原生算子和创建的自定义算子创建图算子完成相应的计算。提供插件Plugin机制用户可以根据自己的需求创建自定义的算子。ATB仓介绍ATB库的目录结构如下ascend-transformer-boost ├── 3rdparty //第三方依赖库文件夹 ├── build //可存放构建生成的文件 ├── ci //持续集成相关的配置文件 ├── docs //文档文件 ├── example //算子调用示例代码包含可直接运行的Demo ├── include //存放公共头文件 ├── ops_configs //存放算子输入输出数据规格约束文件 ├── ops_customize //存放自定义操作相关的文件 ├── output //编译输出文件夹 ├── scripts //脚本文件存放目录 ├── src //主体源代码目录 │ ├── atb │ ├── kernels │ │ ├── configs //支持的配置说明 │ │ ├── include //存放各算子的头文件 │ │ ├── kernels //单算子存放目录 │ │ ├── lcal //通信算子存放目录 │ │ ├── mixkernels //融合算子存放目录 │ │ ├── tbe_adapter //TBE 适配器相关的源代码 │ │ └── CMakeLists.txt │ ├── ops │ │ ├── ops_common │ │ ├── ops_infer //推理OP │ │ └── ops_train //训练OP │ ├── torch_atb //与PyTorch相关的atb库文件 │ └── CMakeLists.txt ├── tests //测试代码 └── torch_atb为什么选择ATB对Transformer模型的高效加速ATB加速库通过优化矩阵乘法等核心算子和注意力机制的实现方式实现了对Transformer模型的高效加速。高性能和效率ATB加速库充分利用了Ascend AI处理器的硬件特性如算力、存储带宽和内存带宽通过硬件加速和数据重用等技术进一步提升了性能和效率。提供了底层基础的高性能算子以及高效的算子组合技术。支持多种模型框架如PyTorch、MindSpore、Paddle。二、环境构建版本兼容性说明ATB的API保证前后一年的ABI兼容能力在不涉及新功能的情况下调用者升级一年内的ATB版本不会出现兼容问题。由于CANN出包目录调整ATB 8.5版本以及主线分支必须匹配8.5或以上版本的toolkit包。快速安装CANN软件本节提供快速安装CANN软件的示例命令更多安装步骤请参考详细安装指南。安装前准备在线安装和离线安装时需确保已具备Python环境及pip3当前CANN支持Python3.7.x至3.11.4版本。 离线安装时请单击获取链接下载CANN软件包并上传到安装环境任意路径。安装CANN由于CANN出包目录调整ATB 8.5版本以及主线分支必须匹配8.5或以上版本的toolkit包。chmod x Ascend-cann-toolkit_${VERSION}_linux-$(arch).run # 其中${VERSION}表示对应的CANN版本如8.2.RC1 ./Ascend-cann-toolkit_${VERSION}_linux-$(arch).run --install安装后配置配置环境变量脚本set_env.sh当前安装路径以${HOME}/Ascend为例。source ${HOME}/Ascend/ascend-toolkit/set_env.sh安装业务运行时依赖的Python第三方库如果使用root用户安装请将命令中的--user删除。pip3 install attrs cython numpy1.19.2,1.24.0 decorator sympy cffi pyyaml pathlib2 psutil protobuf3.20.0 scipy requests absl-py --user安装ops算子包安装前需已安装配套版本的Toolkit并配置环境变量chmod x Ascend-cann-${chip_type}-ops_${VERSION}_linux-$(arch).run #其中${chip_type}表示昇腾产品类型如A3 ./Ascend-cann-${chip_type}-ops_${VERSION}_linux-$(arch).run --installCANN详细安装指南开发者可访问昇腾文档-昇腾社区-CANN社区版-软件安装查看CANN软件安装引导根据机器环境、操作系统和业务场景选择后阅读详细安装步骤。ATB安装部署相关依赖说明在编译加速库之前需访问加速库包安装部署查看加速库相关依赖的版本要求并进行对应依赖的安装部署。基础工具版本要求与安装安装CANN之后您可安装一些工具方便后续开发参见以下内容CANN依赖列表CANN安装后操作三、快速上手ATB编译加速库下载git clone https://gitcode.com/cann/ascend-transformer-boost.git您可自行选择需要的分支。环境变量设置在编译前需要安装nnal软件包安装方法请查看常见问题与回答并根据nnal安装路径设置环境变量ATB_BUILD_DEPENDENCY_PATHexport ATB_BUILD_DEPENDENCY_PATH{nnal install path}/nnal/atb/latest/atb/cxx_abi_{cxx_abi_version}注未设置将使用默认路径/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_{cxx_abi_version}加速库编译编译加速库设置加速库环境变量cd ascend-transformer-boost bash scripts/build.sh source output/atb/set_env.sh注意该编译过程涉及①拉取算子库/MKI并编译②加速库的编译两个过程。更多编译命令说明请参考编译与构建调用示例说明本节示例代码分别展示了如何通过Python和C调用算子。Python在运行python代码前需要导入ATB Python API模块torch_atb该插件运行依赖PyTorch和torch_npu可访问加速库包安装部署查看版本要求和安装指导。安装PyTorch和torch_npu之后需要手动安装torch_atb安装方法包含以下两种通过./Ascend-cann-nnal_${version_info}_linux-aarch64.run --install --torch_atb命令安装在编译步骤中使用bash scripts/build.sh --torch_atb编译命令output/whl文件夹下会生成torch_atb的whl文件您可使用如下命令进行安装:pip3 install torch_atb-{version}-py3-none-any.whl如下代码展示了如何通过Python调用算子注意不要在ATB代码仓同名目录下运行import torch import torch_atb#导入ATB Python API模块 #创建参数对象 linear_param torch_atb.LinearParam() linear_param.has_bias False #创建算子对象 op torch_atb.Operation(linear_param) #准备输入数据 x torch.randn(2, 3, dtypetorch.float16).npu() y torch.randn(2, 3, dtypetorch.float16).npu() #使用forward方法完成操作并获取输出 outputs op.forward([x, y]) torch.npu.synchronize()如需查看调用结果可以打印第一个输出resultoutputs[0].cpu().numpy() print(result)代码编写指导可访问算子使用指导ATB Python API-昇腾社区。C在ATB仓库的example/op_demo目录下存放了多个不依赖测试框架、即编可用的算子调用Demo示例。进入对应目录执行如下命令就可完成一个算子的调用执行。代码完整内容可参考example/op_demo/faupdate/faupdate_demo.cpp下面仅展示其核心内容// 设置卡号、创建context、设置stream atb::Context *context nullptr; void *stream nullptr; CHECK_STATUS(aclInit(nullptr)); CHECK_STATUS(aclrtSetDevice(DEVICE_ID)); CHECK_STATUS(atb::CreateContext(context)); CHECK_STATUS(aclrtCreateStream(stream)); context-SetExecuteStream(stream); // 创建op atb::Operation *faupdateOp nullptr; CHECK_STATUS(CreateFaUpdateOperation(faupdateOp)); // 准备输入tensor atb::VariantPack variantPack; CHECK_STATUS(PrepareInTensor(context, stream, variantPack.inTensors)); // 放入输入tensor // 准备输出tensor atb::Tensor output; CHECK_STATUS(CreateTensor(ACL_FLOAT, aclFormat::ACL_FORMAT_ND, {LOCALOUT_DIM_1, LOCALOUT_DIM_2}, output)); variantPack.outTensors {output}; // 放入输出tensor uint64_t workspaceSize 0; // 计算workspaceSize大小 CHECK_STATUS(faupdateOp-Setup(variantPack, workspaceSize, context)); uint8_t *workspacePtr nullptr; if (workspaceSize 0) { CHECK_STATUS(aclrtMalloc((void **)(workspacePtr), workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST)); } // faupdate执行 CHECK_STATUS(faupdateOp-Execute(variantPack, workspacePtr, workspaceSize, context)); CHECK_STATUS(aclrtSynchronizeStream(stream)); // 流同步等待device侧任务计算完成 // 释放资源 for (atb::Tensor inTensor : variantPack.inTensors) { CHECK_STATUS(aclrtFree(inTensor.deviceData)); } for (atb::Tensor outTensor : variantPack.outTensors) { CHECK_STATUS(aclrtFree(outTensor.deviceData)); } if (workspaceSize 0) { CHECK_STATUS(aclrtFree(workspacePtr)); } CHECK_STATUS(atb::DestroyOperation(faupdateOp)); // operation对象概念先释放 CHECK_STATUS(aclrtDestroyStream(stream)); CHECK_STATUS(DestroyContext(context)); // context全局资源后释放 CHECK_STATUS(aclFinalize());文件编译说明进入example/op_demo/faupdate执行bash build.sh完成编译和执行。出现以下提示即为成功faupdate demo success!代码编写指导可访问单算子-昇腾社区。样例安全声明example目录下的样例旨在提供快速上手、开发和调试ATB特性的最小化实现其核心目标是使用最精简的代码展示ATB核心功能而非提供生产级的安全保障。与成熟的生产级使用方法相比此样例中的安全功能如输入校验、边界校验相对有限。ATB不推荐用户直接将样例作为业务代码也不保证此种做法的安全性。若用户将example中的示例代码应用在自身的真实业务场景中且发生了安全问题则由用户自行承担。日志和环境变量说明加速库日志请参考日志与调试加速库日志现在已经部分适配CANN日志环境变量说明请参考CANN社区版文档/环境变量参考四、自定义算子开发您可参考以下文档进行自定义算子的开发从开发一个简单算子开始以简单的Add算子的增加为例介绍了ATB算子开发的交付件和开发流程适合新入门的选手。开发指南以一个融合算子为例详细介绍了ATB算子开发的流程以及如何对算子进行功能、精度、性能测试。注意您在开发过程中遇到的问题可参考ATB日志与调试尝试解决。五、参与贡献fork仓库修改并提交代码新建 Pull-Request详细步骤可参考贡献指南六、学习资源编译与构建ATB的编译命令说明。从开发一个简单算子开始以简单的Add算子的增加为例介绍了ATB算子开发的交付件和开发流程。开发指南以一个融合算子为例详细介绍了ATB算子开发的流程以及如何对算子进行功能、精度、性能测试。贡献指南介绍了如何向ATB库贡献代码。日志与调试介绍ATB的日志相关环境变量以及调测方法。API文档介绍了ATB库的接口和相关术语。常见问题与回答介绍ATB的编译和安装使用过程中遇到的一些常见问题与解决方案。问题报告通过Issue提交发现的问题。七、参考文档CANN社区版文档ATB社区版文档【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考