Reading

Infra: Flash Attention


Fanyi Pu · GPT-6 Astra · GPT‑5.6 Sol

· Updated

FlashAttention (Dao et al., 2022) 是 dense attention 的 IO-aware 实现:通过分块、online softmax 和重计算,减少中间矩阵在 GPU HBM 中的存储与传输。它保持 attention 的数学定义;浮点运算顺序不同可能产生舍入差异,不要求逐 bit 相同。

Attention 与计算范围

以下讨论单个 attention head,忽略 batch,令 query、key、value 的维度均为 d,序列长度为 N。采用 scaled dot-product attention:

Q,K,V∈RN×d,S=QK⊤d,P=softmax⁡(S),O=PV.

Softmax 沿每一行进行。Sij 是第 i 个 query 对第 j 个 key 的 score,Pij 是归一化后的权重;输出 Oi 是所有 value 的加权和。

分别执行矩阵乘法、softmax 和矩阵乘法,会物化两个 N×N 中间矩阵 S,P,最终输出只有 N×d。FlashAttention 在片上生成 score tile,完成归一化状态更新与加权累加后丢弃该 tile,只保留后续计算需要的状态。

推导默认没有 dropout,每行至少有一个有效 key。

GPU 存储层次与 IO

假设 N=8192、d=128,只看单个 head,张量用 BF16 存储,每个元素 2 bytes:

张量元素数量存储量
单个 Q、K、V 或 ONd2 MiB
单个 S 或 PN2128 MiB

若同时保留 S 和 P,仅这两项就是 256 MiB;32 个 head 则是 8 GiB。这里没有计入其他层的 activation、梯度或临时 workspace,也没有假设真实 kernel 的所有中间量都是 BF16。

GPU 的 HBM 容量大;shared memory 和 registers 位于片上,容量小、访问快。中间结果留在片上并直接供后续计算使用,可以省去写回 HBM 和再次读取的往返。

在上面的三算子实现中,S 要写出、读回,P 也要写出、读回。按相同的 BF16 假设,仅这四次中间矩阵传输就有 512 MiB/head,还没算 Q,K,V,O 的流量。

Kernel fusion 将相邻计算放进同一个 GPU kernel,直接传递中间结果,减少算子之间的 HBM 读写 (Liang, 2026)。算术量相同的实现,可能因此具有不同的运行时间。

Tiling 将计算划分为适合片上容量的小块,并复用已经加载的输入。Fusion 与 tiling 需要配合:完整的 N2 中间矩阵通常无法驻留片上,而分块 softmax 必须维护跨块的归一化状态。

Online softmax 与状态合并

稳定归一化

固定一个 query,将它与所有 key 的 score 记为 x1,…,xN。对应的 softmax 权重为:

pj=exj∑k=1Nexk.

Softmax 的分母依赖整行 score。对各块分别做 softmax,得到的是块内归一化权重,不能直接作为整行的权重。

直接累加 exj 可能溢出。Stable softmax 先减去整行最大值,再求指数和:

m=maxjxj,ℓ=∑jexj−m,pj=exj−mℓ.

分子、分母都乘了 e−m,所以权重不变;同时 xj−m≤0,指数就不会因为 score 太大而溢出。这是 stable softmax,也常称 safe softmax。

分块状态的合并

上述实现需要先求最大值,再计算指数和。Online softmax (Milakov & Gimelshein, 2018) 在扫描 score 时同时维护当前最大值和指数和。最大值改变时,将已有的指数和换算到新基准,无需重新读取之前的 score。

设两段 score 分别为 A,B,各自的状态为 (mA,ℓA) 和 (mB,ℓB),其中 mA 是 A 的最大值,ℓA 是减去 mA 后的指数和,B 同理。合并后的最大值为 m=max(mA,mB)。将 A 的指数和从基准 mA 换算到 m:

∑j∈Aexj−m=emA−m∑j∈Aexj−mA=emA−mℓA.

换算只需要给已有指数和乘以共同因子;最大值不变时,因子为 1。对 A,B 分别换算,再累加:

α=emA−m,β=emB−m,ℓ=αℓA+βℓB.

