Flash Attention 笔记
内容提要
Flash Attention 通过分块计算 Q、K、V，避免存储 seq_len×seq_len 的 softmax 中间矩阵，减少 HBM 访问，使注意力内存随序列长度近似线性增长。虽然 FLOP 更高，但注意力受内存带宽限制，因此速度更快，并支持长序列训练。
延伸解读
内存带宽瓶颈与计算量权衡
文章指出，Flash Attention 的 FLOP 数高于标准注意力，但实际速度更快，原因在于注意力计算受内存带宽限制，而非计算量。它通过减少 HBM 访问次数来提升效率。这提醒我们，优化模型时需先判断瓶颈类型：若为内存受限，减少数据搬运比降低计算量更有效；若为计算受限，则需另寻他法。
分块计算如何降低内存占用
标准注意力需在 HBM 中存储 seq_len×seq_len 的 softmax 中间矩阵，内存占用随序列长度平方增长。Flash Attention 将 Q、K、V 分块，在 SRAM 中计算输出，避免存储完整中间矩阵，使注意力内存近似线性增长。这为长序列训练提供了可行路径，但需注意分块大小与 SRAM 容量的平衡。
与稀疏及低秩近似方法的对比
文章提到 Reformer 和 Performer 分别通过稀疏近似和低秩近似来降低计算与内存需求。Flash Attention 的不同之处在于不近似注意力矩阵，而是通过分块精确计算，在保持模型质量的同时提升效率。对于需要精确注意力的任务，Flash Attention 可能更合适；而稀疏或低秩方法在特定场景下仍有其优势。
实际应用中的考量与局限
Flash Attention 支持块稀疏计算，可跳过注意力掩码中的零块，进一步提升效率。但文章也指出，其反向传播的梯度重计算数学细节尚未完成，且未讨论不同模型或模块中瓶颈的差异。在实际应用中，需根据具体任务和硬件判断是否采用，并关注其与现有框架的兼容性。
Q&A
Flash Attention 为什么能减少注意力机制的内存占用？
Flash Attention 通过将 Q、K、V 分块计算，避免在 HBM 中存储大小为 seq_len×seq_len 的 softmax 中间矩阵，从而减少内存访问，使注意力所需内存随序列长度近似线性增长。
Flash Attention 的 FLOP 更高，为什么反而更快？
因为注意力计算受内存带宽限制（memory-bound），Flash Attention 虽然 FLOP 更高，但减少了 HBM 访问次数，因此速度更快。
Flash Attention 有哪些主要优势？
主要优势包括：1) 更快的模型训练；2) 在长序列任务中模型质量更高；3) 新的基准注意力方法，比现有方法更快且更省内存（2022.5）；4) 支持 Block-Sparse，仅计算注意力掩码中的非零块。
Flash Attention 与标准注意力在计算和内存上有何不同？
Flash Attention 的 FLOP 计数高于标准注意力，但由于注意力是内存访问受限的，且 Flash Attention 的 HBM 访问更少，因此实际更快。同时，它避免了存储大的 softmax 矩阵，内存效率更高。
Flash Attention 适用于哪些场景？
适用于需要长序列训练的任务，可以提升训练速度并支持更高质量的模型。此外，通过 Block-Sparse 技术，也适用于带有注意力掩码的稀疏注意力计算。
Flash Attention 的核心创新点是什么？
核心创新点是通过分块计算 Q、K、V 来避免存储大型 softmax 中间矩阵，从而使注意力内存高效，减少内存瓶颈，并提升训练速度。
