FlashAttention 用 IO 感知重写精确注意力

2022 年春天，斯坦福一带的几位研究者盯着一张并不好看的 profiler 图：显卡的算术单元还空着一截，墙上时间却已经耗尽。罪魁不是“算不动”，而是搬数据搬不动——标准注意力实现把一张巨大的中间分数表写进 GPU 的高带宽内存，序列一长，草稿纸比模型权重更先把显存顶穿。

他们没有改注意力的数学。5 月 27 日，arXiv 上的 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（Tri Dao、Daniel Y. Fu、Stefano Ermon、Atri Rudra、Christopher Ré）标题里的 exact 写得很死：不是稀疏、不是低秩近似，结果要与标准 softmax 注意力一致。变化只发生在实现层——把问题从“FLOPs 够不够”改写成“HBM 读写次数能不能降下来”。

GPU 上有极快却很小的片上 SRAM，以及大得多、相对更慢的 HBM。朴素注意力把巨大分数矩阵物化到 HBM，等于把最贵的带宽花在临时草稿纸上。FlashAttention 把 (Q)、(K)、(V) 切成能装进 SRAM 的块，在片上完成局部打分、归一化与加权，再写回输出；反向用重计算换掉部分需保存的激活。在线 softmax 一类技巧保证分块归一化后仍等于全局 softmax。若只记住“更快”，会漏掉真正的锋利处：瓶颈诊断从 FLOPs 转到 HBM 访问次数。

论文在 BERT、GPT-2 等设定上报告端到端加速，并展示更长上下文在相同硬件上更可行。数字绑在当时硬件与内核上，宜读作数量级论证。二次复杂度的理论阴影仍在；“这块卡上能不能跑起来”却从常常否决，变成可以认真调 kernel 与块大小的对象。开源仓库随后把方法落成可编译的 CUDA 路径。FlashAttention-2、3 会继续拧占用率，但 2022 年文本钉住的是问题重述：在 Transformer 已经统治序列建模之后，关键增益也可以来自对存储层次的诚实建模。精确、可微、对 IO 敏感的融合内核，撑起了长上下文时代的一截基础设施。

对在 2022 年前后试图把上下文从 2K 硬推到更长窗口的人来说，这不是抽象优雅，而是训练任务能否在截止日前跑完的差别。内核不会替你解决数据质量或对齐，它只做一件更窄的事：让精确的二次注意力在商品 GPU 的真实带宽下重新变得可部署。基础设施的进步，常常长这样。

In the spring of 2022, a small group of researchers around Stanford stared at an ugly profiler plot: arithmetic units still had headroom, yet wall time was already gone. The bottleneck was not “not enough compute,” but moving data—standard attention materializes a huge intermediate score table in GPU high-bandwidth memory, and long sequences fill that scratchpad before the model weights do.

They did not change the math of attention. On 27 May, arXiv carried FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness by Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. The word exact is deliberate: not sparse patterns, not low-rank approximation—the result must match standard softmax attention. What changed was the implementation: the question shifted from “are there enough FLOPs?” to “can we cut HBM traffic?”

GPUs have tiny, very fast on-chip SRAM and much larger, slower HBM. Naive attention materializes a huge score matrix in HBM, spending the dearest bandwidth on scratch paper. FlashAttention tiles (Q), (K), and (V) into blocks that fit in SRAM, completes local scoring, normalization, and weighting on-chip, then writes output blocks back; the backward pass recomputes some statistics instead of storing every activation. Online-softmax-style tricks keep tiled normalization equal to global softmax. Remembering only “faster” misses the sharp claim: the bottleneck diagnosis moves from FLOPs to HBM accesses.

On BERT and GPT-2 settings the paper reports end-to-end speedups and longer contexts made feasible on the same hardware. Numbers bind to that era’s chips and kernels; read them as order-of-magnitude argument. Quadratic complexity’s shadow remains; “will it run on this card?” shifts from a frequent no to something engineers can tune with kernels and tile sizes. Open repositories then compiled the method into CUDA paths. FlashAttention-2 and -3 would later squeeze utilization further; the 2022 text nails the reframing: after Transformers already dominate sequence modeling, critical gains can also come from honest modeling of the memory hierarchy. Exact, differentiable, IO-aware fused kernels became a slice of long-context infrastructure.

For people around 2022 trying to push context hard past 2K, this was not abstract elegance but whether a training job would finish before the deadline. A kernel will not solve data quality or alignment; it does one narrower thing: make exact quadratic attention deployable again under commodity GPU bandwidth. Infrastructure progress often looks like that.