Flash Attention

Transformer架構的擴充套件受到自注意力機制的嚴重瓶頸，自注意力機制具有二次時間複雜度和二次記憶體複雜度。加速器硬體的最新發展主要集中在增強計算能力，而不是記憶體以及硬體之間的資料傳輸。這導致注意力操作存在記憶體瓶頸。Flash Attention是一種注意力演算法，用於減少此問題並更有效地擴充套件基於Transformer的模型，從而實現更快的訓練和推理。

標準注意力機制使用高頻寬記憶體（HBM）來儲存、讀取和寫入鍵、查詢和值。HBM記憶體大但處理速度慢，而SRAM記憶體小但操作速度快。在標準注意力實現中，從HBM載入和寫入鍵、查詢和值的成本很高。它將鍵、查詢和值從HBM載入到GPU片上SRAM，執行注意力機制的單一步驟，將其寫回HBM，並對每個注意力步驟重複此過程。相反，Flash Attention一次性載入鍵、查詢和值，融合注意力機制的操作，然後將其寫回。

它已針對支援的模型實現。您可以在此處檢視支援Flash Attention的完整模型列表，其中包含以“flash”為字首的模型。

您可以透過閱讀此連結中的論文來了解更多關於Flash Attention的資訊。

< > 在 GitHub 上更新