上周帮一个团队在 Atlas 800 上部署 LLaMA2-70B 推理服务，要求 4096 上下文、单卡跑通。他们的 PyTorch 模型在昇腾NPU上能跑，但一开 4096 序列就 OOM。

我看了眼 profile 数据，问题很清楚：标准注意力在序列 4096 时，单层中间结果吃掉 4GB 显存。70B 模型 80 层，光注意力中间结果就要 320GB——这不是优化能解决的，得换实现。

方案是 CANN 的 ops-transformer 仓库里的 FlashAttention。说起来一句话，真正跑通花了我们两天。这篇文章把踩过的坑全部记录下来，你可以直接照着走。

---

环境准备：先别急着装东西

很多人一上来就 git clone 然后编译，报错了再排查。别这么干，先确认基础环境。

# 确认 NPU 驱动正常
npu-smi info

# 确认 CANN 版本（FlashAttention 要求 8.0+）
cat /usr/local/Ascend/ascend-toolkit/latest/version.cfg

npu-smi 能正常输出卡信息，CANN 版本 ≥ 8.0，才能继续。低于 8.0 的话 ops-transformer 的 FlashAttention 接口不兼容，得先升级 CANN。

⚠️ 踩坑预警：CANN 升级后要重启机器，不然驱动版本和工具链版本对不上，运行时会报 ACL_ERROR_RT_INITIALIZE。

---

编译 ops-transformer：依赖链别漏

ops-transformer 依赖 opbase，opbase 依赖 CANN 的基础库。编译顺序搞错，link 阶段会找不到符号。

# 第一步：克隆并编译 opbase
git clone https://atomgit.com/cann/opbase.git
cd opbase
export ASCEND_HOME=/usr/local/Ascend/ascend-toolkit/latest
bash build.sh install

# 第二步：编译 ops-transformer
cd ..
git clone https://atomgit.com/cann/ops-transformer.git
cd ops-transformer
bash build.sh

编译成功后，动态库在 output/lib/libascend_ops_transformer.so。

⚠️ 踩坑预警：如果 build.sh 报 cmake not found，先装 cmake 3.16+。CANN 自带了 cmake，路径在 $ASCEND_HOME/toolkit/cmake/bin/cmake，把它加到 PATH 里。

export PATH=$ASCEND_HOME/toolkit/cmake/bin:$PATH

把编译产物拷到 CANN 的库目录：

