訓練大型語言模型時，開發者經常遇到 GPU 記憶體不足嘅瓶頸，尤其係處理長序列輸入嘅 Transformer 注意力層。傳統注意力計算需要大量中間激活值，容易令記憶體爆滿，迫使大家縮減 batch size 或序列長度，嚴重拖慢訓練速度。Flash Attention 正正針對呢個痛點，提供咗一個快速且記憶體高效嘅精確注意力實現，幫 AI 研究員同工程師喺單一 GPU 上處理更長序列，加速模型開發流程。呢個開源項目由 Dao-AILab 維護，專注優化注意力計算，適用於各種深度學習工作負載。
NVIDIA CUDA 支援帶來 2 倍以上速度提升
Flash Attention 喺 NVIDIA GPU 上表現出色，利用 CUDA kernel 重新設計咗注意力計算流程，將原本需要儲存喺 HBM 嘅中間矩陣轉移到更快速嘅 SRAM，減少記憶體讀寫次數。呢個設計唔單止大幅降低峰值記憶體使用量，仲將計算速度提升到傳統實現嘅 2 倍以上。開發者只需簡單替換 PyTorch 嘅 scaled_dot_product_attention，就即時感受到效能躍升，尤其適合訓練長上下文模型如 GPT 或 LLaMA。
實際應用中，呢個 CUDA 優化特別適合高階 AI 實驗室，利用單張 A100 或 H100 GPU 就能處理 64k 序列長度，而唔使分散到多卡訓練。項目仲提供咗詳細嘅安裝指引，支援最新 CUDA 版本，確保相容性。
AMD ROCm 支援擴展到更多硬體平台
唔止 NVIDIA，用家仲可以喺 AMD GPU 上部署 Flash Attention，透過 ROCm 平台實現類似嘅記憶體節省同速度提升。呢個支援令更多開發者受益，尤其係用緊 Instinct MI 系列嘅團隊。ROCm 版本嘅 kernel 經過專門調校，維持咗精確注意力嘅數值穩定性，避開咗浮點誤差問題。
相比純 PyTorch 實現，ROCm 版 Flash Attention 喺記憶體綁定方面更高效，適合成本敏感嘅研究項目。安裝過程簡單，只需跟從 GitHub README 嘅步驟，即可喺 Linux 環境下運行。
2.0 版本完全重寫實現 2 倍加速
Flash Attention 2.0 係一次徹底重構，從頭重新設計咗 kernel 架構，帶來咗 2 倍嘅整體速度提升。呢個版本優化咗 tiling 策略同 I/O 模式，令計算更貼合現代 GPU 硬體特徵。無論係前向還是反向傳播，都明顯更快，特別喺長序列任務上表現突出。
升級到 2.0 後，用家會發現訓練時間大幅縮短，同時保持咗精確性 — 唔係近似方法，而是 100% 等價於標準注意力。呢個改動令 Flash Attention 成為咗 Transformer 模型嘅首選插件。
2.2 版本針對推論階段進一步優化
最新 2.2 版本專注推論優化，調整咗 kernel 以適應生成式任務嘅需求，例如 autoregressive decoding。呢啲改動減少咗不必要嘅記憶體分配，提升咗每秒 token 生成速度。同時，項目仲支援 Transformers 庫，直接整合到 Hugging Face 生態。
對於部署 LLM 服務嘅工程師，2.2 版特別實用，能夠喺相同硬體上支援更長上下文，改善用戶體驗。2.1 版亦有 causal flag 行為調整，確保咗因果遮罩嘅正確性。
產品名稱：Flash Attention
官方網站：https://github.com/Dao-AILab/flash-attention
支援平台：NVIDIA CUDA / AMD ROCm
