CANN-ATB仓库全景-昇腾NPU大模型推理加速库到底干了什么
ATB（Ascend Transformer Boost）是 CANN 生态里最上层的推理加速库。如果你只是想跑 Llama 推理，不需要关心 ops-transformer 怎么融合、ops-blas 怎么 Tiling——ATB 把这些都封装好了。这篇把 ATB 的整体架构和核心能力讲清楚。
定位：CANN 五层架构的最上层
应用层 → 你的模型代码 ATB → 推理加速库（图优化 + 调度 + KV Cache 管理） AOL → ops-transformer / ops-nn / ops-blas / ops-math AOL-Base → opbase（算子基础设施） HAL → 昇腾NPU硬件驱动ATB 不实现任何计算逻辑——它的活是编排。拿到一个 HuggingFace 模型，ATB 做三件事：
- 图优化：把 PyTorch 计算图里的标准算子替换成 ops-transformer 的融合算子
- 内存管理：分配 KV Cache、权重分片、中间 buffer
- 调度执行：决定哪个算子在哪个 AI Core 上跑，什么顺序跑
核心模块
| 模块 | 功能 | 关键代码路径 |
|---|---|---|
| GraphOptimizer | 计算图优化和算子替换 | atb/graph/ |
| KVCacheManager | KV Cache 的分配、复用、淘汰 | atb/kvcache/ |
| BatchScheduler | Continuous batching 调度 | atb/scheduler/ |
| ModelRunner | 模型编译和执行 | atb/runner/ |
| Sampler | Top-K/Top-P 采样 | atb/sampler/ |
推理全流程
1. 模型加载 HuggingFace 权重 → ATB 解析模型结构 → 生成内部计算图 2. 图优化 F.linear → ops-nn MatMul F.scaled_dot_product_attention → FlashAttention 3×F.linear(QKV) → MergedMatMul F.silu + gate*up → fused_gate_up 3. 编译 优化后的计算图 → GE 编译 → NPU 可执行文件（om 文件） 4. 推理执行 请求进来 → BatchScheduler 分配 slot → prefill → decode 循环 → 采样 → 返回步骤 2 和 3 只在首次推理时执行一次，后续请求复用编译结果。
和 vLLM 的关系
vLLM 是 NVIDIA 生态的推理加速框架，ATB 是昇腾NPU生态的对标方案。两者思路相似但实现不同：
| 维度 | vLLM | ATB |
|---|---|---|
| 硬件 | NVIDIA GPU | 昇腾NPU |
| 图优化 | CUDA Graph | GE Graph |
| KV Cache | PagedAttention | BlockManager（类似 Paged） |
| 调度 | Continuous Batching | Continuous Batching |
| 采样 | GPU 采样 | NPU 采样（ops-math） |
| 模型支持 | HuggingFace 原生 | HuggingFace + 自定义转换 |
ATB 的 KV Cache 管理跟 vLLM 的 PagedAttention 思路一样——把 KV Cache 切成固定大小的 block，按需分配，不预留连续显存。这解决了长序列推理时的显存碎片问题。
快速上手
# 安装 ATBpipinstallascend-transformer-boost# 推理python-matb.run--modelmeta-llama/Llama-2-7b-hf--devicenpu:0或者用 Python API：
fromatbimportLLM model=LLM(model="meta-llama/Llama-2-7b-hf",device="npu:0",max_batch_size=32,max_seq_len=4096,dtype=torch.float16)output=model.generate("Hello, world",max_new_tokens=100)print(output)支持的模型
ATB 官方支持列表（CANN 8.5）：
- Llama 2/3/3.1
- Qwen 1.5/2/2.5
- DeepSeek V2/V3
- Mixtral 8x7B
- ChatGLM 3/4
- Baichuan 2
- Yi
不在列表里的模型可以用自定义转换接口接入，但需要自己写图优化规则。
ATB 是昇腾NPU上跑大模型推理的首选方案——不需要手动调算子、不需要操心内存管理，开箱即用。如果你的模型在支持列表里，直接用 ATB 就行。仓库在这里：
https://atomgit.com/cann/ATB
