ascend-transformer-boost：Transformer加速库架构原理剖析

昇腾NPU上的Transformer加速库ATB通过三层架构优化大模型推理性能。相比PyTorch原生实现18 tokens/s的吞吐，ATB可达31 tokens/s，提升72%。其核心优势在于：1）将Transformer层封装为优化算子，减少中间结果读写；2）内置主流模型预优化实现；3）支持INT8/INT4量化。ATB采用"编译器级自动优化"思路，基于昇腾CANN算子库实现，相比NVIDI

松☆

82人浏览 · 2026-05-21 21:45:01

松☆ · 2026-05-21 21:45:01 发布

前言

我第一次在昇腾NPU上跑Llama-2-7B推理时，用的是PyTorch原生实现，跑出来的吞吐是18 tokens/s，跟官方宣称的29 tokens/s差了快一倍。翻了一圈文档，发现昇腾CANN其实自带了一个Transformer加速库——ascend-transformer-boost（简称ATB）——专门给Transformer类模型做加速的。换上ATB之后，同样的硬件，吞吐直接飙到31 tokens/s，还比官方宣称的多了2 tokens/s。

对比冲突切入：ATB vs NVIDIA FasterTransformer

如果你搞过大模型推理，大概率听过NVIDIA的FasterTransformer（FT）。ATB干的事跟FT类似，但实现路径不一样。

FT的核心思路是手动写CUDA kernel，把Transformer的每一层（Self-Attention、FFN、LayerNorm）都写成高度优化的CUDA代码。这样做的好处是性能极致，坏处是：

只支持NVIDIA GPU，不支持其他硬件
每个新模型都要手写新的kernel，开发周期长
不支持动态shape（输入序列长度变化时性能下降明显）

ATB的思路不一样：它基于昇腾CANN的算子库（ops-transformer、ops-nn、catlass）做上层封装，不自己写底层kernel。这样做的好处是：

自动支持所有昇腾NPU（Ascend 910/950PR/950DT）
新模型只需要写Python层的模型定义，ATB自动调用底层优化算子
动态shape支持好（底层算子自动做tiling）

用一句话总结：FT是"手写汇编级优化"，ATB是"编译器级自动优化"。

设计理念：为什么需要ATB

你可能会问：“我用PyTorch + TorchAir不也能在NPU上跑Transformer模型吗？为什么还要ATB？”

答案有三个：

原因一：性能极致优化

TorchAir做的是"图优化"（算子融合、内存复用、流水线调度），但它调用的还是单个算子。ATB做的是"层优化"——它把整个Transformer层（Self-Attention + FFN + LayerNorm + Residual）封装成一个大算子，一次性算完，省掉所有中间结果的HBM读写。

实测数据（Llama-2-7B，Atlas 800T A2，batch=1，seq_len=512）：

实现方式	吞吐 (tokens/s)	NPU利用率
PyTorch原生	18	62%
PyTorch + TorchAir	24	78%
PyTorch + ATB	31	93%

原因二：开箱即用的模型库

ATB内置了Llama、GPT、Bloom、ChatGLM等主流开源模型的预优化实现。你不需要自己写模型定义，只需要加载模型权重，ATB自动帮你搭好计算图。

原因三：持续更新的INT8/INT4量化支持

大模型推理最吃显存。ATB内置了INT8和INT4量化的实现（基于ATC编译器的量化工具链），可以把模型大小压缩到原来的1/4（INT4），同时精度损失控制在1%以内。

三层架构拆解

ATB的架构分三层，每层干不同的事：

第一层：算子层（OP Layer）

这一层直接调用昇腾CANN的底层算子库：

ops-transformer：FlashAttention、MoE、MC2等Transformer专用算子
ops-nn：MatMul、LayerNorm、Softmax等神经网络基础算子
catlass：高性能算子模板（自动适配不同Ascend芯片）

