CANN ATB加速库工作原理

张

张建站

2026/5/9 12:59:38

10分钟阅读

ATB工作原理【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost算子下发原理深度学习模型可以抽象为由一个个算子组合而成计算图节点代表算子边代表张量数据依赖关系。在模型训练和推理时模型主体程序在CPU上执行过程中将算子一个个下发到设备侧Device NPU或GPU上执行并在必要的时候进行同步。整个步骤可以简单抽象为下图两种类型的性能瓶颈由于Host准备算子上下文并下发算子也是需要时间因此这种工作模式下存在两种可能的性能瓶颈Host BoundHost下发较慢设备侧NPU执行算子较快 Host执行效率成为整体性能瓶颈。在profiling图上表现为Stream上的Kernel间存在空泡。此时设备侧的算力没有得到充分利用需要优化Host程序加快算子下发。Device Bound Host下发较快 Device执行较慢 Device执行效率成为性能瓶颈。这种场景下设备侧算力被充分利用如想继续提高性能则需要考虑优化kernel.上图分别给出了这两种性能瓶颈的示例。算子下发过程单个算子的下发过程可以简化为如下步骤合法性检查检查算子输入、输出、参数是否符合算子要求防止错误参数提交到Device后导致错误。输出shape推导Infer Shape通过算子的输入Shape和Data type推导输出Shape和Data Type。例如上述一个简单的Matmul算子左矩阵Shape为M * K右矩阵为K * N 可以推导输出矩阵Shape为M * N。计算Tiling大多数情况下单个AI Core一次能处理的数据有限算子的输入数据无法一次完全载入完成计算需要将输入切分成多块分块完成计算这个过程叫Tiling数据切分的算法称为Tiling算法或者Tiling策略。对于复杂算子来说每个Kernel实现都可能有自己的Tiling算法根据输入输出Tensor的Shape和其他信息计算出进行Kernel执行过程中如何切分。计算的结果一般保存在一个自定义的Tiling数据结构中。上图所示是一个矩阵乘Matmul的Tiling策略首先要进行多核切分: 根据当前核数对M、K、N进行多核切分得到单核内shape大小singleCoreM、singleCoreK、singleCoreN。接着要进行核内切分: 根据Local Memory的大小约束对单核内的Shape大小进一步切分得到A、B、C矩阵参与一次矩阵乘指令的Shape大小baseM、baseN、baseK。ATB会把Tiling策略用一个结构体保存起来后续传给算子核函数使用。在此例中是struct matmulTilingData。struct matmulTilingData { uint singleCoreM uint singleCoreK uint singleCoreN uint baseM; uint basek uint baseN }Tiling策略对复杂算子的性能影响巨大同一个算子在不同Tiling策略下可能有10倍性能差异。获取Workspace大小算子内部有时需要通过额外的HBM内存进行数据交换或者缓存这部分空间称为算子的Workspace。需要在算子实际执行前分配好。上面的示例中先进行矩阵乘法然后进行一次Reduce。需要使用Workspace暂存矩阵乘结果。分配Workspace对于ATB和aclnn这样的两段式算子接口来说这个步骤一般由执行框架如torch-npu进行分配而不是算子内部实现。这样外部框架可以管理整个模型执行过程中间的HBM资源提高分配效率。算子下发将之前准备好的输入输出Tensor地址、Tiling信息、Workspace地址内存空间以及其他参数封装成argument list 调用Launch Kernel接口通知Device侧按照上面的参数执行Kernel。ATB工作原理当模型变得越来越复杂算子越来越多前述的Host Bound会逐渐显现。为了解决这个问题ATB进行了针对性的优化。它提供如下功能定制化融合算子: 提供Transformer结构常用的算子如PageAttention、Linear等。ATB提供的算子通常是针对主流模型经过精心设计的融合算子具有较高的性能。轻量级组图: 支持使用上述算子或第三方算子组图然后像操作单算子一样操作图下面称其为图算子。图算子可以很方便地在不同模型、不同layer之间复用。运行时优化: 采用了多种优化方案提升Host性能并降低Device内存占用。具体总结如下Tiling Cache机制通过缓存计算好的Tiling以存代算减少重复计算。调度优化优化组图模式下算子下发调度方式使设备侧算子运行无间隙解决Host Bound。内存优化通过基于内存Block分裂、合并、尾块优化的内存分配算法实现图算子内部中间Tensor复用平均节省Workspace 50%提升大模型推理Batch Size上限。组图实例详见 tests/framework/c/layer_ops/llama65b/layer/llama65b_layer_mlp_graph_builder.cpp 中CreateLlamaMlpOperationByGraphOpBuilder函数。下面仅包含组图逻辑主体。atb::Status CreateLlamaMlpOperationByGraphOpBuilder(const LlamaMlpParamGb param, atb::Operation **operation) { atb::GraphOpBuilder* graphOpBuilder; CreateGraphOpBuilder(graphOpBuilder); /* 此处省略了参数创建 */ graphOpBuilder-Init( LlamaMlpGraphOp, inferShapeFunc, {hidden_states, weight}, {mlp_out} ); graphOpBuilder-Reshape(hidden_states, reshape_01_2, hidden_states_); graphOpBuilder-AddOperation(Linear(param), {hidden_states_, weight}, {linear_out}); graphOpBuilder-Reshape(linear_out, unsqueueze_0, linear_out_); graphOpBuilder-AddOperation(Split(param), {linear_out_}, {gate_out, up_out}); graphOpBuilder-AddOperation(Swish(param), {gate_out}, {swish_out}); graphOpBuilder-AddOperation(Mul(param), {swish_out, up_out}, {mlp_out}); *operation graphOpBuilder-Build(); DestroyGraphOpBuilder(graphOpBuilder); return atb::NO_ERROR; }上述代码组建了一个由四个算子组成的图算子。逻辑视图如下在ATB的内部使用两个Vector容器分别存放算子节点和算子的输入输出。图算子Setup和Execute流程由于ATB中图算子只是单算子的组合不涉及Kernel融合因此图算子的Setup和Execute过程与单算子类似区别仅在于Setup阶段进行了Workspace优化。Setup和Execute流程分别如下所示运行时优化Setup复用和Cache优化实际推理过程中即使是动态Shape场景下多次推理过程的输入Shape也大概率重复。基于这个特征可以进行如下优化使用一个Cache保存一个算子常用的多份Tiling信息默认每个算子保存10份 Shape相同场景下可以避免重复计算。每个算子执行上下文中保存了上一次执行的Tensor信息、Tiling信息、Worksspace Size信息。如果某次执行的Shape与上次完全相同则可以直接复用上下文跳过整个Setup阶段。上述两种优化对图算子和单算子都适用。HBM内存优化ATB在图算子Setup阶段尽可能复用HBM 使得整个图算子的Workspace size比内部单算子Workspace size的总和要小。具体方式如下一个流中的算子Kernel是顺序执行的所以前一个算子的Workspace可以给后一个算子使用。一个图算子内部的中间Tensor不需要保留到图算子执行完毕只要最后一个使用它的单算子执行完毕后就可以释放空间给其他Tensor使用。下发优化优化前下发调度逐个算子执行Setup和execution容易在NPU上形成空泡基础优化ATB通过图算子批量进行算子Setup和任务下发可有效减少NPU空泡。这一步优化是组图模式自动实现的。不需要用户特殊操作。双线程下发优化推荐使用通过双线程分别进行算子批量Setup和批量任务下发可以同时减少host执行时间和NPU空泡。这种当时需要用户创建两个线程其中一个线程处理Setup 另一个线程处理Execute。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

