【Transformer 与注意力机制】42｜FlashAttention：注意力计算的硬件级重写

很多人第一次听到 FlashAttention，会以为它是一种新的注意力机制：也许像稀疏注意力那样少算一部分 token，或者像线性注意力那样把 softmax 近似掉。这个理解正好反了。FlashAttention 最重要的特点是：它算的仍然是标准 scaled dot-product attention，而且是 exact attention。

若要从 CUDA kernel、在线 softmax 递推与历史实测正确性入手，见 GPU 算子工程系列的 FlashAttention：在线 softmax 与 IO-aware 注意力（算子实现路径）；本文侧重机制边界与学术脉络。

它真正改写的是计算方式。标准 attention 的数学公式很简洁：

\[ \mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V \]

但这个公式如果按朴素方式实现，会在 GPU 内存里产生巨大的中间矩阵。FlashAttention 的问题意识不是“attention 数学上能不能少算”，而是“同样的数学结果，能不能少从高带宽显存里搬数据”。

本篇能让你学会三件事：

标准 attention 的实际瓶颈为什么经常不是 FLOPs，而是内存读写；
tiling 和 online softmax 如何让 FlashAttention 不物化完整注意力矩阵；
为什么 FlashAttention 很重要，但它没有从理论上消灭 $O(n^2)$。

一、标准 attention 到底在内存里做了什么

先看朴素实现。给定 $Q,K,V$，第一步计算 $S=QK^T/\sqrt{d}$。如果序列长度是 $n$，那么 $S$ 是一个 $n \times n$ 矩阵。第二步对 $S$ 的每一行做 softmax，得到注意力概率矩阵 $P$。第三步计算 $O=PV$，得到输出。

问题在于，$S$ 和 $P$ 都很大。序列长度翻倍，矩阵元素数变成四倍。训练时还不只是 forward 要用，backward 也需要中间信息。朴素实现通常会把这些中间矩阵写到 HBM（High Bandwidth Memory，高带宽显存）里，后面再读回来。

从数学公式看，这只是三步矩阵运算；从硬件执行看，它可能变成大量显存读写。GPU 的算力很强，但不同层级内存速度差异很大。片上 SRAM 很快但容量小，HBM 容量大但访问代价高。一个算法如果反复把大矩阵写到 HBM 再读回来，就会被内存 I/O 卡住。

这就是 FlashAttention 论文标题里 “IO-Awareness” 的含义（Dao et al., NeurIPS 2022）。它不是只数乘加次数，而是显式建模 HBM 与片上 SRAM 之间的搬运代价：在 SRAM 容量为 $M$ 的假设下，标准实现需要 $\Theta(N^2)$ 级别的中间矩阵读写，而 IO-aware 分块可以把 HBM 流量压到与 $N^2 d / M$ 相关的量级——FLOP 仍是 $\Theta(N^2 d)$，变的是数据路径。这与 CPU 上 GotoBLAS 用 cache blocking 降低 DRAM 流量是同一类思想，只是层级换成了 GPU 的 HBM / shared / 寄存器。

二、为什么瓶颈常常不是 FLOPs

讨论深度学习性能时，人们习惯说 FLOPs。但在现代 GPU 上，很多算子并不是算术单元不够，而是数据喂不进去。矩阵乘法通常能较好利用 GPU，因为它有高算术强度：同一批数据被加载后可以参与很多乘加。相反，如果一个操作需要频繁读写大张量，但每个元素参与的计算不多，就容易被内存带宽限制。

标准 attention 的麻烦在于中间矩阵太大。$QK^T$ 得到的分数矩阵需要被写出；softmax 读入分数，写出概率；$PV$ 再读入概率。即使这些步骤各自可以调用高效 kernel，kernel 之间的边界仍然迫使中间结果落到 HBM。

这和写程序时的直觉很像：如果你把一个巨大临时数组写到内存，下一步马上再读回来，真正耗时的可能不是某个算术表达式，而是这次写和读。FlashAttention 的目标就是避免这种“写出去又读回来”的中间矩阵。

注意，这里不是说 FLOPs 不重要。attention 的 $O(n^2d)$ 乘加仍然存在。FlashAttention 改进的是实际硬件执行路径：让更多数据停留在更快的片上内存中，减少 HBM traffic，从而提高吞吐并降低显存占用。

三、FlashAttention 的核心思想：不物化完整注意力矩阵

FlashAttention 把 $Q,K,V$ 切成 block。每次只处理一小块 query 和一小块 key/value，把它们加载到片上内存里，计算局部 attention 分数，然后把结果累积到输出。关键是：完整的 $n \times n$ 注意力矩阵从来不需要整体写回 HBM。

