ascend-transformer-boost：昇腾上跑大模型推理，这个加速库到底加速了什么

2601_96139665

13人浏览 · 2026-05-21 23:45:45

2601_96139665 · 2026-05-21 23:45:45 发布

ascend-transformer-boost：昇腾上跑大模型推理，这个加速库到底加速了什么

在昇腾 NPU 上跑大模型，绕不开一个库叫 ATB。它的全称是 Ascend Transformer Boost，翻译过来就是昇腾 Transformer 加速库。

这个名字挺实在的——它不是通用推理框架，就是专门给 Transformer 类模型做推理加速的。Llama、Qwen、ChatGLM 这些常见的大语言模型，用 ATB 加速效果普遍比较明显。但我发现很多人对它的理解还停留在"装上就能快"的层面，对它具体做了什么、怎么配、什么场景效果最好，没有系统认知。今天把这块说清楚。

ATB 是什么，不是什么

先说清楚 ATB 是什么。

ATB 是昇腾官方出的大模型推理加速库，基于 CANN 的底层算子能力，封装了一套高层推理接口。你不需要自己调 AscendCL，不需要自己组合算子，直接把模型扔给 ATB，它帮你搞定算子融合、内存调度、硬件亲和等一堆优化。

ATB 不是推理框架。它不负责模型解析、图优化这些通用工作流，它的核心定位是在"算子层"做加速——PyTorch 模型经过 ATB 加速的那一层，性能能提升多少，取决于你用的是哪一层、模型结构是什么。

换个角度说：ATB 相当于昇腾 NPU 上的"推理加速插件"，底层跑的是 ops-transformer 里的融合算子，但它把这些算子打包成了开箱即用的高层接口。不理解这一层的，会觉得 ATB 很神秘；理解了这层关系，就知道它其实就是 CANN 算子能力的产品化封装。

ATB 加速了什么

ATB 在推理场景里主要做了三件事：算子融合、内存优化、并行调度。

第一，算子融合。Transformer 模型里有很多可以融合的小算子——矩阵乘法 + ReLU、GEMM + Softmax、GEMM + Scale——分开跑的话，中间结果要写回显存再读出来，每次访存都有开销。ATB 会把这些相邻算子合并成单个融合算子，减少显存在读写的次数。

这和昇腾 graph-autofusion 仓库的定位有重叠，但区别在于：ATB 做的是离线融合优化（在模型编译阶段完成），graph-autofusion 做的是运行时动态融合（推理过程中根据实际数据 shape 决定）。如果你跑的是离线量化后的模型，ATB 的融合效果更稳定；如果你需要在运行时适应变长序列，graph-autofusion 更灵活。

第二，内存优化。大模型推理的显存瓶颈主要在两部分：模型权重和 KV Cache。ATB 对 KV Cache 做了动态分块管理，根据序列实际长度分配显存，而不是预分配最大可能的显存空间。另外它支持权重量化压缩，int8 量化后权重显存直接减半，对显存敏感的推理场景效果明显。

from ascend_transformer_boost import AtbPipeline

# 初始化 ATB pipeline
pipeline = AtbPipeline(
    model_path="qwen-7b-npu",          # 模型路径
    dtype="bfloat16",                    # 推理精度
    kv_cache_dtype="float16",            # KV Cache 精度（可单独设置）
)

# ATB 的显存优化策略不需要手动配置，pipeline 会自动：
# 1. 按实际序列长度分配 KV Cache
# 2. 对权重做 int8 量化压缩（如果有量化模型）
# 3. 融合可合并的算子层
result = pipeline.generate(prompt, max_length=512)

第三，并行调度。ATB 支持多级并行策略：算子级并行（一个算子的计算分到多个计算单元）、数据级并行（batch 并行）、流水线并行（Prefill 和 Decode 阶段并行调度）。最后这个是最值得说的——ATB 内置了 PD 分离的调度能力，Prefill 和 Decode 不再是串行的，而是在 ATB 的调度下异步执行，Decode 阶段不用等 Prefill 完全结束就能开始吐出 token。

