閃光注意
Flash Attention 是一種計算 Transformer 內部注意力步驟的巧妙方法，無需編寫巨大的注意力矩陣來減慢記憶體速度。
概述
It makes long-context models far faster and more memory-efficient without changing their math.
深入探討
標準注意力將每個標記與其他每個標記進行比較，產生一個 N×N 分數矩陣，該矩陣隨序列長度呈二次方增長。簡單地說，該矩陣被寫入 GPU 高頻寬記憶體 (HBM) 並從中讀回，而真正的瓶頸是穿梭（而不是乘法）。 Tri Dao 及其同事於 2022 年推出的 Flash Attention 重新組織了計算，因此矩陣永遠不會完全儲存。它處理適合快速片上 SRAM 的小塊中的查詢、鍵和值，計算部分結果，並使用在線運行的 softmax 技巧將它們拼接在一起。輸出在數學上與普通注意力相同，但使用線性記憶體並且運行速度快幾倍，尤其是在長序列上。
技術洞察
關鍵技巧是平鋪加上線上 softmax。 Softmax 通常需要整行分數來計算其分母，但 Flash Attention 在串流傳輸每個圖塊時會保留運行最大值和運行總和，重新縮放早期的部分輸出，以便最終結果準確。由於中間分數保留在 SRAM 中（比 HBM 快幾個數量級），因此該演算法具有 IO 感知能力：它最大限度地減少了記憶體讀取和寫入，而不是原始算術運算。
戰略影響
成本與預算
多年來，架構決策決定著效能和營運成本。
更明確的決策
技術教育幫助團隊選擇正確的堆疊，而不僅僅是最新的堆疊。
品質管控
更好的工程選擇可以減少生產中的可靠性事故。
Flash 注意力的未來
Flash Attention 已成為預設建置區塊，FlashAttention-2 和 FlashAttention-3 透過改進工作分割區並利用低精確度 FP8 路徑，從 H100 等新型 GPU 中獲得更多吞吐量。預計將繼續與硬體進行協同設計，更緊密地整合到訓練和推理框架中，以及針對稀疏、滑動視窗和超長上下文注意力進行調整的變體。隨著上下文視窗擴展到數百萬個令牌，像這樣的 IO 感知核心對於保持記憶體和速度實用仍然至關重要。
現實世界的實施
以更低的記憶體成本訓練大型語言模型，例如 Llama 和 GPT 類系統，具有更長的上下文視窗。
透過加快首次閱讀長提示的預填階段，更快地為聊天助理提供服務。
透過在單一 GPU 上實現長序列注意力，啟用可擷取整本書或程式碼庫的文件分析工具。
為視覺和音頻變壓器提供動力，其中高解析度輸入創建非常長的令牌序列。
風險與防護欄
優化一項基準測試可以隱藏更廣泛的系統弱點。
基礎設施和維護成本常常被低估。
隨著系統變得更加複雜，安全性和可觀察性差距可能會擴大。
實施路線圖
在實施之前定義延遲、品質和成本目標。
在實際負載和資料條件下進行基準測試。
儀器監控錯誤、漂移和使用者影響。
在擴展之前準備回滾和事件回應路徑。
不斷探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the Flash Attention quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常見問題
What is Flash Attention?
Flash Attention 是一種計算 Transformer 內部注意力步驟的巧妙方法，無需編寫巨大的注意力矩陣來減慢記憶體速度。它使長上下文模型速度更快、記憶體效率更高，而無需改變其數學原理。
Flash Attention的主要瓶頸是什麼？
Flash Attention 具有 IO 感知能力：它減少了快速片上 SRAM 和慢速高頻寬記憶體之間傳輸的數據，這是真正的瓶頸，而不是演算法本身。
Flash Attention 如何避免儲存完整的 N×N 注意力矩陣？
它平鋪計算，因此每個區塊都適合快速 SRAM，計算和累積部分輸出，而無需在 HBM 中具體化整個矩陣。
什麼技術可以讓 Flash Attention 正確計算 softmax，而無需一次看到整行？
線上 softmax 在串流瓦片時保持運行最大值和運行分母，重新縮放早期的部分輸出，以便最終標準化是準確的。
為什麼 Flash Attention 對於長序列最有幫助？
樸素注意力矩陣在記憶體中以 N 平方縮放，因此當序列很長時，避免其完全儲存會產生最大的節省。
Flash Attention 的「IO 感知」設計優先考慮最小化什麼？
IO 感知意味著演算法是圍繞記憶體層次結構中移動資料的成本設計的，最大限度地減少 HBM 流量而不是算術運算。