前言

上个月帮一个团队把LLaMA3-8B从GPU迁到昇腾910。模型能跑，但decode吞吐只有GPU上的一半。他们第一反应是"昇腾算力不行"，我跑了一下msprof——Attention占42%、LayerNorm+残差连接占21%、MoE Router占15%。问题不在算力，在算子没融合。每个小操作都单独下发一次ACL调用，几百次调用光调度开销就吃掉了3ms。

装了ATB之后，同样配置吞吐翻了2.7倍。

很多人以为ATB是个"推理框架"。不是。ATB是算子加速库，它干的事就一件：把Transformer里那些反复出现的计算模式，打包成高度优化的融合算子，让你调一行接口就能用。

ATB 在 CANN 架构里的位置

CANN五层架构里，ATB不属于任何一层——它是第2层（计算服务层）之上的加速库，依赖ops-transformer提供的基础算子（FlashAttention、MoE等），往上给MindIE推理引擎和训练框架提供加速接口。

用户代码（PyTorch/MindSpore）
  ↓
MindIE 推理引擎 / 训练框架
  ↓
ATB 加速库（融合算子编排）← 你在这
  ↓
ops-transformer（FlashAttention / MoE / MC2）
  ↓
Ascend C 算子实现
  ↓
达芬奇架构硬件

简单说：ops-transformer提供"砖头"，ATB把砖头搭成"墙"，MindIE拿去盖"房子"。

这个分层很重要。你不需要懂Ascend C，也不需要手写Tiling策略，ATB帮你搞定了。

工程经验： 不复用ATB自己写Transformer优化，要手写Attention+LayerNorm+FFN的融合算子，开发周期2-3周。用ATB的融合算子，改2行配置，半天搞定。不是ATB多神奇，是它把融合逻辑都封装好了。

ATB 解决了什么问题

三个核心痛点：

1. 算子调度开销大

Transformer一层前向传播涉及QKV投影、LayerNorm、Attention、残差连接、FFN、MoE Routing……十几个算子。不融合的话每个算子单独走一次ACL→GE→Runtime调用链，单次开销12-15μs。30层就是400多次调用，纯调度5-6ms。decode阶段每个token的预算可能才10-15ms，调度就吃了快一半。

# 不复用ATB：30层Transformer，每层12个算子，共360次ACL调用
import torch
import torch_npu

model = LLaMA3_8B().npu()

# 每层Forward（伪代码）
for layer in model.layers:
    # QKV投影（3次ACL调用）
    Q = layer.q_proj(x)  # ACL调用1
    K = layer.k_proj(x)  # ACL调用2
    V = layer.v_proj(x)  # ACL调用3
    
    # Attention（3次ACL调用）
    S = Q @ K.T / sqrt(d)  # ACL调用4
    P = torch.softmax(S, dim=-1)  # ACL调用5
    O = P @ V  # ACL调用6
    
    # LayerNorm（2次ACL调用）
    x = layer.ln1(x)  # ACL调用7
    x = x + O  # ACL调用8（逐元素加）
    
    # FFN（3次ACL调用）
    hidden = layer.ffn1(x)  # ACL调用9
    hidden = torch.silu(hidden)  # ACL调用10
    x = layer.ffn2(hidden)  # ACL调用11
    
    # 残差连接（1次ACL调用）
    x = x + O  # ACL调用12

# 30层 × 12次 = 360次ACL调用
# 单次ACL调用开销：12-15μs
# 总调度开销：360 × 15μs = 5.4ms

# 用ATB：30层Transformer，每层1-2次ACL调用
import torch
from atb import TransformerLayer

model = TransformerLayer().npu()

# 每层Forward（融合算子）
for layer in model.layers:
    x = layer(x)  # 1次ACL调用（Attention+FFN+LayerNorm+残差全部融合）

# 30层 × 1次 = 30次ACL调用
# 总调度开销：30 × 15μs = 0.45ms
# 调度开销降低：92%

ATB把一层Transformer的前向/反向计算融合成1-2个Kernel，调度开销从5ms降到0.1ms以下。

2. 显存碎片浪费

每层中间结果（Q、K、V、Attention输出、FFN中间值）各自申请一块显存，用完释放。动态分配导致HBM碎片化，实测带宽利用率只有35%左右。ATB用预分配+buffer复用策略，同一层的前向计算共享中间buffer，带宽利用率拉到80%+。

# 不复用ATB：每层中间结果各自申请显存
def forward_no_fusion(x):
    # QKV投影（3次显存申请）
    Q = torch.empty(M, K, dtype=torch.float16).npu()  # 显存申请1
    K = torch.empty(M, K, dtype=torch.float16).npu()  # 显存申请2
    V = torch.empty(M, K, dtype=torch.float16).npu()  # 显存申请3
    
    # Attention（2次显存申请）
    S = torch.empty(M, N, dtype=torch.float16).npu()  # 显存申请4
    P = torch.empty(M, N, dtype=torch.float16).npu()  # 显存申请5
    
    # FFN（1次显存申请）
    hidden = torch.empty(M, 4*K, dtype=torch.float16).npu()  # 显存申请6
    
    # 共6次显存申请/释放，HBM碎片化严重
    # 带宽利用率：35%