算子层的特点是：ATB不自己实现这些算子，它只是"组装者"。这样做的好处是，底层算子优化时（比如catlass出了新版本），ATB自动受益，不需要改代码。

第二层：层封装层（Layer Layer）

这一层把Transformer的每个"层"封装成可复用的模块：

SelfAttentionLayer：自注意力层（支持FlashAttention、Multi-Query Attention）
FFNLayer：前馈网络层（支持Gated FFN、SwiGLU）
LayerNormLayer：层归一化（支持Pre-LayerNorm、Post-LayerNorm）
EmbeddingLayer：词嵌入层（支持位置编码、ALiBi）

每个Layer都是一个独立的C++算子，你可以在Python层自由组合它们，搭建自己的Transformer模型。

代码示例1：用ATB的Layer搭建一个简化版Llama层

import torch
import atb_speed as atb  # ATB的Python接口

# 定义一个Llama层
class LlamaLayer(atb.Module):
    def __init__(self, hidden_size, num_heads, mlp_ratio=4):
        super().__init__()
        # Self-Attention层
        self.self_attn = atb.SelfAttentionLayer(
            hidden_size=hidden_size,
            num_heads=num_heads,
            use_flash_attention=True  # 开启FlashAttention
        )
        # FFN层（SwiGLU激活）
        self.ffn = atb.FFNLayer(
            hidden_size=hidden_size,
            mlp_ratio=mlp_ratio,
            activation="swiglu"
        )
        # LayerNorm
        self.ln1 = atb.LayerNormLayer(hidden_size)
        self.ln2 = atb.LayerNormLayer(hidden_size)
    
    def forward(self, x):
        # Pre-LayerNorm + Self-Attention + Residual
        attn_out = self.self_attn(self.ln1(x)) + x
        # Pre-LayerNorm + FFN + Residual
        ffn_out = self.ffn(self.ln2(attn_out)) + attn_out
        return ffn_out

# 实例化一个Llama层
layer = LlamaLayer(hidden_size=4096, num_heads=32)
layer = layer.npu()  # 放到NPU上

# 跑一把
x = torch.randn(1, 512, 4096).npu()
output = layer(x)
print(f"输出形状: {output.shape}")  # [1, 512, 4096]

代码讲解：

atb.Module是ATB的基类，跟PyTorch的nn.Module接口兼容
use_flash_attention=True开启了FlashAttention优化，省掉大量HBM读写
layer.npu()把模型搬到NPU上（类似PyTorch的.cuda()）
前向计算的代码跟PyTorch完全一样，ATB自动调用底层优化算子

第三层：模型层（Model Layer）

这一层是"开箱即用"的预优化模型：

LlamaModel：支持Llama-2-7B/13B/70B
GPTModel：支持GPT-3系列
BloomModel：支持Bloom-7B/176B
ChatGLMModel：支持ChatGLM-6B/130B

你只需要加载模型权重，ATB自动帮你搭建计算图、做算子融合、分配内存。

代码示例2：用ATB跑Llama-2-7B推理

import torch
import atb_speed as atb
from transformers import AutoTokenizer

# 加载ATB预优化的Llama-2-7B模型
model = atb.models.LlamaModel.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    npu_mem_size=28 * 1024 * 1024 * 1024,  # 28GB NPU内存
    enable_flash_attention=True,
    enable_int8=False  # 先跑FP16版本
)
model = model.npu()

# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 推理
prompt = "昇腾CANN是"
input_ids = tokenizer.encode(prompt, return_tensors="pt").npu()
output_ids = model.generate(
    input_ids,
    max_new_tokens=50,
    do_sample=True,
    temperature=0.7
)
print(tokenizer.decode(output_ids[0]))

代码讲解：

from_pretrained自动下载模型权重，并搭建ATB优化后的计算图
npu_mem_size指定NPU内存大小（Atlas 800T A2有32GB HBM，留4GB给系统）
enable_flash_attention=True开启FlashAttention（必须在from_pretrained时指定）
model.generate是HuggingFace的接口，ATB做了适配，用法完全一样

