FlashAttention

5.2.1 Flash Attention

FlashAttention 是一种精确 Attention 算法，不是近似 Attention。它的目标不是改变 Transformer 的数学定义，而是改变 Attention 在 GPU 上的计算方式：

普通 Attention 会显式构造巨大的注意力矩阵

S = Q K^{⊤} \in R^{n \times n},

然后再计算

P = softmax (S), O = P V .

FlashAttention 的关键是：

不显式保存完整的 (QK^\top \in \mathbb{R}^{n\times n}) 注意力矩阵，而是分块计算，并在 GPU SRAM 中完成局部计算，只把最终结果写回 HBM。

原始 FlashAttention 论文将其定义为一种 IO-aware exact attention algorithm，核心是通过 tiling 减少 GPU 高带宽显存 HBM 与片上 SRAM 之间的数据读写。

1. 普通 Attention 的主要瓶颈

Attention (Q, K, V) = softmax (\frac{Q K^{⊤}}{\sqrt{d}}) V .

普通 Attention 的理论计算复杂度是： $O (n^{2} d)$ ，显存复杂度主要来自注意力矩阵： $S, P \in R^{n \times n}$

所以显存复杂度是：

O (n^{2})

当序列长度 (n) 很大时，(n^2) 会非常夸张。

例如： $n = 8192$ ，则注意力矩阵大小为 $8192^{2} = 67, 108, 864$ ，如果使用 FP16，每个元素 2 bytes，那么一个矩阵大约需要：

67, 108, 864 \times 2 \approx 134 MB .

但训练时通常不止保存一个矩阵，还要保存 softmax 结果、dropout mask、中间激活，用于反向传播。因此显存消耗会远大于这个数。

更本质的问题是：

GPU 很快，但显存读写相对慢。普通 Attention 的大量时间花在 HBM 与 SRAM 之间搬运 (n\times n) 矩阵。

FlashAttention 要优化的正是这个 IO 瓶颈。

2. GPU 内存层级：为什么 IO 很重要？

可以粗略理解为 GPU 有两类重要存储：

存储位置	说明	容量	速度	特点	用途
HBM（GPU 显存）	离计算核心较远，访问延迟更高	大（几十 GB 到上百 GB）	带宽高，比 CPU 内存快很多，但相对片上缓存较慢	深度学习大张量主要存放处	参数、激活值、梯度等大规模数据
SRAM / Shared Memory / Registers（片上缓存）	离计算单元非常近，访问延迟低	小（远小于 HBM）	极快	适合频繁重复访问的数据	常配合矩阵乘法、卷积、Attention 等 block/tile 级计算

HBM 全称是 High Bandwidth Memory，可以理解为 GPU 上的大容量显存。在深度学习中，HBM 主要存放大规模张量：

模型参数、输入数据、激活值、梯度、优化器状态、中间计算结果

SRAM 全称 Static Random Access Memory ，它是 GPU 芯片内部的高速存储。在深度学习中，它主要用于：

临时存放正在计算的一小块数据

可以总结概括为：

HBM 是 GPU 的大仓库，但每次从 HBW 取数据到计算核心都有代价
SRAM 是计算核心旁边的小工作台

FlashAttention 的核心就是：

尽量在工作台 SRAM 上连续完成更多操作，少访问大仓库 HBM

普通 Attention 的流程大概是：

\underset{HBM}{Q, K} \overset{读入}{\to} \underset{片上 SRAM}{Q, K} \overset{计算核心： Q K^{⊤}}{\to} \underset{片上 SRAM}{S \in R^{N \times N}} \overset{写回}{\to} \underset{HBM}{S} \underset{HBM}{S} \overset{读入}{\to} \underset{片上 SRAM}{S} \overset{计算核心：softmax (S)}{\to} \underset{片上 SRAM}{P \in R^{N \times N}} \overset{写回}{\to} \underset{HBM}{P} \underset{HBM}{P, V} \overset{读入}{\to} \underset{片上 SRAM}{P, V} \overset{计算核心： P V}{\to} \underset{片上 SRAM}{O \in R^{N \times d}} \overset{写回}{\to} \underset{HBM}{O}

问题是 $S$ 和 $P$ 都是 $N \times N$ ，非常大，反复在 HBM 和 SRAM 之间搬运代价极高。

FlashAttention 的想法是：

把 $Q, K, V$ 分块搬进 SRAM，在计算核心中完成局部矩阵乘法和在线 softmax，将输出 $O$ 逐块累积在 SRAM 中，完整的 $S$ 和 $P$ 始终不写回 HBM，最终只将 $O$ 写回 HBM。

3. FlashAttention

（1）分块矩阵乘法

Attention 中的 $Q K V$ ：

Q = (\begin{matrix} q_{1} \\ q_{2} \\ q_{3} \\ ⋮ \\ q_{n} \end{matrix}) K = (\begin{matrix} k_{1} \\ k_{2} \\ k_{3} \\ ⋮ \\ k_{n} \end{matrix}) V = (\begin{matrix} v_{1} \\ v_{2} \\ v_{3} \\ ⋮ \\ v_{n} \end{matrix}), Q K V \in R^{n \times d}

