FlexAttention是由PyTorch團隊開發的一個PyTorch編程模型，於2024年7月公開。它允許用戶通過定義score_mod和mask_mod函式來描述各種注意力變體，並利用torch.compile自動生成高效的FlashAttention核心。
在PyTorch 2.5.0版本中，FlexAttention作為原型特性推出，並針對推理場景進行了最佳化。此後，它獲得了對FlashAttention-4後端的支持，並在PyTorch 2.9版本中在Intel GPU上獲得了原生支持。
基本介紹
- 外文名：FlexAttention
- 產品類型：編程模型
- 上市時間：2024年7月
研發背景,發展歷程,核心技術,核心團隊,重要事件,性能表現,
研發背景
在過去的七年中，注意力機制已成為深度學習最重要的原語之一。最佳化注意力的主要方法是FlashAttention，它通過運算元融合顯著提升了運行效率和記憶體消耗。然而，FlashAttention的重要性與其單體性相結合，給試圖嘗試新注意力變體的研究人員帶來了問題，即“軟體彩票”，新想法的成功取決於是否有配套的高效核心實現。這個問題因編寫高效融合的注意力核心本身具有挑戰性而加劇，傳統的基於編譯器的方法難以處理。為解決這一問題，MetaPyTorch團隊引入了FlexAttention，這是一個新穎的編譯器驅動的編程模型，允許用戶用少量符合PyTorch習慣的代碼定義多種注意力變體，並通過 torch.compile 自動生成融合的FlashAttention核心，其性能可與手工最佳化的核心相媲美。
發展歷程
FlexAttention 由 PyTorch 團隊開發，其發展歷程始於2024年。2024年7月，PyTorch團隊通過官方部落格首次公開介紹了FlexAttention API。該API旨在提供一個靈活且高性能的編程模型，允許用戶通過定義和函式來實現多種注意力變體，並利用自動生成高效的FlashAttention核心。
2024年11月，在PyTorch Nightly版本中被報告存在一個在多GPU環境下的設備同步問題，具體表現為當使用非首個GPU時會出現指針訪問錯誤。
2025年1月，PyTorch 2.5.0版本將FlexAttention作為原型特性正式推出。該版本針對大語言模型推理場景進行了重點最佳化，推出了專門的FlexDecoding後端以支持長上下文解碼，並增加了對分組查詢注意力（GQA）、分頁注意力（PagedAttention）、嵌套鋸齒張量以及可訓練偏差等功能的支持。
2025年3月，PyTorch團隊在英偉達GTC 2025大會上對FlexAttention進行了專題介紹，展示了其在A100和H100GPU上接近甚至部分超越FlashAttention系列的性能表現。
2025年5月，題為《FlexAttention: A Programming Model for Generating Optimized Attention Kernels》的論文在MLSys 2025會議上發表，系統闡述了其編程模型、編譯器驅動實現及性能優勢。
2026年3月，PyTorch官方宣布FlexAttention已支持最新的FlashAttention-4後端。通過集成CuTeDSL並利用JIT編譯，FlexAttention能夠在Hopper和Blackwell架構GPU上自動為自定義注意力變體實例化FlashAttention-4核心，在保持靈活性的同時進一步提升了計算受限工作負載下的性能。
核心技術
FlexAttention是一個由PyTorch團隊開發的編程模型，它提供了一個靈活的API，允許用戶通過定義score_mod和mask_mod函式來描述各種注意力變體及其組合，並利用PyTorch的torch.compile機制自動生成高效的FlashAttention核心。
FlexAttention的核心使用Triton語言實現，其模板包含flex_attention和flex_decoding兩個專門核心。flex_attention核心針對推理和訓練的前填充階段設計，而flex_decoding核心則針對解碼階段進行了最佳化。
FlexAttention支持嵌套鋸齒張量以處理變長序列，並支持分組查詢注意力、分頁注意力以及在score_mod函式中使用可訓練參數等高級特性。
由於用戶邏輯和注意力形狀的多樣性，最佳核心配置因用例而異。通過torch.compile的max-autotune模式，可以自動探索並選擇能提供最佳經驗性能的核心配置。
在PyTorch 2.9及以後版本中，所有FlexAttention場景均在Intel GPU上獲得原生支持，並通過集成Triton XPU編譯器進行專門最佳化以提升性能。
自2026年3月起，FlexAttention支持使用FlashAttention-4作為後端，由PyTorch自動生成CuTeDSL代碼並通過JIT編譯為自定義Attention變體實例化FlashAttention-4。
核心團隊
FlexAttention由Meta Platforms的PyTorch團隊主導開發。
其核心研發團隊包括Juechu Dong、Boyuan Feng、Driss Guessous（Meta Platforms機器學習工程師）、Yanbo Liang和Horace He等人，相關研究成果已以論文形式發表。
重要事件
2025年3月，在英偉達GTC大會上，Meta的機器學習工程師Driss Guessous發表了題為“FlexAttention: The Flexibility of PyTorch With the Performance of FlashAttention”的專題演講，正式向業界介紹了FlexAttention這一兼顧靈活性與高性能的編程模型。
2025年5月，題為《FlexAttention: A Programming Model for Generating Optimized Attention Kernels》的學術論文在MLSys 2025會議上發表並被接受展示。
2026年3月，PyTorch官方宣布FlexAttention正式集成並支持FlashAttention-4後端，此舉使得用戶在使用FlexAttention定義自定義注意力變體時，能夠自動編譯並利用FlashAttention-4的高性能核心，在Hopper和Blackwell GPU上獲得顯著的性能提升。
性能表現
FlexAttention在不同硬體平台和對比基準下展現出與先進方案相媲美的性能。在A100 GPU上，其前向傳播性能可達FlashAttention-2的90%，反向傳播可達85%。在H100 GPU上，前向傳播平均性能是FlashAttention-3的85%，比FlashAttention-2快約25%；反向傳播平均性能是FlashAttention-3的76%，比FlashAttention-2快約3%。在算力受限的工作負載下，使用FlashAttention-4後端相比Triton後端，可獲得1.2倍到3.2倍的性能提升。為實現高性能，FlexAttention核心設計注重在暫存器中執行計算，減少同步和共享記憶體流量，並利用軟體流水線重疊計算與記憶體訪問。