这背后依赖 hixl 做 KV Cache 传输、ascend-boost-comm 做算子复用。ATB 把这些底层能力串起来，提供了一个统一的配置接口。

ATB 的核心配置：几个决定性能的关键参数

ATB 上手容易，但想把性能调到最优，得理解几个关键参数的含义。

compute_dtype：计算精度，影响计算速度和精度之间的权衡。bfloat16 是目前大模型推理的主流，精度比 float16 更好，速度损失不大。如果你在意首 token 延迟，用 float16 可以更快；如果在意生成质量，bfloat16 更稳。

kv_cache_dtype：KV Cache 的精度，可以单独设置，不需要和计算精度一致。比如计算用 bfloat16，KV Cache 用 float16，能省一部分显存，同时对最终输出质量影响有限。这个参数值得多试几组对比。

block_size：KV Cache 的分块大小，默认 128。这个和 ops-transformer 里的 FlashAttention block_size 是同一个概念——太小了 tile 数量多、调度开销大，太大了显存占用高。在昇腾 NPU 上，128 是大多数场景的安全默认值。

parallel_degree：并行度，控制 ATB 同时处理多少个请求。值越大吞吐越高，但对显存的需求也越高。这个参数受限于单卡显存大小，不是越大越好。

from ascend_transformer_boost import AtbPipeline, AutotuneConfig

# 如果默认参数不够用，可以用 Autotune 搜索最优配置
autotune = AutotuneConfig(
    model_path="qwen-7b-npu",
    batch_sizes=[1, 4, 8],
    seq_lengths=[512, 1024, 2048],
    target="throughput"  # 优化目标：吞吐 or latency
)

best_config = autotune.search()  # 搜索过程比较慢，建议离线做一次

pipeline = AtbPipeline(
    model_path="qwen-7b-npu",
    **best_config
)

ATB vs 其他加速方案：什么时候用它

昇腾生态里能做推理加速的方案不止 ATB 一个，我列几个常见的对比一下。

方案	定位	适用场景	门槛
ATB	算子级融合加速	大模型推理吞吐优化	中，API 友好
graph-autofusion	运行时动态融合	变长序列、动态 shape	中
torchtitan-npu	PyTorch 原生适配	训练 + 推理通用	低，会 PyTorch 就能用
PyTorch NPU 后端	底层 NPU 接入	自定义模型、灵活调优	高

ATB 最适合的场景：大模型推理，批量请求，吞吐优先。Llama、Qwen、ChatGLM 这些主流模型，ATB 都有专门的优化路径，直接用就行。

ATB 不适合的场景：模型结构特别定制化，ATB 的融合路径覆盖不到；或者你需要自定义算子行为，ATB 是黑盒封装，调不了底层。这种情况直接用 PyTorch NPU 后端更灵活，或者走 torchtitan-npu 自己调融合策略。

几个踩坑经验

坑一：ATB 对模型格式有要求。 ATB 直接支持的是昇腾离线模型格式（OM），不是原始 PyTorch checkpoint。模型要先用 ATC（AOE 调优引擎附带的模型转换工具）转成 OM 才能喂给 ATB。这个转换过程有坑——量化参数配置不对、shape 填错了，转出来的模型性能会很差。建议参考 CANN 官方文档的转换流程一步一步做，不要自己跳步骤。

坑二：batch_size 和序列长度要平衡。 ATB 的显存分配策略是"先按 max_seq 预分配 KV Cache，再用实际 seq 填充"。如果你把 max_seq 设得很大（比如 8192），但实际跑的大多是 512 以内的请求，显存利用率会很低，batch_size 上不去反而更慢。建议用 profiling 工具实测一下不同配置下的显存占用，再定 max_seq。

坑三：PD 分离模式下 KV Cache 的传输时序要理清楚。 ATB 支持 Prefill-Decode 分离，但这个模式需要你对 hixl 的传输机制有一定了解。Prefill 写完 KV Cache 后、Decode 开始读之前，如果 ATB 的调度没有正确等待传输完成，Decode 可能拿到空数据。这个问题比较隐蔽，建议先用同步模式调通业务逻辑，再切到 PD 分离模式。