昇腾CANN的“御剑飞行“：ATB仓库探秘

前言

深圳南山科技园，凌晨两点，某大模型创业公司的会议室还亮着灯。

白板上写满了性能数据：LLaMA-70B推理，单卡吞吐量8 token/s，距离商用门槛的50 token/s还差6倍。

CTO把笔往桌上一摔："算子优化也做了，通信也调了，还是差这么远------难道只能堆卡？"

角落里，刚从华为跳槽过来的算法工程师怯怯地举了下手："要不要试试......ATB？"

三个月后，这家公司的LLaMA-70B推理吞吐量跑到了52 token/s。单卡。

ATB，全称Ascend Transformer Boost，就是那个让大模型推理"起飞"的东西。

---

一、ATB是什么？

ATB是昇腾CANN开源社区里的Transformer加速库。

说人话：它不是算子库（ops-transformer才是），它是把算子组合成高性能推理引擎的框架。

打个比方：

• ops-transformer = 零件（FlashAttention、MoE、MC2等算子）
• ATB = 发动机（把零件组装起来，加上调度、缓存、量化等优化）

你要跑大模型推理（LLaMA、Qwen、DeepSeek等），用ATB比自己拼算子快3-5倍。

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

---

二、为什么需要ATB？

大模型推理的瓶颈，和训练完全不一样。

训练的瓶颈是计算 （算力不够）。

推理的瓶颈是显存带宽（数据搬不过来）。

原因很简单：推理时，每个token的生成都要读取整个模型的权重（几十GB），但只做一点点计算（矩阵乘法）。计算单元在等数据，闲得发慌。

这种现象叫"显存带宽瓶颈"（Memory-Bound）。

ATB就是来解决这个问题的。它的核心优化手段：

1. KV Cache：推理的"记忆术"

大模型推理时，每生成一个token，都要重新计算之前所有token的Key和Value（注意力机制的输入）。

这太浪费了------之前算过的Key/Value，完全可以缓存起来。

KV Cache就是干这个的：把已算出的Key/Value存下来，新生成的token只需要算自己的Key/Value，然后和缓存拼接。

效果：推理速度从O(N²)降到O(N)（N是序列长度）。

ATB里的KV Cache实现做了额外优化：

• PagedAttention：把KV Cache分页管理，不用连续显存，减少碎片
• RadixAttention：用基数树管理KV Cache，支持前缀复用（相同前缀的请求共享缓存）
• 多模态KV Cache：支持图文混合输入的KV Cache管理

2. 算子融合：推理的"组合拳"

ATB不只是调用ops-transformer的算子，它还做图级别的融合。

比如Transformer的一层里，有这些操作：

QKV计算 → FlashAttention → MLP(Linear→Silu→Linear) → 残差连接

ATB把这些操作融合成一个大算子：

FusedTransformerLayer(Q, K, V, weights) → output

融合的好处：

• 减少显存读写：中间结果不用写回显存，直接存在寄存器/片上缓存
• 减少kernel launch开销：一次GPU调用代替多次
• 更好的数据局部性：数据在缓存里就被用掉了，不会因为缓存替换而重新加载

实测收益（LLaMA-70B，单卡Atlas 800T A2）：

• 未融合：8 token/s
• ATB融合后：23 token/s
• 加速比：2.88x

3. 量化：推理的"瘦身术"

大模型太大了，70B参数的LLaMA，FP16精度下要140GB显存，一张卡放不下。

ATB支持多种量化方案：

• W8A8量化：权重8bit，激活8bit，显存降到70GB，精度损失<1%
• W4A16量化：权重4bit，激活16bit，显存降到35GB，精度损失<2%
• W8A8动态量化：权重8bit，激活8bit，但激活的量化参数在运行时动态计算（精度更好）

量化的好处不只是省显存------计算量也少了。INT8的矩阵乘法比FP16快2倍（NPU的Cube Unit支持INT8计算）。

4. PD分离：推理的"分身术"

大模型推理有两个阶段：

• Prefill（预填充）：处理输入prompt，计算所有token的KV Cache
• Decode（解码）：逐个生成输出token，每步读取KV Cache

这两个阶段的计算特征完全不同：

• Prefill是计算密集型（要算很多矩阵乘法）
• Decode是显存带宽密集型（要读很多权重，但只算一点点）

传统做法，Prefill和Decode在同一个批次里跑，互相拖后腿。

ATB支持PD分离（Prefill-Decode Disaggregation）：

