FlashAttention

FlashAttention 论文里有一组 GPT-2 medium 的数据（A100，序列长 1024，前向加反向）：标准 attention 做 66.6 GFLOPs、读写 HBM 40.3 GB、用时 41.7 ms；FlashAttention 做 75.2 GFLOPs，多算了 13%，HBM 读写却只有 4.4 GB，用时 7.3 ms。多算反而快了 5.7 倍，这说明 attention 的时间根本不由 FLOPs 决定。这一篇讲清楚三件事：标准 attention 的时间花在哪，softmax 这个”需要看完整行”的运算怎么被拆成小块，以及每一代 FlashAttention 在追着什么新瓶颈跑。

前置：GPU 与 Kernel（HBM、shared memory、roofline、tiling）、Transformer 架构（attention 的定义）

标准 attention 的时间花在搬运上

单个 head，Q, K, V 都是 [N, d]。标准实现分三步，每一步都是一个独立的 kernel：

S = Q @ K.T / sqrt(d)     写 N×N 到 HBM
P = softmax(S)            读 N×N，写 N×N 到 HBM
O = P @ V                 读 N×N

读 Q、K、V、写 O 只有 Θ(Nd) 个元素，读写 S 和 P 却是 Θ(N²)。N 远大于 d 时（N = 8192、d = 128），后者占绝对多数。规模感：8K 序列、32 个头、fp16，单层单样本的 S 就是 8192² × 2 字节 × 32 = 4 GiB，要在 HBM 里来回读写好几遍。训练时还得把 P 存下来给反向用，显存也是 O(N²)。

用 GPU 与 Kernel 里的 roofline 看：两个矩阵乘本身强度很高，没问题；夹在中间的 softmax（以及 mask、dropout）每个元素只做几次运算，是典型的 memory-bound，Tensor Core 大部分时间在等数据。所以出路只有一个：让 S 和 P 永远不完整地出现在 HBM 里。矩阵乘可以 tiling，难点全在中间那个 softmax。

为什么 softmax 不能直接分块

softmax 直接按定义写 exp(x) / sum(exp(x)) 是不能用的，fp16 下立刻溢出：

import torch
x = torch.tensor([10.0, 20.0, 30.0], dtype=torch.float16)
print("naive fp16:", (x.exp() / x.exp().sum()).tolist())
print("safe  fp16:", torch.softmax(x, -1).tolist())

naive fp16: [0.0, nan, nan]
safe  fp16: [0.0, 4.5418739318847656e-05, 1.0]

fp16 最大只能表示 65504，e^20 已经是 inf，inf / inf 得 nan（浮点格式的范围见信息的表示）。所以实际用的是 safe softmax：先减去这一行的最大值 m，

m = max_j x_j ,   d = Σ_j exp(x_j − m) ,   y_i = exp(x_i − m) / d

减完以后指数都 ≤ 0，不会溢出。代价是要扫三遍数据：求 m、求 d、归一化，而且前两遍有依赖，求 d 必须先知道全局的 m。在 attention 里，一行就是一个 query 对全部 N 个 key 的分数，“先知道全局最大值”意味着必须先把整行 S 都算出来。分块的路好像被堵死了。

Online softmax：边读边修正

Milakov 和 Gimelshein（2018）的观察是：不必先知道全局最大值，用”目前为止的最大值”先算着，等遇到更大的再回头修正。维护两个数 (m, d)，初值 m = −∞、d = 0，每读一个新元素 x：

m_new = max(m, x)
d     = d · exp(m − m_new) + exp(x − m_new)

修正因子 exp(m − m_new) 的道理只有一行：旧的 d 里每一项是 exp(x_j − m)，乘上 exp(m − m_new) 恰好变成 exp(x_j − m_new)，基准就换到了新的最大值，不用回头重读任何数据。因为 m_new ≥ m，这个因子总是 ≤ 1，不会溢出。走一遍 x = [1, 3, 2]：

import math
x = [1.0, 3.0, 2.0]
m, d = float("-inf"), 0.0
for xi in x:
    m_new = max(m, xi)
    d = d * math.exp(m - m_new) + math.exp(xi - m_new)   # 先把旧和换到新基准，再加新项
    m = m_new
    print(f"read {xi:.0f}: m={m:.0f} d={d:.4f}")
