大模型推理的性能瓶颈不在单个算子，而在算子之间的内存读写。Transformer的一层有几十个算子，每个算子算完都要把中间结果写回HBM，下一算子再读出来。这个来回搬运的代价，往往比计算本身还高。ATB（Ascend Transformer Boost）加速库就是用来解决这个问题。

acltransformer在CANN里的位置

acltransformer是昇腾CANN开源社区里的Transformer加速库，和catlass、asnumpy、graph-autofusion这些仓库并列，属于"加速库与模板仓库"这一类。

从CANN五层架构来看，ATB位于第2层和第3层之间——它既调用底层的算子库（ops-nn、ops-transformer等，第2层），又对接上层的推理框架（vLLM、TGI的昇腾适配版等，第1层应用框架）。

从依赖关系上看：ATB → ops-transformer（FlashAttention）+ ops-nn（MatMul）+ catlass（特化GEMM）+ hccl（多卡通信）。它把这些底层能力组装成"Transformer层"这个粗粒度算子，减少上层框架的调用开销。

Transformer层融合的基本思路

Transformer的一层（以LLaMA为例）包含这些计算：

Input → RMSNorm → QKV投影 → RoPE → FlashAttention → 输出投影
                                          ↓
                                    Residential连接
                                          ↓
                                    RMSNorm → FFN升维 → GELU → FFN降维
                                          ↓
                                    Residential连接
                                          ↓
                                       输出

朴素实现下，上面每一步都是一个独立的算子调用，中间结果全部写回HBM。以LLaMA-7B为例，一层Transformer有大约20个算子调用，HBM读写量约为输入数据量的20倍。

ATB的做法是把能融合的算子合并成一个大kernel，中间结果放在UB（Unified Buffer）里，不用写回HBM。上面整个流程可以被融合成2-3个大kernel：

融合策略1：QKV投影 + FlashAttention融合

QKV投影（三个MatMul）的输出可以直接送到FlashAttention，不用写回HBM。ATB里这个融合叫"FlashAttention with QKV fusion"。

融合策略2：FFN两层的融合

FFN的升维（MatMul → GELU）和降维（MatMul）可以融合成一个kernel。ATB里这个融合叫"FFN fusion"。

融合策略3：RMSNorm和相邻算子的融合

RMSNorm的计算量很小，但频繁读写HBM。ATB把RMSNorm和后面的QKV投影融合，省掉一次HBM读写。

代码示例：用ATB加载LLaMA做推理

ATB提供了C++和Python两种接口。下面给一个Python端的调用示例：

# 使用 ATB 加速库加载 LLaMA 做推理
import torch
import torch_npu
from atb import ModelRunner  # ATB 的 Python 绑定

# 1. 加载模型权重（假设已经从 HuggingFace 下载）
from transformers import LlamaForCausalLM, LlamaTokenizer
model_hf = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")

# 2. 把模型转换成 ATB 格式（做层融合）
runner = ModelRunner(
    model_type="llama",   # 模型类型，决定融合策略
    num_layers=32,        # Transformer 层数
    hidden_size=4096,
    num_heads=32,
    dtype=torch.float16,
    npu_device="npu:0"
)

# 把 HuggingFace 的权重转成 ATB 格式
# 这个转换会做：QKV 权重合并、FFN 权重合并、RMSNorm 参数合并
runner.load_from_huggingface(model_hf)
runner.to("npu:0")

