个人主页ujainu文章目录前言三层架构像搭积木一样组装加速方案三种并行策略让多张卡各干各的EP 并行Expert ParallelismTP 并行Tensor ParallelismPageAttentionKV Cache 的虚拟内存ATB 在 CANN 生态里的位置上游下游ATB 的朋友圈记住这些就够了前言ATB 是昇腾 CANN 生态里专门给 Transformer 模型开快车道的加速库。打个比方CANN 就像一座城市的道路系统ATB 修的是其中最宽的那条高速路——专跑大模型推理普通车不让上。你可能会问大模型推理不就是个输入 token → 输出 token的过程吗搞这么复杂干什么因为一个千亿参数的模型生成一个 token 要做几十上百次矩阵乘法、注意力计算、位置编码……每一层都是计算和搬运。在昇腾 NPU 上算力其实不是瓶颈瓶颈在于数据搬来搬去。ATB 干的活就是把这些零碎的计算步骤打包成流水线作业——能合并的合并能并行的并行能跳过的跳过。那为什么不直接写算子调用而非要搞一层加速库因为大模型推理不是一个算子算完就完事——它是一整条链路prefill 阶段要批量处理 promptdecode 阶段要逐 token 生成KV Cache 要在两个阶段之间无缝衔接多卡之间还要不停交换中间结果。如果每一步都靠推理框架手动调度代码写起来像意大利面条跑起来更像——到处是等待和搬运。ATB 把这些调度逻辑封装成可复用的加速方案让上层框架只管我要跑什么模型不用操心这个算子结果搬到哪里去。三层架构像搭积木一样组装加速方案ATB 的核心设计是三层架构每层解决不同的问题第一层基础原生算子Operation这是最底层的积木块每个 Operation 对应一个具体的计算操作——比如矩阵乘、LayerNorm、RMSNorm、SwiGLU。这些算子大多由 ops-transformer 仓库提供ATB 把它们封装成统一接口方便上层拼接。# Python 层面创建一个 MatMul Operationimportascend_transformer_boostasatb matmul_opatb.MatMul(nameqkv_proj,in_shape[seq_len,hidden_dim],out_shape[seq_len,3*hidden_dim])# 调用输入 hidden_states输出 Q/K/V 投影结果qkvmatmul_op(hidden_states)Operation 层的设计哲学是一个算子只干一件事。好处是灵活——你可以像搭乐高一样自由组合代价是如果每个算子都独立执行中间结果要在 NPU 内存和缓存之间来回搬延迟就上来了。这就引出了下一层。// C 层面定义一个 RMSNorm OperationautormsNormOpatb::Op(RMSNorm).SetInput(x).SetParam(epsilon,1e-6).SetParam(normed_shape,{hidden_dim}).Build();第二层图算子Graph Operator这是 ATB 最关键的一层。图算子把多个 Operation 串成一条流水线中间结果不用写回内存——数据从上一步直接流到下一步省掉了最贵的搬运环节。图算子的本质是一个有向无环图DAG节点是 Operation边是数据流。ATB 的图执行引擎按拓扑序调度这些节点保证每个节点的前置输入就绪后才开始执行并且中间张量尽量留在 NPU 的高速缓冲区里。# 用 ATB 的图算子 API 组装一个 Attention Blockimportascend_transformer_boostasatb graphatb.Graph(AttentionBlock)withgraph:# 第一步RMSNormnormedatb.RMSNorm(hidden_states,gamma,epsilon1e-6)# 第二步QKV 投影qkvatb.MatMul(normed,qkv_weight)q,k,vatb.Split(qkv,chunks3,dim-1)# 第三步旋转位置编码 RoPEqatb.RoPE(q,cos_table,sin_table)katb.RoPE(k,cos_table,sin_table)# 第四步FlashAttentionattn_outatb.FlashAttention(q,k,v,attn_mask)# 第五步输出投影o_projatb.MatMul(attn_out,o_weight)# 整个图编译一次推理时一键执行graph.compile()outputgraph.run(hidden_states)图算子带来的收益不仅是少搬运。ATB 还能对图做全局优化发现两个相邻的 MatMul 可以合并成一个 BatchMatMul发现 RMSNorm 后面紧跟的量化操作可以融合成一个算子——这些优化在单独调用 Operation 时是做不到的因为单算子视角看不到上下文。// C 层面定义图算子的拓扑结构autographatb::Graph(FFNBlock);autonormgraph.AddOp(RMSNorm,{input},{eps1e-6});autogategraph.AddOp(MatMul,{norm},{weightgate_w});autoupgraph.AddOp(MatMul,{norm},{weightup_w});autosilugraph.AddOp(Silu,{gate});automulgraph.AddOp(Mul,{silu,up});autodowngraph.AddOp(MatMul,{mul},{weightdown_w});graph.Freeze();// 冻结图触发编译优化第三层插件机制PluginPlugin 是 ATB 留给开发者扩展的接口。如果你的模型有一些特殊的结构比如新的注意力机制、自定义的激活函数可以通过 Plugin 自定义计算流程不用改 ATB 本身的代码。# 通过 Plugin 接入自定义算子# 就像在自动流水线旁边加了一个手工精加工工位classMyCustomPlugin(atb.Plugin):def__init__(self,config):super().__init__()self.custom_opload_custom_op(my_attention.so)defforward(self,hidden_states,attention_mask):# 你的特殊计算逻辑resultself.custom_op(hidden_states,attention_mask)returnresult# 注册到 ATB 的图算子中graph.register_plugin(my_attention,MyCustomPlugin(config))Plugin 机制让 ATB 具备了开闭原则——对扩展开放对修改关闭。社区贡献的新算子、新策略都可以通过 Plugin 接入ATB 核心代码不用动。比如你用 Ascend C 写了一个高性能的自定义算子编译成.so文件后通过 Plugin 就能无缝嵌入 ATB 的推理图里。# 注册 Ascend C 编写的自定义算子到 PluginclassAscendCPlugin(atb.Plugin):def__init__(self,so_path):self.op_handleatb.load_op(so_path)# 加载 Ascend C 编译的算子defforward(self,*inputs):returnself.op_handle(*inputs)三种并行策略让多张卡各干各的大模型太大了一张卡装不下得拆分到多张卡上。ATB 提供了三种拆法EP 并行Expert ParallelismMoE 模型有几百个专家网络但每次推理只激活其中几个。EP 并行就是把这些专家分配到不同卡上——谁被激活了就去哪张卡算没被激活的专家就安静待着。EP 并行的难点在于路由每个 token 经过 Gate 网络后会得到一个专家编号列表ATB 要把 token 送到对应的卡上算完再收回来。这个发送→计算→收集的过程依赖 hccl 的 AlltoAll 通信原语。# EP 并行8 张卡各管一批专家 # 就像 8 个科室医院病人按挂号去对应科室不用挤在一起 Card 0: Expert 0, 8, 16, 24, 32 Card 1: Expert 1, 9, 17, 25, 33 Card 2: Expert 2, 10, 18, 26, 34 ...# ATB 中配置 EP 并行importascend_transformer_boostasatb ep_configatb.ParallelConfig(parallel_typeep,world_size8,# 8 张卡ep_rank0,# 当前卡编号num_experts64,# 总专家数num_experts_per_token8# 每个 token 激活的专家数)runneratb.Runner(model_config,parallel_configep_config)EP 并行有一个特别好的性质计算量与激活的专家数成正比不激活的专家完全不占计算资源。所以 MoE 模型虽然参数量大实际推理时的计算量比同等参数的稠密模型小很多——EP 并行把这个优势完整地保留了下来。TP 并行Tensor Parallelism把一个大的矩阵乘法切成几块每张卡各算一块最后拼起来。就像把一张大拼图分给几个人各拼各的部分拼完一合就完事了。TP 并行的关键操作是 AllReduce每张卡算完自己那块后要把结果汇总。ATB 里这个通信由 hccl 在后台完成对上层推理逻辑透明。# TP 并行把大矩阵按列切分 # Card 0 算左半边Card 1 算右半边然后拼接 Weight [W_left | W_right] Card 0 Card 1 result concat(Card0_output, Card1_output)# ATB 中配置 TP 并行tp_configatb.ParallelConfig(parallel_typetp,world_size4,# 4 张卡做张量并行tp_rank0,# 当前卡编号backendhccl# 通信后端)runneratb.Runner(model_config,parallel_configtp_config)TP 并行有一个权衡通信量与切分数成正比。2 卡 TP 并行只做一次 AllReduce8 卡 TP 并行就要做更多次。所以在实际部署中TP 并行的卡数通常不超过 8再大就要考虑 EP 并行或者 TPEP 混合并行了。PageAttentionKV Cache 的虚拟内存推理时每个 token 的 KV Cache 大小不一样——有的长有的短。传统的做法是按最大长度预分配内存浪费严重。PageAttention 借用了操作系统虚拟内存的思路把 KV Cache 切成固定大小的页用多少分配多少。# 传统方式按最大长度分配大量浪费 # |████████████████████████████ | → 只用了一半 # PageAttention按需分配像内存分页一样 # |████|████|████|██ | → 用多少占多少 token_0: page [0] → 4K token_1: page [0, 1] → 8K token_2: page [0, 1, 2] → 12K # 释放的时候按页回收碎片自动管理PageAttention 的页表维护在 Host 侧CPU 内存每次 attention 计算时NPU 根据页表索引找到对应的物理页地址。这种间接寻址增加了一点查表开销但换来的是内存利用率的大幅提高——实测在长序列场景下KV Cache 内存占用可以降到原来的 30%-50%。# ATB 中启用 PageAttentionattn_configatb.AttentionConfig(attention_typepaged,page_size16,# 每页 16 个 token 的 KVmax_num_pages4096,# 最大页数dtypefloat16)# 创建带 PageAttention 的图算子paged_attnatb.FlashAttention(attn_config)PageAttention 还有一个好处多请求之间的 KV Cache 页可以共享。比如系统 prompt 对所有请求都一样那它的 KV Cache 只存一份多个请求通过页表指向同一组物理页就行——这在大规模推理服务中能省下可观的内存。ATB 在 CANN 生态里的位置搞清楚 ATB 在哪一层你就知道它该干什么活了┌─────────────────────────────────┐ │ 第1层AscendCL 编程接口 │ ← 上层应用通过这里调用 ├─────────────────────────────────┤ │ 第2层AOL 算子库 │ │ ├─ ops-transformer (底层算子) │ ← ATB 的原材料 │ ├─ ops-nn / ops-blas │ │ └─ hccl (集合通信) │ ← TP/EP 并行的快递公司 ├─────────────────────────────────┤ │ 第3层加速库 │ │ └─ ATB ◄── 就是这里 │ ← 把原材料组装成流水线 ├─────────────────────────────────┤ │ 第4层Runtime 运行时 │ ├─────────────────────────────────┤ │ 第5层硬件驱动 │ └─────────────────────────────────┘ATB 坐在 CANN 五层架构的第 3 层——加速库层。它不上不下恰好是搭桥的角色上面接推理框架的调用下面调 ops-transformer 的底层算子旁边跟 hccl 配合做分布式通信。这个位置决定了 ATB 的职责边界它不管算子怎么在 NPU 上执行那是 Runtime 的事不管底层算子的具体实现那是 ops-transformer 的事也不管模型怎么部署和服务化那是 cann-recipes-infer 的事。ATB 只管一件事把算子编排成高效的执行图让数据在图上尽量少搬运、尽量多复用。# ATB 与上下游的调用关系 cann-recipes-infer → ATB Runner → ops-transformer 算子 ↕ ↕ hccl 通信 graph-autofusion 融合 ↕ ge 图编译上游下游ATB 的朋友圈ATB 不是一个人在战斗它依赖一堆兄弟仓库也被一堆仓库依赖ops-transformer ──→ ATB ──→ cann-recipes-infer ↑ │ │ hccl(通信) graph-autofusion 推理部署 (算子融合) 直接使用 ATB │ ge(图编译)上游ATB 依赖谁ops-transformer提供 FlashAttention、RMSNormRoPECache、SwiGLU 等底层算子——这些是 ATB 流水线上的齿轮。ATB 的图算子编排的每个节点底层基本都调的是 ops-transformer 的算子实现。如果你发现某个算子的行为跟预期不一样大概率要去 ops-transformer 里查。hccl提供 AllReduce、AlltoAll、Broadcast 等集合通信能力——EP 并行和 TP 并行全靠它搬运数据。hccl 底层针对昇腾 NPU 的 HCCS 互联做了深度优化通信带宽利用率远高于通用的 TCP 传输。# 查看当前环境的 hccl 通信拓扑python-cimport hccl; print(hccl.get_comm_topology())# ATB 初始化分布式环境时的 hccl 配置exportHCCL_CONNECT_TIMEOUT7200exportASCEND_RT_VISIBLE_DEVICES0,1,2,3,4,5,6,7下游谁依赖 ATBcann-recipes-infer推理配方直接使用 ATB 部署 DeepSeek、Qwen、GLM 等大模型——这是普通开发者最直接的接触点。cann-recipes-infer 把 ATB 的调用链、模型权重加载、分布式配置都打包好了你只需要改改配置文件就能跑起来。# 用 cann-recipes-infer 快速部署一个模型cdcann-recipes-infer/models/DeepSeek-V3# 修改 run.sh 中的模型路径和并行配置bashrun.sh# 背后自动调用 ATB 的 Runner 和图算子关联仓库graph-autofusion算子自动融合框架能把 ATB 图算子里的多个基础算子进一步合并成单个融合算子。注意graph-autofusion 是融合不是调优——它不调参数只合并算子。少一次算子调度就少一次数据搬运这对于访存密集型的推理场景尤为关键。ge图引擎负责计算图的全局优化和编译下沉。ATB 的图算子在执行前可以下沉到 ge 做进一步优化——比如常量折叠、死代码消除、内存复用分析。ge 跟 ATB 的关系不是谁替代谁而是接力优化ATB 做应用层的算子编排ge 做编译层的图优化两层配合让整个推理链路更顺畅。# 查看 ATB 与 graph-autofusion 的协作状态# graph-autofusion 默认开启会自动扫描 ATB 图中的可融合模式exportATB_AUTOFUSION_ENABLE1# 开启自动融合默认开启exportATB_AUTOFUSION_LOG1# 打印融合日志方便调试记住这些就够了ATB 是什么昇腾 CANN 生态里的大模型加速库干的是把零碎算子打包成流水线的活。三层架构怎么记Operation 是积木块Graph Operator 是拼好的流水线Plugin 是你自己加的扩展工位。并行策略选哪个MoE 模型用 EP 并行大矩阵用 TP 并行KV Cache 省内存用 PageAttention。想动手直接去 cann-recipes-infer 找你用的模型里面已经集成好了 ATB 的完整调用链https://atomgit.com/cann/ascend-transformer-boost
