FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning #
§1 TL;DR #
在 FlashAttention 基础上通过三项优化——减少 non-matmul FLOPs、序列维度并行化、warp 间 split-Q 分工——将 A100 上 attention 前向吞吐从 25–40% 提升至 50–73% of peak，端到端训练达 225 TFLOPs/s（72% MFU）。
§2 痛点 / 方法 / 结果 #
Q1 痛点 #
FlashAttention 虽然将 attention 从 memory-bound 推向 compute-bound，但在 A100 上前向传播仅达理论峰值的 30–50%，反向传播仅 25–35%。而优化后的 GEMM 可达 80–90%。差距来自两个层面：
- Non-matmul FLOPs 开销：GPU Tensor Core 使 matmul 吞吐（312 TFLOPs/s FP16）比 non-matmul（19.5 TFLOPs/s FP32）高 16×。FlashAttention 每步对 $\mathbf{O}_i$ 做 $\text{diag}(\ell)^{-1}$ rescaling，产生不必要的 non-matmul 开销。
- Work partitioning 低效：FlashAttention 外层循环遍历 K/V block（列方向），并行度仅 batch × heads；对长序列小 batch 场景 GPU 占用率低。Warp 间采用 split-K 方案需要 shared memory 同步。
Q2 方法 #
改进 1——减少 non-matmul FLOPs（§3.1）：
- Tweak 1（延迟 rescaling）：不在每步除以 $\ell^{(j)}$，而是维护"未归一化"的输出 $\tilde{\mathbf{O}}$，仅在循环结束后做一次 $\text{diag}(\ell)^{-1}$ 除法：
$$\tilde{\mathbf{O}}^{(j)} = \text{diag}(e^{m^{(j-1)} - m^{(j)}}) \tilde{\mathbf{O}}^{(j-1)} + e^{\mathbf{S}^{(j)} - m^{(j)}} \mathbf{V}^{(j)}$$
$$\mathbf{O} = \text{diag}(\ell^{(T_c)})^{-1} \tilde{\mathbf{O}}^{(T_c)}$$
每次迭代省去一次 element-wise division。
- Tweak 2（logsumexp 融合）：反向传播只需 $L = m + \log(\ell)$ 一个统计量，不需分别存储 $m$ 和 $\ell$，减少内存和 non-matmul 操作。
改进 2——序列维度并行化（§3.2）：
- 前向传播：外层循环改为遍历 Q 行 block（而非 K/V 列 block）——不同 Q block 之间无通信，可跨 thread block 并行。总并行度 = batch × heads × $T_r$。
- 反向传播：外层循环遍历 K/V 列 block（不同 Q block 需累加 $\mathbf{dQ}$），使用 atomic add 通信。
改进 3——Warp 间 split-Q 分工（§3.3）：
FlashAttention 采用 split-K：4 warps 分别处理 K 的不同切片，计算 $\mathbf{Q}\mathbf{K}^\top$ 后需写入 shared memory 同步再合并。FlashAttention-2 改为 split-Q：4 warps 分别处理 Q 的不同切片，各自独立计算 $\mathbf{Q}_{\text{warp}} \mathbf{K}^\top \to \tilde{\mathbf{P}}_{\text{warp}} \mathbf{V}$，无需 warp 间通信。
Causal masking 优化：(1) 跳过列索引 > 行索引的 block（约一半），获得 1.7–1.8× 加速；(2) 对非边界 block 无需 apply mask。
核心技术壁垒：将 FlashAttention 的循环方向从外层-K/V 翻转为外层-Q——这一看似简单的改动解锁了序列维度并行化（Q block 间无依赖），同时使 split-Q warp 分工无需同步。循环翻转最早由 Phil Tillet 在 Triton 实现中提出。
Q3 结果 #
- Kernel 加速：比 FlashAttention 快 1.7–3.0×，前向达 230 TFLOPs/s（73% peak），反向 63% peak。
- 端到端训练：GPT3-2.7B 8K context 达 225 TFLOPs/s / A100（72% MFU），比无 FlashAttention 基线快 2.8×，比 FlashAttention 快 1.3×。
- H100（未优化）：直接运行达 335 TFLOPs/s，不利用 TMA / 4th-gen Tensor Core。
§3 架构 / 方法图 #
flowchart TB
subgraph Forward ["Forward Pass: Outer loop over Q blocks"]
direction LR
Q1["Q block 1
Thread Block A"] --> |"inner loop K,V"| Out1["O block 1"]
Q2["Q block 2
Thread Block B"] --> |"inner loop K,V"| Out2["O block 2"]
Q3["Q block 3
Thread Block C"] --> |"inner loop K,V"| Out3["O block 3"]
end
subgraph Backward ["Backward Pass: Outer loop over K,V blocks"]
direction LR
KV1["K,V block 1
Thread Block X"] --> |"inner loop Q,dO"| dKV1["dK,dV block 1"]
KV2["K,V block 2
Thread Block Y"] --> |"inner loop Q,dO"| dKV2["dK,dV block 2"]
end
subgraph WarpPartition ["Warp Partitioning (Forward)"]
direction TB
Ksm["K in shared memory
(all warps read)"]
Vsm["V in shared memory
(all warps read)"]
W1["Warp 1: Q slice 1"] --> |"QK^T → P → PV"| O1["O slice 1"]
W2["Warp 2: Q slice 2"] --> |"QK^T → P → PV"| O2["O slice 2"]
W3["Warp 3: Q slice 3"] --> |"QK^T → P → PV"| O3["O slice 3"]
W4["Warp 4: Q slice 4"] --> |"QK^T → P → PV"| O4["O slice 4"]
end
Target operation #
- 输入：$\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{N \times d}$，head dim $d \in \{64, 128\}$
- 输出：$\mathbf{O} = \text{softmax}(\mathbf{Q}\mathbf{K}^\top)\mathbf{V}$
- 精度：FP16/BF16 输入，FP32 accumulator
- Mask：causal（可选），block-level skip 优化
Hardware model #
| 参数 | A100 值 |
|---|
| FP16 matmul peak | 312 TFLOPs/s |
| Non-matmul FP32 peak | 19.5 TFLOPs/s |
| Matmul / non-matmul 比 | 16× |
| SMs | 108 |
| Warps per thread block | 4 or 8 |
| Block sizes | $\{64, 128\} \times \{64, 128\}$ |
FLOPs 计算 #
$$\text{FLOPs}_{\text{fwd}} = 4 \times N^2 \times d \times h$$
其中 $h$ = heads。Causal mask 减半。Backward = forward × 2.5（5 matmuls vs 2）。
§4 作者证明 #
符号表 #
| 符号 | 含义 |
|---|
| $B_r, B_c$ | Q 行 block / K,V 列 block 大小 |
| $T_r, T_c$ | Block 数量 |
| $m^{(j)}, \ell^{(j)}$ | 第 $j$ 步的 running max / running sum |
| $\tilde{\mathbf{O}}^{(j)}$ | 未归一化的累积输出 |
| $L$ | logsumexp $= m + \log(\ell)$ |
| $D$ | $\text{rowsum}(\mathbf{dO} \circ \mathbf{O})$，反向传播辅助量 |
方程物理意义 #
延迟 rescaling：原 FlashAttention 每步计算 $\mathbf{O}^{(j)} = \text{diag}(\ell^{(j)})^{-1}(\ldots)$，即每步除以当前归一化因子。FA2 改为维护 $\tilde{\mathbf{O}}^{(j)}$（不除），最后一步才除。数学等价但省去 $T_c - 1$ 次 element-wise division。
logsumexp 融合：$L = m + \log(\ell)$ 将两个标量压缩为一个。反向传播用 $\mathbf{P} = \exp(\mathbf{S} - L)$ 直接恢复 attention weight（vs. FA1 分别用 $m$ 和 $\ell$），减少内存和计算。
无独立性能模型：论文不含 throughput/latency 解析模型，所有效率论证基于 hardware specs（matmul vs non-matmul throughput 16× 差距）和实测 benchmark。
6 项检查 #
- 正确性：Algorithm 1 输出 $\mathbf{O} = \text{softmax}(\mathbf{Q}\mathbf{K}^\top)\mathbf{V}$，证明与 FA1 相同（仅 rescaling 时机不同）。✔
- Non-matmul FLOP 减少：延迟 rescaling 省去每步 $\text{diag}(\ell^{(j)})^{-1}$；logsumexp 省去 $(m, \ell)$ 分存。✔
- 序列并行正确性：前向 Q block 间无依赖（embarrassingly parallel）；反向 dQ 通过 atomic add 累加。✔
- Split-Q 无通信：每 warp 独立计算 $\mathbf{Q}_{\text{warp}} \mathbf{K}^\top \to \tilde{\mathbf{P}} \mathbf{V}$，无 shared memory 写回。✔
- Causal mask 加速：跳过全零 block + 仅边界 block apply mask，理论极限 2×，实测 1.7–1.8×。✔
- GQA/MQA 支持：多个 Q head 共享一个 K/V head 通过隐式索引实现，反向传播对 dK/dV 跨 head 求和。✔
Design space #
| 优化轴 | 选择 | 拒绝的替代方案 | 约束 |
|---|
| 循环方向 | 外层 Q / 内层 K,V | 外层 K,V（FA1） | Q block 间无依赖是并行化前提 |
| Warp 分工 | split-Q | split-K（FA1） | split-K 需 shared memory 同步 |
| Rescaling 时机 | 延迟至循环结束 | 每步 rescale（FA1） | 延迟增加少量寄存器但省 non-matmul |
| Block size | {64,128}² | 更大 block | 寄存器压力 / SRAM 溢出 |
| 统计量存储 | logsumexp $L$ | 分存 $m, \ell$（FA1） | $L$ 合二为一减少 IO 和内存 |
§5 实验与数据 #
Kernel benchmark（A100-80GB） #
| 配置 | FA2 TFLOPs/s | FA1 TFLOPs/s | 加速 |
|---|
| Fwd, hdim 64, no causal, seq 16K | ~230 | ~130 | 1.8× |
| Fwd, hdim 128, causal, seq 16K | ~200 | ~100 | 2.0× |
| Fwd+Bwd, hdim 128, no causal | — | — | 1.7–3.0× |
前向达 73% peak，反向达 63% peak。对比 GEMM 80–90% peak，仍有差距（归因于 non-matmul 操作和反向传播的 SRAM 压力）。
端到端训练（8× A100-80GB） #
| Model | 配置 | 无 FA | FA1 | FA2 |
|---|
| GPT3-1.3B | 2K ctx | 142 | 189 | 196 TFLOPs/s |
| GPT3-1.3B | 8K ctx | 72 | 170 | 220 TFLOPs/s |
| GPT3-2.7B | 2K ctx | 149 | 189 | 205 TFLOPs/s |
| GPT3-2.7B | 8K ctx | 80 | 175 | 225 TFLOPs/s |
GPT3-2.7B 8K：FA2 比无 FA 快 2.8×（80→225），比 FA1 快 1.3×（175→225）。72% MFU。
H100（未针对优化） #
直接运行 FA2 在 H100 上达 335 TFLOPs/s，预计利用 TMA + 4th-gen Tensor Core 可再提升 1.5–2×。
论文承认的弱项 #
- 前向 73% vs GEMM 80–90%：非 matmul 操作无法完全消除（softmax 涉及 exp / sum / max / rescaling）。
- 反向 63% vs 前向 73%：反向 5 个 matmul + 更多 SRAM 压力，优化空间有限。
- Causal mask 实测 1.7–1.8× vs 理论 2×：边界 block 仍需 mask 逻辑。
§6 论证链 #
| 步骤 | 论点 | 证据 |
|---|
| 1 | FlashAttention 在 A100 上仅达 30–50% peak，远低于 GEMM 的 80–90% | §1 para 3（实测数据，Fig. 5 / Fig. 6） |
| 2 | 瓶颈是 non-matmul FLOPs 和 work partitioning，不是 IO | §3.1（matmul vs non-matmul 16× 差距）；§3.2–3.3（low occupancy + shared memory sync） |
| 3 | 延迟 rescaling + logsumexp 减少 non-matmul 开销 | §3.1.1 算法对比（FA1 vs FA2 公式） |
| 4 | 外层-Q 循环解锁序列维度并行 | §3.2 前向 embarrassingly parallel；Phil Tillet Triton 实现先验 |
| 5 | Split-Q warp 分工消除 shared memory 同步 | §3.3 Fig. 3 对比 split-K vs split-Q |
| 6 | 三项改进叠加达 2× kernel 加速，端到端 72% MFU | §4 全部 benchmark 数据 |
§7 实现 cross-reference #
代码仓库：flash-attention
- Kernel：CUDA，手动管理 shared memory 和寄存器分配
- Block size：{64, 128}² 四种组合，按 head dim 手动选择
- Causal mask：编译时分支，跳过全零 block
- GQA/MQA：隐式索引，不复制 K/V
关键实现细节 #
- Split-Q 的寄存器布局：每个 warp 持有 Q slice 的 registers，K/V 在 shared memory 中广播。$\mathbf{Q}_{\text{warp}} \mathbf{K}^\top$ 的结果留在 registers 中直接进入 softmax + $\mathbf{P}\mathbf{V}$ 计算，全程不写 shared memory。
- Backward dQ atomic add：不同 thread block 处理不同 K/V block 但累加到同一 dQ，使用 atomicAdd。论文未讨论非确定性 / 精度影响。
可移植性 #
- 当前仅支持 A100 / H100（NVIDIA Ampere / Hopper）
- Triton 有独立实现（FA2 in Triton），性能略低但更易移植
- 未利用 H100 特有指令（TMA, 4th-gen Tensor Core）——这成为 FA3 的动机
部署上下文 #
- Serving stack：可替代 FA1 成为 vLLM / SGLang / HuggingFace 的默认 attention kernel
- Regime：prefill（大 batch，compute-bound）为主要受益场景；decode（单 token，memory-bound）收益较小
- Launch shape：grid = batch × heads × $T_r$，block = 4 或 8 warps
Software → Hardware 启示 #
- 更高 non-matmul 吞吐：16× 的 matmul/non-matmul 差距是根本限制。若 ISA 提供更快的 exp / sum / max 指令，attention 可进一步接近 GEMM 效率。
- 自动 block size tuning：当前 4 种手动选择，auto-tuning 可扩展到更多硬件配置。
- Warp-level 同步原语：split-Q 成功消除同步，但反向传播仍需部分同步——更灵活的 warp-level 通信原语可进一步优化。