Ascend Transformer Boost 实战：大模型推理性能提升的秘密

高级c

442人浏览 · 2026-05-21 13:06:55

高级c · 2026-05-21 13:06:55 发布

前言

去年帮一个团队把Qwen2.5-72B部署到昇腾910B集群上。模型能跑，但decode吞吐320 TPS，客户说GPU上能到1800。差了5倍多。

跑msprof一看：算子执行占41%，算子调度开销23%，显存申请释放15%，Host-NPU搬运12%，Stream同步9%。NPU的Cube/Vector单元有59%的时间在等——等调度、等显存、等数据。

装上ATB，开FlashAttention + PagedAttention + W8A16量化，吞吐从320拉到1680 TPS。涨了425%。

不是昇腾算力不够，是默认配置没打开加速特性。

Transformer 推理瓶颈到底在哪

大模型推理分两个阶段：Prefill（首token生成，计算密集）和Decode（后续token逐个生成，访存密集）。

Prefill阶段，一次处理整个输入序列，矩阵乘的计算量大，Cube Unit是瓶颈。Decode阶段，每次只处理1个token，但要从HBM读整个KV Cache，HBM带宽是瓶颈。

实测数据（Qwen2.5-7B，910B单卡，FP16）：

阶段	耗时占比	瓶颈
Attention（Prefill）	35%	Cube（QK^T计算量大）
KV Cache读取（Decode）	42%	HBM带宽（每token读57KB）
MoE Router + AllReduce	12%	通信 + 动态路由
算子调度 + 显存管理	11%	Runtime开销

Decode阶段42%的时间在读KV Cache。为什么？7B模型FP16下每个token的KV Cache约57KB，batch=8、seq=2048时就是57KB × 2048 × 8 = 912MB。每生成一个token都要读一遍这912MB，910B的HBM带宽1.2TB/s，光读KV Cache就要0.76ms。

很多人以为推理优化就是"算子写快点"，其实Decode阶段算子根本不是瓶颈——HBM带宽才是。

工程经验： Qwen2.5-72B在910B 4卡上，Decode阶段HBM带宽利用率只有38%。原因？KV Cache按连续地址分配，但推理时不同batch的序列长度不同，Cache里大量空洞。开了PagedAttention之后带宽利用率拉到79%，吞吐直接翻倍。

为什么 GPU 优化方案不能直接复用

三个核心差异：

1. 计算单元架构不同

GPU的SM能同时跑矩阵乘和逐元素运算。FasterTransformer可以把QK^T + Softmax + P×V融成一个巨型Kernel，中间结果全塞L2缓存（几十MB）。

昇腾是Cube干矩阵乘、Vector干逐元素，两个单元之间数据走L1缓存（1MB）。不能融成一个巨型Kernel——L1装不下，必须按Cube/Vector边界切分。

2. 内存层次不同

GPU的L2缓存是全局共享的，所有SM都能访问。昇腾的L1是每个AI Core独立的，跨核通信走HCCL。

这意味着GPU上的Tiling策略可以直接按L2容量规划，昇腾上必须按L1容量（1MB）规划tile大小。同一个FlashAttention，GPU上tile_q可以开到128，昇腾上只能开到64。

3. 调度模型不同

GPU的CUDA Stream调度是硬件级的（SM自动分配），开销<1μs。昇腾的ACL调度走软件调用链（ACL→GE→Runtime），开销12-15μs。7个小算子不融合，光调度就吃3ms。

为什么这样设计？因为达芬奇架构的Cube和Vector是物理独立的执行单元，不像SM是统一调度器管理。Cube/Vector的协作必须软件编排，这既是灵活性的来源（可以精细化控制流水线），也是调度开销的来源。

工程经验： GPU上优化MatMul就够（Cube利用率高），昇腾上要同时优化Cube和Vector——如果Vector太慢（比如softmax算得慢），Cube在等Vector，利用率掉一半。ATB的融合Kernel就是按这个原则设计的：Cube连续的计算塞一个Kernel，Vector操作批量处理，中间靠L1缓存桥接。

ATB 的核心能力

ATB在CANN架构中的定位：ops-transformer提供"砖头"（FlashAttention、MoE等基础算子），ATB把砖头搭成"墙"（融合算子编排），MindIE拿去盖"房子"（完整推理引擎）。

ATB管五件事：

能力	解决的问题	性能收益
Attention优化	QK^T + Softmax + P×V调度开销大	融合成1个Kernel，调度从3ms→0.05ms
KV Cache管理	显存碎片 + 带宽利用率低	PagedAttention消碎片，带宽利用率38%→79%
算子融合	LayerNorm+MatMul+残差反复读写HBM	融合后中间数据走L1，省3次HBM读写
动态Shape适配	不同batch/seq的shape不固定	运行时自动选Tiling策略
MoE加速	Router动态路由 + 通信瓶颈	GatingTopK融合，只算被激活的expert

Attention 优化

