引言

在大模型训练与推理的算力竞赛中，昇腾系列芯片凭借 Ascend 910 的强劲性能崭露头角。然而，硬件算力只是基础，如何让 Transformer 这类核心模型架构在昇腾芯片上跑得更快、更稳、更易用，才是真正的工程挑战。ascend-transformer-boost（简称 ATB）正是为此而生的 Transformer 加速库，它构建在 CANN（昇腾异构计算架构）之上，为 PyTorch、MindSpore、Paddle 等主流框架提供了一层精致而高效的加速抽象。

理解 ATB 的三层架构，不仅是掌握一个加速库的使用方法，更是理解昇腾生态如何将软件抽象与硬件算力无缝衔接的关键。本文将深入剖析 ATB 从原子算子到融合算子，再到底层插件的三层架构设计，通过实际代码演示其在工程实践中的落地方式，并厘清 ATB 与 graph-autofusion、TorchAir 等组件的关系与边界。

---

第一层：基础原生算子——原子能力的直接暴露

ATB 的第一层架构是基础原生算子层，这一层直接封装了 CANN 算子库中的原子算子，通过 ops-transformer 模块对外提供调用接口。所谓"原子算子"，指的是不可再分的最小计算单元，例如矩阵乘法（MatMul）、LayerNorm、Softmax、RoPE（Rotary Position Embedding）等。

这一层的设计哲学是"薄封装、零歧义"。ATB 并没有在原子算子层面做过多的逻辑包裹，而是将 CANN 算子库的能力以最接近硬件的方式暴露出来。对于追求极致性能的高级开发者，可以直接调用这些原子算子，精确控制计算图的每一个节点，避免任何不必要的抽象开销。

以 PyTorch 框架为例，调用 ATB 的原子算子相当直观。以下是一个典型的调用片段：

01: import torch
02: import torch_npu
03: from atb_speed.lib.atb_speed_base import AtbSpeedBase
04: 
05: # 初始化 NPU 设备
06: torch_npu.npu.set_device(0)
07: 
08: # 创建 LayerNorm 算子实例
09: layer_norm_op = torch.ops.atb_ops.LayerNorm(
10:     normalized_shape=4096,
11:     eps=1e-6,
12:     elementwise_affine=True
13: )
14: 
15: # 构造输入张量（batch=2, seq_len=2048, hidden=4096）
16: hidden_states = torch.randn(2, 2048, 4096, device='npu', dtype=torch.float16)
17: 
18: # 直接调用 ATB 算子
19: output = layer_norm_op(hidden_states)

逐段解释：

• 第 01-03 行：导入必要的依赖。torch_npu 是 PyTorch 与昇腾 NPU 的设备适配层，atb_speed 是 ATB 提供的高阶加速封装库，AtbSpeedBase 是其基类。
• 第 06 行：设置当前使用的 NPU 设备编号。set_device(0) 对应第一张昇腾加速卡。
• 第 09-13 行：创建 LayerNorm 算子实例。这里直接通过 torch.ops.atb_ops 命名空间访问 ATB 的原子算子接口。normalized_shape=4096 指定归一化的维度，eps=1e-6 是数值稳定性参数，elementwise_affine=True 表示使用可学习的缩放和偏移参数。
• 第 16 行：在 NPU 设备上创建输入张量。device='npu' 是关键——这告诉 PyTorch 将张量分配到昇腾设备上，而非 CPU 或 CUDA。dtype=torch.float16 使用半精度浮点，这是大模型训练/推理的标准配置。
• 第 19 行：像调用普通 PyTorch 算子一样调用 ATB 算子。这种调用方式对 PyTorch 用户零学习成本，背后却是 ATB 将算子计算下沉到 Ascend 910 的 AICORE 上执行的完整流程。

基础原生算子层的价值在于：它为上层架构提供了不可再分的最小计算原语，同时保留了直接操控硬件的能力。当你需要实现一个 ATB 尚未预定义的特殊算子时，这一层就是你的起点。

---

第二层：图算子机制——预定义融合模式的力量

