上周有个团队问我，为什么同样的昇腾NPU、同样的模型，他们跑出来的吞吐比别人低一倍？我去看了眼代码，发现他们每个算子都单独调用：FlashAttention算完写回HBM，RMSNorm再读出来算，算完又写回去。

我说你们这是"点菜式"调用，每个算子独立执行，中间结果来回搬运。试试ascend-transformer-boost（ATB），它把算子串成流水线，中间结果不写回HBM，直接在片上传递。

改完之后，他们的LLaMA-7B推理吞吐从1200涨到3800 tokens/s。

今天拆一下ATB在大模型推理里的实战用法，看看这个"流水线魔法"是怎么落地的。

场景一：LLaMA推理的算子融合

LLaMA的一个Transformer层包含：

• Self-Attention（FlashAttention）
• RMSNorm
• MLP（两层Linear + SiLU激活）

传统"点菜式"调用：

def llama_layer_naive(x, weights):
 # Self-Attention
 q = x @ weights.wq
 k = x @ weights.wk
 v = x @ weights.wv
 attn_out = flash_attention(q, k, v) # 写HBM
 x = x + attn_out
 
 # Post-Attention RMSNorm
 norm_out = rms_norm(x, weights.norm1) # 读HBM，写HBM
 
 # MLP
 gate = silu(norm_out @ weights.w1) # 读HBM，写HBM
 up = norm_out @ weights.w3 # 读HBM，写HBM
 mlp_out = (gate * up) @ weights.w2 # 读HBM，写HBM
 x = x + mlp_out
 
 return x

问题：每个算子独立调用，HBM读写占了60%的时间。

ATB的"套餐式"融合：

import ascend_transformer_boost as atb

# 创建融合kernel（编译时一次性构建）
fusion_kernel = atb.FusionKernel(
 ops=[
 atb.Op.QKVProjection(wq, wk, wv),
 atb.Op.FlashAttention(causal=True),
 atb.Op.ResidualAdd(),
 atb.Op.RMSNorm(weights.norm1),
 atb.Op.MLP(w1, w2, w3, activation='silu'),
 atb.Op.ResidualAdd()
 ]
)

def llama_layer_fused(x, weights):
 # 一次调用完成整条流水线
 return fusion_kernel.execute(x)

发生了什么？

ATB在编译时把6个算子融合成1个kernel：

传统调用：
QKV投影 → 写HBM → FlashAttention → 写HBM → 残差加 → 写HBM → 
RMSNorm → 写HBM → MLP → 写HBM → 残差加 → 写HBM
访存：12次HBM读写

ATB融合：
QKV投影 → FlashAttention → 残差加 → RMSNorm → MLP → 残差加 → 写HBM
访存：2次HBM读写（输入读一次，输出写一次）

性能对比（昇腾910，LLaMA-7B，batch=8，seq=1024）：

实现	每层延迟	HBM带宽占用	吞吐
ops-transformer单独调用	4.2	85%	1,200
ATB融合调用	1.4	28%	3,800

ATB把HBM带宽占用从85%降到28%，吞吐翻了3.2倍。

场景二：批量推理的动态shape处理

实际推理时，batch size和seq len经常变化。ATB的融合kernel支持动态shape，不用为每个shape重新编译。

# ATB动态shape示例
class DynamicBatchInference:
 def __init__(self, model_weights):
 # 创建动态shape融合kernel
 self.kernel = atb.FusionKernel(
 ops=[...],
 dynamic_dims=['batch_size', 'seq_len'] # 声明动态维度
 )
 
 # 预编译常见shape
 for batch in [1, 2, 4, 8, 16, 32]:
 for seq in [128, 256, 512, 1024, 2048]:
 self.kernel.warmup(batch_size=batch, seq_len=seq)
 
 def infer(self, input_ids):
 batch_size, seq_len = input_ids.shape
 # 自动选择最匹配的编译版本
 return self.kernel.execute(input_ids, batch_size=batch_size, seq_len=seq_len)

为什么动态shape重要？

实际业务中，请求的batch size和seq len千奇百怪：

• 单条推理：batch=1, seq=50
• 批量推理：batch=32, seq=2048
• 长文本推理：batch=1, seq=4096

如果为每个shape都编译一个kernel，编译时间要几十分钟。ATB的动态shape支持让一个kernel处理多种shape，编译一次，到处运行。

场景三：多精度推理的自动切换

ATB支持FP32、FP16、BF16、FP8多种精度，可以根据精度要求自动选择最优实现。

# ATB多精度推理
class MultiPrecisionInference:
 def __init__(self, model_weights, precision='auto'):
 self.precision = precision
 self.kernels = {}
 
 # 为每种精度创建kernel
 for dtype in ['fp32', 'fp16', 'bf16', 'fp8']:
 self.kernels[dtype] = atb.FusionKernel(
 ops=[...],
 dtype=dtype
 )
 
 def infer(self, x, precision=None):
 # 自动选择精度
 if precision is None:
 precision = self._auto_select_precision(x)
 
 # 转换输入精度
 x_converted = x.to(precision)
 
 # 用对应精度的kernel执行
 output = self.kernels[precision].execute(x_converted)
 
 return output
 
 def _auto_select_precision(self, x):
 # 根据输入特征自动选择精度
 if x.shape[0] <= 4: # 小batch用FP16
 return 'fp16'
 elif x.shape[1] >= 2048: # 长seq用FP8
 return 'fp8'
 else:
 return 'bf16'

精度与性能权衡（LLaMA-13B，昇腾910）：