ATB封装了ops-transformer的FlashAttention，并针对达芬奇架构做了Cube/Vector流水线优化。

没有ATB时，Attention的一次前向要走7次ACL调用：

# 不复用ATB：7次ACL调用
Q = torch.matmul(x, Wq)  # ACL调用1
K = torch.matmul(x, Wk)  # ACL调用2
V = torch.matmul(x, Wv)  # ACL调用3
S = torch.matmul(Q, K.transpose(-2, -1))  # ACL调用4
S = S * scale + mask  # Vector运算，ACL调用5
P = torch.softmax(S, dim=-1)  # Vector运算，ACL调用6
O = torch.matmul(P, V)  # ACL调用7

# 7次调用，中间结果（Q、K、V、S、P）全部写回HBM再读出来

ATB的做法：把QKV投影融成1个MatMul（Q/K/V的输入都是x，权重拼接），再把S→P→O融成1个Kernel（FlashAttention）。总共2次ACL调用，中间结果走L1不落HBM。

# 用ATB：2次ATB调用，中间结果走L1
import torch_npu
from atb import AttentionOperation

# QKV投影融合（一次MatMul）
qkv_proj = AttentionOperation(
    type="QKVProjection",
    hidden_size=4096,
    head_num=32,
    head_dim=128,
)
q, k, v = qkv_proj.Run(x, [Wq, Wk, Wv], workspace)

# FlashAttention融合（一次Kernel）
attn_op = AttentionOperation(
    type="FlashAttention",
    head_num=32,
    head_dim=128,
    use_prefill_cache=True,
    sliding_window=4096,
)
output = attn_op.Run(q, k, v, mask, past_kv)

实测性能（Ascend 910B，FP16）：

模型	未融合	ATB融合	提升
Qwen2.5-7B (seq=2048)	34 tokens/s	89 tokens/s	+162%
Qwen2.5-72B (seq=4096, 4卡)	320 TPS	890 TPS	+178%
DeepSeek-V3 (seq=4096)	580 TPS	1420 TPS	+145%

为什么QKV投影能融成一个MatMul？因为Q、K、V的输入都是同一个x，三个权重矩阵 [d×d_q, d×d_k, d×d_v] 可以横向拼接成 [d×(d_q+d_k+d_v)]，一次矩阵乘搞定。这招在GPU上也用，但昇腾上收益更大——省了2次ACL调用（24-30μs），GPU上只省了2次Kernel Launch（~2μs）。

工程经验： DeepSeek-V4有CSA（4倍压缩）和HCA（128倍压缩）两种Attention模式。128K序列下HCA把KV Cache从7.3GB压到57MB，TPOT < 10ms。但短序列（<4K）HCA反而比Full Attention慢12%——Compressor本身有计算开销。ATB的做法是按序列长度动态选策略：seq < 2048用Full Attention，2048-8192用CSA，>8192用HCA。

KV Cache 管理

KV Cache是Decode阶段最大的显存占用者。7B模型FP16，seq=2048、batch=32，KV Cache占3.6GB。

ATB提供三条优化路径：

1. PagedAttention

KV Cache分block（默认16 token一个block），按需分配，消碎片。没有PagedAttention时，不同batch的序列长度不同，连续分配会产生大量空洞。实测显存碎片率30-40%。开了PagedAttention碎片率降到<5%。

# MindIE配置：开PagedAttention
from mindie import MindIERuntime, ModelConfig

config = ModelConfig(
    model_path="/path/to/qwen2.5-72b",
    # PagedAttention：block_size=16
    kv_cache_mode="paged",
    block_size=16,
    # 预分配block池（避免运行时分配开销）
    paged_pool_size_gb=32,
)
runtime = MindIERuntime(config)

2. KV Cache量化

INT8存KV Cache，显存省一半。910B实测：3.6GB → 1.8GB，精度损失 < 0.3%（PPL变化 < 0.01）。

# KV Cache量化配置
config = ModelConfig(
    # KV Cache量化：INT8
    kv_cache_dtype="int8",
    # 可选：前几层FP16、后面INT8（混合策略）
    kv_cache_quant_layers="4:",  # 第4层开始量化
)

工程经验： KV Cache量化有个坑——不同层的KV Cache对量化的敏感度不一样。前几层（靠近embedding）的KV Cache量化后精度影响大，后面层影响小。我们试过"前4层FP16、后面INT8"的混合策略，精度和全INT8一样，但显存多省10%。ATB目前没内置这个策略，得自己改。

3. KV Cache压缩（DeepSeek-V4专用）

HCA 128倍压缩，128K序列KV Cache 57MB。

实测（Qwen2.5-7B，910B单卡，FP16）：

配置	Batch	吞吐(tokens/s)	显存(GB)
无优化	8	72	20.1
+ PagedAttention	16	181	24.3
+ KV Cache INT8 + PagedAttention	32	294	28.1
+ KV Cache INT8 + PagedAttention + W8A16	48	318	31.2

算力没变，显存够用了，batch从8涨到48，吞吐涨341%。

