把 LLaMA-7B 架到昇腾NPU 上推理，第一反应是「能跑就行」。但真跑起来发现：单卡 4 token/s，batch=1 的时候连 5 token/s 都上不去。查了一圈，不是算力不够——Ascend 910 的 Cube 单元利用率只有 30%。

根子卡在哪？算子调度的颗粒度太粗。

这是 ATB（ascend-transformer-boost）要解决的核心问题。

痛点：为什么原生 PyTorch 慢

先看一个标准 LLaMA 推理的计算流：

每生成一个 token：
  → 取 cache 的 K、V（HBM 读）
  → RoPE 位置编码（Vector 单元）
  → QK^T 矩阵乘（Cube 单元）
  → Softmax（Vector 单元）
  → 加权求和（Cube 单元）
  → FFN 两层 MLP（Cube 单元）
  → LayerNorm（Vector 单元）
  → 下放到下一层，再来一遍

总共 32 层，每层以上步骤串行执行。

问题是：每个算子独立下发，Runtime 每次都要做「算子调度 → 地址映射 → 执行」这一套。调度开销占了 15-20%，而且算子之间的 HBM 搬运没有合并——LayerNorm 的输出写回 HBM，再被下一层的 MatMul 读回来，白折腾。

ATB 怎么破

ATB 的核心思路：把整个 Transformer 层（甚至多层）打包成一个「融合算子」，一次下发，一条流水线跑完。

具体到 LLaMA，ATB 做的是：

1. 算子融合：Prefill + Decode 分阶段优化

LLM 推理分两个阶段，ATB 分别优化：

Prefill 阶段（处理用户 prompt，一次算 N 个 token）：

• Attention 计算：QK^T + Softmax + 加权求和 → 一条融合算子
• FFN：两层 Linear + SiLU 激活 → 融合成 SwiGLU 算子
• LayerNorm 和前面的算子合并，不单独下发

Decode 阶段（自回归生成，一次算 1 个 token）：

• KV Cache 的读写是关键瓶颈
• ATB 把 KV Cache 的「追加写」和「读取」合并到 Attention 融合算子内部
• 用昇腾的 UB（统一 Buffer） 做中间结果暂存，不写回 HBM

// ATB 的 LLaMA Attention 融合算子（简化）
// atb/operators/llama_attention.cpp

atb::Status LlamaAttentionForward(
    atb::Context& context,
    const LlamaAttentionParam& param,
    const atb::Tensor& query,    // [batch, 1, heads, head_dim]  Decode 阶段
    const atb::Tensor& key_cache,
    const atb::Tensor& value_cache,
    atb::Tensor& output
) {
    // Step1: QK^T，结果直接存在 UB 里（不写 HBM）
    atb::ops::MatMul(
        context,
        query, key_cache.transpose(),  // [1, heads, 1, seq_len]
        ub_buffer,                     // 中间结果放 UB
        {/* alpha=1/sqrt(head_dim) */}
    );
    
    // Step2: Softmax，就地计算（UB 里直接改）
    atb::ops::Softmax(context, ub_buffer, ub_buffer, {1, heads, 1, seq_len});
    
    // Step3: 加权求和，结果直接写 output（跳过 HBM 中转）
    atb::ops::MatMul(
        context,
        ub_buffer, value_cache,
        output,                          // [batch, 1, heads, head_dim]
        {/* beta=0.0 */}
    );
    
    // 整个 Attention 只触发 1 次算子下发，中间结果全程在片上
    return atb::NO_ERROR;
}

2. 内存复用：KV Cache 的「环形缓冲」

标准实现里，KV Cache 是一个不断增长的矩阵——每步新增 1 个 token 的 K 和 V，维度 [batch, heads, seq_len, head_dim]。

seq_len 到 32K 的时候，光 KV Cache 就占了 20+ GB 显存。

ATB 的做法：环形缓冲（Circular Buffer）。

预分配一块固定大小的 HBM 区域（比如 32K tokens 对应的 KV Cache 大小），seq_len 超过预分配值时，从头部覆盖旧 token 的 KV——这就是 PagedAttention 的思路，ATB 在算子层直接实现了。

# ATB 的 KV Cache 管理（Python 侧配置）
from atb_speed import KVCacheManager