sudo cp output/lib/*.so $ASCEND_HOME/lib64/
sudo cp output/include/* $ASCEND_HOME/include/

---

最小可运行代码：先别碰模型

在接入完整模型之前，先用一个最简单的测试脚本验证 FlashAttention 能正常工作。很多人跳过这步，直接改模型代码，报错了不知道是 FlashAttention 的问题还是模型的问题。

import torch
from opstransformer import flash_attention

# 最小测试：batch=1, heads=2, seq=512, d=64
batch, heads, seq_len, d = 1, 2, 512, 64

q = torch.randn(batch, seq_len, heads, d).npu()
k = torch.randn(batch, seq_len, heads, d).npu()
v = torch.randn(batch, seq_len, heads, d).npu()

# 调用 FlashAttention，BSND 布局
output = flash_attention(
 q, k, v,
 head_num=heads,
 input_layout="BSND",
 scale=0.0
)

print(f"输出shape: {output.shape}") # 应该是 [1, 512, 128]
print(f"输出范围: [{output.min():.4f}, {output.max():.4f}]")

能正常输出，说明编译和安装都没问题。

⚠️ 踩坑预警：输入张量必须在 npu 上。如果在 cpu 上调用，直接段错误，没有任何报错信息。调试了半天才发现。

---

精度验证：必须做的一步

FlashAttention 是分块计算，数值误差会比标准注意力大。必须验证精度是否在可接受范围内。

import torch
import torch.nn.functional as F

# 生成相同输入
torch.npu.manual_seed(42)
q = torch.randn(1, 512, 8, 64).npu()
k = torch.randn(1, 512, 8, 64).npu()
v = torch.randn(1, 512, 8, 64).npu()

# 标准注意力
q_std = q.transpose(1, 2) # [1, 8, 512, 64]
k_std = k.transpose(1, 2)
v_std = v.transpose(1, 2)
scores = torch.matmul(q_std, k_std.transpose(-2, -1)) / (64 ** 0.5)
attn = F.softmax(scores, dim=-1)
output_std = torch.matmul(attn, v_std)
output_std = output_std.transpose(1, 2) # [1, 512, 8, 64]

# FlashAttention
output_flash = flash_attention(q, k, v, head_num=8, input_layout="BSND")

# 对比
diff = (output_std.float() - output_flash.float()).abs()
print(f"最大误差: {diff.max():.6f}")
print(f"平均误差: {diff.mean():.6f}")

float16 下，最大误差 < 0.05 算正常。如果误差超过 0.1，检查 scale 参数是否正确。

---

接入真实模型：改最小代码

以 HuggingFace 的 LLaMA 模型为例，只需要改一个文件。

找到模型里的注意力实现，通常在 modeling_llama.py 的 LlamaAttention.forward 里。

原始代码大概是：

# 原始实现
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
attn_weights = F.softmax(attn_weights, dim=-1)
attn_output = torch.matmul(attn_weights, value_states)

替换为：

# FlashAttention 替换
from opstransformer import flash_attention

# 把 query/key/value 从 [batch, heads, seq, d] 转成 BSND 布局
q = query_states.transpose(1, 2).contiguous() # [batch, seq, heads, d]
k = key_states.transpose(1, 2).contiguous()
v = value_states.transpose(1, 2).contiguous()

attn_output = flash_attention(
 q, k, v,
 head_num=self.num_heads,
 input_layout="BSND",
 scale=1.0 / math.sqrt(self.head_dim)
)

attn_output = attn_output.transpose(1, 2) # 转回 [batch, heads, seq, d]

改动就这么点。剩下的推理脚本、权重加载、tokenizer 全部不用动。

---

显存对比：跑个数据看看效果

import torch

seq_len = 4096
batch = 1
heads = 64
d = 128

torch.npu.reset_peak_memory_stats()

q = torch.randn(batch, seq_len, heads, d).npu()
k = torch.randn(batch, seq_len, heads, d).npu()
v = torch.randn(batch, seq_len, heads, d).npu()

# 标准注意力
q_std = q.transpose(1, 2)
k_std = k.transpose(1, 2)
v_std = v.transpose(1, 2)
scores = torch.matmul(q_std, k_std.transpose(-2, -1)) / (d ** 0.5)
attn = F.softmax(scores, dim=-1)
out_std = torch.matmul(attn, v_std)
mem_standard = torch.npu.max_memory_allocated() / 1024**3

torch.npu.reset_peak_memory_stats()
q2 = q.clone()
k2 = k.clone()
v2 = v.clone()
out_flash = flash_attention(q2, k2, v2, head_num=heads, input_layout="BSND")
mem_flash = torch.npu.max_memory_allocated() / 1024**3

print(f"标准注意力: {mem_standard:.2f} GB")
print(f"FlashAttention: {mem_flash:.2f} GB")
print(f"节省: {(1 - mem_flash/mem_standard)*100:.0f}%")

在 Ascend 910 上，4096 序列、64 头的实测结果：标准注意力 4.1GB，FlashAttention 0.8GB，节省约 80%。

---

常见报错速查

报错信息	原因	解决
Segmentation fault	输入张量不在 npu 上	.npu() 转设备
invalid input layout	张量布局和参数不匹配	BSND/BNSD 对齐
head_num mismatch	head_num 和张量实际维度不一致	检查 num_heads
ACL_ERROR_RT_INITIALIZE	CANN 驱动和工具链版本不一致	重启机器
symbol not found	opbase 没装或路径不对	先编译 opbase
编译报 cmake 错误	cmake 版本太低	用 CANN 自带的 cmake

---

下一步

如果你正在做大模型推理，显存是主要瓶颈，按照上面的步骤从测试脚本开始，验证精度后再接入模型。不要跳步，每步都跑通再往下走。

仓库地址：https://atomgit.com/cann/ops-transformer

ops-transformer 里除了 FlashAttention，还有 MoE、MC2 等大模型算子。如果你在跑 Mixtral 之类的 MoE 模型，可以一起看看，原理类似。