FlashAttention 为什么能 one-pass:在线 softmax(m/l)与 Tiling 的核心思想
从标准注意力的显存 IO 账本出发,解释 FlashAttention 的核心:在线 softmax 维护 (m,l) 并流式累积输出,再配合 tiling 把数据驻留在片上存储,从而避免显式存储 $QK^\top$ 与 softmax 概率矩阵。本文给出可运行的 Numpy 分块注意力实现与数值等价验证,并用可复制的字节算账方法说明它为什么会快。
从标准注意力的显存 IO 账本出发,解释 FlashAttention 的核心:在线 softmax 维护 (m,l) 并流式累积输出,再配合 tiling 把数据驻留在片上存储,从而避免显式存储 $QK^\top$ 与 softmax 概率矩阵。本文给出可运行的 Numpy 分块注意力实现与数值等价验证,并用可复制的字节算账方法说明它为什么会快。
从标准两遍 softmax 的访存模式出发,推导在线 softmax(m,l)更新与正确性;进一步解释在 attention/cross-entropy 中如何通过融合避免落地概率矩阵,并用可运行代码验证等价与估算带宽收益。