Flash Attention

Transformer 架构的扩展受到自注意力机制的严重瓶颈限制，该机制具有二次时间和内存复杂度。加速器硬件的最新发展主要集中在增强计算能力，而不是内存以及硬件之间的数据传输。这导致注意力操作存在内存瓶颈。Flash Attention 是一种用于减少此问题并更有效地扩展基于 Transformer 的模型的注意力算法，从而实现更快的训练和推理。

标准注意力机制使用高带宽内存 (HBM) 来存储、读取和写入键、查询和值。HBM 内存容量大，但处理速度慢，而 SRAM 内存容量较小，但运算速度更快。在标准注意力实现中，从 HBM 加载和写入键、查询和值的成本很高。它从 HBM 将键、查询和值加载到 GPU 片上 SRAM，执行注意力机制的单个步骤，将其写回 HBM，并为每个注意力步骤重复此操作。相反，Flash Attention 加载一次键、查询和值，融合注意力机制的操作，然后将其写回。

它已针对受支持的模型实现。您可以在此处查看支持 Flash Attention 的模型的完整列表，对于带有 flash 前缀的模型。

您可以通过阅读此链接中的论文来了解有关 Flash Attention 的更多信息。

< > 在 GitHub 上更新