单算子优化到头了，推理加速的下一步是算子级融合。ATB 把 Transformer 的一个完整 Decoder Layer 做成一个大 Kernel，从 Attention 到 FFN 一口气跑完，中间不回 HBM。CANN 8.0 之后 ATB 是 LLM 推理的首选加速库。

ATB 在 CANN 五层架构中的位置

ATB（Ascend Transformer Boost）属于昇腾异构计算架构第 2 层（服务层）的 Transformer 加速库，位于 AOL 算子库之上，为 LLM 推理提供端到端的加速方案。

ATB 不直接写 Ascend C 算子，它在现有算子库（ops-transformer、ops-nn）的基础上做图级融合。GE 负责通用算子的融合，ATB 负责 Transformer 特定结构的融合——GE 融合不了的模式，ATB 补充。

从调用链路看：用户代码 → ATB Python/C++ API → 融合图（ATB 内部管理） → GE 编译 → 昇腾NPU执行。

Transformer Decoder Layer 的算子融合图

一个 Decoder Layer 的计算流程拆开是：Input → LayerNorm1 → Attention → Add（残差连接）→ LayerNorm2 → FFN → Add（残差连接）→ Output。

Attention 本身又拆成：QKV 投影 → Scaled Dot-Product → Softmax → 输出投影。

如果每个算子独立跑，中间结果层层写 HBM 再读 HBM，光 LayerNorm 和残差连接之间的 HBM 交互就能吃掉 20% 的延迟。

ATB 的融合策略是把这一整条链合成一个大 Kernel。以 Attention Block 为例：

融合后的 Attention Kernel 包含：QKV 投影（一个 MatMul 融合三个投影的 bias）→ RoPE 旋转 → 分组注意力计算（Q、K 的 topk 选择）→ Softmax → 乘 V → O 投影。

FFN Block 融合：Gate 投影 → SwiGLU（Gate × SiL U激活）→ Up 投影 → Mul → Down 投影。

LayerNorm 通常也融合进去——Attention 的输入 LayerNorm 和 FFN 的输入 LayerNorm 都在对应的 Block 里做掉，不单独出现。

融合 Kernel 之间通过 UB 传递中间结果，不用写 HBM。Block 之间的数据传递通过昇腾NPU的跨核通信机制完成。

# ATB 推理调用示例：LLaMA 模型推理
from atb import TransformerModel, ModelConfig
import torch

# 模型配置
config = ModelConfig()
config.model_type = "llama"           # 支持 llama/chatglm/qwen/baichuan
config.hidden_size = 4096              # LLaMA-2 7B
config.num_layers = 32
config.num_heads = 32
config.head_dim = 128
config.intermediate_size = 11008       # FFN 中间维度
config.vocab_size = 32000

# ATB 推理引擎初始化
model = TransformerModel(config)
model.load_weights("/path/to/llama-2-7b-weights")  # 权重加载

# 推理接口：prefill + decode 分离调用
# Prefill 阶段：处理 prompt，输出第一个 token
input_ids = torch.tensor([prompt_token_ids], dtype=torch.long, device="npu:0")
position_ids = torch.arange(len(prompt_token_ids), dtype=torch.long, device="npu:0")
kv_cache = model.prefill(input_ids, position_ids)  # 返回 KV Cache

# Decode 阶段：自回归生成
generated_ids = []
current_ids = torch.tensor([[last_token]], dtype=torch.long, device="npu:0")
for _ in range(max_new_tokens):
    # kv_cache 会自动更新，不需要手动管理
    logits, kv_cache = model.decode_step(current_ids, kv_cache)
    next_token = torch.argmax(logits, dim=-1)
    generated_ids.append(next_token.item())
    if next_token == eos_token_id:
        break
    current_ids = next_token.unsqueeze(0)

ATB 的好处是不需要手动管理 KV Cache。model.prefill() 和 model.decode_step() 会自动维护 KV Cache 的更新逻辑。在内部，ATB 用了 ops-transformer 仓的 PagedAttention 实现来管理 KV Cache 的内存分配。

ATB 的图执行引擎

ATB 不只是做算子融合，它还有自己的图执行引擎，负责管理融合 Kernel 之间的调度和内存。

ATB 的图执行引擎支持两种模式：

流式模式（Streaming Mode）：各融合 Kernel 串行执行，但通过双缓冲和指令流水隐藏 Kernel 之间的启动开销。适合 Prefill 阶段——算力密集，延迟要求高。

Continuous Batching 模式：多个请求的 decode 步骤并行执行，通过请求级别的调度来最大化吞吐。适合高并发推理服务场景。

ATB 的 Continuous Batching 实现有自己的调度器。调度器维护一个 running requests 队列，每个请求在每个 step 会申请分配资源（KV Cache slot、计算核）。当某个请求完成 decode，调度器立即回收它的资源，插入新请求。调度的调度周期是 token-level 的，理论上可以实现 100% 的 GPU 利用率。

ATB vs 非融合推理

非融合推理走的是这条链路：LayerNorm → Attention（QKV+RoPE+SDPA+Proj）→ Add → LayerNorm → FFN → Add。每个节点之间的数据都要落 HBM。

融合推理走的是：LayerNorm+Attention 融合 Kernel → FFN 融合 Kernel。两个 Kernel 之间通过 UB 或 HCCL 传递。

粗略估算，融合后 HBM 访问次数减少约 60%，对应的延迟收益在 30-50%（取决于 seq_len 和 batch size）。

不过融合也有代价：融合 Kernel 的编译时间比单算子长很多（可能长 10 倍），出错了更难 debug，内存占用也更大（融合 Kernel 需要同时容纳更多中间结果）。

ATB 当前支持的模型

ATB 的 atb/models/ 目录列出了当前官方支持的模型：

LLaMA 系列（LLaMA-2、LLaMA-3、CodeLLaMA）：最完善，Prefill 和 Decode 都支持。

ChatGLM 系列（ChatGLM2、ChatGLM3）：支持，但某些特殊算子（如 GLU 激活）需要额外适配。

Qwen 系列（Qwen-1.5、Qwen-2）：支持。

Baichuan：支持。

其他模型（MOSS、Falcon 等）：可以尝试用 LLaMA 模式近似覆盖，但不一定能跑通。

如果模型不在支持列表里，ATB 提供了自定义模型接口 atb.CustomModel，用户可以自己定义融合图的结构。不过这需要比较深入地理解 ATB 的图构造 API，门槛比直接用预置模型高很多。

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