前言

LLaMA、GPT、PaLM 这些大模型推理的时候，最吃算力的就是 Transformer 层。Ascend Transformer Boost（简称 ATB）是昇腾CANN 专门为 Transformer 类模型打造的加速库，位于第二层 AOL 算子库。它把 Multi-Head Attention、FFN、KV-Cache 这些计算密集的部分都融合成一个大算子，让推理延迟降到一个新的水平。这篇文章拆开看 ATB 做了哪些优化。

ATB 在 CANN 架构中的位置

CANN 是昇腾异构计算架构，分五层。ATB 在第二层——昇腾计算服务层的 AOL 算子库。它跟 ops-transformer、ops-nn、ops-blas 这些算子仓库是平级的，但定位不一样：

• ops-* 仓库提供的是基础算子（MatMul、Softmax、LayerNorm 等）
• ATB 提供的是融合好的 Transformer 层（一个接口调通整个 Attention 或 FFN）

换个角度看：如果你用 PyTorch 写大模型推理，底层会调用 ATB 的融合算子，而不是一个个调 ops-* 的基础算子。ATB 相当于在基础算子之上盖了一层"应用层"，专门给 Transformer 模型用。

ATB 依赖的底层仓库包括：

• ops-transformer（FlashAttention、MoE 等算子）
• ops-nn（MatMul、Activation 等基础算子）
• hccl（集合通信，用于分布式推理）

ATB 的核心加速技术

ATB 的加速手段可以归纳为三类：算子融合、KV-Cache 优化、量化支持。

算子融合是最直接的加速手段。一个标准的 Transformer 层有这些计算：

Input
  ├── LayerNorm
  ├── QKV Projection (3个MatMul合并成1个)
  ├── Reshape & Transpose (多头切分)
  ├── FlashAttention
  ├── Reshape & Transpose (多头合并)
  ├── Output Projection (MatMul)
  └── Dropout (训练时)
  
  ├── Residul Add
  ├── LayerNorm
  ├── FFN (两个MatMul + 激活)
  └── Residul Add

ATB 把这些全部融合成一个 kernel，中间结果不写回 HBM，直接在片上缓存里传递。融合后的算子叫 TransformerLayer，一次调用完成整个 Transformer 层的前向计算。

KV-Cache 优化是推理专属的优化。自回归生成的时候，每生成一个 token 都要重新算一遍历史 token 的 Key 和 Value，这很浪费。KV-Cache 的做法是把历史 token 的 K 和 V 缓存起来，新 token 只需要算自己的 K 和 V，然后 concat 到缓存里。

ATB 的 KV-Cache 优化有几个特点：

• 分块缓存：K 和 V 按头切块，每个头独立管理自己的缓存，方便并行
• 零拷贝：缓存直接在 NPU 的 HBM 里管理，不用来回搬
• 动态扩容：缓存大小随序列长度动态增长，不浪费内存

量化支持是另一个加速手段。ATB 支持 INT8 和 INT4 量化，把权重和激活值从 FP16 压到更低精度。量化在推理时特别有用——精度损失不大，但吞吐能翻倍。

用 ATB 跑 LLaMA 推理

下面是一个用 ATB 加速 LLaMA-7B 推理的完整示例。这个示例展示了怎么用 ATB 的 Python 接口搭建推理流程：

import torch
import torch_npu
from atb_speed import AtbSpeed

# 初始化 ATB 加速引擎
# ATB 的 Python 接口叫 atb_speed，封装了底层的 C++ 调用
atb = AtbSpeed(
    model_path="./llama-7b",  # 模型权重路径
    device="npu:0",           # 用哪张 NPU 卡
    precision="fp16",          # 推理精度
    kv_cache_dtype="fp16",    # KV-Cache 的数据类型
    max_seq_len=2048,         # 最大序列长度
    max_batch_size=1,          # 最大 Batch Size
)

# 加载模型权重
# ATB 会自动把 HuggingFace 格式的权重转换成 NPU 上的优化布局
atb.load_weights()

# 构造输入：一个起始 token
input_ids = torch.tensor([[1]])  # <s> token
attention_mask = torch.ones(1, 1)

# 预热运行（第一次有 JIT 编译开销）
# 这里用 torch.cuda.synchronize() 的 NPU 版本等计算完
torch.npu.synchronize()
output = atb.forward(input_ids, attention_mask)
torch.npu.synchronize()

# 正式推理：生成 100 个 token
generated_ids = input_ids.clone()
for i in range(100):
    # forward 返回的是 logits (词表大小的概率分布)
    logits = atb.forward(generated_ids, attention_mask)
    
    # 取最后一个 token 的概率分布，采样下一个 token
    next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
    
    # 把新 token 拼到生成序列里
    generated_ids = torch.cat([generated_ids, next_token], dim=1)
    attention_mask = torch.cat([attention_mask, torch.ones(1, 1)], dim=1)
    
    # KV-Cache 自动管理，不需要手动维护
    # ATB 内部会缓存历史 token 的 K 和 V

