有一个事实，我第一次读到的时候觉得是在开玩笑：FlashAttention 在数学上跟标准 Attention 完全等价，不算近似，FLOPs 一分没省——反向传播时甚至多算了一些——但它就是能比 PyTorch 的标准实现快 2～4 倍，显存占用从 O(N²) 压到 O(N)。
我的第一反应是：不可能。一个实现不可能凭空快出三倍还不减配。直到我把论文里的 IO 复杂度分析看懂，才发现我之前的整个直觉模型都是错的——Attention 的瓶颈根本不在计算，而在数据搬运。
顺着 FlashAttention 的思路走一遍，你会明白一件事：GPU 内存层次的设计，比浮点运算量的数字更值得操心。
瓶颈不在算，在搬
先看 GPU 的内存层次。离计算单元最近的是寄存器，接着是每个 SM 上的 SRAM（也叫 shared memory），再往外是 L2 缓存，最外面才是显存 HBM。在 A100 上，HBM 带宽约为 2TB/s，而片上 SRAM 总带宽能到 19TB/s 量级——接近 10 倍的差距。这个数据出自 FlashAttention 论文第一页的 Figure 1，GPU 内存层次的完整说明可以参考 CUDA C 编程指南。
你可以把 HBM 想成仓库，SRAM 想成工作台。仓库什么都有，但每取一次货都要走很远的路；工作台就在手边，随取随用，可惜地方就这么大。如果一段程序频繁跑仓库，那么加工手法再快，时间也耗在路上了。
而且这个差距还在拉大：GPU 算力的增长远快于显存带宽的增长，很多算子早就不是算不快，而是数据送不进去。这就是常说的内存墙（memory wall）。对 Attention 来说，内存墙尤其致命。
标准 Attention 的内存账本
Attention 的公式很简单：Attention(Q,K,V)=softmax(QKᵀ)V。但把中间张量的流动路径画出来，问题立刻变得扎眼。
第一步算 S=QKᵀ，得到一个 N×N 矩阵。N=4096 时，S 有 1600 万个元素，fp16 下是 32MB；N=8192 时直接到 128MB。而 A100 整个 GPU 的片上 SRAM 总共也只有约 20MB。
于是标准实现只能这样干活：
- 从 HBM 读 Q、K，算出 S，把整个 S 写回 HBM。
- 再从 HBM 把 S 读出来，逐行做 softmax，得到 P，写回 HBM。
- 再从 HBM 读 P 和 V，算出 O，写回 HBM。
S 和 P 各自写一遍、读一遍，就是 4 次 32MB 的 HBM 流量（N=4096 时）。N 每翻一倍，这个开销翻 4 倍。序列一长，Attention 立刻变成显存和带宽的无底洞。
| 序列长度 N | S 矩阵大小（fp16） |
|---|---|
| 512 | 0.5 MB |
| 1024 | 2 MB |
| 2048 | 8 MB |
| 4096 | 32 MB |
| 8192 | 128 MB |
FlashAttention 论文把这种视角称为 IO-aware（IO 感知）——先看内存，再看计算。
FlashAttention 的戏法：让 S 和 P 不落地
FlashAttention 的核心思想一句话就能说清：把 N×N 的运算切成小块，每一块的 QKᵀ、softmax、加权求和全部在 SRAM 里完成，最终只把 O 写回 HBM。S 和 P 从头到尾没有被完整构造出来，也就没有机会写进 HBM。
这个技术叫 tiling。代价是 K、V 会被反复读取——每个 Q 块都要跟所有 K/V 块配对。但因为中间结果全部留在片上，HBM 总流量从 O(N²) 量级降到了 O(N²d²/M) 量级（d 是头维度，M 是 SRAM 容量）。公式看着抽象，结论却很直接：在 A100 的典型配置下，HBM 流量比标准实现低一个数量级以上。
我们提出 FlashAttention，一种 IO 感知的精确注意力算法，利用 tiling 来减少 GPU HBM 与片上 SRAM 之间的读写次数。
出处是FlashAttention 论文的摘要。注意里面两个关键限定词：IO 感知，精确。
但分块引出一个尴尬问题：softmax 要看到完整的一行才能归一化，分块之后你手上只有局部。这怎么解？
Online softmax：边走边修正
常规 softmax 分三步：找整行最大值，逐元素算 e^x，再除以总和。真正卡住分块的是第一步——最大值还没出现，怎么归一化？
FlashAttention 的回答是：不知道也没关系，边走边修正。它维护两个状态：当前见过的最大值 m，以及指数和 l。每处理一个块，如果新块里出现更大的 m，就把之前部分已算出的输出整体乘上 exp(m_old – m_new)，再做累积。这个技巧有正式名字：online softmax。
所有块处理完后，每个局部的结果都被缩放到了跟一次读完一整行相同的尺度上。数学上等价，差异只有浮点舍入误差。这也是 FlashAttention 敢自称精确（exact）而不是近似算法的原因。
反向传播里的一笔账：重算比存储更便宜
到这儿，前向传播的谜底已经揭开。更反直觉的地方在反向传播。训练 Attention 需要梯度，而梯度计算要用到前向的 S 和 P。标准实现直接把 S、P 存下来；FlashAttention 选择不存，反向时当场重算。
这听着很傻，账却很划算：重算让 FLOPs 增加了约 15%，但省掉了 S、P 写入 HBM 的全部流量。GPU 上的 FLOPs 是廉价的，HBM 的字节是昂贵的。用多一点计算，换少一点搬运，这是 IO 感知算法和普通性能优化最根本的分野。
推理时前向根本不需要存 S/P，所以长序列推理的显存也被大幅压缩。这也是为什么今天提到长上下文，最先想到的一定是 memory-efficient 的 Attention kernel。
IO 感知成了一族算法
FlashAttention 的成功带火了一种设计思路：先看数据流经内存的路径，再看计算量。顺着这条路，后续算法一个接一个冒出来。
- FlashAttention-2：改进了并行策略和线程分工，在序列长度维度上并行，同时重新划分了 warp 的工作，减少跨 warp 同步与非矩阵运算的开销。官方报告在 A100 上能达到理论峰值利用率约 70%（见FlashAttention-2 论文）。
- Flash-Decoding：针对长序列生成场景。生成阶段每个 token 的 Attention 只有一个 query，但要把所有 key 全扫一遍；把 K/V 分块并行算部分输出再合并，避免单核 IO 瓶颈（见FlashAttention-2 博客）。
- PagedAttention：vLLM 背后的核心算法。它解决的是 KV Cache 的存储效率，把连续显存分配改成按 block 分页，减少碎片，让更大的 batch 成为可能（来源：vLLM 论文）。
我后来在 vLLM 的源码里看到 PagedAttention 时，才意识到 IO 感知已经不只是一个 kernel，而是一整套设计哲学：先搞清楚数据从哪来、到哪去、以什么顺序流动，再决定怎么优化。
FlashAttention 的边界在哪
这套思路有明确的收益前提，也有明显的边界。
- 它的收益依赖 GPU 的 HBM/SRAM 带宽差。如果换成内存层次差距不大的设备（比如部分移动 GPU 或 CPU），优化空间会大幅缩水。
- 序列必须要够长。N 很小时，分块和 kernel 启动的额外开销可能盖过收益。
- 它解决的是内存访问问题，不是计算问题。Attention 本身的 O(N²) 计算量还在，序列超过几万时，计算自身会成为新的瓶颈，需要配合稀疏化或线性注意力等近似方法。
- 内核高度依赖硬件。block size、线程布局都要针对具体 GPU 调参；PyTorch 2.x 的 scaled_dot_product_attention 在支持的 GPU 上会自动调度到它，但并非任何环境都能吃到同样的红利。
回到开头那个我不信的事实。现在你大概和我一样想通了：FlashAttention 没有任何魔法，它只是让数据少走了几趟远路——从 HBM 到 SRAM 的那段路。
原创文章，作者：guanweilu，如若转载，请注明出处：https://guanweilu.cn/article/487.html