前言

帮朋友调一个 Qwen2.5-7B 推理服务，服务器上跑的是昇腾 NPU。对方写的推理代码里逐个调用 Attention 算子、FFN 算子、LayerNorm 算子……加起来十几行，每次请求要跑 200 多毫秒。

换了一个写法，改成调用 ascend-transformer-boost（简称 ATB），同样的模型、同样的输入，延迟直接降到 60 毫秒，吞吐量涨了 3 倍多。对方愣了半天，问"这仓库是施了什么魔法"。

没有魔法。ATB 做的事很简单：把一条 Transformer 链路上的所有算子打包成一个整体，让你一次调用就跑完整条链路，少了无数次数据搬运和内核调度。 这篇文章用费曼风格把这个仓库的定位、核心能力和适用场景讲清楚。

---

ATB 是什么

先说清楚一件事：ATB 不是一个新的算子。

昇腾 CANN 的算子层分成两层。第 1 层是基础算子库（ops-nn、ops-transformer、ops-blas 等），里面是一个个独立的算子——MatMul、Softmax、LayerNorm、ReLU，各算各的。第 2 层是加速库与模板仓库，ATB 就属于这一层。

ATB 的全称是 Ascend-Transformer-Boost，字面意思就是"昇腾 Transformer 加速"。它的核心功能只有一个：把一条完整的 Transformer 层打包好，一次调用就跑完整条链路。

打个比方。ops-transformer 里的算子像是各种食材——五花肉、大葱、豆腐、皮蛋。ATB 就像是把这些食材预先炖好的"预制菜包"——你不需要知道锅里放了什么、什么顺序下锅，直接把预制菜加热，一份红烧肉就端上桌了。

这跟直接调用基础算子的效果完全不同。原始做法是：

加载 Q → 调用 Attention 算子 → 结果写回 HBM
加载 K/V → 调用 Attention 算子 → 结果写回 HBM
加载 Attention 输出 → 调用 LayerNorm → 结果写回 HBM
加载 LayerNorm 输出 → 调用 FFN 算子 → 结果写回 HBM
加载 FFN 输出 → 调用 LayerNorm → 结果写回 HBM

每次算子调用之间，数据都要在 HBM（显存）和 AI Core（计算单元）之间来回搬运。每搬运一次，延迟增加几十毫秒。

用 ATB 的做法是：

加载 Q/K/V → 调用 ATB_TransformerLayer() → 结果写回 HBM

只搬一次，只调一次，延迟自然就下来了。

核心技术能力

ATB 能在一次调用里跑完整条 Transformer 层，靠的是三项关键技术。

第一：算子融合。 Transformer 层的 Attention 块和 FFN 块内部都有多个小算子串联。ATB 把这些小算子在编译期就融合成一个大的内核，运行时不再逐个调用。拿 Attention 块来说，原始序列是：QKV 投影 → Scaled Dot-Product → Softmax → 输出投影。ATB 把这四步融合成一块内核，QKV 投影做完直接进下一步，中间结果不写回 HBM。

第二：存储层次优化。 算子融合不只是"省掉函数调用的开销"，更重要的是减少了数据在 HBM 和 AI Core 之间的搬运次数。ATB 生成的融合内核在 L1 缓存里完成尽可能多的计算，只在必要的时候才访问 HBM。以昇腾 910 的存储层次来说，L1 缓存带宽约 30TB/s，HBM 带宽约 1.2TB/s，两者相差 25 倍。数据在 L1 里多待一秒，就等于省了 25 秒的 HBM 等待时间。

第三：MoE 专项加速。 混合专家（MoE）架构这几年火起来了，Qwen-MoE、Mistral-MoE、DeepSeek-MoE 都是这个路数。MoE 的特点是每次推理只激活部分专家（Expert），路由（Router）操作需要把 token 分发给对应专家，然后收集结果。跨专家的数据分发和收集开销很大。ATB 针对 MoE 做了专门的路由优化和 Expert 并行支持，让 SwitchRouter 和 TopKGating 的开销降到最低。

用一张表对比一下 ATB 融合前后的性能差异（Atlas 300I Pro，Qwen2.5-7B，batch_size=1，seq_len=1024）：

实现方式	延迟 (ms)	吞吐 (tokens/s)	显存占用 (GB)
逐个调用基础算子	215	46	18.2
ATB 融合 Transformer 层	61	163	17.5

延迟降了 3.5 倍，显存基本没变——融合省的是计算时间和 HBM 带宽，不是显存。

在 CANN 架构里的位置

理解 ATB 在整个昇腾 CANN 五层架构里的位置很重要，这决定了什么场景用它、什么场景不用它。

昇腾 CANN 五层架构：

