Flash Attention 为什么更快：IO 感知的本质

问题的根源是流量，而非算术

直接按照定义来写，对于长度为 S 的序列、注意力头维度为 d，attention 的计算过程是：计算 S×S 个 score 写入内存，再读回来，对每一行做 softmax，写回去，再读出来，乘以 values。一次矩阵的完整往返就产生了四次——而矩阵大小随序列长度的平方增长。

算术复杂度是 O(S²d)，这部分无法避免。而内存流量是 O(S²)，完完全全是写法带来的额外开销。由于 attention 在适中的头维度下本来就偏低算术强度，因此实际成本就取决于这个流量。

这种写法之所以出现，与其说是粗心，不如说是深思熟虑。把 attention 表达为一系列标准操作的序列——矩阵乘法、softmax、再一个矩阵乘法——每一步都是一个库函数调用，各自都经过了良好的优化，而且每个函数都必然要从内存读取输入、再把输出写回去，因为这是库函数调用的契约。效率损失存在于操作的接缝处，而非任何一个操作的内部，这正是融合实现才能解决的问题。

标准实现，每个注意力头，S = 8192，fp16（2 字节）：

  score 矩阵 S*S = 8192^2 = 67.1e6 个元素 = 134 MB

  写一次，供 softmax 读，写回去，
  再供 value 乘积读出来            -> 约 4 x 134 MB
                                      = 每个注意力头约 537 MB

乘以注意力头数、再乘以层数，
一次前向传播的内存流量就达到数百 GB。

Flash attention，相同的 S：

  S*S 矩阵根本不会被写入设备内存。
  流量是 Q、K、V 和输出的 O(S*d)，
  加上每个 query tile 对 K、V 的一次重读。

峰值内存节省和速度提升同样重要。标准形式必须分配 S×S 的矩阵，因此内存随上下文长度二次增长——这就是为什么在此类 kernel 出现之前，长上下文模型根本不现实，与计算耗时无关。

分块，以及 softmax 的障碍

修复方法是内存受限 kernel 的标准解法：将计算分块，使每个块适合放入快速片上内存，处理完一个块的所有工作后再移动到下一个。加载一批 query，然后让 key 和 value 的块流经它，随之累积输出。一个 tile 的 scores 驻留在片上内存中，tile 处理完毕后就丢弃。

有一件事阻挡了这条路径，而这才是有趣的部分。Softmax 不是局部的。对一行做归一化需要整行指数的和，而要先减去行的最大值——这是数值稳定性的标准技巧——也需要整行的最大值。看起来似乎需要在归一化任意一个值之前先拥有所有的 score，而这正是分块拒绝给予你的东西。

解决方案是维护一个运行最大值和一个运行和，并在最大值变化时对累积结果进行重新缩放。处理一个新的块时：

1. 计算该块的 scores 及其局部最大值
2. 令 m_new = max(m_old, m_block)
3. 通过 exp(m_old − m_new) 对运行和与累积输出进行重新缩放，纠正因它们曾以一个现已过时的最大值为基准进行归一化而产生的问题
4. 加上该块的贡献，以 m_new 为基准进行归一化

每次修正对每一行只是一个标量乘法，相比矩阵运算可以忽略不计。最终运行和就是真正的分母，累积输出就是准确的 attention 输出。这和流式计算均值而不存储流的想法如出一辙，只是应用在了 softmax 上——而这正是整个方法得以成立的原因。

这项技术由 Dao 及其同事于 2022 年作为 FlashAttention 发表，后续版本改进了跨芯片的工作划分而非改变这个核心思想。这个通用原理的名称，以及更值得记住的东西，叫做 IO 感知：数的是到内存的往返次数，而不只是操作次数。

需要明确说明，因为这是阅读相关资料时最常令人困惑的地方：结果是精确的。这不是稀疏或低秩 attention 变体那种对 attention 的近似，也不会以质量换速度。对求和进行重排序并对部分结果重新缩放会改变算术顺序，因此输出可能在最后几位有所不同——就像任何浮点数重结合都会出现的情况——但数学上是同一个函数。这很罕见，也是这项技术被基本上普遍采用、而非成为众多选项之一的原因。

反向传播携带了同样的思想，如果你曾经微调而非仅仅 serving，有一个额外的细节值得了解。梯度需要那些被故意不存储的 attention 权重，所以在反向传播期间要从保存的统计量重新计算。当一个 kernel 受内存限制时，用少量冗余算术换取大量内存流量降低是正确的交易，这个模式远超 attention 范围地反复出现。

它给下游带来了什么

长上下文变成了一个内存容量问题，而非内存缩放问题。二次项被移除后，约束条件移到了 KV cache，它随长度线性增长。这是一个实质上不同的问题，这就是为什么此后关于长上下文的有趣工作都集中在 cache 大小上。

Prefill 变得便宜很多。Prefill 是原本要计算完整 attention 矩阵的地方，所以节省就落在这里，在长 prompt 的首 token 时间上有所体现。

Decode 以不同方式受益。生成一个 token 需要在一个很长的历史上做 attention，但只有一个 query 行，因此没有大的 score 矩阵需要避免。这些 kernel 的 decode 特化变体针对的是 key/value 扫描的并行化。

它成了一个可移植性要求。因为这个 kernel 是针对每个架构手写的，"这个后端是否有 IO 感知的 attention 实现" 是任何非现有计算平台最尖锐的问题之一。