在大模型计算中，计算图往往由数十甚至数百个原子算子串联而成。以 Transformer 的一个标准 Attention 层为例，它通常包含 Query/Key/Value 投影、矩阵乘法、Softmax、Dropout、输出投影等十余个步骤。如果这些原子算子逐个在 Ascend 910 上调度执行，算子启动开销（kernel launch overhead）和片外内存读写将成为性能瓶颈。

ATB 的第二层架构——图算子机制，正是为解决这一问题而生。其核心思想是：将多个频繁连续出现的小算子预先融合为一个"融合算子"，在一次 kernel 调用中完成原本需要多次调度才能完成的计算。

这里需要明确一个关键概念：ATB 的图算子机制采用的是预定义融合模式，而非动态自动融合。所谓"预定义"，是指 ATB 的工程团队根据 Transformer 架构的典型计算模式，人工分析并实现了若干高频融合算子。例如：

• SelfAttention 融合算子：将 Query/Key/Value 投影、Attention 计算、输出投影融合为一个算子。
• MLP 融合算子：将 Feed-Forward Network 的两个线性层与激活函数（如 SiLU 或 GELU）融合。
• RMSNorm 融合算子：将 RMSNorm 的计算与相邻的矩阵乘法融合，减少片外内存访问。

这种预定义融合模式的优势在于可预测性与高性能。由于融合模式是固定的，ATB 可以对每一个融合算子进行精细的 tiling 策略和内存布局优化，充分发挥 Ascend 910 的 AICORE 并行计算能力。代价是灵活性相对受限——如果使用的模型架构偏离了预定义模式，可能无法直接从图算子机制中受益。

以下代码展示了如何通过 ATB 的图算子接口调用一个预定义的融合 Attention 算子：

01: from atb.models.base.attention import Attention
02: from atb.core.operation import Operation
03: 
04: # 定义 Attention 配置
05: attn_config = {
06:     "head_num": 32,
07:     "head_size": 128,
08:     "qkv_permutation": [0, 2, 1, 3],
09:     "kv_dim": 2,
10:     "pre_tokens": 65535,
11:     "next_tokens": 0,
12: }
13: 
14: # 创建 Attention 融合算子
15: attn_op = Attention(attn_config)
16: 
17: # 构建计算图（ATB 的图算子机制核心）
18: graph = Operation()
19: graph.add_node("query", "Linear", {"out_features": 4096})
20: graph.add_node("key", "Linear", {"out_features": 4096})
21: graph.add_node("value", "Linear", {"out_features": 4096})
22: graph.add_node("attention", attn_op, inputs=["query", "key", "value"])
23: graph.build()
24: 
25: # 执行计算图
26: query = torch.randn(2, 2048, 4096, device='npu', dtype=torch.float16)
27: key = torch.randn(2, 2048, 4096, device='npu', dtype=torch.float16)
28: value = torch.randn(2, 2048, 4096, device='npu', dtype=torch.float16)
29: output = graph.run({"query": query, "key": key, "value": value})

逐段解释：

• 第 05-12 行：定义 Attention 算子的配置参数。head_num=32 和 head_size=128 共同决定了注意力头的维度划分。qkv_permutation 指定 Query/Key/Value 的维度重排顺序，这是适配不同模型架构（如 GPT 系列 vs LLaMA 系列）的关键参数。pre_tokens 和 next_tokens 控制注意力窗口的大小，用于实现滑动窗口注意力等变体。
• 第 15 行：创建 Attention 融合算子实例。这里的 Attention 不是一个简单的 PyTorch 模块，而是 ATB 实现的一个融合算子，它在底层将 Query/Key/Value 投影、Attention 分数计算、Softmax、输出加权等步骤融合为一个 kernel。
• 第 18-23 行：这是 ATB 图算子机制的核心——通过计算图（Operation）将多个节点组织起来。add_node 方法添加计算节点，第一个参数是节点名称，第二个是算子类型，第三个是算子配置。build() 方法触发计算图的编译与优化，包括节点融合、内存复用分析、流水线调度等。
• 第 26-29 行：构造输入并执行计算图。graph.run() 启动计算，底层会将整个计算图作为一个执行单元调度到 Ascend 910 上，最大化设备利用率。

---

第三层：插件机制——自定义扩展的终极武器

