通过 tiling + recomputation 将 attention 融合为单个 GPU kernel，避免 $N \times N$ 中间矩阵的 HBM 读写，IO 复杂度从 $\Theta(Nd + N^2)$ 降至最优的 $\Theta(N^2 d^2 / M)$，实现 2–4× 加速和线性内存。
标准 attention 实现将 $\mathbf{S} = \mathbf{Q}\mathbf{K}^\top$ 和 $\mathbf{P} = \text{softmax}(\mathbf{S})$ 两个 $N \times N$ 矩阵完整写入 HBM，产生 $O(N^2)$ 内存占用和大量 HBM 读写。在 A100 上，HBM 带宽（1.5–2.0 TB/s）远低于 on-chip SRAM 带宽（~19 TB/s），使 attention 成为 memory-bound 操作。
近似 attention 方法（Reformer、Performer、Linformer 等）降低了 FLOP 复杂度，但未实现 wall-clock 加速，根因是忽略了 IO 开销。
核心原理——IO-Awareness：将 attention 视为 memory-bound 操作，优化目标从 FLOP 减少转向 HBM 读写减少。
技术 1——Tiling（在线 softmax）：
将 $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ 分成大小为 $B_r \times d$ 和 $B_c \times d$ 的 block，外层循环遍历 K/V block（加载至 SRAM），内层循环遍历 Q block。关键是 softmax 可以增量计算：
$$m(x) = \max_i x_i, \quad f(x) = [e^{x_1 - m(x)}, \ldots, e^{x_B - m(x)}], \quad \ell(x) = \sum_i f(x)_i$$
对拼接向量 $x = [x^{(1)}, x^{(2)}]$：
$$m(x) = \max(m(x^{(1)}), m(x^{(2)}))$$
$$\ell(x) = e^{m(x^{(1)}) - m(x)} \ell(x^{(1)}) + e^{m(x^{(2)}) - m(x)} \ell(x^{(2)})$$
维护运行统计量 $(m, \ell)$ 即可逐 block 计算 softmax，无需物化完整 $N \times N$ 矩阵。
技术 2——Recomputation（反向传播）：
前向传播仅存储 $\mathbf{O}$ 和 softmax 统计量 $(m, \ell)$——$O(N)$ 额外内存。反向传播时从 $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ block 在 SRAM 中重新计算 $\mathbf{S}, \mathbf{P}$。虽然增加 FLOP，但因减少 HBM 读写反而更快。
关键中间量 $D_i = do_i^\top o_i$ 替代了需要 $O(N)$ 向量的 $P_{i:}^\top dP_{i:}$，使反向传播也可 block-wise 进行。
Block-Sparse FlashAttention：在 FlashAttention 基础上增加 block-level mask $\mathbf{M}$，跳过零 block 的计算，IO 复杂度降至 $\Theta(Nd + N^2 d^2 M^{-1} s)$（$s$ 为非零 block 比例）。
核心技术壁垒：在保持 exact attention 输出的前提下，将 softmax 的全局归一化分解为可增量更新的局部统计量——这是一个数学非平凡的观察（algebraic aggregation），它使得整个 attention（matmul + softmax + matmul + dropout + mask）可融合为单个 kernel，消除所有中间 HBM 读写。
数据流：外层循环遍历 K/V block（红色路径），内层遍历 Q block（蓝色路径）。每次在 SRAM 内完成 $\mathbf{S}_{ij}$ 计算 → 在线 softmax → $\tilde{\mathbf{P}}_{ij} \mathbf{V}_j$ → 累加到 $\mathbf{O}_i$ 并更新 $(m, \ell)$。$N \times N$ 矩阵从不被物化到 HBM。
| 参数 | A100 值 |
|---|---|
| HBM 容量 / 带宽 | 40–80 GB / 1.5–2.0 TB/s |
| SRAM 容量 / 带宽 | 192 KB/SM × 108 SMs / ~19 TB/s |
| Tensor Core peak (FP16) | 312 TFLOP/s |
| Block sizes | $B_c = \lceil M/(4d) \rceil$, $B_r = \min(\lceil M/(4d) \rceil, d)$ |
| 符号 | 含义 |
|---|---|
| $N$ | 序列长度 |
| $d$ | Head dimension |
| $M$ | On-chip SRAM 大小 |
| $B_r, B_c$ | Q 行 block 和 K/V 列 block 大小 |
| $T_r, T_c$ | Block 数量 $\lceil N/B_r \rceil$, $\lceil N/B_c \rceil$ |
| $m, \ell$ | 在线 softmax 的 running max 和 running sum |
| $\mathbf{S}, \mathbf{P}$ | Score matrix 和 attention weight matrix |
Theorem 1（正确性）：FlashAttention 输出 $\mathbf{O} = \text{softmax}(\mathbf{Q}\mathbf{K}^\top)\mathbf{V}$，FLOPs $O(N^2 d)$，额外内存 $O(N)$。通过对外层循环变量 $j$ 的归纳证明：每步后 $m^{(j)} = \text{rowmax}(\mathbf{S}_{:,:j})$，$\ell^{(j)} = \text{rowsum}(\exp(\mathbf{S}_{:,:j} - m^{(j)}))$，$\mathbf{O}^{(j)} = \mathbf{P}_{:,:j}\mathbf{V}_{:j}$。
Theorem 2（IO 复杂度）：
$$\text{标准 attention IO} = \Theta(Nd + N^2)$$
$$\text{FlashAttention IO} = \Theta(N^2 d^2 / M)$$
推导：每次加载大小 $\Theta(M)$ 的 K/V block，需要 $T_c = \Theta(Nd/M)$ 次遍历 Q。每次遍历加载 $\Theta(Nd)$，总计 $\Theta(N^2 d^2 / M)$。
对于典型 $d = 64$, $M \approx 100$ KB：$d^2/M \approx 4096/10^5 \approx 0.04 \ll 1$，FlashAttention IO 比标准 attention 少约 $9\times$。
Proposition 3（下界/最优性）：不存在精确 attention 算法在所有 $M \in [d, Nd]$ 上实现 $o(N^2 d^2 / M)$ IO 复杂度。证明：若存在则在 $M = \Theta(Nd)$ 时实现 $o(Nd)$ IO，但输入输出大小为 $Nd$，矛盾。
标准 attention 的 arithmetic intensity = $O(N^2 d) / O(Nd + N^2) \approx O(d)$（当 $N \gg d$），位于 memory-bound 区域。FlashAttention 通过减少 IO 提升 arithmetic intensity 至 $O(M/d)$，推向 compute-bound 边界。
实测数据（A100, GPT-2 medium, seq 1024）：
| GFLOPs | HBM R/W (GB) | Runtime (ms) | |
|---|---|---|---|
| Standard | 66.6 | 40.3 | 41.7 |
| FlashAttention | 75.2 | 4.4 | 7.3 |
FlashAttention FLOPs 更高（75.2 vs 66.6，因 recomputation）但 HBM 读写少 9.2×，运行快 5.7×——直观证明 FLOPs 是误导性指标。
| 模型 | 基线 | FlashAttention | 加速 |
|---|---|---|---|
| BERT-large (seq 512) | 20.0 ± 1.5 min (MLPerf) | 17.4 ± 1.4 min | 15% |
| GPT-2 small (seq 1K) | 9.5 days (HF) | 2.7 days | 3.5× |
| GPT-2 medium (seq 1K) | 21.0 days (HF) | 6.9 days | 3.0× |
| LRA (seq 1K–4K) | — | — | 2.4× |
| GPT-2 small | Context | Perplexity | Time |
|---|---|---|---|
| Megatron-LM | 1K | 18.2 | 4.7 days |
| FlashAttention | 1K | 18.2 | 2.7 days |
| FlashAttention | 4K | 17.5 | 3.6 days |
4K 上下文：比 Megatron 1K 快 30% 且 perplexity 低 0.7。
FlashAttention 是首个在 Path-X（seq 16K）上达到 61.4% 准确率的 Transformer（所有其他方法均为随机水平）。Block-sparse FlashAttention 在 Path-256（seq 64K）上达到 63.1%。
| GPU | 典型加速范围 |
|---|---|
| A100 (hdim 64) | 2–4× |
| A100 (hdim 128) | 略低（block 更大，SRAM 压力更高） |
| RTX 3090 | 2.5–4.5×（内存带宽更低，IO 优化收益更大） |
| T4 | 加速较低（SRAM 更小 → block 更小） |
| 步骤 | 论点 | 证据 |
|---|---|---|
| 1 | 标准 attention 是 memory-bound：FLOPs 不是瓶颈，HBM IO 才是 | §2.1 A100 HBM vs SRAM 带宽差 10×；Fig. 2 table（75.2 vs 66.6 GFLOPs 但 9.2× fewer HBM R/W → 5.7× faster） |
| 2 | IO-awareness 原则：优化 HBM 读写次数而非 FLOP 数 | §1 引用 IO complexity 文献 [1]；§3.2 Theorem 2 推导 FlashAttention IO = $\Theta(N^2 d^2/M)$ |
| 3 | Tiling + online softmax 使 attention 可逐 block 计算 | §3.1 softmax 增量分解公式；Theorem 1 归纳证明正确性 |
| 4 | Recomputation 以少量额外 FLOP 换取大量 HBM IO 节省 | §3.1 反向传播仅存 $(m, \ell)$；Fig. 2 证明 IO 减少主导 runtime |
| 5 | IO 复杂度渐近最优 | Proposition 3 下界证明 |
| 6 | 实验验证速度和质量双赢 | §4 BERT/GPT-2/LRA 加速 + Path-X 首次解决 + GPT-2 perplexity 提升 |
代码仓库：flash-attention
论文承认 CUDA kernel 方法不可跨 GPU 架构移植（§5），呼吁类似 Halide 的编译器框架。当前仅支持 NVIDIA Turing/Ampere。