FlashAttention

FlashAttention 是 IO 感知的精确注意力实现:把 $Q K^T$ 分块 (tiling) 放进 SRAM 计算,配合 online softmax 维护运行最大值 $m_i$ 与归一化因子 $l_i$(合并新块时递归缩放:$m_{new} = max(m, m_{local})$),前向不把 $O(n^2)$ 注意力矩阵写回 HBM,反向用重算而非存储;动机是 HBM 与 SRAM 的带宽差距约 10-20 倍(A100: 40GB HBM vs 约 108KB SRAM)。

大模型训练与推理的事实标准实现(PyTorch SDPA、NVIDIA FA2/FA3、DeepSeek/Meta 训练栈),序列越长收益越大,长上下文训练几乎必用。

标准实现把注意力矩阵反复写读 HBM,IO 成为主要瓶颈;分块 + 重算把 HBM 访问从 $O(n^2 d)$ 降到 $O(\frac{n^2 d^2}{M})$($M$ 为 SRAM 大小),速度提升 2-4 倍、显存从 $O(n^2)$ 降到 $O(n)$,且结果与标准 softmax 逐位一致(不是近似)。

检验攻克程度：针对「FlashAttention」专属刷题排雷

做单选排雷题、推导选项机制，答错自动收录进专属错题本。