flagos-ai/FlagAttention

解決的問題

FlagAttention 提供記憶體效率高的注意力運算子，相比標準的 CUDA 實作（如 FlashAttention）更易修改與自訂。它解決語言建模中對非標準注意力機制的需求，例如需要針對注意力分數計算、KV 快取布局或推論路徑進行專案特定的客製化。

工作原理

使用 Triton 語言實作，本專案透過分塊（tiling）與重計算（re-computation）避免產生完整的注意力分數矩陣，從而大幅降低記憶體用量與記憶體流量。它提供多個專用運算子：

• flash_attention：支援 MQA/GQA 與 dropout 的 Triton 版 FlashAttention v2 實作。
• piecewise_attention：用於非線性化位置嵌入（NLPE）的擴充，使用兩組查詢與鍵，根據距離閾值計算注意力分數。
• flash_attention_split_kv：專為長 KV 序列設計的分塊 KV 快速解碼運算子。
• paged_attention：用於推論期間分頁 KV 快取管理的運算子。

適用對象

適用於需要高效率注意力運算子但又需要比現成套件更高彈性之 AI 研究人員與開發者，或希望對注意力分數應用自訂轉換的使用者。

亮點

• 基於 Triton：相比原始 CUDA 程式碼更易理解與修改。
• 記憶體高效：透過分塊與重計算減少記憶體流量與用量。
• 可自訂：特別支援非標準運算子，如用於 NLPE 的分段注意力（piecewise_attention）。
• 硬體支援：已在 Nvidia Ampere GPU（A100、RTX-3090）與 Iluvatar CoreX GPU 上完成測試。
• 功能完整：支援 float16、bfloat16、因果/非因果模式以及 MQA/GQA。