直觉上，可以把标准 attention 想成先铺开一整张大表，再对这张表做 softmax 和加权求和。FlashAttention 则像一边扫描表的一小块，一边维护每一行 softmax 所需的统计量和输出累积。扫描结束后，得到的结果和完整铺表再计算一样，但中间大表没有被保存。

这听起来简单，难点在 softmax。softmax 不是局部线性运算，一行里的每个元素都依赖整行最大值和整行指数和：

\[ \mathrm{softmax}(x_i)=\frac{e^{x_i-m}}{\sum_j e^{x_j-m}} \]

其中 $m=\max_j x_j$ 用于数值稳定。如果一行被分成很多块，必须在不知道整行所有元素的情况下逐块更新最大值和归一化因子。这就是 online softmax 要解决的问题。

四、online softmax 为什么可行

softmax 的稳定计算通常先找整行最大值 $m$，再计算 $\sum_j e^{x_j-m}$。如果分块处理，第一块有自己的最大值 $m_1$ 和和 $l_1$，第二块有自己的最大值 $m_2$ 和和 $l_2$。合并时不能简单相加，因为两个和是以不同最大值为基准计算的。

解决办法是把它们转换到共同基准。新的最大值是：

\[ m=\max(m_1,m_2) \]

新的归一化和是：

\[ l=e^{m_1-m}l_1+e^{m_2-m}l_2 \]

这样就能逐块合并 softmax 的统计量。输出 $PV$ 也可以用类似方式增量更新：当新的 block 改变全局最大值和归一化因子时，旧的输出累积要按比例重新缩放，再加上新 block 的贡献。

这个思想非常关键。它让 FlashAttention 可以在不保存完整 $S$ 和 $P$ 的情况下，仍然得到和标准 attention 相同的结果。它不是近似，不是随机采样，也不是稀疏化，而是重排了精确计算。

五、FlashAttention 是 exact attention，不是 sparse attention

这一点值得单独强调。Longformer、BigBird、Sparse Transformer 这类方法改变了 attention pattern：不是每个 token 都看每个 token，而是看局部窗口、全局 token 或某种稀疏连接。它们降低复杂度的方式是少算一部分注意力。

线性注意力则尝试改变 softmax attention 的形式，用核技巧或其他近似把复杂度从二次降到线性或近似线性。它改变的是数学形式或近似目标。

FlashAttention 不这样做。它仍然让每个 query attend 到所有 key，仍然计算标准 softmax attention。它降低的是内存 I/O 和中间状态保存成本，而不是 attention 关系本身。

三条路径的对比如下（复杂度指序列长度 $N$、头维 $d$ 时的渐近阶，常数与硬件代际另论）：

路径	代表 work	是否改变 attention 语义	复杂度目标	典型代价
Exact + IO-aware	FlashAttention 系列	否	FLOP 仍 $\Theta(N^2 d)$；HBM 中间态 $\Theta(N)$	kernel 复杂、需跟硬件演进
Sparse / 局部 pattern	Longformer, BigBird	是（限定可见 key 集合）	可降至 $\Theta(N \cdot w)$（$w$ 为窗口等）	需按任务设计 pattern、可能损精度
Linear / 核近似	Performer, Linear Transformer	是（改写 softmax 或用特征映射）	常目标 $\Theta(N d^2)$ 或 $\Theta(N)$	近似误差、与预训练权重不直接兼容

FlashAttention 落在第一行：语义与 Vaswani et al. (2017) 的标准 attention 一致，只是把“先物化 $N\times N$ 再 softmax”换成“分块流式 + 在线归一化”。这也是为什么它可以被 PyTorch / Hugging Face 等栈当作 drop-in kernel 替换，而不必重训模型。稀疏与线性方法则在问另一个问题：当 $N$ 大到 exact 路径即使 IO 优化仍不可承受时，能否用更少的 token 对或不同的数学形式换可训练性——二者互补，不是谁取代谁。

六、显存收益来自哪里

训练时，显存不只被参数占用，还被激活和中间结果占用。标准 attention 如果保存完整 attention matrix，长序列下显存压力非常大。FlashAttention 避免把完整 $S$ 和 $P$ 写回显存，因此 forward 阶段需要保存的中间状态显著减少。

backward 阶段需要梯度。朴素实现可以直接使用保存下来的 $P$；FlashAttention 则倾向于保存更少统计量，在 backward 中重算部分 attention。这里有一个经典 trade-off：用少量额外计算换显存和 I/O。对现代 GPU 来说，这往往是划算的，因为 HBM traffic 比片上重算更贵。

所以 FlashAttention 的收益不是来自“少做数学”，而是来自“少保存、少搬运、必要时重算”。这个思想在深度学习系统里很常见：activation checkpointing 也是用计算换显存。FlashAttention 的特别之处在于，它针对 attention 的 softmax 结构做了精细的 I/O-aware 设计；反向传播不存完整 $P$，而存 $O(N)$ 的 running max / sum 并在 backward 重算分块概率——与 forward 的 online softmax 是同一笔“算力换 HBM”的交易。

