当 Transformer 遇上昇腾NPU：FlashAttention 算子的"记忆瘦身"秘诀

刚接触昇腾CANN那会，我被大模型推理时的显存占用吓了一跳。跑个 70B 参数模型，光是注意力机制的临时显存就能吃掉几十GB，像是给NPU喂了一头大象。后来在 ops-transformer 仓库里翻到 FlashAttention 算子，才发现这玩意儿本质上是给显存做了一次"断舍离"。

注意力机制为什么这么"记仇"？

要理解 FlashAttention，得先搞明白标准 Attention 到底干了什么蠢事。

Transformer 的核心就是 Attention 公式：Attention(Q,K,V) = softmax(QK^T / √d_k)V。看起来挺简洁，但实际跑起来是个显存怪兽。标准实现会先把 QK^T 算出来存成矩阵（大小是 seq_len × seq_len），再做 softmax，再乘 V。每一步都要把完整矩阵写在昇腾NPU的显存上。

打个比方：你要统计全校学生的身高分布，标准做法是先把所有人的身高写在一张巨大的纸上，再统计。纸太大，桌面（显存）放不下，只能折起来放（放到HBM高带宽内存），用的时候再展开。折痕（数据搬运）费时间，桌面越大（显存越大）也架不住这么折腾。

FlashAttention 的核心思路：不展开整张纸，分批处理，边算边扔。

算子是怎么"分批"的？

ops-transformer 仓库里的 FlashAttention 算子，关键技术是 Tiling（分块计算） 和 Recomputation（重计算）。

Tiling：把大矩阵切成小方块

算子不再一次性算完整的 QK^T 矩阵，而是把 Q、K、V 都切成小块（tile），比如 128×128 的子矩阵。每次只取一小块放到昇腾NPU的片上内存（L1 Buffer / Unified Buffer），算完立刻写回显存，再取下一个小块。

这一步最关键的是 在线 softmax（online softmax）。标准 softmax 要把整行算完才知道分母（sum of exp），但分块后你只能看到一部分。FlashAttention 用了一个数学技巧：先算每块的局部最大值和求和，最后再修正。像是你边逛街边记"目前看到最贵的衣服"，最后离店时再统一比较，不用把每件衣服都拍照存下来。

Recomputation：用计算换显存

反向传播时需要用到正向的注意力矩阵，标准做法是要存下来（显存占用 O(N²)）。FlashAttention 选择不存，反向时重新算一遍正向结果。这是用"时间（多算一次）“换"空间（少存矩阵）”。

在昇腾NPU上，这个取舍很划算。达芬奇架构的算力够猛，但显存（尤其是片上内存）是稀缺资源。重计算的开销远小于显存带宽的压力。

在昇腾NPU上跑有什么不一样？

ops-transformer 是 CANN 的 Transformer 类算子库，专门为昇腾NPU 的硬件特性做了适配。FlashAttention 算子用 Ascend C 编程语言写成，直接调用达芬奇架构的矩阵计算单元（Cube Unit）和向量计算单元（Vector Unit）。

实际跑起来，你能看到几个明显收益：

显存占用暴降。标准 Attention 是 O(N²)，FlashAttention 降到 O(N)。同样是 4K 序列长度，标准实现可能要 16GB 显存放注意力矩阵，FlashAttention 只需要几百MB。

速度更快。因为数据不用在 HBM 和片上内存之间反复横跳，带宽利用率高了，相同 batch size 下吞吐能提升 2-3 倍（具体看序列长度和头数）。

支持更长序列。显存不再是瓶颈后，70B 模型跑 32K 上下文不再是梦，只要其他部分（比如 RoPE 位置编码）也跟着适配。

怎么用它？

如果你在写 PyTorch 模型，想调用 ops-transformer 的 FlashAttention，大致流程是这样：

import torch
from ops_transformer import flash_attention

# 假设 Q/K/V 已经是昇腾NPU上的 tensor
# 形状：[batch, seq_len, num_heads, head_dim]
Q = torch.randn(2, 4096, 32, 128, device='npu')
K = torch.randn(2, 4096, 32, 128, device='npu')
V = torch.randn(2, 4096, 32, 128, device='npu')

# 直接调 FlashAttention，不用自己管分块
output = flash_attention(Q, K, V, causal=True)  # causal=True 用于自回归模型

# output 形状跟 Q 一样，直接下游用

代码里的坑：第一次跑会有 JIT 编译开销（大概几秒），后面就快了。如果报 OOM，先检查 causal 参数有没有乱设，再看看序列长度是不是真的需要这么长（4K 和 32K 的显存占用差 64 倍）。

跟其他方案比怎么样？

方案	显存占用	速度	适用场景
标准 Attention	O(N²)，大	慢（带宽瓶颈）	短序列（≤512）
FlashAttention（ops-transformer）	O(N)，小	快（计算瓶颈）	长序列（≥2K）
多卡并行（序列并行）	分散到多卡	取决于通信	超长序列（≥32K）

如果你们的模型序列长度超过 8K，不用 FlashAttention 基本跑不动。用 ops-transformer 的好处是它跟 CANN 深度绑定，调优是原生团队做的，踩坑概率低一些。

下一步可以做什么？

如果你还没试过，建议先在自己机器上跑一个 4K 序列的小模型，对比标准 Attention 和 FlashAttention 的显存占用（用 npu_mem_get_info() 打出来看看）。差距大到你会怀疑之前是怎么忍过来的。

如果想深入，可以去 ops-transformer 仓库翻 FlashAttention 的 Ascend C 源码，看看 Tiling 策略是怎么定的（哪个参数决定 tile 大小，为什么是 128 而不是 256）。这部分写清楚了，你对达芬奇架构的理解能上一个台阶。