長序列模型的實用化關鍵：FlashAttention-2 如何將硬體效率推向極限

近年來，大型語言模型（LLM）的長序列（long-context）能力突飛猛進，從數千 token 擴展到數十萬甚至百萬。我們常將此歸因於模型架構創新，但真正的突破點，其實發生在硬體與軟體之間的底層運算層。我認為，長序列能力之所以能從學術論文走向實際應用，關鍵在於 FlashAttention-2 這樣的工程傑作。它透過極致的平行化與記憶體存取優化，將注意力機制的運算效率逼近 GPU 硬體理論上限，讓過去因成本過高而難以實現的長序列架構變得可行。

標準注意力機制的瓶頸：為何記憶體 I/O 比運算更致命？

要理解 FlashAttention 的重要性，我們必須先回到注意力機制（Attention Mechanism）的核心挑戰。傳統的注意力機制在計算上具有二次方複雜度（O(N²)），其中 N 是序列長度。當序列長度從 2,048 擴展到 32,768 時，所需的運算量與記憶體會呈平方級增長，很快就變得令人望而卻步。

然而，真正的瓶頸往往不是浮點運算（FLOPs）本身，而是記憶體的讀寫頻寬。在標準實作中，為了計算注意力分數，GPU 需要生成一個巨大的 N×N 注意力矩陣（Attention Matrix），並將其完整寫入高頻寬記憶體（HBM）。HBM 雖然快，但相較於 GPU 核心內部的超高速快取（SRAM），其速度與延遲仍然是天壤之別。反覆在 SRAM 與 HBM 之間讀寫這個龐大矩陣，成了拖慢整體效能的主因。

換句話說，在處理長序列時，GPU 大部分時間不是在「計算」，而是在「等待」資料從相對慢速的 HBM 傳輸過來。這使得硬體的理論運算能力大量閒置，效率極低。

FlashAttention-2 如何進一步榨乾 GPU 潛力？

為了解決這個 I/O 瓶頸，Tri Dao 在 2022 年提出了第一代 FlashAttention。其核心思想是「核心融合（Kernel Fusion）」與「區塊計算（Tiling）」。它將整個注意力計算過程（縮放、遮罩、Softmax、輸出加權）合併成單一的 CUDA 核心，並將巨大的注意力矩陣拆分成小區塊（tiles）。這些小區塊可以在更快的 SRAM 中完成計算，無需將整個中介矩陣寫回 HBM，從而大幅減少了記憶體讀寫次數。

而在 2023 年發表的FlashAttention-2，則是在這個基礎上進行了更細膩、更極致的工程優化。它並未改變核心演算法，而是專注於如何讓運算與硬體特性更完美地對齊，主要改進包括：

• 減少非矩陣乘法運算： 調整演算法以減少非矩陣乘法（non-matmul）的 FLOPs，讓更多運算能被專為矩陣運算設計的 Tensor Cores 高效執行。
• 更佳的平行化策略： 在前向傳播（forward pass）中，除了在批次大小（batch size）和頭（head）維度上平行化，FlashAttention-2 還巧妙地在序列長度（sequence length）維度上進行了平行化，進一步提升了 GPU 的利用率。
• 優化的工作排程： 在 GPU 的執行緒區塊（thread blocks）與 warp（一組 32 個執行緒）之間更細緻地分配工作，減少了不同 warp 之間的等待與同步，確保運算單元時刻保持忙碌。

這些看似微小的改動，卻帶來了驚人的成果。根據論文數據，FlashAttention-2 的執行速度約為第一代的 2 倍。更重要的是，在 NVIDIA A100 GPU 上的測試顯示，其運算吞吐量達到了硬體理論峰值的 50-73%，這個效率已經非常接近高度優化的矩陣乘法（GEMM）運算（通常可達 60-85%）。這意味著它幾乎將 GPU 的潛力壓榨到了極限。

從理論到實踐：工程優化如何解鎖長序列應用的未來？

FlashAttention-2 的成功，完美詮釋了底層工程優化對於上層應用的巨大價值。它沒有提出全新的注意力變體，而是透過對硬體架構的深刻理解，解決了長久以來的工程瓶頸。這使得許多過去僅存在於理論中的長序列模型，如 Ring Attention 或基於區塊的注意力，變得在工程上可行且高效。

當我們看到 Gemini 1.5 Pro 能夠處理 100 萬 token 的上下文，或是開源社群能夠訓練出具有 32K 甚至 128K 序列長度的模型時，其背後都有 FlashAttention 這類高效能運算核心的影子。它將訓練和推論長序列模型的成本與時間大幅降低，為需要處理長篇文件、多輪複雜對話、或分析整段程式碼庫的應用鋪平了道路。

最終，AI 系統的進步不僅來自於演算法的靈光一閃，也同樣仰賴於那些在數百萬次運算迴圈中，為我們節省幾個時脈週期、減少幾次記憶體存取的無名工程英雄。FlashAttention-2 正是這樣一個典範，它提醒我們，真正的突破，往往發生在理論與實踐交會的最深處。

延伸閱讀

• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (arXiv)
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv)
• The Illustrated Transformer by Jay Alammar

我是江中喬，一位具有 TPM 與產品管理背景的 AI 系統建構者，目前專注於 AI 認知增強系統與多 Agent 協作架構的設計與實踐。