预定义融合模式覆盖了大部分主流场景，但大模型领域的创新速度意味着总会有新的算子或融合模式超出 ATB 的预定义范围。为了应对这种不确定性，ATB 设计了第三层架构——插件机制（Plugin Mechanism），允许开发者以标准化的方式注入自定义算子实现。

插件机制的核心抽象是 Plugin 基类。任何继承自 Plugin 的自定义算子，只需要实现几个标准接口，就能无缝接入 ATB 的计算图体系，获得与内置算子同等的调度优化与内存管理能力。

以下是一个自定义 Plugin 的典型实现框架：

01: from atb.core.plugin import Plugin
02: from atb.core.context import Context
03: 
04: class CustomRoPEPlugin(Plugin):
05:     """自定义 RoPE（Rotary Position Embedding）插件"""
06:     
07:     def __init__(self, max_seq_len=2048, base=10000):
08:         super().__init__()
09:         self.max_seq_len = max_seq_len
10:         self.base = base
11:         self.rope_embed = None
12:     
13:     def initialize(self, context: Context):
14:         """插件初始化——在计算图编译阶段调用"""
15:         hidden_size = context.get_input_shape(0)[-1]
16:         self.rope_embed = self._build_rope_cache(hidden_size, self.max_seq_len)
17:         # 将 RoPE 缓存注册到上下文，供后续算子复用
18:         context.set_attr("rope_cache", self.rope_embed)
19:     
20:     def execute(self, inputs, context: Context):
21:         """算子执行逻辑——在推理/训练阶段调用"""
22:         hidden_states = inputs[0]
23:         seq_len = hidden_states.shape[1]
24:         rope_cache = context.get_attr("rope_cache")
25:         # 调用底层 CANN 算子完成 RoPE 计算
26:         output = torch.ops.cann.RoPE(hidden_states, rope_cache[:seq_len])
27:         return output
28:     
29:     def _build_rope_cache(self, dim, max_len):
30:         """预计算 RoPE 频率缓存"""
31:         inv_freq = 1.0 / (self.base ** (torch.arange(0, dim, 2).float() / dim))
32:         t = torch.arange(max_len).type_as(inv_freq)
33:         freqs = torch.outer(t, inv_freq)
34:         return torch.cat((freqs, freqs), dim=-1)
35: 
36: # 注册插件到 ATB 算子库
37: from atb.core.registry import register_plugin
38: register_plugin("CustomRoPE", CustomRoPEPlugin)

逐段解释：

• 第 04 行：自定义插件类继承自 Plugin 基类。这是 ATB 插件机制的标准做法。
• 第 13-18 行：initialize 方法在计算图编译阶段调用，用于一次性初始化操作。这里预计算了 RoPE 的位置编码缓存，并通过 context.set_attr 将其注册到计算图上下文中，避免每次推理时重复计算。Context 是 ATB 提供的跨算子状态共享机制，类似于一个计算图级别的全局字典。
• 第 20-27 行：execute 方法是插件的核心，在每次推理或训练时调用。inputs 是输入张量列表，context 提供跨算子状态访问。第 26 行调用了 CANN 的底层 RoPE 算子——这展示了插件机制的真正力量：你可以在自定义插件中直接调用 CANN 算子库的任何接口，包括那些尚未被 ATB 封装的实验性算子。
• 第 37-38 行：通过 register_plugin 将自定义插件注册到 ATB 的算子库中。注册之后，这个插件就可以像内置算子一样在计算图中通过名称引用。

插件机制的精妙之处在于它打破了 ATB 预定义算子与自定义算子之间的边界。一旦注册，自定义插件就成为了 ATB 计算图体系的一等公民，可以参与图算子融合、内存复用优化、流水线调度等所有高层优化。这种设计让 ATB 具备了近乎无限的扩展能力——无论未来的模型架构如何演进，开发者都可以通过插件机制快速适配。

---

ATB 与 graph-autofusion 的区别——预定义融合 vs 自动发现

在昇腾的软件栈中，除了 ATB 的图算子机制，还存在另一个与算子融合相关的组件——graph-autofusion。这两个概念经常被混淆，但它们的设计哲学与技术路径有着本质区别。