(m,ℓ) 是后续合并所需的状态。如果需要输出完整的 softmax 向量,最终仍需保留或重新读取 score,计算各个 pj。

Attention 加权和的累加

Attention 输出为 Oi=∑jpjVj,不需要输出各个 pj。因此可以同时维护未归一化的加权和 (Jain et al., 2026):

uA=∑j∈Aexj−mAVj.

它与 ℓA 使用相同的指数权重,换算到新基准时也乘以 α。对 B 使用 β,得到

u=αuA+βuB,Oi=uℓ.

每个 query 的持久状态为两个标量 m,ℓ 和一个 d 维向量 u。各块更新分子和分母,扫描结束后做一次除法;已处理块的 score 和指数权重可以丢弃。

以指数值为 1,2,4,8 的情况为例,把 value 简化为标量 V=(1,2,3,4)。按完整公式计算,结果是

Oi=1⋅1+2⋅2+4⋅3+8⋅41+2+4+8=4915.
Online attention: merge two blocks in a common exponential scale左块的指数质量为 1、2,value 为 1、2;右块的指数质量为 4、8,value 为 3、4。左块状态按四分之一缩放,右块状态保持原样,合并得到分母 1.875、加权和 6.125,输出等于 49/15。Block A
exp⁡(xA)=(1,2)
VA=(1,2)
mA=log⁡2,ℓA=1.5,uA=2.5
×14
Block B
exp⁡(xB)=(4,8)
VB=(3,4)
mB=log⁡8,ℓB=1.5,uB=5.5
×1
m=max(mA,mB)=log⁡8
ℓ=1.5×14+1.5×1=1.875
u=2.5×14+5.5×1=6.125
Oi=6.1251.875=4915≈3.2667
左右滑动查看完整示意图
这里的 exp⁡(x) 是便于手算的未平移指数值,实际算法使用减去最大值后的指数。两个块的分母和加权和必须按同一基准合并;value 用标量展示,向量时逐分量执行相同运算。

图中两段的基准分别是 mA=log⁡2、mB=log⁡8。合并时,把左块的 ℓA,uA 同时乘 elog⁡2−log⁡8=1/4,得到 u/ℓ=6.125/1.875=49/15,与完整计算一致。不能直接平均两段各自的 attention 输出,因为两段在完整分母中的份额不同。

分块计算与 Forward

GPU 实现将 Br 个 query 组成一个 tile,复用同一块 KV,并通过矩阵乘法计算 score 和加权和。以下 i,j 表示块编号,每个 KV tile 包含 Bc 个 token:

Qi∈RBr×d,Kj,Vj∈RBc×d,Sij=QiKj⊤d∈RBr×Bc.

实现时,一个 thread block 可以负责一个 query tile,在 GPU 的一个 SM 计算单元上执行,用 shared memory 暂存输入,把累加状态留在 registers 中。两个矩阵乘法可以交给 Tensor Cores;最大值、分母和重缩放仍然对每个 query 独立进行。

将各 query 的 u 排成矩阵 U,令 W 表示当前 score tile 相对运行中最大值计算的指数权重。W 尚未除以分母,通过 WVj 累加到 U:

Attention intermediates: HBM tensors versus on-chip tiles上方三个独立算子将完整 score S 和概率 P 分别写入并读出 HBM。下方将一个 score tile 的矩阵乘法、online softmax 和 value 加权求和融合在片上,保留运行状态并复用临时空间,最终输出 O。
Separate kernels · Q,K,V start in HBM
QK⊤d
matmul
S
HBM · N×N
softmax
row reduction
P
HBM · N×N
PV
matmul
O
HBM · N×d
S and P each make a write → read round trip
Fused tiled kernel · one query tile at a timeOn chip · shared memory + registers
QiKj⊤d
score tile
online softmax
update m,ℓ
WVj
rescale + update U
Oi=Uℓ
write to HBM
Repeat with the next Kj,Vj · reuse tile storage
左右滑动查看完整示意图
上:箭头穿过 S,P 时分别发生 HBM 写入与读回。下:虚线框内只容纳当前 tile 与运行状态,S,P 不形成完整的 HBM 张量。示意数据流,不按容量或耗时比例绘制。