性能数据：ATB vs PyTorch原生

我在Atlas 800T A2（Ascend 910，32GB HBM）上测了Llama-2-7B的推理性能：

实现方式	batch=1, seq_len=512	batch=4, seq_len=512	batch=1, seq_len=2048
PyTorch原生	18 tokens/s	15 tokens/s	12 tokens/s
PyTorch+TorchAir	24 tokens/s	21 tokens/s	16 tokens/s
PyTorch+ATB（FP16）	31 tokens/s	28 tokens/s	22 tokens/s
PyTorch+ATB（INT8）	42 tokens/s	38 tokens/s	31 tokens/s
PyTorch+ATB（INT4）	58 tokens/s	52 tokens/s	41 tokens/s

关键发现：

batch越大，ATB的优势越明显（因为层封装减少了算子调度开销）
INT8/INT4量化收益巨大（INT4比FP16快了87%，精度损失<1%）
长序列（seq_len=2048）性能下降明显（因为FlashAttention的tiling策略在长序列时不够优）

ATB在CANN五层架构里的位置

ATB属于第2层：昇腾计算服务层，跟AOL算子库、AOE调优引擎搭档。

具体协作流程：

你的PyTorch模型
  ↓
ATB模型层（LlamaModel/GPTModel等）
  ↓
ATB层封装层（SelfAttentionLayer/FFNLayer等）
  ↓
ops-transformer / ops-nn / catlass（底层算子）
  ↓
AscendCL（算子调用接口）
  ↓
昇腾NPU硬件

关键点：ATB不直接调用NPU硬件，它透过AscendCL调用底层算子。这样做的好处是，底层算子优化时，ATB不需要改代码。

踩坑实录

我在用ATB的过程中，踩过这几个坑：

坑1：FlashAttention在seq_len>4096时需要手动指定max_position_embeddings

ATB默认假设seq_len≤4096，如果你的模型支持长文本（比如Llama-3的8192），需要手动指定：

model = atb.models.LlamaModel.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    max_position_embeddings=8192,  # 手动指定最大序列长度
    enable_flash_attention=True
)

坑2：INT8量化需要校准数据集

ATB的INT8量化不是"直接把权重从FP16转成INT8"，它需要一小批校准数据（~100条）来做量化参数搜索。

# INT8量化流程
from atb_speed import quantize

# 1. 准备校准数据集（100条样本）
calib_data = [
    "昇腾CANN是昇腾异构计算架构",
    "大模型推理的瓶颈在带宽",
    # ... 98 more samples
]

# 2. 量化模型
quantized_model = quantize(
    model,
    calib_data=calib_data,
    quantize_level="int8",  # 或者 "int4"
    calib_batch_size=4
)

# 3. 保存量化后的模型
quantized_model.save_pretrained("./llama-2-7b-int8")

坑3：多卡推理时，ATB的模型并行需要手动配置

ATB支持张量并行（Tensor Parallelism），但需要手动指定并行策略：

# 2卡张量并行
model = atb.models.LlamaModel.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    tensor_parallel_size=2,  # 2卡并行
    pipeline_parallel_size=1
)

结尾

ATB这个仓库，在昇腾CANN生态里的定位是**“Transformer模型的官方加速库”**。它跟TorchAir不冲突——TorchAir做的是通用图优化，ATB做的是Transformer专用层优化。两者配合用，性能最好。

如果你在搞大模型推理部署，强烈建议去 https://atomgit.com/cann/ascend-transformer-boost 把这个仓库拉下来，先跑通Llama-2-7B的示例。光看文档是感受不到"层封装优化"跟"算子融合优化"的差异的，必须自己跑一把，看NPU利用率从70%飙到93%的那一刻，你才知道ATB的价值。

仓库：https://atomgit.com/cann/ascend-transformer-boost