ATB 的图算子机制采用的是预定义融合模式。ATB 的工程师根据 Transformer 架构的典型计算模式，人工分析并实现了若干高频融合算子。这些融合算子是"写死"在 ATB 代码库中的，每一种融合模式都经过了精细的 hand-tuned 优化。

graph-autofusion 则走的是一条截然不同的路——自动发现融合机会。它不依赖人工预定义的融合模式，而是在计算图层面通过图模式匹配算法，自动识别可以融合的算子子图。例如，如果计算图中出现"Conv2D → BiasAdd → ReLU"这样的连续算子序列，graph-autofusion 会自动将其识别为一个可融合模式，并生成对应的融合 kernel。

两种方案各有优劣：

维度	ATB 图算子机制	graph-autofusion
融合模式来源	人工预定义	自动发现
优化程度	极高（hand-tuned）	较高（依赖编译器优化能力）
覆盖范围	聚焦 Transformer 核心计算模式	通用计算图模式
灵活性	受限于预定义模式	能自动适配新型计算图结构
典型应用场景	大模型训练/推理的关键路径	通用深度学习模型的算子融合优化

在实际的昇腾软件栈中，这两者并非竞争关系，而是互补协作。ATB 的预定义融合算子覆盖了 Transformer 模型中计算最密集、对性能影响最大的核心路径；graph-autofusion 则负责处理模型中非 Transformer 部分的通用算子融合，两者共同构成了昇腾模型的完整融合优化体系。

---

ATB 与 TorchAir 的关系——两条不同的桥接路径

另一个容易混淆的概念是 ATB 与 TorchAir 的关系。两者的定位都是 PyTorch 与昇腾之间的桥接层，但它们所处的抽象层级和优化策略有显著差异。

TorchAir 位于 PyTorch 和 GE（图引擎）之间，它的核心工作是将 PyTorch 的动态计算图"下沉"为 CANN 的静态计算图，然后交给 GE 进行图优化。这意味着 TorchAir 对 PyTorch 的用户代码几乎不做任何修改——你写的 torch.nn.Linear 和 torch.nn.LayerNorm，TorchAir 会自动将它们映射到 CANN 的对应算子上。TorchAir 的优势是零侵入性，劣势是无法利用融合算子的极致性能，因为融合发生在 GE 层，无法感知 ATB 预定义的高性能融合模式。

ATB 则直接位于 PyTorch 算子调用层之下。当你在 PyTorch 中调用 torch.ops.atb_ops.LayerNorm 时，这个调用完全绕过了 GE，直接在 CANN Runtime 上调度 ATB 优化过的算子。这意味着 ATB 可以提供比 TorchAir 更极致的性能，但代价是需要显式调用 ATB 接口，PyTorch 原生的 torch.nn.LayerNorm 不会被自动替换为 ATB 版本。

这两者的协作逻辑是：对于非核心路径的算子，使用 TorchAir 自动映射，享受零侵入的便利；对于 Transformer 的核心计算路径（Attention、MLP、Norm 等），显式使用 ATB 接口，获得极致的融合算子性能。在实际的大模型部署中，常见的做法是 TorchAir 作为默认桥接层，ATB 作为关键算子的性能补丁。

---

三层架构的实际工程价值

ATB 三层架构的设计，体现了昇腾工程团队对"性能"与"灵活性"这对经典矛盾的深刻理解。

第一层给了极致玩家直接操控硬件的能力；第二层让大多数开发者能够以零成本的方式享受融合算子带来的性能红利；第三层确保了 ATB 不会因为预定义模式的局限性而成为创新路上的绊脚石。

对于在昇腾上做推理服务部署的工程师而言，理解了这三层架构，就等于拥有了一把万能钥匙：你可以从第一层出发，精确定位性能瓶颈；可以用第二层快速构建高性能计算图；可以在第三层按需扩展，而不必等待 ATB 官方的算子更新。

这才是好的软件架构该有的样子——既给了足够深的底层能力，又不强迫每个用户都去学底层。

如果你在实际项目中用到了 ATB 的某一层，不妨对比思考一下：这一层解决的是哪类问题？如果让你自己实现同样的能力，需要多少工作量？这种对比会帮助你更深刻地理解昇腾软件栈的工程哲学。