从零搭建 ops-transformer 开发环境：在昇腾NPU上跑通第一个算子

之前帮一个团队在新到的 Atlas 训练服务器上搭建 CANN 开发环境，从装驱动到跑通 ops-transformer 里的 FlashAttention 算子，整条流程我完整走了一遍。把每一步操作记录下来，省得下次再踩同样的坑。

解局易否结局

13人浏览 · 2026-05-20 21:23:10

解局易否结局 · 2026-05-20 21:23:10 发布

从零搭建 ops-transformer 开发环境：在昇腾NPU上跑通第一个算子

前置条件：确认你的硬件和系统

在动手之前，先确认这几件事：

npu-smi info

如果这条命令能输出 NPU 的型号、驱动版本、固件版本等信息，说明驱动已经装好了。如果报 command not found，你需要先去昇腾社区下载对应型号的驱动和固件。

⚠️ 踩坑预警：Atlas A3 服务器和 Atlas 800 训练服务器的镜像包不一样，别下错了。按你的设备型号去下载页面对应选。

我的环境是这样的，你可以对照一下：

项目	版本
NPU	Atlas 800T A2
驱动	23.0.x
固件	23.0.x
CANN	8.0.RCx
OS	EulerOS 2.10 / Ubuntu 22.04 都行
Python	3.10

驱动和固件的版本号要一致，CANN 版本要 ≥ 驱动版本。这三者的版本兼容表在昇腾社区文档里有，查一下再往下走。

第一步：安装 CANN 工具包

从昇腾社区下载 CANN 软件包（.run 安装文件），上传到服务器后执行：

# 建议装在默认路径，自定义路径后续环境变量要自己配
chmod +x Ascend-cann-toolkit_8.0.RCx_linux-x86_64.run
./Ascend-cann-toolkit_8.0.RCx_linux-x86_64.run --install

# 装完加载环境变量
source /usr/local/Ascend/ascend-toolkit/set_env.sh

装完验证一下：

python3 -c "import acl; print(acl.get_version())"

能输出版本号就说明 CANN 装好了。如果 import 报错，大概率是环境变量没加载，检查一下 set_env.sh 是否 source 了。

第二步：安装昇腾版 PyTorch

CANN 自带针对昇腾NPU 优化的 PyTorch，别自己从源码编译，直接用官方的 wheel 包：

pip install torch==2.1.0  # 版本号以 CANN 对应的文档为准
pip install torch-npu==2.1.0.xxx  # 昇腾 NPU 后端

具体的 wheel 包下载地址在 cann-learning-hub 的「环境搭建」章节里有链接，按你的 CANN 版本选对应的包。

验证 NPU 是否可用：

import torch
import torch_npu  # 这行是关键，注册 npu 设备

# 确认 NPU 设备被识别
print(torch.npu.device_count())  # 应该输出你的 NPU 卡数

# 跑个简单矩阵乘验证算力
a = torch.randn(1024, 1024).npu()
b = torch.randn(1024, 1024).npu()
c = torch.mm(a, b)  # 在 NPU 上执行
print(c.shape)  # torch.Size([1024, 1024])

如果 torch.npu.device_count() 输出 0，检查驱动是否正常、torch-npu 版本是否和 PyTorch 版本匹配。

第三步：拉取 ops-transformer 代码

git clone https://atomgit.com/cann/ops-transformer.git
cd ops-transformer

先看一下仓库结构，不用全看懂，关注这几个目录就行：

ops-transformer/
├── flash_attention/     # FlashAttention 算子实现
├── moe/                # MoE（混合专家）算子
├── mc2/                # MC2 算子
├── cmake/              # 编译配置
└── samples/            # 调用示例（重要！）

samples/ 目录是你上手最快的地方，里面有 Python 调用示例和 C++ 调用示例。

第四步：编译算子库

ops-transformer 的算子是 Ascend C 写的，需要先编译才能被 PyTorch 调用：