print(f"two-pass d={sum(math.exp(v - max(x)) for v in x):.4f}")
print("softmax", [round(math.exp(v - m) / d, 3) for v in x])

read 1: m=1 d=1.0000
read 3: m=3 d=1.1353
read 2: m=3 d=1.5032
two-pass d=1.5032
softmax [0.09, 0.665, 0.245]

读到 3 时最大值变了，旧的 d = 1 乘上 e^(1−3) = 0.1353 再加 1。结果和两遍算法完全一致。

更重要的是，这个状态可以两两合并。两段数据各自算出 (m_a, d_a) 和 (m_b, d_b)，合并为 m = max(m_a, m_b)、d = d_a·exp(m_a − m) + d_b·exp(m_b − m)。把 [1, 3] 和 [2] 分开算再合并，得到的还是 1.5032。合并满足结合律，于是数据可以任意切块、按任意顺序、甚至由不同的线程块并行处理后再合并。这正是分块需要的性质。

停一下，检查理解单选

对 x = [2, 1, 4] 做 online softmax。读完前两个数后 m = 2、d = 1 + e^(-1) ≈ 1.368。读到 4 之后 d 是多少？

1.368 + 1 = 2.3681.368 × e^(-2) + 1 ≈ 1.1851.368 × e^(2) + 1 ≈ 11.11e^(-2) + 1 ≈ 1.135

把输出也在线累加：分块 attention

attention 最终要的不是 softmax 本身，而是 O = P V。观察到 O 的每一行是 V 的行按 softmax 权重的加权和，它的分子 Σ_j exp(s_j − m) v_j 和分母 d 用的是同一个基准 m，所以可以用同一个修正因子一起在线更新。把 K、V 沿序列切成大小为 B_c 的块，对每个 Q 块，依次处理第 j 个 KV 块：

S_j    = Q_i K_j.T / sqrt(d)                       [B_r, B_c]，只在片上
m_new  = max(m, rowmax(S_j))
P_j    = exp(S_j − m_new)                          未归一化的分子
l      = l · exp(m − m_new) + rowsum(P_j)
O      = O · exp(m − m_new) + P_j V_j
最后    O = O / l

下面用 float64 在 CPU 上验证它和标准实现一致，序列长 1024，每次只看 128 个 key：

import torch
torch.manual_seed(0)
N, d, Bc = 1024, 64, 128
q, k, v = (torch.randn(N, d, dtype=torch.float64) for _ in range(3))
 
def standard(q, k, v):
    S = q @ k.T / d**0.5                     # 完整的 N×N
    return torch.softmax(S, -1) @ v
 
def blocked(q, k, v, Bc):
    m = torch.full((N, 1), float("-inf"), dtype=q.dtype)
    l = torch.zeros(N, 1, dtype=q.dtype)
    o = torch.zeros_like(q)
    for j in range(0, N, Bc):
        S = q @ k[j:j+Bc].T / d**0.5         # 任一时刻只有 N×Bc 的分数
        m_new = torch.maximum(m, S.amax(-1, keepdim=True))
        alpha = torch.exp(m - m_new)          # 旧状态换到新基准
        P = torch.exp(S - m_new)
        l = l * alpha + P.sum(-1, keepdim=True)
        o = o * alpha + P @ v[j:j+Bc]
        m = m_new
    return o / l, m + torch.log(l)            # 顺便返回每行的 logsumexp，反向要用
 
out, lse = blocked(q, k, v, Bc)
print("max |diff| =", (out - standard(q, k, v)).abs().max().item())

max |diff| = 7.771561172376096e-16

差异在 1e-16 量级，只来自浮点运算顺序不同。**FlashAttention 是精确 attention，不是近似。**它没有改变数学，只改变了计算顺序。