将 $Q K V$ 分块：

Q = [\begin{matrix} Q_{1} \\ Q_{2} \\ ⋮ \\ Q_{T_{Q}} \end{matrix}], K = [\begin{matrix} K_{1} \\ K_{2} \\ ⋮ \\ K_{T_{K}} \end{matrix}], V = [\begin{matrix} V_{1} \\ V_{2} \\ ⋮ \\ V_{T_{K}} \end{matrix}] Q_{i} \in R^{B_{q} \times d}, K_{j}, V_{j} \in R^{B_{k} \times d} .

其中， $B_{q}$ 和 $B_{k}$ 分别是每个分块 $Q_{i}$ ， $K_{j} V_{j}$ 的词向量的个数。

对于每一个 query block (Q_i)，FlashAttention 逐个遍历 key/value block。这里以 $B_{q} = 2, B_{k} = 3$ 为例：

Q_{i} = [\begin{matrix} q_{i 1} \\ q_{i 2} \end{matrix}], K_{j} = [\begin{matrix} k_{j 1} \\ k_{j 2} \\ k_{j 3} \end{matrix}], V_{j} = [\begin{matrix} v_{j 1} \\ v_{j 2} \\ v_{j 3} \end{matrix}]

计算小块 attention score：

S_{i j} = \frac{Q_{i} K_{j}^{T}}{\sqrt{d}} = \frac{1}{\sqrt{d}} (\begin{matrix} q_{i 1} k_{j 1}^{T} & q_{i 1} k_{j 2}^{T} & q_{i 1} k_{j 3}^{T} \\ q_{i 2} k_{j 1}^{T} & q_{i 2} k_{j 2}^{T} & q_{i 2} k_{j 3}^{T} \end{matrix}) \in R^{B_{q} \times B_{k}}

这个小矩阵 (S_{ij}) 会放在 SRAM / shared memory 中，用完就丢，不写回 HBM。

接下来会经过局部 Softmax 更新，得到 $softmax (S_{i j})$ ，具体如何处理后面再讨论。

最后，局部 Value 加权：

o_{i j} = softmax (S_{i j}) V_{j} \in R^{B_{q} \times d}

对于 query block $Q_{i}$ ，这样逐个遍历 key/value block $(K_{1}, V_{1}), (K_{2}, V_{2}), \dots, (K_{T_{k}}, V_{T_{k}})$ ：

o_{i} = {o_{i 1}, o_{i 2}, \dots, o_{i T_{k}}}

将这些结果相加，得到最终 attention 对应 $Q_{i}$ 位置上的输出：

O_{i} = (softmax (\frac{Q K^{⊤}}{\sqrt{d}}) V)_{i} = \sum_{j = 1}^{T_{k}} o_{i j} \in R^{B_{q} \times d}

将每个 query block 遍历的输出拼接起来，就得到了最终的 Attention 模块的输出：

O = Concat (O_{1}, O_{2}, \dots, O_{T_{Q}}) \in R^{n \times d}

（2）Online Softmax

回到之前留下的问题，在得到小块 attention score $S_{i j}$ 之后，我们需要将其进行局部 softmax 更新。分块计算的难点是：虽然每次只看到一个 KV block，但最终必须得到全局 softmax 的结果。

对第 (i) 个 query token，标准 Attention 输出是：

o_{i} = \sum_{j = 1}^{n} \frac{\exp (s_{i j})}{\sum_{t = 1}^{n} \exp (s_{i t})} v_{j}, s_{i j} = \frac{q_{i} k_{j}^{T}}{\sqrt{d}}

注意 softmax 的分母是整行求和 $\sum_{t = 1}^{n} \exp (s_{i t})$ ，但分块运算时我们只拿到一个 block，例如只看到 $s_{i 1}, s_{i 2}, \dots, s_{i B_{k}}$ ，无法直接计算出 softmax 的分母。因此，我们就无法直接对每个 block 单独做 softmax。

Online Softmax 的核心思想是：维护运行统计量（running statistics），每来一个新的 block，就更新这些统计量，最终得到与全局 softmax 完全一致的结果。

Online Softmax 的核心

先考虑一个单独的 query token：

q_{i} \in R^{1 \times d}, K \in R^{n \times d}, V \in R^{n \times d}

现在受限与 SRAM 内存，只能先进入

K_{1} = [\begin{matrix} k_{1}, \\ ⋮, \\ k_{N_{1}} \end{matrix}], V_{1} = [\begin{matrix} v_{1}, \\ ⋮, \\ v_{N_{1}} \end{matrix}]

对于 $K_{1}, V_{1}$ ，计算局部 attention 输出：

\begin{array}{r} s_{i j} = \frac{q_{i} k_{j}^{⊤}}{\sqrt{d}} \in R^{1}, softmax (s_{i j}) = \frac{\exp (s_{i j})}{\sum_{r = 1}^{N_{1}} \exp (s_{i r})} \in R^{1}, \\ o_{i}^{N_{1}} = softmax (s_{i \cdot}) V_{1} = \sum_{j = 1}^{N_{1}} softmax (s_{i j}) v_{j} \in R^{1 \times d} \end{array}

