FlashAttention

FlashAttention是由史丹福大學Tri Dao等人於2022年提出的一種快速且記憶體高效的精確注意力算法。它旨在解決Transformer模型自注意力機制在處理長序列時存在的計算速度慢、記憶體占用高（O(N)）的問題。該算法核心在於具備I/O感知能力，通過平鋪（Tiling）和重計算（Recomputation）技術，最佳化GPU不同層級記憶體間的數據訪問，從而在不改變計算結果的前提下，將記憶體複雜度降至線性（O(N)），並將計算速度提升數倍。該算法及其後續疊代版本已被集成到PyTorch等主流深度學習框架中，成為加速大模型訓練的關鍵技術之一。

基本介紹

• 外文名：FlashAttention

研發背景

Transformer模型的自注意力機制在處理長序列時，其計算複雜度和記憶體占用均與序列長度呈平方關係（O(N)），導致模型訓練和推理速度慢、記憶體消耗巨大，成為制約模型處理更長上下文的關鍵瓶頸。

為解決這一瓶頸，史丹福大學計算機系與紐約州立大學布法羅分校的研究團隊（Tri Dao、Dan Fu等）於2022年提出了FlashAttention算法。該算法是一種具有I/O（輸入/輸出）感知能力的精確注意力算法，旨在通過算法層面的最佳化，顯著減少對慢速高頻寬記憶體（HBM）的訪問，從而在保持計算精度的同時，實現更快的計算速度和更低的記憶體占用。

發展歷程

FlashAttention由Tri Dao等人於2022年提出，是一種IO感知的精確注意力算法。該算法通過分塊（Tiling）和重計算（Recomputation）技術，將注意力計算的記憶體複雜度從O(N)降低至O(N)。實驗表明，其運行速度比PyTorch標準實現快2-4倍，記憶體占用僅為後者的5%-20%。

FlashAttention-2於2023年發布，在算法、並行化和工作分區方面進行了顯著改進。該版本最佳化了GPU執行緒塊和warps之間的工作分配，減少了非矩陣乘法運算。在A100 GPU上，其吞吐量約為初代版本的2倍，端到端訓練速度可達225 TFLOPs/s，模型FLOPs利用率達72%。同時，它支持更高的頭維數（最高256）及多查詢注意力（MQA）和分組查詢注意力（GQA）。

FlashAttention-3於2024年發布，專門針對英偉達Hopper架構（如H100 GPU）最佳化。該版本利用了異步執行（重疊GEMM和softmax計算）、低精度（FP8）計算等硬體特性。在H100 GPU上，FP16精度下性能最高可達740TFLOPS（理論峰值利用率達75%），FP8精度下性能接近1.2 PFLOPS，速度是FlashAttention-2的1.5-2.0倍。

FlashAttention-4於2026年3月正式發布，針對Blackwell GPU架構（如B200）進行了算法與核心的協同設計。該版本通過新型軟體流水線、多項式近似指數計算、條件式softmax重縮放以及利用2-CTA MMA指令等技術，進一步最佳化性能。在B200 GPU（BF16精度）上，其性能最高可達1605 TFLOPs/s（利用率71%），比cuDNN 9.13快1.3倍。

核心技術

FlashAttention是一種快速且記憶體高效的精確注意力算法，其核心思想是通過平鋪（tiling）和重計算（recomputation）技術，將大型注意力矩陣分塊處理並利用GPU的片上SRAM進行計算，從而避免存儲完整的中間注意力矩陣，將記憶體使用量從序列長度的二次方O(N)降低到線性O(N)。該算法通過減少對高頻寬記憶體（HBM）的訪問次數，實現了比標準注意力實現快2-4倍的速度，並將記憶體使用量減少至標準方法的5%-20%。

FlashAttention-2在初始版本基礎上進行了算法最佳化，通過改進工作分區（work partitioning）和並行化策略，減少了非矩陣乘法運算的需求，在A100 GPU上實現了更高的計算吞吐量，達到理論FLOPs的50%-73%。FlashAttention-3則針對Hopper架構（如H100 GPU）的硬體特性進行了深度最佳化，採用了異步執行技術（如重疊GEMM與Softmax計算）並利用硬體對FP8低精度計算的支持，在H100上實現了高達740 TFLOPS（理論峰值利用率75%）的性能，速度是FlashAttention-2的1.5-2.0倍。FlashAttention-4於2026年3月發布，針對Blackwell架構（如B200 GPU）重新設計了算法流水線，引入了多項式指數模擬、條件式Softmax重縮放等新機制，以最大化硬體資源利用率。

該算法採用安全Softmax（或稱線上Softmax計算技術）來處理數值穩定性問題。然而，在BF16等低精度訓練場景下，標準的FlashAttention實現可能因特定數值誤差積累導致訓練後期損失函式（Loss）突然爆炸，研究表明這源於安全Softmax計算中產生的系統性數值偏置。一種修複方法是對安全Softmax算法進行極小修改，以穩定低精度訓練。

運營模式

FlashAttention系列算法以開源形式發布，並迅速被集成到PyTorch等主流深度學習框架和庫中。PyTorch從2.0版本開始集成FlashAttention核心，並在2.2版本中集成了FlashAttention-2。社區和廠商也致力於將FlashAttention移植適配到其他硬體平台，例如通過TileLang在AMD GPU上實現，或通過Ascend C在華為昇騰AI處理器上進行最佳化。在行業套用方面，FlashAttention已成為LLaMA、Falcon等知名大模型訓練和推理的基礎設施之一。

核心團隊

FlashAttention算法由史丹福大學的研究者提出。其主要提出者為Tri Dao，他是史丹福大學計算機科學博士，後任職於普林斯頓大學計算機科學系助理教授，並擔任Together AI首席科學家。共同作者包括Dan Fu，他是史丹福大學計算機科學博士生，是FlashAttention初代論文的共同一作。

隨著算法版本的疊代，後續貢獻者團隊不斷擴展。FlashAttention-3由來自Meta、NVIDIA、Together AI、喬治亞理工學院、普林斯頓大學等多機構的研究者共同推出。

重要事件

2022年6月，FlashAttention算法論文正式發布。2023年，FlashAttention被集成至PyTorch2.0版本。2026年3月，清華大學研究人員發表論文，首次從機制上解釋了FlashAttention在BF16低精度訓練下可能導致Loss爆炸的問題，並提出了修複方案。同月，FlashAttention-4正式發布，並宣布PyTorch FlexAttention已支持其作為後端。