七、学术谱系：从标准 attention 到 IO-aware kernel

今天生产栈里的 FlashAttention 可以沿下面几条线读（奠基 → 工程化，不代表时间上的严格先后）：

标准 scaled dot-product attention（Vaswani et al., NeurIPS 2017）：定义 $ (QK^/)V $ 与 $O(N^2)$ 的 token 全连接关系；朴素实现物化 $N\times N$ 分数/概率矩阵。
在线 softmax / 流式归一化（Milakov & Gimelshein, 2018；Rabe & Staar, 2021 等在长序列推理中的独立推导）：说明 softmax 的分块最大值与归一化和可增量合并——FlashAttention 论文 §3.1 直接建立在这一思想上，使“不看到整行也能得到相同 softmax”成为可证明的 exact 步骤，而非启发式近似。
IO-aware exact attention（Dao et al., NeurIPS 2022）：把 $QK^\top$ 与 $PV$ 两阶段 GEMM 与 online softmax 融合进 tiling 循环，给出 HBM 访问的渐近分析；这是“算法论文 + GPU 内存模型”的分叉点，区别于仅改 attention pattern 的稀疏 work。
并行与硬件代际（Dao, arXiv:2307.08691, 2023，未经 peer review；Shah et al., 2024 FA3 材料）：FA2 优化 warpgroup 划分与 work partitioning，在 A100 类硬件上逼近 GEMM 利用率；FA3 面向 Hopper 的 TMA、异步 wgmma 与 FP8 等——公式仍 exact，差异在能否吃满新一代 Tensor Core 与异步拷贝。

谱系上的分界很清楚：改内存访问 schedule 的 exact kernel（FlashAttention）与 改 attention 定义或近似目标（稀疏、线性、状态空间等）是两条线。长上下文工程里常见组合是：先用 FlashAttention 把标准 attention 训到尽可能长的 $N$，再在仍不够时引入稀疏/线性/压缩 KV 等架构手段。

八、FlashAttention-2 / FlashAttention-3 的演进方向

FlashAttention-2（Dao, arXiv:2307.08691, 2023）的重点是更好的并行性与 work partitioning。第一版已减少 HBM I/O，但 GPU 实际吞吐还取决于：一个 thread block 负责多少 query/key 行、warp 间如何分摊 $QK^\top$ 与 $PV$、head 维与 batch 维如何映射到 grid——FA2 在这些维度上重新划分，报告在 A100 上前向可接近同形状 GEMM 的硬件利用率，相对 FA1 约 2× 量级加速（引用数据，具体 shape 与精度见原文 table）。

FlashAttention-3（Shah et al., 2024）进一步面向 Hopper（H100）：利用 TMA 异步搬 tile、warpgroup 级异步 MMA，并探索 FP8 等低精度路径。它仍然输出 exact attention（在声明的数值语义下），不是新的稀疏或线性机制；收益高度依赖 sm_90 特性，与 Ampere 消费级卡上的 FA1/FA2 不可直接比吞吐。

版本演进背后的共同方向很清楚：数学目标不变，实现必须跟硬件代际走。随着模型变大、上下文变长，单纯写出 $\mathrm{Attention}(Q,K,V)$ 已经不足以判断能否训练——还要问中间矩阵会不会把 HBM 打满、kernel 能否融合 GEMM 与 softmax、Tensor Core 能否持续喂饱。算法与 kernel 的边界在变薄；具体 CUDA 与 CUTLASS 实现见 gpu-kernel/14。

九、它改变了什么，没有改变什么

FlashAttention 改变了 Transformer 训练和推理的实际可用边界。更少显存、更高吞吐，意味着同样硬件上可以使用更长序列、更大 batch，或者把显存留给更大的模型和 optimizer state。它让很多原本“理论上能训、实际上太贵”的配置变得可行。

但它没有改变 attention 的理论复杂度。每个 query 仍然要和每个 key 交互，元素数量仍然随序列长度平方增长。序列长度从 8K 到 16K，attention 关系数量仍然变成四倍。FlashAttention 能让这件事更高效，但不能把二次关系变成线性关系。

这就是为什么 FlashAttention 和稀疏注意力、线性注意力、状态空间模型并不是互相替代。FlashAttention 是“把标准 attention 做得更好”；后几类方法是在问“能不能不用完整标准 attention”。前者是工程路径，后者是架构路径。

理解这个边界很重要。否则很容易产生误解：用了 FlashAttention，就没有长上下文瓶颈了。事实是，瓶颈被推远了，但没有消失。