下面采用 FlashAttention-2 (Dao, 2023) 的 query tile 外循环与延迟归一化形式,展示数据依赖。rowmax、rowsum 沿列归约,行向量按行 broadcast;这是片上计算的伪代码,逐行写成普通 PyTorch 操作不会自动得到一个 fused kernel。

伪代码中,m 是旧的最大值,m′ 是合并后的最大值。新块直接相对 m′ 求指数,前面公式里的 β 就已经包含在 W 中,只需用 α 重缩放旧状态。

Algorithm 1FlashAttention forward

Input Q,K,V in HBM,按 Br,Bc 分块

Output O,以及 backward 所需的 L

  1. parallel for each query tile Qi do
  2. load Qi
  3. m←−∞,ℓ←0,U←0
  4. for each KV tile Kj,Vj do
  5. load Kj,Vj
  6. S←QiKj⊤/d
  7. mask S全被 mask 的行跳过本轮更新
  8. m′←max(m,rowmax⁡(S))
  9. α←exp⁡(m−m′)
  10. W←exp⁡(S−m′)
  11. ℓ←αℓ+rowsum⁡(W)
  12. U←αU+WVj
  13. m←m′
  14. end for
  15. Oi←U/ℓ最后才归一化
  16. Li←m+log⁡ℓ
  17. store Oi,Li
  18. end for

标记行:计算重缩放系数,并用它同时调整分母与加权和。

交互演示一条长 sequence,怎样逐块算完?

点左侧选择 Q 块,观察它怎样扫过整条 KV 序列;也可以改变 KV 块大小。

沿 sequence 维度切成小块伪代码 01、04 行

Q 分成 6 块,每块 4 行;K、V 分成 6 块,每块 4 行。粗线围出 4 × 4 的 score tile。

Q ↓K,V sequence · 24 tokens →输出
1–4
5–8
9–12
13–16
17–20
21–24
把 24 × 24 的 attention 计算划分成 4 × 4 的 tile行是 query,列是 key。固定一条 query 横带,逐块向右扫描 KV;斜线只表示该范围已合并,不表示保存了 score。
O1
O2
O3
O4
O5
O6
当前计算已合并,仅标记进度
临时 S,W
4×4
尚未读取
m,ℓ,U 跨块保留已覆盖 0 / 24 个 key
Oi=U/ℓ扫描完才能归一化

每个 Q 块扫描 6 次 KV;每次只算一个 4×4 tile。

图中 24 个 token,Q 每 4 行一块;KV 每 4 或 8 行一块,分别扫描 6 或 3 次。各段依次合入同一份运行状态,扫完后得到对应的 Oi。大网格只标记计算范围,不保存完整 score 矩阵;本例不加 mask。

符号与边界条件

这里的 S,W 只活在当前 tile 内,累加完就能复用空间。Li 是逐行的 log-sum-exp,训练时留给 backward 使用。

初始化时,首个有效 tile 的 m′ 有限,因此 α=0,旧的零状态不会贡献结果。对于全被 mask 的 tile 行,需要显式跳过更新,避免计算 −∞−(−∞);若整个 query 行都无有效 key,则要由算子另行约定输出,不能直接套用这里的除法。

Causal attention: traverse KV tiles for one query tile序列长 12,每块 2 个 query 和 2 个 key。当前 query 块为第 3 块,已处理前 2 个 KV 块,正在计算第 2 块,然后处理对角块。未来的整块跳过,对角块只计算有效的下三角。
K,V tiles →
0
Q0
1
Q1
2
Q2
3
Q3
✓✓●
4
Q4
5
Q5
Qi stays on chip
Br×d
m,ℓ,U
Br, Br, Br×d · running state
Oi=Uℓ
after all valid KV tiles
Conceptual score grid · never stored in HBM
左右滑动查看完整示意图
每格代表 2×2 个 score,块编号从 0 开始。粗框表示当前 query tile 的有效范围,✓ 表示已累加,● 表示当前 tile,斜线右上方是被 mask 的位置。右侧状态每处理一块就更新一次。

