项目 15：IO 感知的艺术：FlashAttention 的硬件级优化

在 GPU 体系结构中，计算速度（FLOPS）通常远高于显存带宽（Memory Bandwidth）。对于标准的注意力计算，GPU 需要频繁地在 HBM（高带宽显存）和 SRAM（片上高速缓存）之间搬运巨大的 $L \times L$ 矩阵。

FlashAttention 不再一次性计算全量的注意力矩阵。

分块加载：将 $Q$、$K$、$V$ 划分为小的 Block（例如 64x64），使其能够完全放入 GPU 只有几十 MB 的 SRAM 中。
流水线处理：在 SRAM 内部完成局部点积、局部 Softmax 和局部加权和。
重新计算：在反向传播时，FlashAttention 不存储中间的注意力矩阵（这节省了巨大的 $O(L^2)$ 空间），而是根据前向传播的信息实时重新计算。

分块计算最大的数学难题在于 Softmax。传统的 Softmax 公式 $P_{i} = \frac{e^{x_{i}}}{\sum e^{x_{j}}}$ 要求必须知道所有 $x_{j}$ 之后才能计算分母。这与分块（只看局部数据）是矛盾的。

FlashAttention 引入了 Online Softmax 算法：

我们维持两个局部统计量：当前块的最大值 $m$ 和局部指数和 $d$。
当新块到来时，通过缩放因子实时更新 $m$ 和 $d$。 $m_{new} = \max(m_{old}, m_{block})$ $d_{new} = d_{old} \times e^{m_{old}-m_{new}} + d_{block} \times e^{m_{block}-m_{new}}$ 这种增量更新机制使得我们可以在只看到局部数据的情况下，最终计算出数学上完全一致的全局 Softmax 结果。

引入 FlashAttention 后，你会观察到惊人的变化：

虽然在 Capstone 项目中你通常直接调用 torch.nn.functional.scaled_dot_product_attention（其后端就是 FlashAttention），但你应该深入研究其 triton 实现：

FlashAttention 是过去三年 LLM 工程领域最重要的突破之一。它深刻揭示了一个真理：在现代硬件架构下，减少数据搬运比减少计算量更有意义。理解了项目 15，你就掌握了高性能算子开发的灵魂——即如何与硬件的存储层级进行共舞。

（采用 CC BY-NC-SA 4.0 许可协议进行授权）

本文标题:项目 15：IO 感知的艺术：FlashAttention 的硬件级优化

本文最后一次更新为天前，文章中的某些内容可能已过时！

跨越显存墙：让算力不再等待数据