之前的几篇把 Self-Attention 的公式拆得比较清楚了，从 QKV 的计算含义到 Multi-Head 的分工。另一条线是 GPU 性能——Arithmetic Intensity、memory bound、HBM 带宽这些概念也已经聊过。关于 HBM 和片上存储的物理区别，可以参考GPU 存储层级那篇。
这两条线最近在我脑子里碰到了一起。
Attention 公式里有一步 \(QK^T\)，产生一个 \(S \times S\) 的 Score Matrix。之前拆公式的时候，我关心的是每一项的计算含义；但从 GPU 实现的角度看，这个矩阵本身就是一个问题。
FlashAttention 的核心主张是：这个矩阵根本不需要完整存下来。
第一次看到这个说法的时候我的反应是——Softmax 需要全局归一化，不把整行 Score 算完，怎么做 Softmax？不做完 Softmax，怎么和 Value 加权求和？
这篇想把这个问题从头想通。
那个 S×S 的中间矩阵#
先把公式摆出来（推导过程在之前那篇）：
$$ O = \operatorname{Softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$Q、K 的 shape 都是 \(S \times d_k\)，\(QK^T\) 的结果是 \(S \times S\)。
取一个具体尺寸感受一下：\(S = 8192\)，\(d_k = 128\)。Score Matrix 有 \(8192 \times 8192 \approx 6700\) 万个元素。FP32 下需要约 256 MB，BF16 也要 128 MB——而这只是一个 attention head。如果有 32 个 head，FP32 下单层 Attention 的 Score Matrix 就要 8 GB。
然而这个矩阵是个中间产物。最终需要的只是 \(O\)——shape 为 \(S \times d_k\) 的 Attention Output，对应这个例子只有 \(8192 \times 128 \approx 100\) 万个元素，4 MB。
传统实现的问题在于：完整的 \(S \times S\) Score Matrix 被写进 HBM，后续再读回来做 Softmax，Softmax 的结果又写回 HBM，然后再读回来和 V 做矩阵乘。中间矩阵在 HBM 里反复来回搬运。之前讨论过，HBM 带宽是 GPU 性能的关键瓶颈——这些中间矩阵的读写代价不小。
既然最终需要的只是 \(O\)，一个自然的想法就是：能不能不把完整的 Score Matrix 存下来，算完一小块就丢掉？
这里先澄清一点。FlashAttention 不是不计算 Score。小块 Score 仍然会在计算过程中短暂存在于 GPU 的片上存储里。它省掉的是把完整 \(S \times S\) 矩阵物化到 HBM 这个动作。
算一行丢一行#
用一个 \(16 \times 16\) 的 Score Matrix 当教学示例。
最简单的做法：一次处理一个 Query。取 \(q_1\)，和全部 16 个 Key 做 dot product，得到一行 Score（\(1 \times 16\)）。对这一行做 Softmax，和对应的 Value 加权求和，得到 \(q_1\) 的输出。写出结果，丢掉 Score。然后处理 \(q_2\)，依此类推。
这个算法完全正确，而且已经避免了物化完整的 \(16 \times 16\) Score Matrix——任何时刻只需要保存一行。
但下一个问题是：为什么不一次算多行？比如 4 个 Query 一起处理（\(4 \times 16\) 的分块）？
原因很直接。4 个 Query 和同一组 K 做 dot product 的时候，K 的数据可以被复用——4 个 Query 一起算，K 读一次就够了，不用反复搬运。而且 \(4 \times 16\) 的矩阵乘是一个形状合理的 GEMM，比逐行的 GEMV 更能发挥 GPU 矩阵计算单元的吞吐。
分块方式不仅要考虑"能不能算"，还要考虑硬件是否擅长这种计算模式。
4×16 还是 4×4#
到这里，沿 Query 方向分块的动机清楚了。但 FlashAttention 还在 Key 方向继续分块。为什么？
先别急着回答。对于刚才那个 \(16 \times 16\) 的小矩阵，\(4 \times 16\) 的分块——每次处理 4 个 Query，看完全部 16 个 Key——完全可行，甚至很可能比 \(4 \times 4\) 更高效。因为每个 Query 的 Softmax 可以在一轮内完成，不需要任何跨块的修正。
小矩阵只是用来演示分块思想的。不能因为教学示例画了 \(4 \times 4\) 就认为 \(4 \times 4\) 天然更好。
但现在把 Sequence Length 放大。设 \(S = 8192\)，\(d_k = 128\)，Query 方向取 Block Size \(B_q = 64\)。如果 Key 方向不分块，一个 Block 需要计算 \(64 \times 8192\) 的 Score，同时还要在片上持有对应的 K、V 数据和累加状态。
GPU 的片上存储——Shared Memory 和 Registers——容量有限。A100 每个 SM 的 Shared Memory 可配置最大 228 KB，Register File 是 256 KB。当需要同时保留大量 K/V tile、Score 分块以及各种累加状态时，片上资源很快就不够用了。
这里有个容易犯的错误：不要简单地把所有 tensor 大小加起来，然后声称它们必须全部同时存在于 Shared Memory。真实的 CUDA kernel 会把数据分配到 Registers、Shared Memory 等不同层级，也可能不物化某些临时矩阵。但核心约束是真实存在的——片上资源有限，同时能持有的数据规模是受控的。
Block 太大，Register Pressure 上升，Occupancy 下降，某些中间结果可能溢出到慢速存储。Block 太小，循环次数增加，每次 GEMM 规模太小，计算效率和数据复用都打折扣。Block Size 不是越小越好也不是越大越好，而是硬件资源与计算吞吐之间的权衡。
总结一下这两节的思路：
- Query 方向分块：让多个 Query 共享 K/V 数据，提高矩阵计算效率。
- Key 方向分块：限制每次处理的数据规模，让片上资源需求保持可控。
- Block Size 的选择：是一个具体硬件上的工程权衡问题。
分块之后 Softmax 不对了#
Query 方向的分块没有数学问题——每组 Query 彼此独立，分开算完结果一样。
但 Key 方向的分块引入了一个真正的困难。
Softmax 的定义：
$$ p_i = \frac{e^{s_i - m}}{\sum_j e^{s_j - m}}, \quad m = \max_j s_j $$分母里的求和覆盖了整行所有 Score。如果 Key 被拆成了多个 Block，每个 Block 只看到部分 Score，直接对各个 Block 分别做 Softmax 再拼接，结果是错的。
用一个数值例子。设一行 Score 是：
$$ S = [0, \; 0, \; \ln 2, \; \ln 2] $$分成两个 Block：\(B_1 = [0, 0]\)，\(B_2 = [\ln 2, \ln 2]\)。
如果各自独立做 Softmax：
- \(B_1\)：\(m = 0\)，\(\sum = 2\)，概率 \([1/2, \; 1/2]\)
- \(B_2\)：\(m = \ln 2\)，\(\sum = 2\)，概率 \([1/2, \; 1/2]\)
拼起来得到 \([1/2, \; 1/2, \; 1/2, \; 1/2]\)。概率之和是 2，明显不对。
正确的全局 Softmax：\(m = \ln 2\)，分母 \(= e^{0 - \ln 2} + e^{0 - \ln 2} + e^{\ln 2 - \ln 2} + e^{\ln 2 - \ln 2} = 1/2 + 1/2 + 1 + 1 = 3\)，正确概率是 \([1/6, \; 1/6, \; 1/3, \; 1/3]\)。
问题很明确：每个 Block 只看到局部 Score，无法得到全局的 max 和 sum，因此不能独立完成归一化。
只维护 m 和 l#
这是 Online Softmax 解决的问题。
核心想法：不需要存下所有 Score，只需要在线维护两个标量——
- \(m\)：到目前为止见过的最大 Score
- \(l\)：在当前 \(m\) 基准下，所有已处理 Score 的指数和
当新的 Key Block 到来时，更新规则是：
$$ m_{\text{new}} = \max(m_{\text{old}}, \; \max(S_{\text{block}})) $$$$ l_{\text{new}} = e^{m_{\text{old}} - m_{\text{new}}} \cdot l_{\text{old}} + \sum_{s \in S_{\text{block}}} e^{s - m_{\text{new}}} $$第二个公式是关键。它的直觉是：如果新 Block 出现了更大的 Score，之前累计的分母只需要乘以一个缩放系数 \(e^{m_{\text{old}} - m_{\text{new}}}\)，就能转换到新的归一化基准。 不需要回头重算之前的 Score——旧的分母乘一个数就行了。
用刚才的例子验证。
处理 \(B_1 = [0, 0]\)：
$$ m_1 = 0, \quad l_1 = e^{0-0} + e^{0-0} = 2 $$处理 \(B_2 = [\ln 2, \ln 2]\)：
$$ m_2 = \max(0, \; \ln 2) = \ln 2 $$$$ l_2 = e^{0 - \ln 2} \cdot 2 + e^{\ln 2 - \ln 2} + e^{\ln 2 - \ln 2} = \tfrac{1}{2} \cdot 2 + 1 + 1 = 3 $$
最终 \(m = \ln 2\)，\(l = 3\)，和全局 Softmax 的分母完全一致。
有一点要说清楚。\(m\) 和 \(l\) 足以维护 Softmax 的归一化分母，但如果要恢复每个位置的概率值，还需要对应位置的原始 Score。Online Softmax 的意义不是"用两个标量替代所有概率"，而是让归一化统计量能跨 Block 在线维护，使得分块处理成为可能。
（在实数数学上，Online Softmax 和普通 Softmax 严格等价。浮点实现下由于运算顺序不同，结果不保证逐位一致，但精度差异通常极小。）
既然 Value 就在手边#
Online Softmax 解决了归一化分母的跨块维护。但 Attention 的最终输出不是概率，而是 \(O = PV\)——概率加权的 Value 之和。
把 \(O\) 展开成求和形式：
$$ O = \frac{\sum_i e^{s_i - m} \cdot V_i}{l} $$分子部分同样可以在线累计。定义一个累加器 \(acc\)，每处理一个 Key Block：
$$ acc_{\text{new}} = e^{m_{\text{old}} - m_{\text{new}}} \cdot acc_{\text{old}} + P_{\text{block}} \cdot V_{\text{block}} $$其中 \(P_{\text{block}}\) 是当前 Block 的未归一化指数权重，\(P_{\text{block},j} = e^{s_j - m_{\text{new}}}\)——注意这不是最终的 Softmax 概率，只是分子部分的局部贡献。
全部 Key Block 处理完后：
$$ O = acc \;/\; l $$两件事值得注意。第一，当前 Block 的 Score 和 Value 已经在片上了，顺手做一次矩阵乘并累加到 \(acc\)，不需要额外的 HBM 读写。第二，处理完当前 Block 后，这个 Block 的 Score 就可以丢弃了——不需要为了后续步骤保留任何历史 Score。
这一步没有引入新的数学思想。它只是把 Online Softmax “在线维护归一化统计量"的思路扩展到"在线累积输出分子”。同一个缩放系数 \(e^{m_{\text{old}} - m_{\text{new}}}\)，同时应用于 \(l\) 和 \(acc\)。
想到了 Halide#
写到这里我忽然想到以前接触过的 Halide。
Halide 有一个核心设计理念：Algorithm 和 Schedule 分离。 Algorithm 定义"算什么"——像素级的计算逻辑。Schedule 定义"怎么算"——循环顺序、分块大小、在哪里存储中间结果、何时计算何时消费。同一个 Algorithm，换一套 Schedule，性能可能差十倍。
FlashAttention 和这个思路有很强的对应关系。Tiling 对应把循环维度分块，Compute Placement 对应在消费点就地计算（类似 Halide 的 compute_at，Score 在需要的位置计算而不是 materialize 到全局存储），Producer-Consumer Fusion 对应 Score 生产出来后立刻参与 Softmax 和累计（不写回 HBM 再读），Storage Locality 对应中间结果尽量停留在片上。
但有一个关键区别。
Halide 的 Schedule 优化不改变 Algorithm 的数学语义。图像处理里的 elementwise 或 stencil 操作，改循环顺序和 tiling 策略就够了，不需要碰算法本身。
FlashAttention 不只是 Schedule 优化。 Softmax 存在跨 Key Block 的全局归一化依赖——必须拿到整行的 max 和 sum 才能算出概率。这个依赖不是改循环顺序就能消除的。Online Softmax 是一次数学上的等价重构：把"先算完全部 Score，再一次性归一化"改写成"逐块处理，在线维护 \(m\) 和 \(l\)"。重构之后，全局依赖被替换成了一个可增量更新的状态，后续的 Tiling、Fusion 才有可能执行。
也就是说：
- Online Softmax 是算法层面的重构——改变了计算的数学形式（结果等价），让 Softmax 能分块执行。
- Tiling、Fusion、数据局部性优化是调度层面的优化——不改变数学，但改变执行方式和数据搬运模式。
FlashAttention 同时做了这两件事。先用 Online Softmax 打破全局归一化的依赖，才让后续的分块和融合成为可能。这一点让我觉得它比单纯的 kernel fusion 要有意思得多。
写在最后#
回过头看，FlashAttention 在三个层面给了我启发。
数学层面：标准 Attention 的公式没有改变，计算复杂度仍然是 \(O(S^2 d)\)。改变的是计算的等价形式——Online Softmax 让 Softmax 从全局操作变成了可分块的在线操作。
GPU 工程层面：完整的 \(S \times S\) Score / Probability Matrix 不再需要写入 HBM。中间结果在片上计算、片上消费、片上丢弃，显存流量和中间存储开销大幅降低。
一般化的工程思想：一个计算图 \(A \to B \to C\)，不意味着必须完整物化 \(A\)、\(B\) 后才能开始算 \(C\)。通过 Tiling、Fusion、Streaming、Data Locality 等手段，可以重新组织执行过程——让中间结果不落地到慢速存储。
很多高性能计算优化并不是发明更少的数学运算，而是避免让数据在不同存储层次之间无谓搬运。
参考资料#
- Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS, 2022.
- Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR, 2024.
- Milakov, Gimelshein. Online normalizer calculation for softmax. arXiv, 2018.
- Halide: a language for fast, portable computation on images and tensors.
- 从动态读取到 QKV：Self-Attention 的公式是怎么长出来的
- 从 CS336 到 ML Infra：理解 Arithmetic Intensity