当前位置：首页 > news >正文

FlashAttention 在昇腾NPU上的极致优化

news 2026/5/31 1:00:41

刚接触 FlashAttention 那会，我被一个困惑砸懵了：明明 Attention 机制的计算量已经是 O(n²) 了，业界还在拼命优化它，图什么？

直到我看见一组数据才明白——训练一个 1750 亿参数的 GPT-3，光是 Attention 计算就要消耗 60% 的算力。这东西要是跑得慢，整个模型就是摆设。

为什么标准 Attention 是个"内存吞金兽"

传统 Attention 的问题不在计算量，在于它来来回回读写 HBM（高带宽内存）的次数太多。

算一次 Self-Attention，标准流程是这样的：

Q、K、V 三个矩阵从 HBM 读进来
计算 QK^T，得到 n×n 的注意力分数矩阵
这个矩阵要 softmax，softmax 要取指数、取和，光这一步就涉及多次矩阵运算
最后乘以 V，结果写回 HBM

问题出在哪？中间那个 n×n 的矩阵。对于一个 4096 长度的序列，这个矩阵是 4096×4096 = 1600 万个元素，单精度浮点数就是 64MB。跑一次前向传播，这个矩阵要进进出出 HBM 至少 3-4 次。光这一项，内存带宽就被吃干净了，GPU/NPU 计算单元反而在"等米下锅"。

FlashAttention 的核心思路很简单：让数据在 SRAM 里多转几圈，少回 HBM 串门。

昇腾NPU上怎么"省内存"

ops-transformer 仓里的 FlashAttention 算子，是基于昇腾异构计算架构（昇腾CANN）实现的。它的优化策略可以总结为三个字：分块计算。

具体来说，FlashAttention 把 Q、K、V 切成小块（Tile），每次只把一个小块加载到加速器的片上缓存，计算出这一块的 Attention 结果，然后和已计算的部分做融合。

这么做有两个好处：

第一，峰值内存从 O(n²) 降到 O(n)。不需要一次性把完整的注意力分数矩阵存下来了。拿 4096 序列长度来说，标准实现需要约 64MB 中间buffer，FlashAttention 只需要几百 KB 的片上缓存，差距是几百倍。

第二，计算量和标准实现完全等价。没有因为省内存就牺牲精度，数学上严格等价。

实测数据：省内存不省速度

我拿到一组在 Ascend 910 上的实测数据（来自 cann-recipes-infer 仓库的 Benchmark）：

配置	序列长度	显存占用	吞吐量
标准 Attention	4096	16.8 GB	1,250 tokens/s
FlashAttention（融合版）	4096	2.1 GB	3,870 tokens/s

显存降到原来的八分之一，吞吐量反而提升了 2 倍多。这才是真正的"降本增效"。

为什么会这样？显存带宽省下来之后，数据搬运的瓶颈没了，计算单元可以满载跑。

在昇腾NPU上怎么用

代码比想象中简单：

import torch from cann import ops # Q/K/V: [batch, heads, seq_len, head_dim] q = torch.randn(1, 32, 4096, 64, device='npu') k = torch.randn(1, 32, 4096, 64, device='npu') v = torch.randn(1, 32, 4096, 64, device='npu') # 直接调用融合算子，一次搞定 output = ops.flash_attention(q, k, v, head_dim=64)

这里没有手写 attention_mask、没有手动做 softmax 归一化，算子内部全给你融合好了。开发团队在注释里写了句大实话：