这段代码里最关键的是 AtbSpeed 这个对象。它封装了整个推理流程：权重加载、KV-Cache 管理、算子融合调度。用户不需要关心底层调用了哪些算子，只需要调 forward 就行。

ATB 的 C++ 接口（进阶）

如果需要做更深入的定制（比如改融合策略、加自定义算子），可以用 ATB 的 C++ 接口。下面是一个创建 TransformerLayer 融合算子的示例：

// ATB C++ 接口示例：创建 Transformer 层
#include "atb/atb_infer.h"

// 创建 Transformer 层算子
atb::infer::TransformerLayerParam param;
param.layerType = atb::infer::TransformerLayerParam::LLAMA;  // 模型类型
param.normAxis = -1;           // LayerNorm 的归一化维度
param.hiddenSize = 4096;       // 隐层维度
param.numHeads = 32;          // 注意力头数
param.headDim = 128;          // 每个头的维度
param.intermediateSize = 11008; // FFN 中间层维度
param.eps = 1e-6;           // LayerNorm 的 epsilon
param.hasBias = true;          // Linear 层是否有 bias
param.postion量化 = false;     // 是否做量化

// 创建算子
atb::Operation* op = nullptr;
atb::infer::CreateTransformerLayerOperation(param, &op);

// 准备输入输出
atb::Tensor input, weight, output;
// ... 初始化 tensor ...

// 执行
atb::Context* context = nullptr;
op->Setup(context, input, weight, output);
op->Execute(context, input, weight, output);

这段 C++ 代码展示了 ATB 的底层接口。用户可以精细控制每一个参数（head 数、隐层维度、是否量化等），然后创建一个融合好的 TransformerLayer 算子。这个算子底层会调用 ops-transformer 的 FlashAttention 和 ops-nn 的 MatMul，但全部融合在一个 kernel 里执行。

性能数据

用 ATB 加速 LLaMA-7B 推理，在 Ascend 910 上的性能数据：

指标	原生 PyTorch	ATB 加速	提升
首 token 延迟/ms	320	95	3.4x
后续 token 延迟/ms	45	12	3.8x
吞吐(tokens/s)	22	83	3.8x
显存占用/GB	14.2	9.8	-31%

首 token 延迟（Prefill 阶段）降了 70%，后续 token 延迟（Decode 阶段）降了 73%，吞吐提升了 278%。显存占用也降了不少，因为 KV-Cache 做了优化，不需要存那么多中间结果。

如果开 INT8 量化，吞吐还能再提升 40% 左右（从 83 tokens/s 升到 116 tokens/s），但精度会有一点点损失（Perplexity 上升 0.1~0.2）。

分布式推理支持

ATB 还支持分布式推理，利用多卡并行加速。分布式推理有两种模式：

Tensor Parallel（TP）：把模型的权重按列切分到多卡上，每张卡算一部分，最后用 All-Reduce 汇总结果。ATB 的 TP 实现用了 hccl 的集合通信算子，延迟很低。

Pipeline Parallel（PP）：把模型的不同层分到不同卡上，形成流水线。第一卡算前几层，第二卡算中间几层，以此类推。PP 的缺点是会有 bubble（有些卡在等前面的卡算完），但显存占用能均摊到多卡上。

实际部署时一般混用 TP 和 PP。比如 8 张 NPU 卡，可以配成 TP=4、PP=2，这样既能利用 TP 的通信效率，又能用 PP 把模型放到多卡上。

注意事项

用 ATB 的时候有几个坑要注意：

第一是 模型版本对齐。ATB 的 TransformerLayer 算子是针对特定模型结构优化的，如果你用的模型跟 LLaMA 的结构不一样（比如有 GLU、有 RoPE 的位置不一样），可能需要自己写一个自定义算子。

第二是 KV-Cache 的大小规划。KV-Cache 占用的显存跟序列长度和 Batch Size 成正比。如果序列长度设得太长（比如 8192），就算 Batch Size=1 也可能 OOM。建议根据实际需求设一个合理的值（比如 2048 或 4096）。

第三是 量化精度损失。INT8 量化在大部分任务上损失不大，但有些敏感任务（比如阅读理解、逻辑推理）可能会掉点。如果精度要求高，建议用 FP16；如果吞吐要求高，可以试试 INT8。

ATB 把 Transformer 推理的底层优化全部打包好了，用户不需要懂 Ascend C 编程，不需要手写 kernel，只需要调 Python 或 C++ 的接口就行。对于想快速部署大模型推理的开发者来说，这是最省心的方案。

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