CANN ATB：Transformer 推理加速库的融合策略

单算子优化到头了推理加速的下一步是算子级融合。ATB 把 Transformer 的一个完整 Decoder Layer 做成一个大 Kernel从 Attention 到 FFN 一口气跑完中间不回 HBM。CANN 8.0 之后 ATB 是 LLM 推理的首选加速库。ATB 在 CANN 五层架构中的位置ATBAscend 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 投影 → SwiGLUGate × SiL U激活→ Up 投影 → Mul → Down 投影。LayerNorm 通常也融合进去——Attention 的输入 LayerNorm 和 FFN 的输入 LayerNorm 都在对应的 Block 里做掉不单独出现。融合 Kernel 之间通过 UB 传递中间结果不用写 HBM。Block 之间的数据传递通过昇腾NPU的跨核通信机制完成。# ATB 推理调用示例LLaMA 模型推理fromatbimportTransformerModel,ModelConfigimporttorch# 模型配置configModelConfig()config.model_typellama# 支持 llama/chatglm/qwen/baichuanconfig.hidden_size4096# LLaMA-2 7Bconfig.num_layers32config.num_heads32config.head_dim128config.intermediate_size11008# FFN 中间维度config.vocab_size32000# ATB 推理引擎初始化modelTransformerModel(config)model.load_weights(/path/to/llama-2-7b-weights)# 权重加载# 推理接口prefill decode 分离调用# Prefill 阶段处理 prompt输出第一个 tokeninput_idstorch.tensor([prompt_token_ids],dtypetorch.long,devicenpu:0)position_idstorch.arange(len(prompt_token_ids),dtypetorch.long,devicenpu:0)kv_cachemodel.prefill(input_ids,position_ids)# 返回 KV Cache# Decode 阶段自回归生成generated_ids[]current_idstorch.tensor([[last_token]],dtypetorch.long,devicenpu:0)for_inrange(max_new_tokens):# kv_cache 会自动更新不需要手动管理logits,kv_cachemodel.decode_step(current_ids,kv_cache)next_tokentorch.argmax(logits,dim-1)generated_ids.append(next_token.item())ifnext_tokeneos_token_id:breakcurrent_idsnext_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 → AttentionQKVRoPESDPAProj→ Add → LayerNorm → FFN → Add。每个节点之间的数据都要落 HBM。融合推理走的是LayerNormAttention 融合 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