算子融合

昇腾上每次算子下发走ACL→GE→Runtime调用链，开销12-15μs。30层Transformer几百次调用，光调度开销3-5ms。

ATB的融合策略：

LayerNorm + MatMul + 残差

三层算子融成一个Kernel。LayerNorm输出直接喂MatMul，MatMul输出加残差，中间数据走L1不落HBM。省2次HBM读写 + 2次ACL调用。

GatingTopK（MoE专用）

Router输出 + TopK + Expert分发融成一个Kernel。不融合时Router输出落HBM再给TopK读，融合后直接L1传数据。单层省18μs，40层省720μs。

SAS（DeepSeek-V4专用）

统一Window/Sparse/Compress Attention的融合Kernel。不同Attention类型共用同一个Kernel框架，按参数切换模式，省去Kernel切换开销。

// graph-autofusion配置（GE环境变量）
export GE_FUSION_PASS_MODE=aggressive  // 激进融合策略
export GE_ENABLE_OP_FUSION=1          // 开算子融合
export ATB_FUSION_MODE=aggressive     // ATB也开激进融合

// 融合效果验证：看GE编译日志
// 搜索"Fusion success"，看哪些算子融了
grep "Fusion success" ge_compile.log | wc -l

实测融合收益（Qwen2.5-7B，seq=2048）：

融合方案	省掉的HBM读写	省掉的ACL调用	性能提升
QKV投影融合	2次	2次	+8%
FlashAttention	3次	4次	+162%
LayerNorm+MatMul+残差	2次	2次	+15%
GatingTopK（MoE）	1次	2次	+22%

为什么不能把所有算子都融成一个巨型Kernel？两个原因：一是L1只有1MB，整层Transformer的中间结果装不下；二是Cube和Vector是物理独立的，跨单元的数据传递必须显式编排流水线，不像GPU的SM内部寄存器直接传。

动态 Shape

大模型推理的输入shape不是固定的——不同请求的序列长度不同，MoE每层的expert激活模式也不同。

ATB的做法：编译时生成多套Tiling策略，运行时按shape选最优的一套。

# 动态Shape配置（ATB）
from atb import DynamicShapeConfig

dynamic_config = DynamicShapeConfig(
    # 定义shape范围
    shape_ranges={
        "seq_len": [(1, 128), (129, 512), (513, 2048), (2049, 8192)],
        "batch": [(1, 1), (2, 4), (5, 16), (17, 64)],
    },
    # 为每个range生成对应的Tiling参数
    gen_tiling_for_ranges=True,
    # 运行时自动选最优Tiling
    auto_select_tiling=True,
)

编译阶段：

shape_range = [(1,128), (129,512), (513,2048), (2049,8192)]
为每个range生成对应的Tiling参数和Kernel二进制

运行阶段：

实际shape = (1, 1024)
→ 落在(513, 2048)这个range
→ 用第3套Tiling策略

为什么这样设计？因为Tiling参数影响Cube/Vector的利用率和L1缓存命中率。seq=128时tile_q=32就够了，seq=8192时tile_q要开到128才能填满Cube的MAC阵列。一套Tiling策略适配所有shape，要么短序列浪费算力，要么长序列L1溢出。

工程经验： 动态Shape有个隐藏开销——Tiling策略切换需要重新加载Kernel二进制（~50μs）。如果请求的序列长度频繁跳档（比如一会儿100 token，一会儿8000 token），切换开销会吃掉5-8%的吞吐。我们用"分桶策略"缓解：把请求按序列长度分到固定几个桶（0-256, 256-1024, 1024-4096），桶内统一padding到桶的上限，避免频繁切换。

Batch 调优

BatchSize对吞吐的影响不是线性的。

实测数据（Qwen2.5-7B，910B单卡，FP16，ATB全开）：

BatchSize	吞吐(tokens/s)	TTFT(ms)	显存(GB)	HBM带宽利用率
1	38	4	5.2	21%
4	127	8	16.8	51%
8	178	12	20.1	68%
16	231	20	29.3	79%
32	268	38	39.2	85%
64	244	72	OOM	81%

BatchSize=32是拐点。再大，KV Cache占显存太多，开始swap到Host内存，HBM带宽利用率反而降了。

为什么不是越大越好？Decode阶段每个token都要读整个KV Cache。batch越大，要读的KV Cache越多，HBM带宽撑不住。batch=32时HBM带宽利用率85%，接近理论上限；batch=64时KV Cache太大，部分swap到Host内存，带宽利用率反而掉到81%。

场景选型：

在线对话（TTFT < 200ms）：batch=4-8
离线批处理（吞吐优先）：batch=16-32

工程经验： Batch调优不是只看吞吐和延迟。batch=32时吞吐最高，但如果请求的序列长度方差大（有的50 token，有的4000 token），短请求等长请求的时间会拉高P99延迟。我们实际部署时用batch=16，牺牲13%吞吐换P99延迟降40%。调优要看业务指标，不是看NPU指标。