Flash Attention

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

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

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

您可以通过阅读这篇论文了解更多关于 Flash Attention 的信息链接。

< > 在 GitHub 上更新