在 GPU 上，每个线程块负责一个 Q 块 Q_i，它和 O_i、m、l 一起待在 shared memory 和寄存器里；K、V 的块一块一块地读进来，用完即丢。块多大由片上容量决定：一次内层迭代要同时放下 Q、O、K、V 四个块，每块每行 d 个元素，约 4·B·d ≤ M（M 是一个 SM 能用的 SRAM 元素数）。FA1 论文取 B_c = ceil(M / 4d)、B_r = min(B_c, d)。比如 M = 49152 个 fp32（192 KB），d = 64 时 B_c = 192、B_r = 64。head dim 越大，块越小。实际 kernel 还要算上寄存器和双缓冲，块大小是手调的，FA2 在 {64, 128} × {64, 128} 里按 head dim 挑。

实现时有几个坑。某一行被整行 mask（比如 padding）时 m_new = −∞，exp(−∞ − (−∞)) 是 nan，要先特判。累加器 O 和 l 必须用 fp32，即使输入是 bf16。GPU 上 exp2 比 exp 快，所以把 log2(e) 预先乘进缩放系数。

停一下，检查理解单选

分块 attention 处理到新的 KV 块时，这一块的最大值比之前的 m 更大。对已经累加好的 l 和 O 应该怎么做？

丢弃之前的结果，从第一个块开始按新最大值重算两者都乘上 exp(m_old − m_new)，再加上新块的贡献不用处理，减最大值只是为了数值稳定，用哪个最大值都一样只修正 l，O 在最后除以 l 时会自动修正

省了多少 HBM 读写

FA1 论文给出的 IO 复杂度（N 序列长，d head dim，M 是 SRAM 大小，d ≤ M ≤ Nd）：

	HBM 访问量	额外显存
标准 attention	Θ(Nd + N²)	O(N²)（存 S、P）
FlashAttention	Θ(N²d² / M)	O(N)（每行一个 logsumexp）

推导按 FA1 的循环顺序（外层 KV 块、内层 Q 块）：K、V 每块只读一次，共 Θ(Nd)；但每处理一个 KV 块都要把全部 Q、O 扫一遍，共有 N / B_c = Θ(Nd / M) 个 KV 块，每遍 Θ(Nd)，合计 Θ(N²d²/M)。d = 64–128、M 约 100 KB 时 d²/M 远小于 1，比 N² 小很多倍。论文还证明了下界：对所有 M ∈ [d, Nd]，不存在 HBM 访问量为 o(N²d²/M) 的精确 attention 算法，这个量级已经最优。

开头那组数据就是这张表的实测：FLOPs 从 66.6 涨到 75.2 GFLOPs（多出的是反向重算），HBM 读写从 40.3 GB 降到 4.4 GB（约 1/9），时间从 41.7 ms 降到 7.3 ms。它用到了 GPU 与 Kernel 里的三招：tiling（S 按块算）、fusion（matmul、mask、softmax、dropout、matmul 融成一个 kernel）、重算（下一节）。

有一个常见误解要澄清：FlashAttention 没有降低计算复杂度，FLOPs 仍是 O(N²d)，长序列下 attention 依然是平方的。它降的是 HBM 访问和显存。想要降低计算复杂度，得换注意力的形式（稀疏、线性、滑动窗口，见注意力变体与 MoE）。

停一下，检查理解多选

关于 FlashAttention，哪些说法正确？

它的 FLOPs 比标准 attention 更多它算的是精确 attention，与标准实现只有浮点运算顺序带来的差异它把 attention 的计算量降到了序列长度的线性额外显存从 O(N²) 降到 O(N)，这是长上下文能训练的关键之一它的块越大越好，块大小只受序列长度 N 限制

反向：只存一个数，现算 P

标准反向要用 P 来算 dV、dP 再到 dS，所以前向得把 N×N 的 P 存进 HBM。FlashAttention 前向只存输出 O 和每行一个数：logsumexp L = m + log(l)。反向时重新读入 Q、K 的块，在片上用 P = exp(S − L) 现算出 P 的对应小块，再算 dQ、dK、dV：

S = q @ k.T / d**0.5
P_recomp = torch.exp(S - lse)                 # 只用前向存下的 lse，不需要 m、l 分开存
print("recomputed P matches softmax:", torch.allclose(P_recomp, torch.softmax(S, -1)))
print("lse shape", tuple(lse.shape))