Causal attention 中,整块位于未来的 tile 可以直接跳过;与对角线相交的 tile 在片上把无效 score 设为 −∞,其指数权重就是 0。这样无需单独存一个完整的三角 mask。图中各个 query tile 可以独立计算,但每个 tile 内仍要沿着有效的 KV 块合并状态。

Tile 越大,输入复用通常越充分;但 Qi、Kj、Vj、U、Sij 和 softmax 状态都要占片上资源。Registers 或 shared memory 占得太多,会减少同一 SM 能同时驻留的工作,甚至导致 register spilling。实际选块要同时考虑数据类型、head dimension 和硬件资源。

FlashAttention 节省的是中间结果的存储和搬运,dense attention 的乘法仍要做:

指标(单个 head)三算子、显式中间矩阵FlashAttention
Dense attention 算术量Θ(N2d)仍为 Θ(N2d)
HBM 中保存的 S,PΘ(N2)不保存完整矩阵
除输入输出外的归一化状态取决于实现Θ(N)
HBM 访问量有 Θ(N2) 的中间矩阵读写依赖 tile、循环顺序和片上容量

输入和输出本身仍占 Θ(Nd)。所以,显存占用随序列长度线性增长,不代表计算量或总 HBM 访问量也变成了线性。

原论文的 IO 访问量分析

原始 FlashAttention 论文 (Dao et al., 2022) 用一个两级存储模型分析 IO。设片上容量为 M 个标量元素,取 d≤M≤Nd。其 KV tile 外循环令 Bc=Θ(M/d),每个 KV 块扫描一次 Q 和输出状态;共有约 Nd/M 个 KV 块,每轮搬运 Θ(Nd) 个元素,于是 HBM 访问量为

Θ(N2d2M).

这是原论文循环与存储模型下的界,不是上面伪代码的逐项流量统计,也没有把不同精度、cache 命中或硬件调度算进去。它展示了片上空间如何换取数据复用:在 M≫d2 的适用区间,相比显式中间矩阵带来的 Θ(N2) 流量,收益明显。

固定 M,d 时,上面的 IO 界仍随 N2 增长,KV 也可能被不同 query tile 多次读取。可以用一个粗略的性能下界帮助判断优化方向:

t≳max(FLOPscompute throughput,HBM bytesHBM bandwidth).

它忽略了同步、调度以及不能完全重叠的操作,不能直接当作耗时预测。但它说明,如果瓶颈是 HBM 带宽,减少中间矩阵的读写会直接降低这一项成本;优化后,还需要检查其他执行单元是否成为新的瓶颈。

Backward 与重计算

训练时,backward 需要 attention 概率或足以恢复它的信息。保存完整 P 需要平方级空间。FlashAttention 使用 recomputation (Dao et al., 2022):保留 Q,K,V,O,并为每行保存 log-sum-exp:

Li=mi+log⁡ℓi=log⁡∑jeSij.

Backward 重新计算当前 tile 的 score,用 Pij=eSij−Li 恢复概率,计算局部梯度后丢弃中间结果。L 占 N 个标量;训练 activation 还包括 Q,K,V,O。这是 attention 内部的 activation checkpointing。

重计算增加 score 和指数运算,减少大矩阵的 HBM 存储与传输,因此可能同时降低显存占用和耗时。梯度按 tile 累加,forward 与 backward 均不物化完整的 N×N 中间矩阵。

分块梯度计算

记上游梯度为 G=∂L/∂O,则

∂L∂V=P⊤G,∂L∂P=GV⊤.

Softmax 的 backward 是逐行的。令 H=GV⊤,有

∂L∂Sij=Pij(Hij−Di),Di=∑jPijHij.

利用 Oi=∑jPijVj,将 Di 改写为

Di=∑jPij(Gi⋅Vj)=Gi⋅(∑jPijVj)=Gi⋅Oi.

因此可先由 G,O 计算 D,无需扫描完整 P。每个 tile 恢复局部 P,H 后,计算局部 ∂L/∂S,再累加

∂L∂Q=1d∂L∂SK,∂L∂K=1d(∂L∂S)⊤Q.

