昇腾NPU小模型推理性能调优实战：从1.5s到0.7s的优化之路

不要相信第一感觉最开始觉得是.to('cpu')的问题，结果profiling后发现完全不是。性能问题一定要用工具定位，不能靠猜。Profiling是必备技能MindStudio的Timeline视图非常直观，能精确到每个算子的微秒级耗时。建议每个做昇腾开发的同学都学会用。Warmup很重要昇腾NPU首次推理会慢很多（算子编译+缓存预热），线上服务一定要做充分的warmup。版本很关键很多优化特性

chaser

260人浏览 · 2026-01-21 16:31:53

chaser · 2026-01-21 16:31:53 发布

本文目录：

一、问题背景

最近做了个模型迁移的项目，遇到了个典型的性能问题。原本在英伟达GPU上跑得好好的小模型，推理时间稳定在1秒左右，结果迁移到华为昇腾300I Duo卡后，整个流程耗时飙到了1.5秒。

具体场景是这样的：

原环境：NVIDIA GPU + vLLM框架，端到端推理时间 ~1.0s
新环境：昇腾300I Duo + PyTorch迁移方案，推理+数据下发总耗时 ~1.5s
核心差异：昇腾推理结果默认留在NPU设备上，需要额外调用.to('cpu')把数据搬到主机内存

50%的性能倒退显然不能接受，于是开始了这次调优之旅。

二、调优全流程

1.初步问题定位

一开始的思路很直接——既然多了个.to('cpu')操作，那性能损失肯定在这儿。于是我分别统计了两段耗时：

import time

# 推理阶段
start = time.time()
output = model(input_ids)
infer_time = time.time() - start

# 数据下发阶段
start = time.time()
output_cpu = output.to('cpu')
transfer_time = time.time() - start

print(f"推理耗时: {infer_time:.3f}s")
print(f"下发耗时: {transfer_time:.3f}s")

窗口上显示的结果让我更确信了这个判断：transfer_time占了大头。

基于这个"错误的定位"，我尝试了几种优化手段：

用异步拷贝减少等待时间
调整数据传输的batch size
换了几种不同的tensor格式

结果全都效果不佳，甚至有的方案反而更慢了。这时候意识到，可能方向搞错了。

2.采集Profiling数据

既然凭感觉不靠谱，那就上工具。昇腾提供了PyTorch Profiler接口，可以精确记录每个算子的执行时间。

采集方法

在推理代码里插入profiling代码（参考官方文档）：

import torch
import torch_npu
from torch_npu.profiler import profile

# 在推理代码外层包裹profiler
with profile(
    activities=[torch_npu.profiler.ProfilerActivity.CPU,
                torch_npu.profiler.ProfilerActivity.NPU],
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    
    # 你的推理代码
    with torch.no_grad():
        output = model(input_ids)
        output_cpu = output.to('cpu')
    
    # 确保数据同步（重要！）
    torch_npu.npu.synchronize()

# 导出profiling结果
prof.export_chrome_trace("./profiling_result.json")

在这里插入图片描述

采集时的两个坑：

采集不到数据？
- 原因：CPU和NPU是异步执行的，你打印时间戳的时候NPU可能还在跑
- 解决：在prof.stop()之前加上torch_npu.npu.synchronize()强制同步
循环场景下数据丢失？
- 原因：prof.step()是用来标记每轮迭代的，如果只跑一次推理可能触发不了采集
- 解决：单次推理场景下直接删掉prof.step()

3.用MindStudio分析数据

采集完数据后，用MindStudio打开profiling_result.json（下载地址见官方文档）。

打开Timeline视图后，有了关键的发型：
.to('cpu')算子本身只花了约20ms，并且真正的时间黑洞在推理阶段的几个核心算子上，之前的"下发耗时"其实包含了NPU推理的等待时间——因为Python层面打印时间戳的时候，NPU还没跑完，也就是说问题不是数据传输慢，而是推理本身变慢了。

4.根因分析

通过对比Timeline上的算子执行情况，定位到几个性能瓶颈：

算子调度开销大：小模型的单个算子执行时间短（微秒级），但PyTorch调度overhead相对明显
内存拷贝碎片化：频繁的小tensor拷贝导致带宽利用率低
动态编译损耗：首次推理时算子需要JIT编译，即使是第二次推理也有编译缓存查询的开销

5.针对性优化方案

基于上面的分析，有两条路可以走：

5.1换框架

昇腾针对小模型场景优化了专用推理框架：

1. TorchAIR框架

仓库地址：https://gitee.com/ascend/torchair
核心优势：图编译优化，将PyTorch动态图编译成静态执行图，减少调度开销
适用场景：模型结构固定，batch size变化不大

2. MindIE-Torch框架

文档地址：https://www.hiascend.com/document/detail/zh/mindie/20RC2/mindietorch/Torchdev/mindie_torch0002.html
核心优势：算子融合+内存优化，针对Transformer类小模型有专项优化
适用场景：Encoder-only或小型生成模型

这两个框架的改造成本都不高，基本只需要替换推理入口即可。

5.2PyTorch原地优化

如果业务限制不能换框架，可以试试下面几个调优开关：

1. 启用流水优化（针对Host-bound场景）

# 在启动脚本中设置环境变量
export TASK_QUEUE_ENABLE=2

这个参数让NPU的任务队列管理更激进，减少CPU-NPU之间的同步等待。实测在小batch场景下能提速15%-20%。

2. 禁用算子在线编译

import torch_npu

# 在模型加载后、推理前设置
torch_npu.npu.set_compile_mode(jit_compile=False)
torch_npu.npu.config.allow_internal_format = False

原理解释：

第一行关闭JIT编译，强制使用预编译算子库
第二行禁止算子内部格式转换（比如自动转NZ格式），减少不必要的layout变换

版本要求：

驱动固件：>=23.0.3
CANN工具包：>=8.0.RC1

⚠️ 老版本设置这两个参数无效，必须先升级！

3. 完整优化代码示例

import torch
import torch_npu
import os

# 1. 设置环境变量
os.environ['TASK_QUEUE_ENABLE'] = '2'

# 2. 模型加载
model = YourModel().to('npu:0')
model.eval()

# 3. 编译优化
torch_npu.npu.set_compile_mode(jit_compile=False)
torch_npu.npu.config.allow_internal_format = False

# 4. Warmup（重要！让算子预热充分）
with torch.no_grad():
    dummy_input = torch.randn(batch_size, seq_len).to('npu:0')
    for _ in range(10):
        _ = model(dummy_input)
    torch_npu.npu.synchronize()

# 5. 正式推理
with torch.no_grad():
    output = model(input_ids)
    output_cpu = output.to('cpu')