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 :
𝑄 , 𝐾 , 𝑉 ∈ ℝ 𝑁 × 𝑑 , 𝑆 = 𝑄 𝐾 ⊤ √ 𝑑 , 𝑃 = s o f t m a x ( 𝑆 ) , 𝑂 = 𝑃 𝑉 . Q , K , V ∈ R N × d , S = Q K ⊤ d , P = softmax ( S ) , O = P V .
Softmax 沿每一行进行。𝑆 𝑖 𝑗 S i j 是第 𝑖 i 个 query 对第 𝑗 j 个 key 的 score,𝑃 𝑖 𝑗 P i j 是归一化后的权重;输出 𝑂 𝑖 O i 是所有 value 的加权和。
分别执行矩阵乘法、softmax 和矩阵乘法,会物化两个 𝑁 × 𝑁 N × N 中间矩阵 𝑆 , 𝑃 S , P ,最终输出只有 𝑁 × 𝑑 N × d 。FlashAttention 在片上生成 score tile,完成归一化状态更新与加权累加后丢弃该 tile,只保留后续计算需要的状态。
推导默认没有 dropout,每行至少有一个有效 key。
GPU 存储层次与 IO
假设 𝑁 = 8 1 9 2 N = 8192 、𝑑 = 1 2 8 d = 128 ,只看单个 head,张量用 BF16 存储,每个元素 2 bytes:
张量 元素数量 存储量 单个 𝑄 Q 、𝐾 K 、𝑉 V 或 𝑂 O 𝑁 𝑑 N d 2 MiB 单个 𝑆 S 或 𝑃 P 𝑁 2 N 2 128 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 需要配合:完整的 𝑁 2 N 2 中间矩阵通常无法驻留片上,而分块 softmax 必须维护跨块的归一化状态。
Online softmax 与状态合并
稳定归一化
固定一个 query,将它与所有 key 的 score 记为 𝑥 1 , … , 𝑥 𝑁 x 1 , … , x N 。对应的 softmax 权重为:
𝑝 𝑗 = 𝑒 𝑥 𝑗 ∑ 𝑁 𝑘 = 1 𝑒 𝑥 𝑘 . p j = e x j ∑ k = 1 N e x k .
Softmax 的分母依赖整行 score。对各块分别做 softmax,得到的是块内归一化权重,不能直接作为整行的权重。
直接累加 𝑒 𝑥 𝑗 e x j 可能溢出。Stable softmax 先减去整行最大值,再求指数和:
𝑚 = m a x 𝑗 𝑥 𝑗 , ℓ = ∑ 𝑗 𝑒 𝑥 𝑗 − 𝑚 , 𝑝 𝑗 = 𝑒 𝑥 𝑗 − 𝑚 ℓ . m = max j x j , ℓ = ∑ j e x j − m , p j = e x j − m ℓ .
分子、分母都乘了 𝑒 − 𝑚 e − m ,所以权重不变;同时 𝑥 𝑗 − 𝑚 ≤ 0 x j − m ≤ 0 ,指数就不会因为 score 太大而溢出。这是 stable softmax,也常称 safe softmax。
分块状态的合并
上述实现需要先求最大值,再计算指数和。Online softmax (Milakov & Gimelshein, 2018 ) 在扫描 score 时同时维护当前最大值和指数和。最大值改变时,将已有的指数和换算到新基准,无需重新读取之前的 score。
设两段 score 分别为 𝐴 , 𝐵 A , B ,各自的状态为 ( 𝑚 𝐴 , ℓ 𝐴 ) ( m A , ℓ A ) 和 ( 𝑚 𝐵 , ℓ 𝐵 ) ( m B , ℓ B ) ,其中 𝑚 𝐴 m A 是 𝐴 A 的最大值,ℓ 𝐴 ℓ A 是减去 𝑚 𝐴 m A 后的指数和,𝐵 B 同理。合并后的最大值为 𝑚 = m a x ( 𝑚 𝐴 , 𝑚 𝐵 ) m = max ( m A , m B ) 。将 𝐴 A 的指数和从基准 𝑚 𝐴 m A 换算到 𝑚 m :
∑ 𝑗 ∈ 𝐴 𝑒 𝑥 𝑗 − 𝑚 = 𝑒 𝑚 𝐴 − 𝑚 ∑ 𝑗 ∈ 𝐴 𝑒 𝑥 𝑗 − 𝑚 𝐴 = 𝑒 𝑚 𝐴 − 𝑚 ℓ 𝐴 . ∑ j ∈ A e x j − m = e m A − m ∑ j ∈ A e x j − m A = e m A − m ℓ A .
换算只需要给已有指数和乘以共同因子;最大值不变时,因子为 1。对 𝐴 , 𝐵 A , B 分别换算,再累加:
𝛼 = 𝑒 𝑚 𝐴 − 𝑚 , 𝛽 = 𝑒 𝑚 𝐵 − 𝑚 , ℓ = 𝛼 ℓ 𝐴 + 𝛽 ℓ 𝐵 . α = e m A − m , β = e m B − m , ℓ = α ℓ A + β ℓ B .
( 𝑚 , ℓ ) ( m , ℓ ) 是后续合并所需的状态。如果需要输出完整的 softmax 向量,最终仍需保留或重新读取 score,计算各个 𝑝 𝑗 p j 。
Attention 加权和的累加
Attention 输出为 𝑂 𝑖 = ∑ 𝑗 𝑝 𝑗 𝑉 𝑗 O i = ∑ j p j V j ,不需要输出各个 𝑝 𝑗 p j 。因此可以同时维护未归一化的加权和 (Jain et al., 2026 ) :
𝑢 𝐴 = ∑ 𝑗 ∈ 𝐴 𝑒 𝑥 𝑗 − 𝑚 𝐴 𝑉 𝑗 . u A = ∑ j ∈ A e x j − m A V j .
它与 ℓ 𝐴 ℓ A 使用相同的指数权重,换算到新基准时也乘以 𝛼 α 。对 𝐵 B 使用 𝛽 β ,得到
𝑢 = 𝛼 𝑢 𝐴 + 𝛽 𝑢 𝐵 , 𝑂 𝑖 = 𝑢 ℓ . u = α u A + β u B , O i = u ℓ .
每个 query 的持久状态为两个标量 𝑚 , ℓ m , ℓ 和一个 𝑑 d 维向量 𝑢 u 。各块更新分子和分母,扫描结束后做一次除法;已处理块的 score 和指数权重可以丢弃。
以指数值为 1 , 2 , 4 , 8 1 , 2 , 4 , 8 的情况为例,把 value 简化为标量 𝑉 = ( 1 , 2 , 3 , 4 ) V = ( 1 , 2 , 3 , 4 ) 。按完整公式计算,结果是
𝑂 𝑖 = 1 ⋅ 1 + 2 ⋅ 2 + 4 ⋅ 3 + 8 ⋅ 4 1 + 2 + 4 + 8 = 4 9 1 5 . O i = 1 ⋅ 1 + 2 ⋅ 2 + 4 ⋅ 3 + 8 ⋅ 4 1 + 2 + 4 + 8 = 49 15 .
左右滑动查看完整示意图
这里的 e x p ( 𝑥 ) exp ( x ) 是便于手算的未平移指数值,实际算法使用减去最大值后的指数。两个块的分母和加权和必须按同一基准合并;value 用标量展示,向量时逐分量执行相同运算。
图中两段的基准分别是 𝑚 𝐴 = l o g 2 m A = log 2 、𝑚 𝐵 = l o g 8 m B = log 8 。合并时,把左块的 ℓ 𝐴 , 𝑢 𝐴 ℓ A , u A 同时乘 𝑒 l o g 2 − l o g 8 = 1 / 4 e log 2 − log 8 = 1 / 4 ,得到 𝑢 / ℓ = 6 . 1 2 5 / 1 . 8 7 5 = 4 9 / 1 5 u / ℓ = 6.125 / 1.875 = 49 / 15 ,与完整计算一致。不能直接平均两段各自的 attention 输出,因为两段在完整分母中的份额不同。
分块计算与 Forward
GPU 实现将 𝐵 𝑟 B r 个 query 组成一个 tile,复用同一块 KV,并通过矩阵乘法计算 score 和加权和。以下 𝑖 , 𝑗 i , j 表示块编号,每个 KV tile 包含 𝐵 𝑐 B c 个 token:
𝑄 𝑖 ∈ ℝ 𝐵 𝑟 × 𝑑 , 𝐾 𝑗 , 𝑉 𝑗 ∈ ℝ 𝐵 𝑐 × 𝑑 , 𝑆 𝑖 𝑗 = 𝑄 𝑖 𝐾 ⊤ 𝑗 √ 𝑑 ∈ ℝ 𝐵 𝑟 × 𝐵 𝑐 . Q i ∈ R B r × d , K j , V j ∈ R B c × d , S i j = Q i K j ⊤ d ∈ R B r × B c .
实现时,一个 thread block 可以负责一个 query tile,在 GPU 的一个 SM 计算单元上执行,用 shared memory 暂存输入,把累加状态留在 registers 中。两个矩阵乘法可以交给 Tensor Cores;最大值、分母和重缩放仍然对每个 query 独立进行。
将各 query 的 𝑢 u 排成矩阵 𝑈 U ,令 𝑊 W 表示当前 score tile 相对运行中最大值计算的指数权重。𝑊 W 尚未除以分母,通过 𝑊 𝑉 𝑗 W V j 累加到 𝑈 U :
左右滑动查看完整示意图
上:箭头穿过 𝑆 , 𝑃 S , P 时分别发生 HBM 写入与读回。下:虚线框内只容纳当前 tile 与运行状态,𝑆 , 𝑃 S , P 不形成完整的 HBM 张量。示意数据流,不按容量或耗时比例绘制。
下面采用 FlashAttention-2 (Dao, 2023 ) 的 query tile 外循环与延迟归一化形式,展示数据依赖。rowmax、rowsum 沿列归约,行向量按行 broadcast;这是片上计算的伪代码,逐行写成普通 PyTorch 操作不会自动得到一个 fused kernel。
伪代码中,𝑚 m 是旧的最大值,𝑚 ′ m ′ 是合并后的最大值。新块直接相对 𝑚 ′ m ′ 求指数,前面公式里的 𝛽 β 就已经包含在 𝑊 W 中,只需用 𝛼 α 重缩放旧状态。
Algorithm 1 FlashAttention forward Input 𝑄 , 𝐾 , 𝑉 Q , K , V in HBM,按 𝐵 𝑟 , 𝐵 𝑐 B r , B c 分块
Output 𝑂 O ,以及 backward 所需的 𝐿 L
parallel for each query tile 𝑄 𝑖 Q i do
load 𝑄 𝑖 Q i
𝑚 ← − ∞ , m ← − ∞ , ℓ ← 0 , ℓ ← 0 , 𝑈 ← 0 U ← 0
for each KV tile 𝐾 𝑗 , 𝑉 𝑗 K j , V j do
load 𝐾 𝑗 , 𝑉 𝑗 K j , V j
𝑆 ← S ← 𝑄 𝑖 𝐾 ⊤ 𝑗 / √ 𝑑 Q i K j ⊤ / d
mask 𝑆 S 全被 mask 的行跳过本轮更新
𝑚 ′ ← m ′ ← m a x ( 𝑚 , r o w m a x ( 𝑆 ) ) max ( m , rowmax ( S ) )
𝛼 ← α ← e x p ( 𝑚 − 𝑚 ′ ) exp ( m − m ′ )
𝑊 ← W ← e x p ( 𝑆 − 𝑚 ′ ) exp ( S − m ′ )
ℓ ← ℓ ← 𝛼 ℓ + α ℓ + r o w s u m ( 𝑊 ) rowsum ( W )
𝑈 ← U ← 𝛼 𝑈 + 𝑊 𝑉 𝑗 α U + W V j
𝑚 ← 𝑚 ′ m ← m ′
end for
𝑂 𝑖 ← O i ← 𝑈 / ℓ U / ℓ 最后才归一化
𝐿 𝑖 ← L i ← 𝑚 + l o g ℓ m + log ℓ
store 𝑂 𝑖 , 𝐿 𝑖 O i , L i
end for
标记行:计算重缩放系数,并用它同时调整分母与加权和。
交互演示 一条长 sequence,怎样逐块算完? 点左侧选择 Q 块,观察它怎样扫过整条 KV 序列;也可以改变 KV 块大小。
Q 分成 6 块,每块 4 行;K、V 分成 6 块,每块 4 行。粗线围出 4 × 4 的 score tile。
𝑄 Q ↓𝐾 , 𝑉 K , V sequence · 24 tokens →输出
𝑄 1 Q 1 1–4 𝑄 2 Q 2 5–8 𝑄 3 Q 3 9–12 𝑄 4 Q 4 13–16 𝑄 5 Q 5 17–20 𝑄 6 Q 6 21–24
把 24 × 24 的 attention 计算划分成 4 × 4 的 tile 行是 query,列是 key。固定一条 query 横带,逐块向右扫描 KV;斜线只表示该范围已合并,不表示保存了 score。 𝑂 1 O 1 ·
𝑂 2 O 2 ·
𝑂 3 O 3 ·
𝑂 4 O 4 ·
𝑂 5 O 5 ·
𝑂 6 O 6 ·
当前计算 已合并,仅标记进度
↓ 合入同一份状态
𝑚 , ℓ , 𝑈 m , ℓ , U 跨块保留
已覆盖 0 / 24 个 key ↓ 全部扫完后
𝑂 𝑖 = 𝑈 / ℓ O i = U / ℓ 扫描完才能归一化
每个 Q 块扫描 6 次 KV;每次只算一个 4 × 4 4 × 4 tile。
𝑄 Q ↓𝐾 , 𝑉 K , V sequence · 24 tokens →输出
𝑄 1 Q 1 1–4 𝑄 2 Q 2 5–8 𝑄 3 Q 3 9–12 𝑄 4 Q 4 13–16 𝑄 5 Q 5 17–20 𝑄 6 Q 6 21–24
把 24 × 24 的 attention 计算划分成 4 × 8 的 tile 行是 query,列是 key。固定一条 query 横带,逐块向右扫描 KV;斜线只表示该范围已合并,不表示保存了 score。 𝑂 1 O 1 ·
𝑂 2 O 2 ·
𝑂 3 O 3 ·
𝑂 4 O 4 ·
𝑂 5 O 5 ·
𝑂 6 O 6 ·
当前计算 已合并,仅标记进度
↓ 合入同一份状态
𝑚 , ℓ , 𝑈 m , ℓ , U 跨块保留
已覆盖 0 / 24 个 key ↓ 全部扫完后
𝑂 𝑖 = 𝑈 / ℓ O i = U / ℓ 扫描完才能归一化
每个 Q 块扫描 3 次 KV;每次只算一个 4 × 8 4 × 8 tile。
图中 24 个 token,Q 每 4 行一块;KV 每 4 或 8 行一块,分别扫描 6 或 3 次。各段依次合入同一份运行状态,扫完后得到对应的 𝑂 𝑖 O i 。大网格只标记计算范围,不保存完整 score 矩阵;本例不加 mask。
符号与边界条件 这里的 𝑆 , 𝑊 S , W 只活在当前 tile 内,累加完就能复用空间。𝐿 𝑖 L i 是逐行的 log-sum-exp,训练时留给 backward 使用。
初始化时,首个有效 tile 的 𝑚 ′ m ′ 有限,因此 𝛼 = 0 α = 0 ,旧的零状态不会贡献结果。对于全被 mask 的 tile 行,需要显式跳过更新,避免计算 − ∞ − ( − ∞ ) − ∞ − ( − ∞ ) ;若整个 query 行都无有效 key,则要由算子另行约定输出,不能直接套用这里的除法。
左右滑动查看完整示意图
每格代表 2 × 2 2 × 2 个 score,块编号从 0 开始。粗框表示当前 query tile 的有效范围,✓ 表示已累加,● 表示当前 tile,斜线右上方是被 mask 的位置。右侧状态每处理一块就更新一次。
Causal attention 中,整块位于未来的 tile 可以直接跳过;与对角线相交的 tile 在片上把无效 score 设为 − ∞ − ∞ ,其指数权重就是 0。这样无需单独存一个完整的三角 mask。图中各个 query tile 可以独立计算,但每个 tile 内仍要沿着有效的 KV 块合并状态。
Tile 越大,输入复用通常越充分;但 𝑄 𝑖 Q i 、𝐾 𝑗 K j 、𝑉 𝑗 V j 、𝑈 U 、𝑆 𝑖 𝑗 S i j 和 softmax 状态都要占片上资源。Registers 或 shared memory 占得太多,会减少同一 SM 能同时驻留的工作,甚至导致 register spilling。实际选块要同时考虑数据类型、head dimension 和硬件资源。
FlashAttention 节省的是中间结果的存储和搬运,dense attention 的乘法仍要做:
指标(单个 head) 三算子、显式中间矩阵 FlashAttention Dense attention 算术量 Θ ( 𝑁 2 𝑑 ) Θ ( N 2 d ) 仍为 Θ ( 𝑁 2 𝑑 ) Θ ( N 2 d ) HBM 中保存的 𝑆 , 𝑃 S , P Θ ( 𝑁 2 ) Θ ( N 2 ) 不保存完整矩阵 除输入输出外的归一化状态 取决于实现 Θ ( 𝑁 ) Θ ( N ) HBM 访问量 有 Θ ( 𝑁 2 ) Θ ( N 2 ) 的中间矩阵读写 依赖 tile、循环顺序和片上容量
输入和输出本身仍占 Θ ( 𝑁 𝑑 ) Θ ( N d ) 。所以,显存占用随序列长度线性增长,不代表计算量或总 HBM 访问量也变成了线性。
原论文的 IO 访问量分析 原始 FlashAttention 论文 (Dao et al., 2022 ) 用一个两级存储模型分析 IO。设片上容量为 𝑀 M 个标量元素 ,取 𝑑 ≤ 𝑀 ≤ 𝑁 𝑑 d ≤ M ≤ N d 。其 KV tile 外循环令 𝐵 𝑐 = Θ ( 𝑀 / 𝑑 ) B c = Θ ( M / d ) ,每个 KV 块扫描一次 𝑄 Q 和输出状态;共有约 𝑁 𝑑 / 𝑀 N d / M 个 KV 块,每轮搬运 Θ ( 𝑁 𝑑 ) Θ ( N d ) 个元素,于是 HBM 访问量为
Θ ( 𝑁 2 𝑑 2 𝑀 ) . Θ ( N 2 d 2 M ) .
这是原论文循环与存储模型下的界,不是上面伪代码的逐项流量统计,也没有把不同精度、cache 命中或硬件调度算进去。它展示了片上空间如何换取数据复用:在 𝑀 ≫ 𝑑 2 M ≫ d 2 的适用区间,相比显式中间矩阵带来的 Θ ( 𝑁 2 ) Θ ( N 2 ) 流量,收益明显。
固定 𝑀 , 𝑑 M , d 时,上面的 IO 界仍随 𝑁 2 N 2 增长,KV 也可能被不同 query tile 多次读取。可以用一个粗略的性能下界帮助判断优化方向:
𝑡 ≳ m a x ( F L O P s c o m p u t e t h r o u g h p u t , H B M b y t e s H B M b a n d w i d t h ) . t ≳ max ( FLOPs compute throughput , HBM bytes HBM bandwidth ) .
它忽略了同步、调度以及不能完全重叠的操作,不能直接当作耗时预测。但它说明,如果瓶颈是 HBM 带宽,减少中间矩阵的读写会直接降低这一项成本;优化后,还需要检查其他执行单元是否成为新的瓶颈。
Backward 与重计算
训练时,backward 需要 attention 概率或足以恢复它的信息。保存完整 𝑃 P 需要平方级空间。FlashAttention 使用 recomputation (Dao et al., 2022 ) :保留 𝑄 , 𝐾 , 𝑉 , 𝑂 Q , K , V , O ,并为每行保存 log-sum-exp:
𝐿 𝑖 = 𝑚 𝑖 + l o g ℓ 𝑖 = l o g ∑ 𝑗 𝑒 𝑆 𝑖 𝑗 . L i = m i + log ℓ i = log ∑ j e S i j .
Backward 重新计算当前 tile 的 score,用 𝑃 𝑖 𝑗 = 𝑒 𝑆 𝑖 𝑗 − 𝐿 𝑖 P i j = e S i j − L i 恢复概率,计算局部梯度后丢弃中间结果。𝐿 L 占 𝑁 N 个标量;训练 activation 还包括 𝑄 , 𝐾 , 𝑉 , 𝑂 Q , K , V , O 。这是 attention 内部的 activation checkpointing。
重计算增加 score 和指数运算,减少大矩阵的 HBM 存储与传输,因此可能同时降低显存占用和耗时。梯度按 tile 累加,forward 与 backward 均不物化完整的 𝑁 × 𝑁 N × N 中间矩阵。
分块梯度计算 记上游梯度为 𝐺 = 𝜕 L / 𝜕 𝑂 G = ∂ L / ∂ O ,则
𝜕 L 𝜕 𝑉 = 𝑃 ⊤ 𝐺 , 𝜕 L 𝜕 𝑃 = 𝐺 𝑉 ⊤ . ∂ L ∂ V = P ⊤ G , ∂ L ∂ P = G V ⊤ .
Softmax 的 backward 是逐行的。令 𝐻 = 𝐺 𝑉 ⊤ H = G V ⊤ ,有
𝜕 L 𝜕 𝑆 𝑖 𝑗 = 𝑃 𝑖 𝑗 ( 𝐻 𝑖 𝑗 − 𝐷 𝑖 ) , 𝐷 𝑖 = ∑ 𝑗 𝑃 𝑖 𝑗 𝐻 𝑖 𝑗 . ∂ L ∂ S i j = P i j ( H i j − D i ) , D i = ∑ j P i j H i j .
利用 𝑂 𝑖 = ∑ 𝑗 𝑃 𝑖 𝑗 𝑉 𝑗 O i = ∑ j P i j V j ,将 𝐷 𝑖 D i 改写为
𝐷 𝑖 = ∑ 𝑗 𝑃 𝑖 𝑗 ( 𝐺 𝑖 ⋅ 𝑉 𝑗 ) = 𝐺 𝑖 ⋅ ( ∑ 𝑗 𝑃 𝑖 𝑗 𝑉 𝑗 ) = 𝐺 𝑖 ⋅ 𝑂 𝑖 . D i = ∑ j P i j ( G i ⋅ V j ) = G i ⋅ ( ∑ j P i j V j ) = G i ⋅ O i .
因此可先由 𝐺 , 𝑂 G , O 计算 𝐷 D ,无需扫描完整 𝑃 P 。每个 tile 恢复局部 𝑃 , 𝐻 P , H 后,计算局部 𝜕 L / 𝜕 𝑆 ∂ L / ∂ S ,再累加
𝜕 L 𝜕 𝑄 = 1 √ 𝑑 𝜕 L 𝜕 𝑆 𝐾 , 𝜕 L 𝜕 𝐾 = 1 √ 𝑑 ( 𝜕 L 𝜕 𝑆 ) ⊤ 𝑄 . ∂ L ∂ Q = 1 d ∂ L ∂ S K , ∂ L ∂ K = 1 d ( ∂ 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 ) 的主要优化包括:
少做非矩阵运算。 保留未归一化的 𝑈 U ,最后才除以 ℓ ℓ ,省掉反复归一化输出的工作。Tensor Cores 的矩阵吞吐很高,指数、归约、除法等操作不能按同样的 FLOPs 吞吐估算。
增加独立的 query 工作。 只按 batch 和 head 分配 thread block 时,小 batch、长序列可能留下许多空闲 SM。再沿 query 序列分块,就有更多互不依赖的输出行可供调度。
减少 warp 之间的部分和归约。 一个 warp 是一起执行的 32 个 GPU 线程。在 thread block 内,让各 warp 负责不同 query 行并共享 KV,可以让它们各自产生完整的输出行。
左右滑动查看完整示意图
以两个 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 的 𝑃 𝑉 P V 要等待对应 softmax 权重就绪。重叠执行需要额外 buffer 和 registers,需同时评估资源占用与吞吐收益。FA3 还探索 FP8,量化误差需要独立评估。
FlashAttention-4:Blackwell 资源分配
FlashAttention-4 (Zadouri et al., 2026 ) 于 2026 年 3 月发布,面向 Blackwell attention。以 B200 为例,Tensor Cores 吞吐的增长快于指数运算吞吐和 shared memory 带宽的增长,softmax 与片上搬运因而可能成为瓶颈。
其实现说明 介绍了以下优化:
用 TMEM 重排流水线。 Blackwell 的异步 MMA 将累加结果写入 tensor memory(TMEM),减轻 registers 压力。FA4 交错处理两个 query tile,让一块的矩阵乘法与另一块的 softmax 重叠,并把输出重缩放交给单独的 warpgroup。
让其他计算单元分担指数。 部分指数通过乘加单元上的多项式近似计算,与硬件指数单元并行。多项式也需要指令和 registers,因此要选择合适的分担比例,并评估近似误差。
继续减少 backward 的片上流量。 将部分中间结果保留在 TMEM,直接供矩阵单元读取;再用 2-CTA MMA 让两个 thread block 协作,配合梯度 tile 的重排,减少 shared memory 的重复读取和梯度的全局原子累加。
FA4 还通过 conditional rescaling 减少对累加向量的反复缩放:最大值只增长一点时,可以暂时沿用旧的指数基准,让分子、分母和新 tile 都在同一尺度下继续累加;增长超过阈值时再一起调整。
Conditional rescaling 的不变量 标准 online softmax 每块将 𝑚 m 更新为当前最大值。归一化在代数上只要求分子与分母使用相同基准,不要求该基准始终等于最大值。用 𝑟 r 表示采用的基准,对已处理的位置集合 𝐴 A ,定义
ℓ 𝐴 ( 𝑟 ) = ∑ 𝑗 ∈ 𝐴 𝑒 𝑥 𝑗 − 𝑟 , 𝑢 𝐴 ( 𝑟 ) = ∑ 𝑗 ∈ 𝐴 𝑒 𝑥 𝑗 − 𝑟 𝑉 𝑗 . ℓ A ( r ) = ∑ j ∈ A e x j − r , u A ( r ) = ∑ j ∈ A e x j − r V j .
分子、分母都有共同因子 𝑒 − 𝑟 e − r ,所以 𝑢 𝐴 ( 𝑟 ) / ℓ 𝐴 ( 𝑟 ) u A ( 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 × 𝑁 1 × N ,计算量已经是 Θ ( 𝑁 𝑑 ) Θ ( N d ) 。
因此 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