这些是完整矩阵写法;实现时仍然逐块累加,不物化 P,H 或 score gradient。与 forward 的独立 query 行不同,多个块可能贡献给同一份梯度,backward 的工作划分还要处理这种归约。

启用 attention dropout 时,还必须重现 forward 的随机 mask,例如保存可重放的随机数状态;上面无 dropout 的梯度推导不能原样忽略这一步。

FA2 至 FA4 的硬件优化

减少 HBM 中间矩阵读写后,并行度、线程间通信、非矩阵运算和片上带宽仍可能限制性能。FA2 至 FA4 分别优化工作划分、异步流水线和硬件资源分配。

FlashAttention-2:并行度与工作划分

FlashAttention-2 (Dao, 2023) 的主要优化包括:

  1. 少做非矩阵运算。 保留未归一化的 U,最后才除以 ℓ,省掉反复归一化输出的工作。Tensor Cores 的矩阵吞吐很高,指数、归约、除法等操作不能按同样的 FLOPs 吞吐估算。
  2. 增加独立的 query 工作。 只按 batch 和 head 分配 thread block 时,小 batch、长序列可能留下许多空闲 SM。再沿 query 序列分块,就有更多互不依赖的输出行可供调度。
  3. 减少 warp 之间的部分和归约。 一个 warp 是一起执行的 32 个 GPU 线程。在 thread block 内,让各 warp 负责不同 query 行并共享 KV,可以让它们各自产生完整的输出行。
Within a thread block: split KV versus split Q左侧两个 warp 处理相同的 query 和不同的 KV 范围,输出是同一行的部分结果,需要归约。右侧两个 warp 处理不同的 query 行和共享的 KV,各自得到完整的输出行,无需在 warp 之间归约这些输出。
Split K,V · partial outputs
Split Q · independent output rows
warp 0
Qi; K(0),V(0)
m(0),ℓ(0),U(0)
warp 0
Qi(0) · shared K,V
Oi(0)
warp 1
Qi; K(1),V(1)
m(1),ℓ(1),U(1)
warp 1
Qi(1) · shared K,V
Oi(1)
merge → Oi
shared memory + sync
Each warp owns its output rows
左右滑动查看完整示意图
以两个 warp 示意 forward 的工作划分,依据 FlashAttention-2 的 split-K 与 split-Q 思路重绘。上标 (w) 标记 warp w 对应的分片。右侧省掉输出部分和的跨 warp 合并,加载共享 K,V 时仍需必要的同步。

图中若沿 KV 拆,同一个 query 的结果散落在不同 warp,需要合并部分状态;沿 Q 拆,则每个 warp 负责自己的输出行。这减少的是输出累加所需的 warp 间通信,共享 KV 的加载仍需协调。这里的划分发生在一个 thread block 内,与把多个 query tile 分配给不同 thread block 是两个层次。

FlashAttention-3:异步流水线

FlashAttention-3 (Dao, 2024) 利用 Hopper 的异步数据搬运 TMA 和矩阵乘法 WGMMA,将搬运和计算分配给不同的 warp。在当前 tile 计算期间预取下一块 KV,并使不同工作组的 softmax 与矩阵乘法重叠。

流水线必须满足数据依赖:当前 tile 的 PV 要等待对应 softmax 权重就绪。重叠执行需要额外 buffer 和 registers,需同时评估资源占用与吞吐收益。FA3 还探索 FP8,量化误差需要独立评估。

FlashAttention-4:Blackwell 资源分配

FlashAttention-4 (Zadouri et al., 2026) 于 2026 年 3 月发布,面向 Blackwell attention。以 B200 为例,Tensor Cores 吞吐的增长快于指数运算吞吐和 shared memory 带宽的增长,softmax 与片上搬运因而可能成为瓶颈。

