在昇腾NPU上部署LLaMA模型：从环境搭建到推理加速

本文分享了在昇腾NPU上部署LLaMA模型的经验。首先介绍了CANN异构计算架构的核心定位和五层架构，重点讲解了环境搭建、模型转换和性能优化三个关键环节。在性能优化部分，详细说明了FlashAttention算子和量化压缩的应用效果，并记录了首次推理慢和序列长度优化两个常见问题的解决方法。最后提供了实测数据对比，展示了不同优化方案在延迟、吞吐和显存占用方面的提升效果。文章指出昇腾NPU部署门槛虽低

菠萝吹雪272

357人浏览 · 2026-05-20 22:48:51

菠萝吹雪272 · 2026-05-20 22:48:51 发布

从一个实际问题开始

朋友手上有块昇腾910，想跑个7B参数的LLaMA模型做推理。他在GPU上用PyTorch跑得很顺，但换到NPU上一脸懵：驱动装了，CANN也装了，接下来呢？

这个问题其实触及了CANN的核心定位。昇腾CANN不是编译器，也不是某个框架，而是一套完整的昇腾异构计算架构。它把硬件能力抽象出来，向上提供统一的编程接口。理解这点很重要，后面很多概念都建立在这个基础上。

CANN的架构分五层，从下往上分别是基础层、执行层、编译层、服务层、语言层。我们今天要用的，主要是第2层的算子库和第4层的运行时。LLaMA模型推理需要的算子，大部分已经在ops-nn和ops-transformer这两个仓库里实现了。

环境搭建：比想象中简单

先说环境。CANN 8.0之后的版本安装简化了很多，不用像以前那样折腾依赖。

1. 驱动和固件

这一步没什么好说的，按官方文档来就行。装完后用npu-smi info看一眼，能显示设备信息就说明驱动正常。

2. CANN软件栈

CANN软件栈包括toolkit、kernel、算子库几个部分。对于推理场景，装toolkit就够了：

# 下载CANN toolkit，版本选8.0或更高
chmod +x Ascend-cann-toolkit_8.0.0_linux-x86_64.run
./Ascend-cann-toolkit_8.0.0_linux-x86_64.run --install

# 设置环境变量（必须，否则后续编译找不到头文件）
source /usr/local/Ascend/ascend-toolkit/setenv.sh

这里有个坑：环境变量一定要source，不然后面编译算子时会找不到头文件。我第一次装的时候忘了这一步，折腾了半天才发现。

3. PyTorch适配

昇腾对PyTorch的适配做得不错，安装也很简单：

pip install torch-npu

装完后测试一下：

import torch
import torch_npu

# 检查NPU是否可用
print(torch_npu.npu.is_available()) # 应该返回True
print(torch_npu.npu.device_count()) # 显示NPU数量

如果输出正常，说明环境搭建完成。

模型转换：从PyTorch到昇腾

环境有了，下一步是把模型转成昇腾能跑的格式。LLaMA模型本身是PyTorch格式，理论上可以直接加载，但为了性能，还是建议转换一下。

昇腾NPU的达芬奇架构和GPU的CUDA架构不同，算子的实现方式也不一样。PyTorch原生算子在NPU上跑，需要通过适配层转换，会有额外开销。如果提前把模型转成昇腾优化的格式，能省掉这部分开销。

直接加载PyTorch模型的代码：

import torch
import torch_npu
from transformers import LlamaForCausalLM, LlamaTokenizer

# 加载模型到NPU
model = LlamaForCausalLM.from_pretrained(
 "meta-llama/Llama-2-7b-hf",
 torch_dtype=torch.float16,
 device_map="npu:0" # 指定NPU设备
)
tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 推理
input_text = "你好，请介绍一下昇腾NPU"
inputs = tokenizer(input_text, return_tensors="pt").to("npu:0")
outputs = model.generate(**inputs, max_length=100)
print(tokenizer.decode(outputs[0]))

这个方案最简单，适合快速验证。但要注意，第一次运行会比较慢，因为算子需要编译。

性能优化：让推理更快

模型能跑了，下一步是优化性能。这部分是重头戏，也是最容易踩坑的地方。

1. FlashAttention加速

LLaMA是Transformer架构，注意力计算是性能瓶颈。CANN的ops-transformer仓库里实现了FlashAttention算子，能把注意力计算的复杂度从O(N²)降到O(N)。

import torch_npu
from ops_transformer import FlashAttention

# 使用FlashAttention替代原始attention
flash_attn = FlashAttention(head_dim=128, num_heads=32, causal=True)
output = flash_attn(q, k, v) # q, k, v: [batch, seq_len, num_heads, head_dim]

# 显存占用从约16GB降到4GB左右（seq_len=4096时）
# 原因：不需要存储完整的N×N注意力矩阵

FlashAttention的收益不只是速度，显存占用也大幅降低。这个优化对长序列场景特别有价值，比如LLaMA推理时batch_size=4、seq_len=4096，显存占用能从16GB降到4GB左右。

2. 量化压缩

如果显存还是不够，可以考虑量化。CANN内置了量化工具，能把FP16模型转成INT8，显存减半，速度提升。

from amct_pytorch import Calibrator

# 量化校准
calibrator = Calibrator(model, quant_config)
calibrator.run_calibration(calibration_data)
quantized_model = calibrator.convert()

量化会有精度损失，但LLaMA这类大模型对INT8量化比较鲁棒，实测困惑度下降不到1%。

踩坑实录

坑1：第一次推理很慢

第一次跑模型时，推理延迟特别长。后来才知道，这是Ascend C算子的编译缓存过程。昇腾NPU上的算子是用Ascend C语言写的，第一次调用时需要编译成二进制，后面就会快了。

解决方法：跑个预热请求，触发算子编译。

# 预热，触发算子编译
warmup_input = tokenizer("warmup", return_tensors="pt").to("npu:0")
_ = model.generate(**warmup_input, max_length=10)
torch.npu.synchronize() # 等待编译完成

# 正式推理（这次就快了）
start_time = time.time()
outputs = model.generate(**inputs, max_length=100)
print(f"推理耗时: {time.time() - start_time:.2f}秒")

坑2：序列长度不是block size整数倍

FlashAttention对序列长度有要求，最好是block size的整数倍。如果序列长度不满足，会有padding开销。

# 不好的做法：直接用原始长度
inputs = tokenizer(text, return_tensors="pt")

# 好的做法：pad到block size整数倍
block_size = 64 # FlashAttention的block size通常是64或128
seq_len = inputs["input_ids"].shape[1]
pad_len = (block_size - seq_len % block_size) % block_size
if pad_len > 0:
 inputs["input_ids"] = torch.nn.functional.pad(
 inputs["input_ids"], (0, pad_len), value=tokenizer.pad_token_id
 )