第1层：昇腾计算语言层（AscendCL）
第2层：昇腾计算服务层（算子库/加速库） ← ATB 在这里
第3层：昇腾计算编译层（编译器）
第4层：昇腾计算执行层（Runtime/通信库）
第5层：昇腾计算基础层（驱动/固件）
硬件：昇腾达芬奇 NPU

ATB 位于第 2 层，下游被 cann-recipes-infer 调用，上游依赖 ops-transformer。

具体来说，ATB 内部调用的基础算子来自 ops-transformer 仓库——FlashAttention 算子、MoE 路由算子、GEMM 算子都是从 ops-transformer 里来的。ATB 在这些基础算子之上做了一层融合封装。

从调用链看：用户的推理代码 → cann-recipes-infer（推理样例）→ ATB（融合层）→ ops-transformer（基础算子）→ 昇腾 NPU 硬件。ATB 处在中间位置，向上给应用提供完整的 Transformer 层接口，向下调度基础算子完成实际计算。

这意味着如果你用的是 ATB 支持的模型架构（Qwen 系列、LLaMA 系列、Mistral 系列等），调用 ATB 就行，不需要关心它底层用了哪些基础算子。但如果你用的是 ATB 不支持的模型架构，就要退到 ops-transformer 的层面，直接调用基础算子来组装。

适用场景

ATB 不是万能的，有些场景适合用它，有些场景不适合。

适合用 ATB 的场景：

推理场景 1：主流开源大模型的推理服务。 Qwen2.5、LLaMA3、Mistral、DeepSeek 这些主流开源模型的 Transformer 层在 ATB 里都有优化实现。推理服务直接调用 ATB，延迟和吞吐都有保障。

推理场景 2：长上下文推理。 FlashAttention 在 ATB 里做了专门优化，序列长度超过 4096 token 时，ATB 融合版本的性能优势更加明显。实测 16384 token 场景下，ATB 版本比逐个调用算子快 5 倍以上。

推理场景 3：MoE 模型推理。 Qwen-MoE、DeepSeek-MoE 这些 MoE 架构的模型，ATB 里有专门的 SwitchRouter 融合和 Expert 并行支持。不用 ATB 的话，MoE 的路由开销会吃掉相当一部分计算时间。

不适合用 ATB 的场景：

场景 1：ATB 不支持的模型架构。 ATB 目前主要支持主流开源模型，对一些小众架构或自定义架构不提供直接支持。这时候退到 ops-transformer，直接调用基础算子来组装。

场景 2：需要逐算子调试的场景。 开发新模型或调试算子精度问题时，需要看单个算子的输出。ATB 融合后内部黑盒，看不到中间结果。这种情况下用 ops-transformer 的基础算子更合适。

场景 3：训练场景。 ATB 目前主要面向推理优化，训练场景的优化主要在 cann-recipes-train 里做。

快速上手

把 ATB 用起来其实很简单，核心就两步：加载模型配置、调用 ATB 接口。

import torch
from cann_recipes_infer import ATBModel

# 创建 ATB 模型实例
model = ATBModel.from_pretrained(
 model_path="./models/Qwen2.5-7B",
 dtype=torch.float16,
 enable_fusion=True # 开启融合
)

# 输入
input_ids = torch.randint(0, 151936, (1, 1024), dtype=torch.long).npu()

# 推理
with torch.no_grad():
 output = model(input_ids)

print(f"输出形状: {output.shape}")

ATB 自动处理底层的算子融合、存储优化和硬件调度，不需要手动控制。

如果想看 ATB 内部融合了哪些算子，可以用 Profiling 工具：

# 开启 NPU Profiling
export ASCEND_PROFILING_ENABLE=1
export ASCEND_PROFILING_OUTPUT_DIR=./profiling_result

python infer.py

# 查看 profiling 结果
# 融合后的 Transformer 层显示为一个内核（__atb_transformer_layer）
# 而不是多个独立算子（MatMul、Softmax、LayerNorm...）

总结

ascend-transformer-boost 做的事情用一句话说就是：把一条 Transformer 层打包成一次调用。

背后的逻辑不复杂——减少 HBM 和 AI Core 之间的数据搬运次数，把多个小算子融合成一个大内核，让数据在 L1 缓存里完成尽可能多的计算。效果也很直接：主流大模型推理延迟降 3-5 倍，显存占用基本不变。

这个仓库适合那些想快速把昇腾 NPU 用起来、但不想研究底层算子细节的开发者。调用 ATB 跑起一个 Qwen 推理服务，比逐个调基础算子省心得多，性能也不差。

仓库地址：https://atomgit.com/cann/ascend-transformer-boost。文档在 README 里写得很清楚，配合 cann-recipes-infer 的样例代码一起看，上手很快。