其实现说明介绍了以下优化:

  1. 用 TMEM 重排流水线。 Blackwell 的异步 MMA 将累加结果写入 tensor memory(TMEM),减轻 registers 压力。FA4 交错处理两个 query tile,让一块的矩阵乘法与另一块的 softmax 重叠,并把输出重缩放交给单独的 warpgroup。
  2. 让其他计算单元分担指数。 部分指数通过乘加单元上的多项式近似计算,与硬件指数单元并行。多项式也需要指令和 registers,因此要选择合适的分担比例,并评估近似误差。
  3. 继续减少 backward 的片上流量。 将部分中间结果保留在 TMEM,直接供矩阵单元读取;再用 2-CTA MMA 让两个 thread block 协作,配合梯度 tile 的重排,减少 shared memory 的重复读取和梯度的全局原子累加。

FA4 还通过 conditional rescaling 减少对累加向量的反复缩放:最大值只增长一点时,可以暂时沿用旧的指数基准,让分子、分母和新 tile 都在同一尺度下继续累加;增长超过阈值时再一起调整。

Conditional rescaling 的不变量

标准 online softmax 每块将 m 更新为当前最大值。归一化在代数上只要求分子与分母使用相同基准,不要求该基准始终等于最大值。用 r 表示采用的基准,对已处理的位置集合 A,定义

ℓA(r)=∑j∈Aexj−r,uA(r)=∑j∈Aexj−rVj.

分子、分母都有共同因子 e−r,所以 uA(r)/ℓA(r) 不随 r 改变。初始化之后,只要指数与累加值仍在安全的数值范围内,就可以暂时保留旧的 r,让新 tile 也相对这个基准累加。这样,旧的向量 u 就不必每次都乘一个缩放因子。

FA4 的官方 softmax 实现 用阈值控制这种更新:最大值相对当前基准跳得足够大时,才更新基准并重缩放状态;否则保留旧基准,缩放因子取 1。关键是 ℓ,u 和新 tile 始终使用同一尺度,最后再归一化。只更新基准却省掉旧状态的重缩放,会破坏这个不变量。

FA4 的官方实现 使用 CuTe DSL,以 Python 表达并编译 GPU kernel。

Prefill 与 Decode

训练与 prefill 通常同时处理很多 query,容易形成较大的矩阵乘法,也有大量 N×N 中间结果可以省掉。普通单 token decode 中,query 长度是 1,KV 长度是 N;每个 head 当前步骤的 score 只有 1×N,计算量已经是 Θ(Nd)。

因此 decode 仍能受益于融合与 online reduction,但常见难点转向读取长 KV cache,以及如何为少量 query 提供足够的并行工作。前面的状态合并公式不限定扫描顺序:可以将 KV 分成多段分别计算,再按树形合并结果。实数运算下结果相同,浮点下可能有舍入差异;额外的并行度也会带来部分结果存储和归约成本。不能直接把长序列训练的加速比套到 decode 上。

PagedAttention 管理跨生成步骤保留的 KV 布局和共享;FlashAttention 优化单次 attention 运算内部的数据流。两者可以结合,作用范围不同:融合计算不会消除 KV cache,也不会减少 dense attention 需要处理的有效 token。

性能评估需要分别测量计算量、HBM 流量、片上资源占用、并行度和同步开销。训练、prefill 与 decode 的形状和瓶颈不同,应分别选择 tile、工作划分和归约方式。

References

Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. doi.org
Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. tridao.me
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems, 35, 16344–16359. arxiv.org
Jain, M., Kandar, T., Srivastava, S., Zhu, J., Chaudhary, J., Srivastava, D., & Kong, J. (2026). CSE 291A/DSC 291, Lecture 15: KV Cache Management and Flash Attention. haoailab.com
Liang, P. (2026). CS336: Language Modeling from Scratch, Lecture 6: Kernels. github.com
Milakov, M., & Gimelshein, N. (2018). Online Normalizer Calculation for Softmax. doi.org
Zadouri, T., Hoehnerbach, M., Shah, J., Liu, T., Thakkar, V., & Dao, T. (2026). FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. doi.org

Cite this post

@misc{pu2026mlmlrevisitinfraflashattention,
  author = {Pu, Fanyi and {GPT-6 Astra} and {GPT‑5.6 Sol}},
  title  = {Infra: Flash Attention},
  year   = {2026},
  month  = {9},
  url    = {https://pufanyi.com/blog/ml/ml-revisit/infra/flash-attention}
}