为避免数值溢出，使用稳定 softmax：

o_{i}^{N_{1}} = \sum_{j = 1}^{N_{1}} \frac{\exp (s_{i j} - m_{i}^{N_{1}}) v_{j}}{\sum_{r = 1}^{N_{1}} \exp (s_{i r} - m_{i}^{N_{1}})}, m_{i}^{N_{1}} = max_{1 \leq j \leq N_{1}} s_{i j} .

但上面的 $softmax$ 方法是不对的，我们还有 $K_{2}, V_{2}$ 没有纳入计算。

对于 $q_{i}, K_{1}, V_{1}$ ，维护三个运行统计量：

\begin{array}{r} m_{i}^{N_{1}} = max_{1 \leq j \leq N_{1}} s_{i j} \end{array} ℓ_{i}^{N_{1}} = \sum_{r = 1}^{N_{1}} \exp (s_{i r} - m_{i}^{N_{1}}) A_{i}^{N_{1}} = \sum_{r = 1}^{N_{1}} \exp (s_{i r} - m_{i}^{N_{1}}) v_{r}

现在处理 $K_{2}, V_{2}$ ：

m_{i}^{N_{2}} = max (m_{i}^{N_{1}}, max_{N_{1} + 1 \leq j \leq N_{2}} s_{i j}) l_{i}^{N_{2}} = (m_{i}^{N_{1}} - m_{i}^{N_{2}}) l_{i}^{N_{1}} + \sum_{r = N_{1} + 1}^{N_{2}} \exp (s_{i r} - m_{i}^{N_{2}}) A_{i}^{N_{2}} = (m_{i}^{N_{1}} - m_{i}^{N_{2}}) A_{i}^{N_{1}} + \sum_{r = N_{1} + 1}^{N_{2}} \exp (s_{i r} - m_{i}^{N_{2}}) v_{r}

处理完所有 $K, V$ 之后，最终输出为：

o_{i} = \frac{A_{i}^{N_{2}}}{l_{i}^{N_{2}}}

现在推广到 query block $Q_{i}$ 。

对 $Q K V$ 进行分块：

Q = [\begin{matrix} Q_{1} \\ Q_{2} \\ ⋮ \\ Q_{T_{Q}} \end{matrix}], K = [\begin{matrix} K_{1} \\ K_{2} \\ ⋮ \\ K_{T_{K}} \end{matrix}], V = [\begin{matrix} V_{1} \\ V_{2} \\ ⋮ \\ V_{T_{K}} \end{matrix}] Q_{i} \in R^{B_{q} \times d}, K_{j}, V_{j} \in R^{B_{k} \times d}

局部 attention score 为：

S_{i j} = \frac{Q_{i} K_{j}^{⊤}}{\sqrt{d}} \in R^{B_{q} \times B_{k}} .

此时 (Q_i) 中有 (B_q) 个 query token，处理第 $j$ 个 KV block 时，需要对每一行分别维护：

m_{i}^{(j)} \in R^{B_{q}}, ℓ_{i}^{(j)} \in R^{B_{q}}, A_{i}^{(j)} \in R^{B_{q} \times d} .

初始化：

m_{i}^{(0)} = - \infty, ℓ_{i}^{(0)} = 0, A_{i}^{(0)} = 0.

处理第 (j) 个 KV block 时，我们已经有 $m_{i}^{(j - 1)}, ℓ_{i}^{(j - 1)}, A_{i}^{(j - 1)}$ 。

先计算：

S_{i j} = \frac{Q_{i} K_{j}^{T}}{\sqrt{d}} \in R^{B_{q} \times B_{k}} ‘

对每一行求最大值：

m_{i}^{{block}_{j}} = rowmax (S_{i j}) \in R^{B_{q} \times 1}

更新全局最大值：

m_{i}^{(j)} = rowmax (m_{i}^{(j - 1)}, m_{i}^{{block}_{j}})

分母更新：

ℓ_{i}^{(j)} = \exp (m_{i}^{(j - 1)} - m_{i}^{(j)}) ⊙ ℓ_{i}^{(j - 1)} + rowsum (\exp (S_{i j} - m_{i}^{(j)})) \in R^{B_{q} \times 1}

加权分子更新：

A_{i}^{(j)} = \exp (m_{i}^{(j - 1)} - m_{i}^{(j)}) ⊙ A_{i}^{(j - 1)} + \exp (S_{i j} - m_{i}^{(j)}) V_{j} \in R^{B_{q} \times d}

处理完所有 KV block 后，得到最后的 attention 输出：

O_{i} = \frac{A_{i}^{T_{k}}}{ℓ_{i}^{T_{k}}} \in R^{B_{q} \times d}

最后把所有 query block 的输出拼接起来：

O = Concat (O_{1}, O_{2}, \dots, O_{T_{q}}) \in R^{n \times d}

FlashAttention ​

5.2.1 Flash Attention ​

1. 普通 Attention 的主要瓶颈 ​

2. GPU 内存层级：为什么 IO 很重要？ ​

3. FlashAttention ​

（1）分块矩阵乘法 ​