# 用ATB：预分配+buffer复用
def forward_with_atb(x):
    # 预分配中间buffer（1次）
    buffer = torch.empty(max_M * max_K * 4, dtype=torch.float16).npu()
    
    # 所有中间结果复用同一个buffer
    Q = buffer[0:M*K].view(M, K)
    K = buffer[M*K:2*M*K].view(M, K)
    V = buffer[2*M*K:3*M*K].view(M, K)
    # ...
    
    # 零显存申请/释放，HBM无碎片
    # 带宽利用率：80%+

3. MoE 模型支持差

DeepSeek、Mixtral这类MoE模型，Router决定哪些expert被激活是动态的。传统方案要么全量算所有expert（浪费算力），要么用if-else分支（分支预测失败代价高）。ATB内置GatingTopK融合算子，直接在NPU上做动态路由，只算被激活的expert。

# 不复用ATB：MoE Router（逐expert算）
def moe_router_no_fusion(x, experts, top_k=2):
    # Router算logits（1次ACL调用）
    logits = router(x)  # ACL调用1
    
    # TopK（1次ACL调用）
    topk_idx = torch.topk(logits, top_k)[1]  # ACL调用2
    
    # 逐expert算（expert_count次ACL调用）
    outputs = []
    for i in topk_idx:
        outputs.append(experts[i](x))  # ACL调用3~3+expert_count
    
    # 合并（1次ACL调用）
    output = torch.stack(outputs).sum(dim=0)  # ACL调用4+expert_count
    
    # 总ACL调用：4 + expert_count（8个expert就是12次）
    # 分支预测失败率高（动态路由）

# 用ATB：GatingTopK融合（1次ACL调用）
from atb import MoELayer

def moe_router_with_atb(x, experts, top_k=2):
    # GatingTopK融合算子（Router+TopK+Dispatch一次完成）
    output = MoELayer(experts, top_k)(x)  # ACL调用1
    
    # 只算被激活的expert，省掉未激活expert的计算
    # 总ACL调用：1

工程经验： Mixtral 8×7B在910B上，不开ATB时MoE部分吞吐280 TPS；开ATB的MoE融合后到920 TPS。省的不是算力，是那些没被激活的expert根本没进计算流水线。

核心模块

ATB主要管四件事：

1. Attention 加速

封装ops-transformer的FlashAttention，自动处理mask、位置编码、KV Cache格式。支持Full Attention、Window Attention、以及DeepSeek-V4的CSA/HCA压缩Attention。

# ATB Attention配置（PyTorch）
from atb import AttentionConfig, AttentionLayer

# 配置Attention（自动选最优实现）
attn_config = AttentionConfig(
    head_num=32,
    head_dim=128,
    seq_len=2048,
    use_flash_attention=True,  # 开FlashAttention
    use_window_attention=False,  # 不开Window Attention
    use_csa=False,  # 不开DeepSeek-V4的CSA
)

# 创建Attention层（自动融合QKV投影+Attention+输出投影）
attn_layer = AttentionLayer(attn_config).npu()

# 前向（1次ACL调用）
output = attn_layer(Q, K, V, mask)

2. MoE 加速

GatingTopK融合（router+topk+dispatch一次完成）、Expert并行调度、All-to-All通信优化。MoE模型在昇腾上的性能瓶颈往往不在计算而在通信，ATB对这块做了专门优化。

# ATB MoE配置（PyTorch）
from atb import MoEConfig, MoELayer

# 配置MoE（自动融合Router+TopK+Dispatch）
moe_config = MoEConfig(
    expert_num=8,
    top_k=2,
    expert_capacity=4,
    use_all_to_all=True,  # 开All-to-All通信优化
)

# 创建MoE层（自动只算被激活的expert）
moe_layer = MoELayer(experts, moe_config).npu()

# 前向（1次ACL调用）
output = moe_layer(x)

3. FFN / LayerNorm 融合

把LayerNorm+线性投影+激活函数+残差连接融成一个Kernel。别小看这个——30层Transformer，每层省4次HBM读写，总共省120次。

# 不复用ATB：LayerNorm+线性投影+激活+残差（4次ACL调用）
def ffn_no_fusion(x):
    # LayerNorm（1次ACL调用）
    x = layer_norm(x)  # ACL调用1
    
    # 线性投影（1次ACL调用）
    hidden = linear1(x)  # ACL调用2
    
    # 激活（1次ACL调用）
    hidden = silu(hidden)  # ACL调用3
    
    # 残差连接（1次ACL调用）
    output = linear2(hidden) + x  # ACL调用4

# 用ATB：FFN+LayerNorm融合（1次ACL调用）
from atb import FFNLayer

def ffn_with_atb(x):
    # FFN+LayerNorm融合算子（LayerNorm+线性投影+SiLU+残差一次完成）
    ffn_layer = FFNLayer(hidden_dim=4096, act_fn='silu').npu()
    output = ffn_layer(x)  # ACL调用1

4. 量化支持

W8A16、W8A8C16量化推理，INT8权重存HBM，计算时反量化或直接低精度计算。910B实测W8A16吞吐比FP16高45%，显存省30%。