mkdir build && cd build
cmake .. \
  -DCMAKE_INSTALL_PREFIX=$HOME/ops-transformer-install \
  -Dtorch_path=$(python3 -c "import torch; print(torch.utils.cmake_prefix_path)") \
  # 这行把 torch 的 cmake 路径传进去，算子编译时能找到 torch 的头文件
cmake --build . -j$(nproc)
cmake --install .

⚠️ 踩坑预警：cmake 找不到 torch 的 cmake 文件时，手动指定 -Dtorch_path。$(python3 -c "import torch.utils.cmake_prefix_path") 这个写法在某些 shell 里可能不 work，直接 python3 -c "import torch; print(torch.utils.cmake_prefix_path)" 拿到路径后硬填就行。

编译完之后，算子库会安装到你指定的 CMAKE_INSTALL_PREFIX 目录。把它加到环境变量里：

export PYTHONPATH=$HOME/ops-transformer-install/lib:$PYTHONPATH

建议把这一行加到 ~/.bashrc 里，免得每次重新 source。

第五步：跑通 FlashAttention 调用示例

进入 samples 目录，找一个 FlashAttention 的示例：

import torch
import torch_npu
from flash_attention import flash_attention_npu  # 编译好的算子

# 构造输入
batch, heads, seq_len, dim = 1, 32, 512, 64
q = torch.randn(batch, heads, seq_len, dim, dtype=torch.float16).npu()
k = torch.randn(batch, heads, seq_len, dim, dtype=torch.float16).npu()
v = torch.randn(batch, heads, seq_len, dim, dtype=torch.float16).npu()

# 调用 FlashAttention 算子
# causal=True 用于自回归模型（GPT/LLaMA），因果语言模型必备
out = flash_attention_npu(q, k, v, causal=True)

print(out.shape)  # 应该是 (1, 32, 512, 64)
print(out.dtype)  # torch.float16

# 和 PyTorch 原生实现做一次结果对比，验证正确性
# 两者的输出不会逐元素相等（浮点精度），但误差应该在 1e-3 以内

能跑出正确 shape 的输出就说明算子链路打通了。对比结果如果误差超过 1e-3，检查一下 dtype 是否一致——Ascend C 的 BF16 和 PyTorch 的 float16 计算精度有细微差异，这是正常的。

第六步：用 cann-learning-hub 系统化学习

跑通单个算子只是起点。想系统化地理解整个 CANN 算子开发体系，cann-learning-hub 是最好的入口。

打开 https://atomgit.com/cann/cann-learning-hub ，按这个顺序看：

CANN 架构概览——先搞懂五层架构，知道 ops-transformer 在第二层的哪个位置
Ascend C 编程入门——学会 Kernel 函数怎么写，Cube/Vector 指令怎么调
算子开发实战——从零写一个自定义算子，走完注册→编译→调用的全流程
cann-samples 示例集合——里面有几十个算子的调用示例，当参考手册用

cann-learning-hub 不是仓库文档那种干巴巴的 API 列表，它按学习阶段做了分层，每层都有前置知识说明。如果你是第一次接触昇腾NPU 开发，别跳步，从头过一遍。

验证清单：怎么确认环境没问题

按顺序检查，全部通过就可以开始正式开发了：

检查项	验证命令	预期结果
NPU 驱动	`npu-smi info`	输出 NPU 型号和状态
CANN 工具包	`python3 -c "import acl"`	无报错
PyTorch+NPU	`torch.npu.device_count()`	≥1
算子编译	`ls ~/ops-transformer-install/lib/`	有 .so 文件
算子调用	运行 FlashAttention 示例	输出正确 shape

下一步行动建议

环境搭好、算子跑通之后，建议按这个路径继续深入：

先去 cann-recipes-train 仓库看一个完整的 LLaMA 训练流程，理解 ops-transformer 的算子在真实模型中是怎么被 ATB（Transformer 加速库）调度使用的。然后回到 ops-transformer 源码，挑一个你感兴趣的算子（MoE 或者 MC2），对照 cann-learning-hub 里的 Ascend C 教程读代码。最后试试在 cann-competitions 里参加一次算子优化竞赛，实战出真知。