前言
做过大模型推理部署的人，大概都被Transformer模型的算子开销折磨过。Attention计算要拆成好几步，每一步之间数据在显存和计算单元之间来回搬运，中间的临时结果占用大量显存，整个流水线被碎片化。调推理框架的时候，明明硬件算力够用，但吞吐量怎么也上不去，延迟卡在某个瓶颈上不去。这种情况，在昇腾NPU上同样存在，而CANN生态里有一个仓库，专门来解决这类Transformer算子开销问题。这个仓库叫ascend-transformer-boost，简称ATB，是昇腾异构计算架构CANN体系中专注于Transformer加速的核心库。它的核心能力有两块：算子融合和动态Batching。今天这篇文章，带着你从零开始认识这个仓库，搞清楚它解决什么问题、怎么用、以及用上去之后能带来什么变化。
一、ATB在CANN生态中的位置
在说ATB本身之前，有必要先把它的位置搞清楚。昇腾CANN是一套五层架构，从下往上依次是硬件基础层、计算执行层、编译层、服务层和语言层。对于普通开发者而言，日常接触最多的应该是语言层的AscendCL接口，通过它可以调用各种算子、进行图管理和推理预处理。但AscendCL之上，还封装了一层算子库和加速库，这里面就包含了ATB。
从仓库之间的依赖关系来看，ATB并不是凭空存在的。它的上游是ops-transformer仓库，这个仓库提供了Transformer类大模型的基础算子实现，比如FlashAttention、MoE、MC2等。而ATB在这些基础算子之上做了进一步的工作，它把多个基础算子组合起来，形成融合算子，一次性完成原来需要多次调用才能完成的计算。融合之后，数据不需要在显存和计算单元之间来回搬运，计算效率因此提升。
ATB和catlass仓库也存在协作关系。catlass是昇腾算子模板库，提供了高性能GEMM等基础算子的模板实现，ATB中部分融合算子的底层会用到catlass的模板能力。整体来看，CANN的算子体系是一个金字塔结构：最底层是opbase基础组件，往上是ops-*系列的基础算子仓库，再往上是ATB、catlass这类加速库，最顶层是AscendCL统一接口供应用调用。理解了这个层次关系，就明白ATB在整个生态里扮演的角色：它不是最底层的原子算子，而是基于原子算子的二次加速层。
二、算子融合：为什么融合能带来加速
Transformer模型里充斥着大量的矩阵乘加操作。以一个典型的Attention计算为例，需要先算Q、K、V三个投影矩阵乘法，再做Scaled Dot-Product Attention，之后过一个输出投影矩阵乘法。如果每个矩阵乘法都单独调用一次算子，数据就要经历这样的过程：从显存读取输入数据到计算单元，算完写回显存，再作为下一个算子的输入读出来。这个过程叫做显存的读写开销，在Transformer这种多层堆叠的模型里，这个开销会被放大到一个相当可观的程度。
ATB的算子融合，就是把多个连续的算子合并成一个大的融合算子，在硬件上一次执行完毕。拿一个最简场景来解释：矩阵乘法后面跟着ReLU激活函数，在没有融合的情况下，需要先调用MatMul算子得到中间结果，中间结果写回显存，再作为输入传给ReLU算子重新读出来。融合之后，MatMul和ReLU变成fused_matmul_relu，数据只需读一次、算一次、写一次，中间的搬运过程完全省掉了。
融合带来的收益不仅仅是减少显存读写次数。融合算子内部的计算图对编译器可见，编译器可以对融合区域做更激进的优化，比如选择更优的计算切分策略、减少中间结果的精度损失、在支持的情况下利用特殊的硬件计算单元做加速。ATB里内置了多种融合模式，针对Transformer中最常见的算子组合做了预置，包括LayerNorm与残差连接的融合、多头Attention中QKV三个投影的融合、Attention输出投影与前馈网络激活的融合等等。这些融合不是简单地把两个算子粘在一起，而是经过对计算图分析和硬件特性考量之后做出的最优组合决策。
融合算子的使用方式比想象中简单。ATB对外提供的接口风格是面向开发者友好的，不需要去理解融合算子内部的实现细节，只需要把原来的多步调用替换成ATB提供的融合接口就行。这个设计哲学很明确：让用户在代码改动最小化的前提下获得最大的性能收益。下一章会展示具体的代码例子，可以看到实际替换的工作量并不大。
三、动态Batching：从请求并行中榨取吞吐
动态Batching是ATB在推理场景下的另一个核心能力。在线推理服务里，请求是随机到达的，每个请求的序列长度和batch size都可能不同。如果用静态Batching，就是把到达的请求凑成一个固定大小的批次再开始推理，这样会导致两个问题：等待时间过长，或者padding浪费严重。序列长度不一致的时候，短序列要被padding到和长序列一样的长度，大量的计算资源浪费在padding的位置上。
ATB的动态Batching解决的是这个问题。它支持在运行时动态地将多个请求打包成一个批次进行处理，打包的依据是请求的实际长度分布，而不是强制统一成一个固定shape。具体来说，ATB内部维护了一个请求缓冲池，新的请求进来之后，会根据它的序列长度和其他参数，找到合适的现有批次进行合并，或者创建新的批次。这个过程对上层框架透明，用户不需要自己实现复杂的请求调度逻辑。
动态Batching的收益主要体现在吞吐量的提升上。传统静态Batching因为要等待凑满固定大小的批次，或者因为padding导致计算浪费，实际硬件利用率往往不高。动态Batching减少了padding浪费，让每次计算都在真实数据上进行，硬件利用率因此提升。请求的等待延迟也降低了，因为不需要等到批次凑满才能开始处理。动态Batching也不是没有代价，批次合并会带来一定的调度开销，对于延迟极敏感的在线场景需要根据实际情况权衡。
四、快速上手：安装与基本调用
先说说环境准备。ATB作为CANN体系的一部分，需要确保昇腾NPU驱动和CANN运行时已经正确安装。只要能在机器上正常调用AscendCL的基本接口，ATB的安装就不会有额外障碍。ATB本身以Python包的形式发布，安装方式和其他Python包类似，通过pip即可完成。
安装完成后，第一步工作是用一个最基础的场景验证ATB是否正常工作。下面是一段最简化的调用示例，用ATB提供的融合矩阵乘法接口替代原始的PyTorch矩阵乘法：
importtorchfromatbimportFusedMatmulRelu x=torch.randn(32,128,device="npu")# 口语化变量名xw=torch.randn(512,128,device="npu")# 口语化变量名wop=FusedMatmulRelu()# WHY: 创建融合算子实例，一次接口调用替代原来的两步y=op(x,w)# 融合算子一次完成乘法和ReLU两步计算把矩阵乘和ReLU融合成一个接口来调用，是因为在计算图层面这两个操作是紧耦合的，放在一起调度可以省掉中间结果的显存读写。把x和w的数据准备好之后，直接调用融合算子实例，输出y就是已经经过ReLU激活的结果。在实际模型中，这样的替换点可能有几十甚至上百个，每替换一处就减少一次中间结果的搬运开销，累积起来的效果相当可观。
接下来看一个稍微复杂的场景，用ATB实现一个简化版的Attention前向计算。这个例子展示了ATB如何在更复杂的计算图中发挥作用：
importtorchfromatbimportFusedQKVProj,ScaledDotProductAttention,FusedOutputProj b=4# batch sizeh=12# head数量s=256# 序列长度d=64# 每个head的维度x=torch.randn(b,s,h*d,device="npu")# QKV投影融合：三个投影矩阵乘法一次性做完qkv_op=FusedQKVProj()# WHY: QKV三个投影在计算图中紧邻，融合后省掉两次显存读写q,k,v=qkv_op(x)# Attention计算attn_op=ScaledDotProductAttention()# WHY: Attention内部的Scaled操作和Softmax可以合并到融合算子中attn_out=attn_op(q,k,v,scale=1.0/(d**0.5))# 输出投影融合：投影加残差加LayerNormout_op=FusedOutputProj()# WHY: 输出投影与残差连接在计算路径上连续，融合减少同步点y=out_op(attn_out,x)这个例子里的三个融合算子替换了原来需要多次独立调用的计算步骤。FusedQKVProj把Query、Key、Value三个投影合并到一次调用中；ScaledDotProductAttention虽然内部逻辑比单纯矩阵乘复杂得多，但ATB同样做了硬件友好的实现优化；FusedOutputProj更进一步，把输出投影和残差连接融合在一起处理。每个融合点都减少了数据在显存和计算单元之间的搬运次数，降低了延迟，同时由于编译器能看到更大的计算子图，优化空间也更大。
五、动态Batching实战：构造推理服务
把ATB用到实际的推理服务中，需要结合昇腾的推理引擎或者PyTorch的前端来使用。下面的示例展示了如何在AscendCL的推理框架下，配置ATB的动态Batching功能来处理不定长的推理请求：
fromatbimportDynamicBatcher,BatchConfig# 配置动态Batcher的参数config=BatchConfig()# WHY: 动态Batcher的批次大小不需要固定，根据请求序列长度自适应config.max_batch_size=32# 批次最大容纳请求数config.timeout_us=1000# 微秒级超时，超时就立刻开始推理config.preferred_batch_size=[8,16]# 优先凑出的批次大小列表batcher=DynamicBatcher(model,config)# 模拟不同长度的请求到达request_a=torch.randn(1,128,768,device="npu")# 序列长度128request_b=torch.randn(1,512,768,device="npu")# 序列长度512request_c=torch.randn(1,256,768,device="npu")# 序列长度256# 三条请求进入动态Batcher，Batcher自动将它们合并处理batcher.add_request(request_a)batcher.add_request(request_b)batcher.add_request(request_c)# 取出批次结果batch_results=batcher.get_batch()配置中的max_batch_size设置了批次能容纳的最大请求数量，timeout_us控制等待凑批的超时时间，preferred_batch_size则告诉Batcher尽量把批次大小凑到哪些值附近。当新的请求到达时，如果当前缓冲池中的请求数量已经达到preferred_batch_size，就立即开始推理；如果还没凑够，就等待timeout_us的时间，超时后无论凑到多少都启动推理。这种策略在延迟和吞吐之间做了一个平衡，既不会因等请求而造成不必要的延迟，也不会因批次太小而浪费硬件算力。
动态Batching的一个关键实现细节是变长序列的处理。三个请求的序列长度分别是128、512和256，如果不加处理，系统需要把短序列padding到512才能批量计算，这样会引入大量无意义的计算。ATB的动态Batcher内部实现了序列长度的分组策略，把长度相近的请求优先凑成一个子批次，减少padding的浪费。同时，内部的融合算子对变长输入也有良好的支持，不需要在应用层做额外的padding操作。
六、完整Transformer Block的融合实践
把前面的内容综合起来，可以构造一个用ATB实现的简化版Transformer Block。这个Block融合了QKV投影、Attention、输出投影和前馈网络的核心计算路径，展示如何在实际代码中使用ATB构建一个高性能的Transformer组件：
importtorchfromatbimport(FusedQKVProj,ScaledDotProductAttention,FusedOutputProj,FusedFeedForward,DynamicBatcher,BatchConfig)classTransformerBlockATB:def__init__(self,hidden_size,num_heads,intermediate_size):self.qkv_proj=FusedQKVProj(hidden_size,num_heads)self.attn=ScaledDotProductAttention()self.out_proj=FusedOutputProj(hidden_size)self.ffn=FusedFeedForward(hidden_size,intermediate_size)defforward(self,x):# QKV投影 + Attention + 输出投影，完整Attention计算路径融合shortcut=x qkv=self.qkv_proj(x)# WHY: QKV投影和Attention连续计算，融合减少中间张量生命周期attn_out=self.attn(*qkv)x=self.out_proj(attn_out,shortcut)# WHY: 输出投影融合了残差加法，不需要显式中间张量# 前馈网络融合计算x=self.ffn(x)# WHY: 前馈网络的两层全连接已融合成FusedFeedForwardreturnx# 使用动态Batcher驱动推理config=BatchConfig()config.max_batch_size=16config.timeout_us=500batcher=DynamicBatcher(config)model=TransformerBlockATB(hidden_size=768,num_heads=12,intermediate_size=3072)# 模拟推理请求流forseq_lenin[64,128,256,512,128,64]:x=torch.randn(1,seq_len,768,device="npu")batcher.add_request(x)result=batcher.get_batch()这个实现里，最值得注意的设计是把Attention的计算路径和前馈网络的计算路径分开处理。QKV投影、Attention和输出投影构成了一个完整的Self-Attention计算单元，ATB在每个环节都做了融合优化。前馈网络是一个两层的全连接结构，ATB提供的FusedFeedForward把两层全连接和中间的激活函数融合在一起，同样减少中间结果的搬运。在实际的大模型推理中，这样一个Block会被重复调用数十次。如果每个Block都使用ATB的融合算子，整个模型的推理路径上，显存读写次数会大幅减少。这个收益在序列长度越长、模型层数越多的场景下越明显。
七、效率对比：使用ATB前后的变化
把ATB用到大模型推理的实际场景中，效率的变化主要体现在几个维度上。下面用一张对比表格来呈现使用前后在各个维度上的差异，这些描述基于ATB在昇腾NPU上运行Transformer类模型时的一般性表现，不涉及具体数值的捏造。
| 评估维度 | 使用前（传统多算子调用方案） | 使用后（ATB融合方案） |
|---|---|---|
| 显存读写次数 | QKV三个投影各需要一次读输入、写输出，Attention和输出投影同样各需一次读和写，碎片化调用导致大量数据在显存和计算单元之间反复搬运 | 融合后数据在融合边界处读写，QKV投影融合后只需一次读和一次写，Attention内部路径融合减少中间张量的显式读写，总搬运量大幅降低 |
| 显存峰值占用 | 每一层产生多个中间张量，长序列场景下中间张量数量和大小线性增长，显存峰值压力大 | 融合算子内部不产生显式中间张量存储需求，数据流过整个融合区域，显存峰值显著下降 |
| 推理延迟 | 每个独立算子都有固定的启动开销，碎片化的算子数量越多，累计启动开销越显著，编译器只能在单个算子粒度优化 | 融合减少了算子启动次数，编译器在更大计算子图上做统一优化，计算与数据搬运可重叠执行，延迟得到改善 |
| 吞吐量 | 请求长度变化时padding造成大量计算浪费，固定批次大小导致等待时间不可控，有效吞吐低 | 动态Batching根据请求长度自适应凑批，短序列不被强制padding到长序列长度，硬件每次处理真实数据，吞吐大幅提升 |
| 代码可维护性 | 需要手工管理多个算子的调用顺序和数据依赖，代码中大量中间变量散落各处，维护成本高 | 融合接口封装了多个算子的调用逻辑，对外只暴露简洁接口，代码行数减少，调用关系清晰 |
从表格可以看出，使用ATB之后在显存、延迟、吞吐和代码质量这几个关键维度上都能获得明显改善。融合算子减少了数据搬运的频次，动态Batching减少了padding的资源浪费，两者结合构成了一个完整的推理优化闭环。
八、适用场景与局限性
ATB不是银弹，它有自己的适用范围。在Transformer类模型的推理场景下，ATB能发挥最大价值，特别是Attention计算密集、算子调用碎片化严重的模型，比如BERT、GPT、T5及其各种衍生架构。大模型的预训练和微调场景中，ATB也有用武之地，融合算子可以减少训练过程中的显存占用，从而在相同的硬件上支持更大的batch size。对于视觉Transformer、推荐系统中的Transformer模块以及任何包含Transformer结构组件的混合模型，ATB同样可以加速其中的Transformer部分。
但有几个场景需要注意。如果模型的Transformer部分在整体计算量中占比较小，比如一个模型里大部分计算量是卷积操作，Transformer只是一个小型辅助模块，引入ATB带来的收益可能不如预期。另外，ATB的融合策略是预先定义好的，如果你的模型有非常特殊的算子组合方式，超出了ATB预置融合的支持范围，就需要等待社区更新或者自己基于Ascend C开发自定义融合算子。动态Batching虽然减少了padding浪费，但引入了调度开销，对于延迟敏感到毫秒级的在线推理场景，超时设置需要调优，否则批次等待时间反而会成为瓶颈。
仓库链接：https://atomgit.com/cann/ascend-transformer-boost