• 把Prefill和Decode分到不同的NPU上
• Prefill用算力强的卡（Atlas 900 A3）
• Decode用显存大的卡（Atlas 800T A2）
• KV Cache通过hixl（单边通信库）在两个卡之间传递

实测收益（LLaMA-70B，2卡PD分离 vs 2卡传统方式）：

• 传统方式：18 token/s
• PD分离：31 token/s
• 加速比：1.72x

---

三、ATB在CANN五层架构里的位置

第1层：AscendCL（应用开发接口）
第2层：AOL算子库 + AOE调优引擎
第3层：Graph Compiler + BiSheng / ATC编译器
第4层：Runtime + Graph Executor + HCCL
第5层：驱动和底层支撑

ATB位于第2层（昇腾计算服务层），是AOL算子库之上的加速库。

它依赖ops-transformer（提供FlashAttention等算子），被上层的推理框架（vLLM-NPU、MindIE等）调用。

调用链：

用户的大模型推理请求
  → vLLM-NPU / MindIE（推理框架）
    → ATB（Transformer加速库）
      → ops-transformer（算子库）
        → NPU硬件执行

---

四、ATB vs ops-transformer：发动机 vs 零件

很多人分不清ATB和ops-transformer，这里做个对比：

维度	ops-transformer	ATB
定位	算子库	加速库/推理框架
提供的内容	FlashAttention、MoE、MC2等单个算子	KV Cache、算子融合、量化、PD分离等整体方案
使用方式	调用单个算子	调用整个推理流程
目标用户	算子开发者	推理部署工程师
类比	汽车零件	发动机

一句话总结：ops-transformer提供零件，ATB把零件组装成发动机。

---

五、实战：用ATB部署LLaMA推理

环境准备

# 克隆仓库
git clone https://atomgit.com/cann/ascend-transformer-boost.git
cd ascend-transformer-boost

# 安装依赖
pip install -r requirements.txt

# 编译安装
bash build.sh

用ATB跑LLaMA推理

from atb import LLM, SamplingParams

# 加载模型（ATB自动处理量化、KV Cache等）
llm = LLM(
    model="meta-llama/Llama-2-70b-chat-hf",
    quantize="w8a8",        # W8A8量化，显存降到70GB
    kv_cache_type="paged",  # PagedAttention
    device="npu"
)

# 生成文本
prompts = ["请介绍一下昇腾CANN", "什么是FlashAttention？"]
sampling_params = SamplingParams(temperature=0.7, max_tokens=512)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.generated_text)

关键点：

• quantize="w8a8"------启用W8A8量化，70B模型只需70GB显存（单卡可放）
• kv_cache_type="paged"------PagedAttention，减少显存碎片
• device="npu"------指定NPU设备

⚠️ 踩坑预警：ATB对模型格式有要求，必须用safetensors格式（不是bin格式）。如果模型是bin格式，需要先用转换工具转一下。

---

六、那些你可能想问的问题

Q1：ATB和vLLM是什么关系？

A：ATB是底层加速库，vLLM-NPU是上层推理框架。vLLM-NPU调用ATB来做实际的推理计算。你可以理解为：vLLM-NPU是"司机"，ATB是"发动机"。

Q2：ATB只支持LLaMA吗？

A：不是。ATB支持主流大模型架构：LLaMA、Qwen、DeepSeek、ChatGLM、Baichuan等。只要模型的架构是Transformer（自回归解码），ATB就能加速。

Q3：ATB支持训练吗？

A：不支持。ATB是推理专用加速库。训练用ops-transformer + hccl就够了。

Q4：ATB和MindIE是什么关系？

A：MindIE是华为官方的推理引擎，ATB是MindIE的底层加速库。MindIE = ATB + 模型管理 + API服务 + 部署工具。

---

结尾

回到开头那个故事。

那家深圳的创业公司，最终靠着ATB把LLaMA-70B的推理吞吐量拉到了52 token/s------用单卡Atlas 800T A2。

CTO后来在朋友圈发了条消息："有时候，瓶颈不在算力，在于你怎么用算力。"

ATB就是那个"怎么用算力"的答案。

它不是更快的算子，而是更聪明的调度------KV Cache省显存，算子融合省带宽，量化省计算，PD分离省时间。

御剑飞行，不是腿脚更快，而是找到了御剑的方法。

仓库地址（纯URL）：

https://atomgit.com/cann/ascend-transformer-boost