recomputed P matches softmax: True
lse shape (1024, 1)

（接着上一节的代码运行。为了验证这里一次性算了完整 S，真实 kernel 是逐块算的。）一个 [N, 1] 的向量就概括了每一行的归一化信息。

为什么”多算一遍”划算？多出来的是一次 QK^T 和 exp，矩阵乘跑在 Tensor Core 上，是 compute-bound 的，额外 FLOPs 只占一小部分；省掉的是 Θ(N²) 的 HBM 写出和读回，而 attention 恰好卡在访存上。这和 activation checkpointing 不一样：checkpointing 用时间换显存，会变慢（见并行训练）；FlashAttention 的重算既省显存又更快，因为它重算的是便宜的东西、省下的是贵的东西。

反向的并行方式也和前向相反：每个线程块固定一个 KV 块 K_j, V_j，遍历所有 Q 块，dK_j、dV_j 在寄存器里累加、最后写一次；而 dQ_i 会被所有 KV 块贡献，只能用 fp32 的 atomic add 累加到 HBM 里的缓冲区。

停一下，检查理解单选

FlashAttention 反向需要注意力概率 P，它是怎么得到的？

前向把完整的 P 存进 HBM，反向直接读前向把 P 压成 FP8 存下来，反向解压前向只存每行的 logsumexp L，反向重新读 Q、K 块，用 exp(S − L) 在片上现算反向不需要 P，梯度只用 O 和 dO 就能算出来

FA1 到 FA4：每一代追着新的瓶颈

FA1 解决了 HBM 读写，但在 A100 上只跑到峰值的 25–40%。之后每一代的数学都没变，变的是执行方式，追的是硬件上新冒出来的瓶颈。

**FA2（2023，A100）：减少非矩阵乘运算，改并行划分。**A100 的 BF16 矩阵乘峰值 312 TFLOP/s，非矩阵乘的 FP32 只有 19.5 TFLOP/s，一次非矩阵乘 FLOP 约等于 16 次矩阵乘的代价。所以 FA2 不再每步把 O 除以 l，只乘修正因子、最后除一次；反向只存 L 一个数。并行上，FA1 一个线程块负责一个 (batch, head)，batch 为 1、16 个头时只有 16 个 block，108 个 SM 大多闲着；FA2 把 grid 改成 (Q 块数, B·H)，沿序列方向也并行。这之所以可行，是因为 softmax 按行进行，O_r 只依赖 Q_r 和全部 K、V，Q 块之间不需要通信。循环顺序也换了：外层 Q、内层 KV，O_i 从头到尾留在寄存器里，只写一次 HBM；block 内的 4 个 warp 按 Q 的行切分，不再像 FA1 那样切 K、再经 shared memory 归约。结果是比 FA1 快约 2 倍，A100 上前向到峰值的约 73%。

decode 阶段只有一个 query，序列方向切不出多个 Q 块，FA2 的并行又失效了。Flash-Decoding 改为沿 KV 切分，每块算出局部的 (m, l, O)，再用上面那个可结合的合并公式拼起来，细节在推理系统。

**FA3（2024，H100）：让 Tensor Core 和 exp 同时工作。**H100 SXM5 的 BF16 矩阵乘是 989 TFLOP/s，而算 exp 的 SFU 只有约 3.9 TFLOP/s（每 SM 每周期 16 次 × 132 个 SM × 1.83 GHz），差约 250 倍。head dim 128 时，矩阵乘的 FLOPs 约是 exp 次数的 512 倍，算下来 exp 的耗时能达到矩阵乘的一半。FA2 的写法是”读数据 → GEMM → softmax → GEMM”依次执行，softmax 时 Tensor Core 闲着，GEMM 时 SFU 闲着。FA3 用上 Hopper 的异步硬件：TMA 单元负责 HBM 和 shared memory 之间的块搬运，WGMMA 异步发出矩阵乘；warp 特化成生产者（只发 TMA）和消费者（只计算）；两个消费者 warpgroup 用 ping-pong 错开，一个做 softmax 时另一个正好在做 GEMM。它还支持 FP8，用 Hadamard 变换打散离群值降低量化误差。论文报告 BF16 约 740 TFLOP/s（约 75% 利用率），比 FA2 快 1.5–2 倍。

