在昇腾NPU上从零跑通FlashAttention：五天实操记录

Day 1：环境装了一整天，torch_npu版本配错两次。Day 2：标准attention跑通了，显存炸了。Day 3：切FlashAttention，layout传错排查了三小时。Day 4：数值验证和性能测试。Day 5：嵌入完整模型，端到端跑通。

五天时间，从零到FlashAttention在昇腾NPU上跑起来。每一步的具体操作和踩坑记录都在下面。

Day 1：环境搭建

硬件是Atlas 800训练服务器，里面是Ascend 910。操作系统EulerOS 2.10。

# CANN 8.0安装（从昇腾官网下对应版本） tar -zxvf Ascend-cann-toolkit_8.0_linux-x86_64.tar.gz cd Ascend-cann-toolkit_8.0_linux-x86_64 ./install.sh # 环境变量，加到~/.bashrc里 source /usr/local/Ascend/ascend-toolkit/set_env.sh # 验证 npu-smi info # 应该能看到910设备列表

然后装torch_npu。这一步我翻车了两次：

# ❌ 第一次翻车：CANN 8.0配了torch_npu 2.0的包 pip install torch_npu==2.0.1 # 结果import报错，API对不上 # ✅ 正确做法：查版本对应表再装 # CANN 8.0 对应 torch_npu 2.1.0 pip install torch==2.1.0 pip install torch_npu==2.1.0.post3

验证torch_npu装好了：

import torch import torch_npu print(torch.npu.is_available()) # True print(torch.npu.device_count()) # 看到卡数 print(torch.npu.get_device_name(0)) # Ascend 910

如果is_available()返回False，大概率是torch_npu和CANN版本不匹配。cann-learning-hub的入门教程里有版本对应表，先查再装。

Day 2：标准attention跑通，搞清楚瓶颈

先不碰FlashAttention，把标准attention跑一遍，亲眼看到显存问题：

import torch import torch_npu from torch_npu.contrib import transfer_to_npu # .cuda()自动重定向到.npu() import time def bench_standard_attn(batch, heads, seq_len, dim): """标准attention性能测试""" q = torch.randn(batch, heads, seq_len, dim, device='npu', dtype=torch.float16) k = torch.randn(batch, heads, seq_len, dim, device='npu', dtype=torch.float16) v = torch.randn(batch, heads, seq_len, dim, device='npu', dtype=torch.float16) # 预热3轮 for _ in range(3): scores = torch.matmul(q, k.transpose(-2, -1)) / (dim ** 0.5) attn = torch.softmax(scores, dim=-1) out = torch.matmul(attn, v) torch.npu.synchronize() # 计时20轮 t0 = time.time() for _ in range(20): scores = torch.matmul(q, k.transpose(-2, -1)) / (dim ** 0.5) attn = torch.softmax(scores, dim=-1) out = torch.matmul(attn, v) torch.npu.synchronize() latency = (time.time() - t0) / 20 * 1000 mem = torch.npu.memory_allocated() / 1024**3 return latency, mem # 从短到长测 for seq in [512, 2048, 4096]: ms, mem = bench_standard_attn(4, 32, seq, 128) print(f"seq={seq}: {ms:.1f}ms, 显存={mem:.1f}GB")

输出：

seq=512: 8.2ms, 显存=2.1GB seq=2048: 48.6ms, 显存=9.8GB seq=4096: 187.3ms, 显存=34.2GB

seq=8192直接OOM。原因很简单：scores矩阵大小是4×32×8192×8192×2字节≈16GB，再加上softmax的中间结果，单层attention就要几十GB。标准attention的显存是O(N²)，序列一长就炸。

Day 3：切FlashAttention，踩最大一个坑

第一版代码（有bug）

import torch_npu q = torch.randn(4, 32, 4096, 128, device='npu', dtype=torch.float16) k = torch.randn(4, 32, 4096, 128, device='npu', dtype=torch.float16) v = torch.randn(4, 32, 4096, 128, device='npu', dtype=torch.float16) out = torch_npu.npu_flash_attention( q, k, v, head_num=32, input_layout="BSND", scale=1.0 / (128 ** 0.5), keep_prob=1.0, )

报错：RuntimeError: shape mismatch。

排查：layout搞反了

我的tensor shape是[4, 32, 4096, 128]，意思是[batch, heads, seq, dim]，即BNSD格式。但我传了input_layout="BSND"，接口按[batch, seq, heads, dim]理解，把32当成了序列长度，4096当成了头数——当然对不上。

# ✅ 修正：layout跟实际shape匹配 out = torch_npu.npu_flash_attention( q, k, v, head_num=32, input_layout="BNSD", # 改成BNSD scale=1.0 / (128 ** 0.5), keep_prob=1.0, ) print(out.shape) # [4, 32, 4096, 128] ✅

这个坑cann-learning-hub的FlashAttention教程里专门有一节讲，我当初跳过了，代价是三小时排查。