# ATB量化配置（PyTorch）
from atb import QuantConfig, QuantLinear

# 配置量化（W8A16）
quant_config = QuantConfig(
    weight_bits=8,  # INT8权重
    activation_bits=16,  # FP16激活
    compute_bits=16,  # FP16计算
    symmetric=True,  # 对称量化
)

# 创建量化线性层
linear = QuantLinear(in_features=4096, out_features=4096, 
                     quant_config=quant_config).npu()

# 量化权重（离线）
linear.quantize_weights()

# 前向（自动反量化）
output = linear(x)

和 CUDA 生态对比

维度	CUDA 生态	ATB
定位	对标FasterTransformer + TransformerEngine	同左
底层算子	CUTLASS + 自定义CUDA Kernel	ops-transformer + Ascend C
融合方式	手写fused kernel或nvFuser	graph-autofusion自动融合 + 手动融合算子
框架接入	通过plugin接入TensorRT	通过Framework Adaptor接入MindSpore/PyTorch

最大的区别在底层。CUDA的SM能同时跑矩阵乘和逐元素运算，所以NVIDIA可以把整层Transformer融合成一个巨型kernel。昇腾是Cube干矩阵乘、Vector干逐元素，两个单元之间数据走L1传递。ATB的做法是按Cube/Vector边界切分融合粒度——Cube连续的计算尽量塞到一个kernel，Vector操作批量处理，中间靠L1缓存桥接。

这导致ATB的融合kernel数量比FasterTransformer多，但每个kernel内部的效率更高。trade-off不同，但最终性能在同一水平线。

工程经验： ATB的融合粒度比FasterTransformer细（按Cube/Vector边界切分），调试更容易（哪个kernel慢一目了然）。FasterTransformer的巨型kernel调试很难（一个kernel几万行PTX）。

怎么上手

ATB的使用门槛不高。如果你用MindSpore训练/推理，ATB已经内置在里面了，开启对应配置就行。PyTorch用户通过torchtitan-npu接入。

最直接的用法是看cann-recipes-infer仓库里的示例——里面有完整的模型加载、ATB配置、推理跑通的流程。从"能跑"到"跑满"，中间差的往往就是一个配置项的事。

# 完整示例：LLaMA3-8B推理（ATB加速）
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from atb import TransformerConfig, TransformerModel

# 1. 加载模型（HuggingFace格式）
model = AutoModelForCausalLM.from_pretrained("llama3-8b")
tokenizer = AutoTokenizer.from_pretrained("llama3-8b")

# 2. 配置ATB（开融合）
atb_config = TransformerConfig(
    use_flash_attention=True,  # 开FlashAttention
    use_moe_fusion=True,  # 开MoE融合（如果是MoE模型）
    use_ffn_layernorm_fusion=True,  # 开FFN+LayerNorm融合
    use_quantization=False,  # 不开量化
    torch_compile_backend="npu",  # 开GE图编译
)

# 3. 创建ATB模型（自动融合）
atb_model = TransformerModel(model, atb_config).npu()

# 4. 编译（GE图编译，优化算子调度）
atb_model = torch.compile(atb_model, backend="npu")

# 5. 推理
inputs = tokenizer("Hello, ", return_tensors="pt").input_ids.npu()
outputs = atb_model.generate(inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# 运行
cd cann-recipes-infer/examples/llama3-8b/
python infer_atb.py

# 输出：
# [INFO] Load model: llama3-8b, params: 8.0B
# [INFO] ATB config: FlashAttention=ON, MoE=OFF, Quant=OFF
# [INFO] GE graph compile time: 12.3s
# [INFO] Generate: "Hello, how are you doing today? I hope..."
# [INFO] Throughput: 89 tokens/s（vs 不复用ATB的34 tokens/s，+162%）

踩坑实录

坑1：ATB跟torch.compile冲突（GE图编译失败）

ATB内部已经做了GE图编译，再调torch.compile(backend="npu")会重复编译，报错GE_ERROR_DUPLICATE_COMPILE。

解决：ATB模型不要再调torch.compile。ATB配置里设torch_compile_backend="npu"就行。

坑2：MoE模型不开All-to-All通信优化，多卡吞吐反而降

MoE模型多卡推理，All-to-All通信占60%时间。不开use_all_to_all=True，多卡吞吐反而比单卡低。

解决：MoE模型多卡推理，必开use_all_to_all=True。

坑3：量化配置不对（W8A8C16，activation_bits=8），精度掉15%

W8A8C16的激活量化成INT8，精度损失大（FP16→INT8，动态范围压缩）。

解决：用W8A16（权重INT8，激活FP16），精度损失<1%，吞吐比FP16高45%。

坑4：ATB的FlashAttention不支持Window Attention（报错）

ATB的FlashAttention实现只支持Full Attention和CSA/HCA，不支持Window Attention（Sliding Window）。

解决：Window Attention用标准Attention（不复用FlashAttention）。设use_flash_attention=False, use_window_attention=True。

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

https://atomgit.com/cann/cann-recipes-infer

https://atomgit.com/cann/ops-transformer