Flash Attention

擴展 Transformer 架構的主要瓶頸在於自我注意力（self-attention）機制，其時間與記憶體複雜度呈二次方成長。近期硬體加速器的發展主要集中於提升計算能力，而非記憶體容量及硬體間的資料傳輸效率，這導致注意力運算受到記憶體頻寬的限制。Flash Attention 是一種旨在緩解此問題並更有效率地擴展 Transformer 模型架構的注意力演算法，能加速模型的訓練與推論過程。

標準的注意力機制使用高頻寬記憶體（HBM）來儲存、讀取與寫入鍵（Keys）、查詢（Queries）與值（Values）。HBM 擁有較大的儲存空間，但處理速度較慢；相對地，SRAM 的容量較小，但運算速度極快。在標準的注意力實作中，從 HBM 讀取與寫入鍵、查詢與值會消耗大量成本。它需要將鍵、查詢與值從 HBM 載入至 GPU 的晶片內 SRAM，執行單一步驟的注意力機制，再寫回 HBM，並對每個注意力步驟重複此過程。相反地，Flash Attention 只需載入鍵、查詢與值一次，將注意力機制的運算進行融合（fuse），最後再寫回記憶體。

此功能已實作於支援的模型中。您可以點擊此處查看完整支援 Flash Attention 的模型清單（名稱包含 flash 前綴的模型）。

若想深入了解 Flash Attention，您可以閱讀此連結中的研究論文。

在 GitHub 上更新