序列长度对齐

跑通之后换seq=3000试试，又报错。原因：FlashAttention在昇腾NPU上要求序列长度是16的倍数。

def pad_seq(tensor, align=16): """把序列长度padding到align的倍数""" seq = tensor.size(2) if tensor.dim() == 4 else tensor.size(1) if seq % align == 0: return tensor, seq padded = (seq // align + 1) * align diff = padded - seq if tensor.dim() == 4 and tensor.shape[1] < tensor.shape[2]: # BNSD格式 pad = torch.zeros(tensor.size(0), tensor.size(1), diff, tensor.size(3), device=tensor.device, dtype=tensor.dtype) return torch.cat([tensor, pad], dim=2), seq else: # BSND格式 pad = torch.zeros(tensor.size(0), diff, tensor.size(2), tensor.size(3), device=tensor.device, dtype=tensor.dtype) return torch.cat([tensor, pad], dim=1), seq # 使用 q, orig_len = pad_seq(q, 16) k, _ = pad_seq(k, 16) v, _ = pad_seq(v, 16) out = torch_npu.npu_flash_attention(q, k, v, head_num=32, input_layout="BNSD", scale=1.0/(128**0.5), keep_prob=1.0) # 截回原始长度 out = out[:, :, :orig_len, :]

Day 4：数值验证和性能测试

跟标准attention对比

def verify_flash_vs_standard(): # 小规模数据，方便在CPU上跑FP32标准版做参考 batch, heads, seq, dim = 1, 8, 256, 64 q = torch.randn(batch, heads, seq, dim, device='npu', dtype=torch.float16) k = torch.randn(batch, heads, seq, dim, device='npu', dtype=torch.float16) v = torch.randn(batch, heads, seq, dim, device='npu', dtype=torch.float16) # NPU FlashAttention out_flash = torch_npu.npu_flash_attention( q, k, v, head_num=heads, input_layout="BNSD", scale=1.0/(dim**0.5), keep_prob=1.0, ) # CPU标准attention（FP32精度） q32 = q.cpu().float() k32 = k.cpu().float() v32 = v.cpu().float() scores = torch.matmul(q32, k32.transpose(-2, -1)) / (dim**0.5) attn = torch.softmax(scores, dim=-1) out_ref = torch.matmul(attn, v32) diff = (out_flash.cpu().float() - out_ref).abs() print(f"最大误差: {diff.max().item():.6f}") print(f"平均误差: {diff.mean().item():.6f}") # FP16下最大误差 < 0.02正常 assert diff.max().item() < 0.05 verify_flash_vs_standard()

我第一次跑出来误差0.28，以为FlashAttention有bug。排查发现是scale忘传了——默认的scale不是1/√d。加上之后误差降到0.009。

完整性能对比

def bench_flash_attn(batch, heads, seq_len, dim): q = torch.randn(batch, heads, seq_len, dim, device='npu', dtype=torch.float16) k = torch.randn(batch, heads, seq_len, dim, device='npu', dtype=torch.float16) v = torch.randn(batch, heads, seq_len, dim, device='npu', dtype=torch.float16) # 预热5轮（第一次有算子编译开销） for _ in range(5): torch_npu.npu_flash_attention(q, k, v, head_num=heads, input_layout="BNSD", scale=1.0/(dim**0.5), keep_prob=1.0) torch.npu.synchronize() t0 = time.time() for _ in range(50): torch_npu.npu_flash_attention(q, k, v, head_num=heads, input_layout="BNSD", scale=1.0/(dim**0.5), keep_prob=1.0) torch.npu.synchronize() return (time.time() - t0) / 50 * 1000 for seq in [512, 2048, 4096, 8192]: ms = bench_flash_attn(4, 32, seq, 128) print(f"seq={seq}: {ms:.1f}ms")

完整对比（Ascend 910，batch=4）：

Day 5：嵌入完整模型

单算子跑通只是验证，真正要用的地方是LLM推理。替换一个7B LLaMA的attention层：

class FlashAttnLayer(torch.nn.Module): def __init__(self, hidden_size, num_heads): super().__init__() self.num_heads = num_heads self.head_dim = hidden_size // num_heads self.qkv = torch.nn.Linear(hidden_size, 3 * hidden_size, bias=False) self.o_proj = torch.nn.Linear(hidden_size, hidden_size, bias=False) def forward(self, x): bsz, seq_len, _ = x.shape # Q/K/V一起算，省两次矩阵乘 qkv = self.qkv(x) q, k, v = qkv.chunk(3, dim=-1) # reshape成BNSD格式 q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k = k.view(bsz, seq_len, s ...(truncated)...

打开cann-learning-hub，从FlashAttention入门教程开始走一遍，重点看layout参数和scale参数的说明。在昇腾NPU上跑通单算子后，用文中的验证脚本对比数值一致性。社区博客搜"FlashAttention踩坑"找别人的排查经验。

https://atomgit.com/cann/cann-learning-hub