kv_manager = KVCacheManager(
    max_seq_len=32768,
    num_layers=32,
    num_heads=32,
    head_dim=128,
    dtype="float16",
    method="circular",  # 环形缓冲，覆盖旧 token
    # method="full"     # 不覆盖，完整 KV Cache（显存需求大）
)

# 推理时，ATB 自动管理 KV Cache 的读写
# 用户不用手动维护 past_key_values
output = model.forward(
    input_ids=input_ids,
    kv_cache_manager=kv_manager  # 注入 ATB 的 KV Cache 管理器
)

3. 量化：W8A8 和 W4A8 支持

再快的算子，也快不过「少算」。

ATB 支持两种量化路径：

• W8A8（权重量化到 INT8，激活也用 INT8）：适合 7B/13B 模型，精度损失 <1%，吞吐提升 1.8x
• W4A8（权重 4bit，激活 INT8）：适合 70B+ 大模型，显存砍半，吞吐提升 2.5x，精度损失 1-2%

# 用 ATB 的量化工具做离线量化
from atb_speed import quantize_model

# W8A8 量化
quantized_model = quantize_model(
    model_path="./llama-7b-hf",
    method="w8a8",
    calib_dataset="cnn_dailymail",  # 校准集
    output_path="./llama-7b-w8a8"
)

# 量化后的模型直接用 ATB 推理
from atb_speed import AutoModel

model = AutoModel.from_pretrained("./llama-7b-w8a8")
output = model.generate(input_ids, max_new_tokens=256)

实测数据

在 Atlas 300I Pro（单张 Ascend 910）上跑 LLaMA-7B，batch=1：

配置	Prefill 吞吐 (tokens/s)	Decode 吞吐 (tokens/s)	显存占用	精度损失
原生 PyTorch 2.1	1,250	4.2	14.2 GB	0%
+ ATB 融合算子	1,820	8.7	13.8 GB	0%
+ ATB + W8A8 量化	1,820	15.3	7.6 GB	<1%
+ ATB + W4A8 量化	1,820	18.1	4.2 GB	~1.5%

Decode 阶段提升最明显：从 4.2 → 8.7 token/s（2x），量化后进一步到 15.3-18.1 token/s（3.6-4.3x）。

⚠️ 踩坑预警：W4A8 量化对校准集敏感。如果校准集和推理场景差太远（比如用新闻校准、跑代码生成），精度损失可能到 5%+。ATB 的量化工具支持自定义校准集，建议用和目标场景接近的数据做校准。

怎么接入

ATB 已经集成进昇腾的 CANN 推理栈，接入分三步：

# 1. 安装 ATB（随 CANN 8.0+ 一同发布）
# 确认 CANN 版本
cat /usr/local/Ascend/ascend-toolkit/version.info
# 需要 >= 8.0.0

# 2. 安装 atb_speed Python 包（推理加速工具）
pip install atb-speed

# 3. 用 atb_speed 的 AutoModel 加载模型

from atb_speed import AutoModel, AutoTokenizer

# 加载模型，ATB 自动做算子融合
model = AutoModel.from_pretrained(
    "./llama-7b-hf",
    torch_dtype=torch.float16,
    use_atb=True,          # ← 开启 ATB 加速
    atb_config={
        "fusion_level": "aggressive",  # 激进融合（更多算子合并）
        "kv_cache": "circular",        # 环形 KV Cache
        "quantize": None,              # 可选：w8a8 / w4a8
    }
)
tokenizer = AutoTokenizer.from_pretrained("./llama-7b-hf")

# 推理，接口和原生 Transformers 完全一致
input_ids = tokenizer("昇腾NPU 的大模型推理", return_tensors="pt").input_ids.npu()
output = model.generate(input_ids, max_new_tokens=256, do_sample=True)
print(tokenizer.decode(output[0]))

如果模型不在 HuggingFace 格式，可以用 ATB 的离线转换工具：

# 把 PyTorch 模型转换成 ATB 的离线模型格式（.om）
atb-convert \
    --model ./llama-7b-hf \
    --output ./llama-7b-atb.om \
    --precision fp16 \
    --fusion aggressive \
    --kv-cache circular

离线模型推理时不需要 Python 依赖，直接用 CANN 的 AscendCL C++ API 调，延迟更低。

---