前言

刚把项目从 PyTorch 原生 attention 迁移到昇腾融合算子那周，我踩的坑能写一本小册子——编译报错、路径不对、算子精度对不上、融合没生效……全遇上了。

后来静下心来把 ops-transformer 仓库的 FlashAttention 部分从头梳理了一遍，终于搞清楚问题出在哪。

ops-transformer 是 CANN 算子服务层的核心仓库，FlashAttention 融合算子的底层实现就在里面。如果你之前用的是 ATB 的 Python 接口，想往更深的地方走，或者想自己定制融合策略，这篇文章就是给你准备的。

---

一、环境准备

跑 ops-transformer 的 FlashAttention 之前，基础环境必须先到位。

依赖清单：

• 昇腾 NPU（Atlas A2/A3 系列）
• CANN ≥ 8.0（低于这个版本没有 FlashAttention 融合算子支持）
• Python ≥ 3.8
• CMake ≥ 3.16
• Git

⚠️ 踩坑预警： CANN 版本一定要对。之前遇到最常见的问题就是"我明明装的 CANN 8.0，怎么算子编译还是报错"——后来一问，他机器上同时跑了两个 CANN 版本的环境，编译器调用的不是 8.0 那个。跑之前先确认一下：

# 确认当前生效的 CANN 版本
cat /usr/local/Ascend/ascend-toolkit/version.ini

看到 8.0.x 或者更高版本才继续往下走。

检查 NPU 是否可用：

npu-smi info

能正常看到设备信息就行。

---

二、clone ops-transformer 仓库

ops-transformer 仓库在 AtomGit 上，clone 下来：

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

仓库结构如下：

ops-transformer/
├── flash_attention/ # FlashAttention 算子核心实现
│ ├── src/ # Ascend C 源码
│ ├── CMakeLists.txt
│ └── README.md # 先读这个！
├── examples/ # 调用示例
├── test/ # 单元测试
└── docs/ # 架构文档

⚠️ 踩坑预警 2： 不要直接进 src/ 目录开始看代码。flash_attention/README.md 里有编译顺序说明——ops-transformer 依赖 opbase（算子基础组件仓库），必须先编译 opbase 再编译 flash_attention，顺序反了编译会失败。

先看一下 README 的编译顺序：

cat flash_attention/README.md | grep -A 5 "build order"

---

三、编译 FlashAttention 算子

第一步：编译 opbase

ops-transformer 的所有算子都依赖 opbase，先把它拿下：

# 先 clone opbase（README 里会告诉你仓库地址）
git clone https://atomgit.com/cann/opbase.git
cd opbase

# 编译
mkdir build && cd build
cmake ..
make -j16

⚠️ 踩坑预警 3： opbase 的 CMakeLists.txt 里默认没有开 ENABLE_TESTING，先开了方便后面验证：

cmake .. -DENABLE_TESTING=ON

第二步：编译 flash_attention

opbase 编译成功后，回到 ops-transformer 目录：

cd ops-transformer/flash_attention
mkdir build && cd build

# 指定 opbase 的安装路径
cmake .. -DOPBASE_ROOT=/path/to/opbase/install
make -j16

编译成功会生成 libops_transformer_flash_attention.so 动态库文件。

---

四、跑通第一个样例

编译成功后，进 examples 目录看调用示例：

cd ops-transformer/examples/flash_attention/

你会看到两个关键文件：

examples/
├── flash_attention_main.cpp # 主程序
├── run_example.sh # 一键运行脚本
└── config/
 └── fa_config.json # 算子配置（tiling 参数、精度等）

先不急着改代码，直接跑通：

./run_example.sh

正常输出：

[INFO] Loading FlashAttention kernel from libops_transformer_flash_attention.so
[INFO] Kernel loaded successfully.
[INFO] Input shape: [batch=1, heads=12, seq_len=512, head_dim=64]
[INFO] Output shape: [batch=1, heads=12, seq_len=512, head_dim=64]
[INFO] Execution time: 2.31 ms
[INFO] All checks passed.

看到 All checks passed 就说明算子已经在 NPU 上跑通了。

⚠️ 踩坑预警 4： 第一次跑如果报 kernel not found，大概率是 LD_LIBRARY_PATH 没设对：

export LD_LIBRARY_PATH=/path/to/ops-transformer/flash_attention/build:$LD_LIBRARY_PATH

加完再跑一次。

---

五、代码里真正要看懂的地方

跑通之后，回到 src/ 目录看源码。有三处是 FlashAttention 在昇腾上实现的关键：

1. Tiling 参数为什么是 128

// 不是随便选的，是按 Ub 大小 / 数据精度算出来的
// FP16 下，128x128 tiles 能塞进 Ub，全不溢出来回 HBM
constexpr int TILE_SIZE_M = 128;
constexpr int TILE_SIZE_N = 128;

2. 在线 Softmax 在哪

// 传统 Softmax 要两遍，这里一遍搞定
// WHY：Ub 空间不够两遍扫描，只能用在线算法省空间
void online_softmax(...) {
 float max_val = -INFINITY;
 float sum_val = 0.0f;
 // 一遍扫描，同时算 max 和 sum
}

3. 融合策略怎么配

// fa_config.json 里可以配融合策略
{
 "enable_fusion": true,
 "fusion_mode": "flash_attention_v2",
 "causal_mask": true,
 "softmax_scale": 1.0
}

---

六、和 ATB 的关系：怎么选

ops-transformer 是底层，ATB 是封装层。具体怎么选：

场景	选哪个	理由
快速验证想法	ATB Python 接口	一行代码搞定，不用编译
定制融合策略	ops-transformer	直接改 Ascend C 源码
生产部署	ATB C++ 接口	稳定性有保障
调试/学习	ops-transformer	能看到完整计算逻辑

两者可以配合用：先用 ATB 快速验证效果，确认方向对了，再用 ops-transformer 做深度调优。

---

总结：一句话说就是

ops-transformer 的 FlashAttention 编译就三步：先编译 opbase → 再编译 flash_attention → 最后跑 examples 验证。

坑主要集中在两点：编译顺序不能乱（opbase 必须先编译）、LD_LIBRARY_PATH 要设对（第一次加载 so 容易找不到）。

验证跑通之后，想快速出效果接 ATB 的 Python 接口；想深度调优直接改 fa_config.json 里的融合参数，或者进源码改 Ascend C。