**FA4（B200）：exp 本身成了最大瓶颈。**Blackwell 的 Tensor Core 吞吐又翻了一倍，SFU 和 shared memory 带宽却没有同步增长。FA4 把一部分 exp2 交给普通的 FMA 单元用多项式近似算（整数部分直接写进浮点指数位，小数部分用低次多项式）；并且只在最大值增长超过阈值时才 rescale O。数学上这仍是同一个 softmax：O 和 l 始终用同一个指数基准累加，最后相除时基准抵消；但浮点实现里近似 exp 和舍入仍会带来误差。阈值是在 rescale 开销和指数的数值范围之间折中。累加器放进 Blackwell 新增的 TMEM（专给 Tensor Core 用的片上内存）。论文报告 B200 上 BF16 最高约 1613 TFLOP/s（约 71% 利用率）。

版本	硬件	新瓶颈	核心改动
FA1	A100 及更早	N×N 矩阵往返 HBM	分块 + online softmax + 重算
FA2	A100	非矩阵乘运算贵、SM 没用满	推迟除法、沿序列并行、外层 Q
FA3	H100	exp 和数据搬运拖住 Tensor Core	TMA、WGMMA、warp 特化、ping-pong、FP8
FA4	B200	exp 和 shared memory 带宽	软件模拟 exp2、有条件 rescale、TMEM

停一下，检查理解单选

FA2 论文强调减少 rescale、除法这类非矩阵乘运算，哪怕它们只占总 FLOPs 的很小比例。原因是什么？

非矩阵乘运算要写 HBM，矩阵乘不用A100 上非矩阵乘 FP32 吞吐只有矩阵乘的约 1/16，少量这类运算也会占掉可观的时间除法在 GPU 上结果不精确非矩阵乘运算只能由单个线程执行

实际怎么用

大多数时候不用自己写。PyTorch 的 F.scaled_dot_product_attention(q, k, v, is_causal=True) 会按条件自动选 Flash 后端；也可以直接用 flash_attn 包的 flash_attn_func(q, k, v, causal=True)。两个坑：布局不同，SDPA 要 (B, H, N, d)，flash_attn_func 要 (B, N, H, d)，形状碰巧对得上时传错不报错但结果是错的；SDPA 的 Flash 后端只接受 CUDA 上的 fp16 / bf16，输入是 fp32 时会悄悄换成别的后端，不报错也不加速。

GQA（多个 Q 头共享一组 KV 头，见注意力变体与 MoE）在 FlashAttention 里几乎不费事：前向时第 h 个 Q 头用 h / g 找到对应的 KV 头，只改一个指针偏移，K、V 不复制，同组的线程块读的是同一块显存，重复读大多命中 L2。

自己写的场景是改 attention 的形状：新的 mask 模式、新的位置编码、要在 attention 里融合别的运算。这时上面那段 blocked 函数就是 kernel 的骨架，换成 Triton 后外层循环变成 grid，q @ k.T 变成 tl.dot，其他几乎一行对一行。

记住这三件事

标准 attention 慢在把 N×N 的 S、P 往返 HBM，FlashAttention 让它们只以小块出现在片上，FLOPs 不降反升，时间和显存却大幅下降。
online softmax 用修正因子 exp(m_old − m_new) 把旧累加量换到新基准，状态 (m, l, O) 可以两两合并，所以能分块、能并行，结果是精确的。
反向只存每行的 logsumexp、现算 P，这是用便宜的矩阵乘换昂贵的显存访问；FA2 到 FA4 的数学不变，改的是执行方式，追的是每代硬件上新的瓶颈。

CS 与 LLM 学习知识库

探索

FlashAttention

标准 attention 的时间花在搬运上

为什么 softmax 不能直接分块

Online softmax：边读边修正

把输出也在线累加：分块 attention

省了多少 HBM 读写

反向：只存一个数，现算 P

FA1 到 FA4：每一代追着新的瓶颈

实际怎么用

记住这三件事

延伸阅读

目录

反向链接