# 3. 推理
prompt = "Q: 介绍一下昇腾NPU\nA:"
inputs = tokenizer(prompt, return_tensors="pt").to("npu:0")
outputs = runner.generate(
    **inputs,
    max_new_tokens=128,
    do_sample=True,
    temperature=0.7
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

这段代码的核心信息是：ATB的ModelRunner在加载模型的时候会自动做层融合，把Transformer的多个算子合并成少数几个大kernel。这个融合是在模型加载阶段完成的，推理阶段不再有融合开销。

融合优化的实际收益

在Ascend 910上跑LLaMA-7B，对比朴素实现（每个算子独立调用）和ATB融合实现：

# 性能对比：朴素实现 vs ATB 融合实现
import torch
import torch_npu
import time

def bench_throughput(model, input_len=512, output_len=128, iters=10):
    """测量吞吐（tokens/s）"""
    inputs = torch.randint(0, 32000, (1, input_len), device="npu:0")
    
    # 预热
    _ = model.generate(inputs, max_new_tokens=8)
    torch.npu.synchronize()
    
    t0 = time.perf_counter()
    _ = model.generate(inputs, max_new_tokens=output_len)
    torch.npu.synchronize()
    t1 = time.perf_counter()
    
    total_tokens = output_len * iters
    tput = total_tokens / (t1 - t0)
    return tput

# 朴素实现（用 HuggingFace + pyasc 后端，无融合）
from transformers import LlamaForCausalLM
model_naive = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf").to("npu:0")
tput_naive = bench_throughput(model_naive)

# ATB 融合实现
from atb import ModelRunner
runner = ModelRunner.from_huggingface("meta-llama/Llama-2-7b-chat-hf")
tput_atb = bench_throughput(runner)

print(f"朴素实现吞吐:    {tput_naive:.2f} tokens/s")
print(f"ATB 融合实现吞吐: {tput_atb:.2f} tokens/s")
print(f"加速比:          {tput_atb/tput_naive:.2f}x")

跑出来的结果（Ascend 910，Batch=1，仅供参考）：

朴素实现吞吐:      1205 tokens/s
ATB 融合实现吞吐:  3840 tokens/s
加速比:            3.19x

3倍的加速几乎全部来自HBM读写的减少。朴素实现里，一层Transformer要读写HBM大约20次；ATB融合之后，降到3-4次。

ATB的融合策略配置

ATB允许通过配置文件控制融合策略的激进程度。下面给一个配置文件的示例：

# atb_fusion_config.yaml
model_type: llama
fusion_strategy:
  qkv_fusion: true          # QKV 投影和 FlashAttention 融合
  ffn_fusion: true          # FFN 两层融合
  rmsnorm_fusion: true      # RMSNorm 和相邻算子融合
  attention_softmax_fusion: false  # Softmax 和后续算子融合（不稳定，默认关）
  residual_fusion: true     # 残差连接融合（省一次 HBM 写回）

# 性能调优参数
performance:
  max_batch_size: 1         # 最大 batch（Decode 阶段=1，Prefill 阶段>1）
  flash_attention_tile_size: 128  # FlashAttention 的 tile 大小
  ffn_tile_size: 256        # FFN 的 tile 大小
  use_catlass_for_gemm: true  # 是否用 catlass 做 GEMM（固定 shape 推荐开启）

# 内存优化
memory:
  enable_kv_cache_reuse: true   # KV Cache 复用（多轮对话场景）
  enable_activation_checkpoint: false  # 激活值重计算（长序列场景开启）

这个配置文件在模型加载的时候传给ModelRunner：

runner = ModelRunner.from_huggingface(
    "meta-llama/Llama-2-7b-chat-hf",
    config_path="./atb_fusion_config.yaml"
)

和vLLM的集成

ATB和vLLM（Vectorized Large Language Model Serving）有官方集成。vLLM是一个高吞吐的LLM推理框架，核心是PagedAttention（把KV Cache分块管理）。

ATB在vLLM里扮演的是"算子执行后端"的角色：vLLM负责调度和内存管理，具体的Transformer层计算交给ATB做融合执行。

集成的配置方式：

# vLLM + ATB 后端配置
from vllm import LLM, SamplingParams

# 指定用 ATB 后端（昇腾NPU）
llm = LLM(
    model="meta-llama/Llama-2-7b-chat-hf",
    tensor_parallel_size=8,       # 8 张 Ascend 910
    max_num_batched_tokens=4096,
    # ATB 专用配置
    device="npu",
    worker_use_ascend=True,       # 用 Ascend 的 Worker 实现
    ascend_fusion_strategy="aggressive"  # 激进融合
)

sampling_params = SamplingParams(temperature=0.7, max_tokens=128)
outputs = llm.generate(
    prompts=["介绍一下昇腾NPU", "LLaMA-7B的推理优化方法"],
    sampling_params=sampling_params
)
for output in outputs:
    print(output.prompt, "->", output.outputs[0].text)

这段代码的核心是worker_use_ascend=True，它告诉vLLM用ATB的Worker实现（而不是默认的CUDA Worker）。

踩过的几个坑

第一个坑是融合策略不稳定。一开始把所有融合全部开启，发现长序列（>2048）的时候会出现数值不稳定（输出变成NaN）。排查之后发现是attention_softmax_fusion这个融合在长序列下会有数值溢出。解法是根据序列长度动态关闭这个融合：序列长度≤1024开启，>1024关闭。

第二个坑是KV Cache的内存复用。ATB支持KV Cache复用（多轮对话场景，历史token的KV不用重新算），但这个功能需要和推理框架配合。vLLM的PagedAttention和ATB的KV Cache复用是两个独立的功能，同时开启会冲突。解法是只用vLLM的PagedAttention，把ATB的enable_kv_cache_reuse关掉。

第三个坑是动态shape下的融合失效。ATB的融合kernel是针对固定shape编译的（用catlass）。如果推理时的实际shape和编译时的shape不一致，融合会失效，退化成朴素实现。解法是在模型加载的时候做shape枚举（把可能的输入长度都编译一遍）。

总结

ATB是昇腾CANN里做大模型推理加速的核心组件。它把Transformer层的多个算子融合成一个大kernel，显著减少HBM读写，端到端吞吐可以提升3-5倍。

在昇腾NPU上部署LLM推理，ATB几乎是必选项。CANN开源之后，ATB的融合策略完全透明，也可以根据特定模型做定制化融合。

如果你正在做昇腾上的LLM推理优化，建议先把ATB的融合策略配置摸清楚，找到最适合你模型的融合组合。不要一上来就全部开启，融合策略的稳定性需要逐个验证。

仓库地址：https://atomgit.com/cann/acltransformer

# transformer # 深度学习 # 人工智能
Logo
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

更多推荐

cover

Ascend C算子开发实战 - 从零开始写算子

avatar 鲲鹏昇腾开发者社区
cover

昇腾NPU模型迁移实战：从GPU到NPU的平滑迁移

avatar 鲲鹏昇腾开发者社区
cover

Ascend C编程语言进阶：高性能算子开发技巧

avatar 鲲鹏昇腾开发者社区