十、关键概念回顾

HBM：GPU 高带宽显存，容量大但访问代价高于片上 SRAM。
I/O-aware：把内存读写成本纳入算法设计，而不只统计 FLOPs。
tiling：把矩阵分成 block，分块加载到片上内存中计算。
online softmax：逐块维护 softmax 的最大值和归一化和，从而不需要一次看到完整行。
exact attention：输出与标准 softmax attention 一致，不改变数学目标。
materialization：把中间矩阵完整写入内存。FlashAttention 的核心收益之一是避免物化完整 attention matrix。

十一、常见误解

11.1 “FlashAttention 是近似注意力”

不是。FlashAttention 计算的是标准 softmax attention 的精确结果，只是改变了计算和内存访问顺序。

11.2 “用了 FlashAttention 就没有长上下文瓶颈”

不对。FlashAttention 降低了实际显存和 I/O 成本，但 attention 关系数量仍然是二次增长。长上下文瓶颈被缓解，不是被消灭。

11.3 “性能只看 FLOPs”

深度学习系统里，FLOPs 只是一个维度。对 attention 这类操作，HBM 读写、中间矩阵保存和 kernel 调度都可能成为决定因素。

11.4 “FlashAttention 只对训练有用”

训练收益很明显，但推理中的 prefill 阶段同样可以受益。decode 阶段的瓶颈又会转向 KV Cache 读写和自回归串行性，这会在 KV Cache 一篇里展开。

十二、争论与开放问题

争论：有了 FlashAttention，还要不要稀疏/线性 attention？ 一派认为，exact 路径把 HBM 与激活显存压到 $\Theta(N)$ 后，多数预训练与 SFT 应继续用标准 attention，换 kernel 即可（FlashAttention 在主流框架中的默认集成是这一立场的工程证据）。另一派指出，$O(N^2)$ FLOP 与 quadratic 的 token 对关系仍在：极长上下文（100K+）、多模态长序列或算力预算紧时，Longformer 类稀疏、Performer 类线性或 Mamba/SSM 等 改关系或改状态更新 的方法仍有文献与产品路线（参见本系列 43｜稀疏与局部注意力）。共识边界是：FlashAttention 不消灭二次 FLOP，只让 exact 路径在更长 $N$ 上仍可行；是否进一步改架构取决于任务与预算，不存在单一最优。

开放问题（可检验，非展望式口号）：

Prefill vs decode 的分工：FlashAttention 主要优化 prefill 与训练 forward 的 $QK^\top$ / $PV$ 融合；decode 阶段瓶颈常转向 KV Cache 带宽与自回归串行（见本系列 KV Cache 篇）。如何把 IO-aware 思想延伸到 paged KV、speculative decode 与多请求 batching，仍是 serving 栈活跃工程题。
数值与低精度：FA3 在 Hopper 上探索 FP8 等路径；在更低位宽下 exact 的语义（相对 FP32 baseline 的误差界）与训练稳定性如何权衡，peer-reviewed 与生产报告仍在积累。
与编译器/DSL 的边界：Triton、CUTLASS、FlashInfer 等把 tiling 自动化；手写 CUDA 何时仍必要（特殊 mask、融合 layernorm、变长序列）——见 gpu-kernel 系列与 Triton 篇，无统一答案，取决于 shape 分布与维护成本。

十三、下一步

FlashAttention 说明了一件事：同样的注意力公式，换一种硬件友好的计算路径，就能显著改变可训练边界。但推理阶段还有另一套问题：历史 token 的 Key/Value 能不能复用？为什么训练和推理像两种不同程序？这就是后面 KV Cache 要讲的主题。

十四、参考文献

核心论文

Vaswani, A. et al. “Attention Is All You Need.” NeurIPS 2017. 标准 scaled dot-product attention。
Milakov, M. & Gimelshein, N. “Online Normalizer Calculation for Softmax.” 2018. 分块 softmax 归一化的基础。
Dao, T. et al. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS 2022. IO-aware exact attention 原始论文。
Dao, T. “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.” arXiv:2307.08691, 2023. 未经 peer review；FA2 并行划分与 A100 引用数据。
Shah, J. et al. “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision.” 2024. Hopper / FP8 方向。

规范与对照

NVIDIA. “CUDA C++ Programming Guide.” GPU 内存层级与编程模型。
Beltagy, I. et al. “Longformer.” 2020；Tay, Y. et al. “Efficient Transformers: A Survey.” 2020. 稀疏/高效 attention 对照（与 exact IO-aware 路径对比时引用）。

← 上一篇：41｜位置编码演进　|　下一篇：43｜稀疏与局部注意力 →

同主题继续阅读

把当前热点继续串成多页阅读，而不是停在单篇消费。

2026-04-15 · transformer