精度	吞吐	内存占用	精度损失
FP32	1,800	100%	0%
FP16	3,600	50%	0.01%
BF16	3,800	50%	0.02%
FP8	5,200	25%	0.3%

FP8吞吐比FP16高44%，精度损失可控在0.3%以内。

ATB与ops-transformer的协作

ATB不是替代ops-transformer，而是基于ops-transformer的融合封装。

用户代码
 ↓
ATB高层API（融合调度）
 ↓
ops-transformer算子（底层实现）
 ↓
GE图编译（算子融合优化）
 ↓
Runtime执行（调度到NPU）

分工明确：

• ops-transformer：提供FlashAttention、RMSNorm等单算子实现
• ATB：把多个算子打包成融合kernel，优化调用顺序
• GE：编译时进一步优化融合策略
• Runtime：运行时调度执行

什么时候直接用ops-transformer？

• 研究新算子、调试算子行为
• 非标准模型结构（ATB没覆盖的组合）
• 需要细粒度控制算子执行

什么时候用ATB？

• 标准Transformer推理（LLaMA、Qwen、Baichuan等）
• 生产环境追求极致性能
• 多卡推理需要通信优化

实战踩坑记录

踩坑一：融合kernel编译慢

第一次创建融合kernel时，ATB要编译，可能要几秒到几十秒。

解决：启动时预热，把编译时间藏起来。

# 启动时预热
def warmup(model):
 dummy_input = torch.zeros(1, 128, 4096)
 for _ in range(3):
 model.infer(dummy_input)
 print("Warmup done, kernel compiled")

# 服务启动时预热
warmup(model)
# 然后开始接收请求

踩坑二：显存不够

融合kernel要分配L1 Buffer存放中间结果。模型太大时，L1 Buffer不够，融合失败。

解决：调整融合粒度，分多个融合kernel。

# 大模型：分两个融合kernel
# Kernel1：Attention部分
kernel_attn = atb.FusionKernel(ops=[
 atb.Op.QKVProjection(...),
 atb.Op.FlashAttention(...),
 atb.Op.RMSNorm(...)
])

# Kernel2：MLP部分
kernel_mlp = atb.FusionKernel(ops=[
 atb.Op.MLP(...),
 atb.Op.ResidualAdd()
])

def forward(x):
 x = kernel_attn.execute(x)
 x = kernel_mlp.execute(x)
 return x

踩坑三：精度不对齐

ATB的融合kernel内部精度和ops-transformer单独调用不完全一致，导致数值差异。

排查方法：逐层对比。

# 逐层对比调试
def debug_precision(x):
 # ATB融合输出
 output_atb = kernel_atb.execute(x)
 
 # ops-transformer单独调用输出
 q, k, v = compute_qkv(x)
 attn_out = flash_attention(q, k, v)
 norm_out = rms_norm(attn_out)
 output_ops = mlp(norm_out)
 
 # 对比差异
 diff = (output_atb - output_ops).abs().max()
 print(f"Max diff: {diff}")

性能调优技巧

技巧一：调整融合粒度

融合粒度影响L1 Buffer占用和并行度。

# 粗粒度融合：整层融合，L1占用大，但访存最少
kernel_coarse = atb.FusionKernel(ops=[
 atb.Op.QKVProjection(...),
 atb.Op.FlashAttention(...),
 atb.Op.RMSNorm(...),
 atb.Op.MLP(...)
])

# 细粒度融合：分多个kernel，L1占用小，但并行度更高
kernel_fine_1 = atb.FusionKernel(ops=[atb.Op.QKVProjection(...), atb.Op.FlashAttention(...)])
kernel_fine_2 = atb.FusionKernel(ops=[atb.Op.RMSNorm(...), atb.Op.MLP(...)])

经验：小模型（7B以下）用粗粒度融合，大模型（13B以上）用细粒度融合。

技巧二：启用异步执行

ATB支持异步执行，配合Python的asyncio实现高并发推理。

import asyncio

async def async_infer(model, input_ids):
 # 异步执行，不阻塞事件循环
 output = await model.kernel.execute_async(input_ids)
 return output

async def batch_infer(model, batch_inputs):
 # 并发处理多个请求
 tasks = [async_infer(model, inp) for inp in batch_inputs]
 results = await asyncio.gather(*tasks)
 return results

技巧三：监控HBM带宽

ATB提供性能分析工具，可以监控HBM带宽利用率。

# 启用性能分析
with atb.Profiler() as prof:
 output = kernel.execute(x)

# 查看HBM带宽利用率
print(f"HBM bandwidth: {prof.hbm_bandwidth}%")
print(f"Compute utilization: {prof.compute_util}%")

如果HBM带宽利用率很高（>70%），说明访存瓶颈，应该增大融合粒度。如果计算利用率很高（>80%），说明计算瓶颈，融合已经足够。

总结

ATB在大模型推理里的核心价值：把"点菜式"调用变成"套餐式"融合，减少HBM读写。

三个实战场景：

• LLaMA推理：整层融合，吞吐翻3倍
• 动态shape：一次编译，多种shape复用
• 多精度推理：自动选择最优精度，平衡性能和精度

一句话说清楚：ops-transformer是"原材料"，ATB是"预制菜"。预制菜不是原料更好，而是加工流程优化了——中间步骤不写回HBM，直接在片上传递。

昇腾NPU上跑大模型推理，ATB是生产环境首选。ops-transformer适合研究和调试，ATB适合追求极致性能的部署。

意外收获：ATB不只是推理能用，训练也能用。ATB的训练模式提供了反向传播融合，下次有机会可以拆一下训练场景的实战。