Attention 复杂度与 FlashAttention#

1. Self-Attention 的复杂度瓶颈#

对于序列长度 $n$、隐藏维度 $d$：

1.1 为什么 Attention 是 $O(n^2 d)$？#

假设序列有 4 个 token，$d_k = 3$，$Q$ 和 $K$ 各是一个 $4 \times 3$ 的矩阵。$QK^T$ 就是让每个 token 的 Q 向量和每个 token 的 K 向量两两做点积：

         token1  token2  token3  token4
token1 [  Q1·K1   Q1·K2   Q1·K3   Q1·K4 ]
token2 [  Q2·K1   Q2·K2   Q2·K3   Q2·K4 ]
token3 [  Q3·K1   Q3·K2   Q3·K3   Q3·K4 ]
token4 [  Q4·K1   Q4·K2   Q4·K3   Q4·K4 ]

结果是 $4 \times 4 = 16$ 个格子，每个格子算一次 $d_k$ 维点积（$d_k$ 次乘法）。总运算量 $= n^2 \times d_k = O(n^2 d)$。

token 数翻倍（$4 \to 8$）：格子数从 16 变成 64，运算量变 4 倍。这就是平方增长的含义。$n = 128000$ 时，这张矩阵有 128000 × 128000 ≈ 160 亿个格子。

1.2 为什么 FFN 是 $O(n d^2)$？#

FFN 对每个 token 独立做两层线性变换（$d \to 4d \to d$），每个 token 约 $O(d^2)$ 次运算，$n$ 个 token 串联，总量 $= O(nd^2)$。token 数翻倍，运算量只翻倍——线性增长，和 $n$ 无关的平方项。

谁更大？取决于 $n$ 和 $d$ 的关系。

• $n = 512, d = 4096$：FFN 更大（$nd^2 \gg n^2d$）。这是早期 BERT/GPT 的情况。
• $n = 32768, d = 4096$：Attention 更大（$n^2d \gg nd^2$）。这是现在长上下文模型的情况。

所以长序列处理是一个重要研究方向——Attention 的 $O(n^2)$ 在长文本上成为严重瓶颈。

2. Flash Attention（2022 → 现在几乎所有训练都在用）#

Flash Attention 不改变注意力的数学结果，只改变计算顺序——它是一个 IO-aware 算法。

核心问题：标准注意力实现需要把完整的 $n \times n$ 注意力矩阵写入 GPU 的 HBM（高带宽内存），然后读出来做 softmax，再写回去。GPU 的计算速度很快，但内存读写（IO）是瓶颈。

Flash Attention 的做法：把 $Q, K, V$ 分块（tiling），每次只处理一小块，在 SRAM（片上缓存，速度是 HBM 的 10-20 倍）中完成注意力计算，避免把完整的 $n \times n$ 矩阵写入 HBM。

2.1 标准 Attention 的 IO 问题#

标准实现要走三趟 HBM：

第一步：HBM → 读 Q, K → 算 S = QKᵀ → 写 S 回 HBM
第二步：HBM → 读 S   → 算 softmax → 写 A 回 HBM
第三步：HBM → 读 A, V → 算 AV     → 写结果回 HBM

$n = 8192$ 时 $S$ 矩阵有 $8192^2 \approx 6700$ 万个元素。来回搬运的时间比计算本身还长。

2.2 关键：online softmax 修正#

分块的难点是 softmax 需要看到全局才能算分母 $\sum e^{x_j}$。Flash Attention 的解法是延迟除法——每块只算分子，记住局部最大值 $m$ 和局部 exp 之和 $\ell$；新块来了之后用修正系数 $e^{m_\text{old} - m_\text{new}}$ 统一基准，最后一次性除以总 $\ell$。

以分数 $[3, 1, 4, 2]$、分两块处理为例：

第一块 [3, 1]

$m_1 = 3,\quad \ell_1 = e^0 + e^{-2} = 1.135$

$\tilde{O}_1 = e^0 V_0 + e^{-2} V_1$（先不除以 $\ell_1$）

SRAM 记住 $m_1, \ell_1$，写 $\tilde{O}_1$ 回 HBM

第二块 [4, 2]，更新基准

$m_2 = 4,\quad \ell_2 = e^0 + e^{-2} = 1.135$

修正旧结果：$\tilde{O}_1 \times e^{m_1 - m_2} = \tilde{O}_1 \times e^{-1}$

总 $\ell = e^{-1}\ell_1 + \ell_2 = 1.553$，最后除以 $\ell$ 得正确 softmax

结果和一次性算完整行完全一样——数学等价，不是近似。

**效果：**速度提升 2-4 倍，显存从 $O(n^2)$ 降到 $O(n)$（不需要存完整注意力矩阵）。

💡 为什么 Flash Attention 重要？

它让 Transformer 可以处理更长的序列、用更大的 batch size，而不改变模型本身。几乎是"免费的午餐"。现在 PyTorch 2.0+ 内置了 scaled_dot_product_attention，底层就是 Flash Attention。