大模型数学速成（17）：Flash Attention——不物化注意力矩阵

第 05 篇的公式是 (\mathrm{softmax}(QK^\top/\sqrt{d}),V)。第 11 篇又看到长序列下 (T\times T) 的代价。Flash Attention 不改这个数学定义，却尽量不把整张分数矩阵 $S$ 写进高带宽显存——用分块在更快的片上存储里完成 Softmax 归约。

本篇只建立 IO 与分块直觉，不写 CUDA kernel。

这是「大模型数学速成」续篇第 17 篇。建议先读第 05、11 篇。下一篇讲 MoE。

一、标准 Attention 的痛点

单头、序列长 $T$、头维 $d$：

\[S = \frac{QK^\top}{\sqrt{d}} \in \mathbb{R}^{T\times T}, \quad P = \mathrm{softmax}(S),\quad O = PV\]

对象	规模	问题
$S$ 或 $P$	$O(T^2)$	显存随长度平方涨
读写 $S$	多次经过 HBM（大而慢的显存）	带宽墙，往往比算力先爆

长上下文时，很多实现的瓶颈不是「乘加次数」，而是 反复搬运那张大矩阵。

二、数学不变，变的是计算路径

Flash Attention 的目标输出仍是：

\[O_i = \sum_j P_{ij} V_j,\quad P_{ij} = \frac{e^{S_{ij}}}{\sum_k e^{S_{ik}}}\]

与 naive 实现 数值上追求等价（实现里会有稳定化细节）。
它不是一种新的注意力「定义」，而是 同一公式的 IO 友好算法。

一句话：

算的还是 Softmax-Attention；尽量别物化完整的 $S$。

三、分块 + 在线 Softmax 直觉

把 $Q,K,V$ 切成小块（tile），每次只把能塞进 SRAM（小而快）的块搬上来：

取一块 $Q$ 与一块 $K$，算局部分数
用局部结果更新该 $Q$ 行 Softmax 所需的 统计量
再结合对应 $V$ 块累加输出
换下一块 $K,V$，继续更新

关键是 Softmax 可以 在线（online） 维护。对一行 logits，记：

\[m = \max_j S_{j},\qquad \ell = \sum_j e^{S_j - m}\]

若先只看到前半块，再看到后半块，可以用新的最大值 $m'$ 把旧的 $\ell$ 与部分输出 按指数差 rescale 后合并——不必先拿到整行再 Softmax。

迷你数值直觉（一行 4 个数，分两块）

设一行分数（已含缩放）为 $[1,\ 2,\ |\ 0,\ 3]$，竖线表示分块。

块 1：$[1,2]$
$m_1=2$，$\ell_1=e^{1-2}+e^{0}=e^{-1}+1\approx 1.368$

块 2：$[0,3]$ 与块 1 合并：
$m_2=\max(2,3)=3$
$\ell_2 = \ell_1\cdot e^{m_1-m_2} + e^{0-3}+e^{3-3}$
$\approx 1.368\cdot e^{-1} + e^{-3}+1 \approx 1.553$

与直接对 $[1,2,0,3]$ 做 Softmax 的归一化项一致（同一 $m=3$）。
输出混合 $V$ 时用同样的 rescale 系数对齐——这就是「不物化整行也能正确归一化」的代数来源。

四、数据流示意

HBM（大慢）:  Q, K, V, O
                 │ 小块搬运
                 ▼
SRAM（小快）:  算局部 S 块 → 更新 m,ℓ → 累加 O 块
                 │
                 ▼
              写回 O（通常不写完整 S）

存储	特点	Flash 用法
HBM	容量大、带宽相对慢	驻留完整 Q/K/V/O
SRAM	容量小、带宽高	驻留当前 tile 与统计量

因果 mask、GQA 等，是在分块循环里 跳过或屏蔽非法 $(i,j)$，与「是否 Flash」正交。

五、和本系列其它篇的关系

篇	关系
05 Attention	定义 $S,P,O$；Flash 实现同一式
11 KV Cache	Decode 步长 1 时 $S$ 是「1×历史」；Flash 更多惠及 Prefill / 训练长序列
13 因果 mask	分块时仍要下三角可见性
15 量化	另一条省显存轴；可与高效 Attention 叠用

六、读者收益与边界

读论文/引擎文档时：看到 Flash / memory-efficient attention，优先理解成 IO 与分块算法，不是换了一种注意力分数公式。

本篇不讲：具体 CUDA 占用、版本差异、是否 backward 也 fused——那些属实现手册。

七、小结

概念	要点
痛点	$T\times T$ 的 $S$ 吃显存与 HBM 带宽
不变	仍是 $\mathrm{softmax}(QK^\top)V$
手段	tile + 在线 Softmax 统计量 $(m,\ell)$
收益	少物化 $S$，长序列更可行
边界	算法优化，不是新模型结构