强化学习在工业控制中的Sim-to-Real迁移挑战与优化

1. 强化学习在工业控制中的模拟到现实挑战在工业过程控制领域，强化学习（Reinforcement Learning, RL）正逐渐展现出其独特的价值。作为一名长期从事工业自动化与智能控制的研究者，我亲眼见证了RL技术从实验室走向实际产线的艰难历…...

2026/5/9 12:54:44 阅读更多 →

Armv8-M TrustZone安全隔离与IDAU硬件实现解析

1. Armv8-M TrustZone架构与安全隔离机制解析在嵌入式安全领域，Arm TrustZone技术已经成为硬件级隔离的事实标准。作为在Armv8-M架构中实现安全隔离的核心机制，它通过物理硬件信号将处理器状态划分为安全(Secure)和非安全(Non-secure)两个世界。这种隔离…...

2026/5/9 12:54:41 阅读更多 →

K-12人工智能教育：达格斯特三角框架下的课程设计与教学实践

1. 项目概述：为什么K-12阶段需要人工智能教育？最近几年，但凡关注点科技新闻或者教育动态的朋友，恐怕都绕不开“人工智能”这个词。从能写诗作画的AI，到能辅助诊断的医疗系统，再到我们手机里越来越“懂你”的…...

2026/5/9 12:53:34 阅读更多 →

ColorControl：一键掌控多设备显示与智能控制的终极方案

ColorControl：一键掌控多设备显示与智能控制的终极方案【免费下载链接】ColorControl Easily change NVIDIA display settings and/or control LG TVs 项目地址: https://gitcode.com/gh_mirrors/co/ColorControl ColorControl 是一个专注于显示参数优化与智…...

2026/5/8 2:36:29 阅读更多 →

使用Taotoken CLI工具一键配置开发环境与多工具API密钥的教程

使用Taotoken CLI工具一键配置开发环境与多工具API密钥的教程 1. 安装Taotoken CLI工具 Taotoken CLI工具提供两种安装方式。对于需要频繁使用CLI的场景，推荐全局安装： npm install -g taotoken/taotoken若只需临时使用或避免全局依赖，可通…...

2026/5/9 12:07:00 阅读更多 →

C语言固件完整性保护全栈方案（含国密SM4+可信执行环境TEE落地代码）

更多请点击： https://intelliparadigm.com 第一章：Shell脚本的基本语法和命令 Shebang 与执行方式每个可执行 Shell 脚本的第一行应以 Shebang（ #!/bin/bash）开头，用于指定解释器路径。保存为 hello.sh 后&#xf…...

2026/5/7 11:30:49 阅读更多 →

在 Node.js 后端服务中集成 Taotoken 实现多模型对话路由

在 Node.js 后端服务中集成 Taotoken 实现多模型对话路由 1. 准备工作在开始集成 Taotoken 之前，需要确保您的开发环境已满足以下条件。Node.js 版本建议使用 18.x 或更高 LTS 版本。通过运行 node -v 可以检查当前版本。如果尚未安装 openai 包，可以…...

2026/5/8 6:34:49 阅读更多 →