FlashAttention-2

FlashAttention 計算「精確」注意力，卻完全不把 N 乘 N 的注意力矩陣寫進 GPU 記憶體。樸素做法會把每對 token 的分數都實體化，耗費二次方記憶體，並讓緩慢的高頻寬記憶體充斥讀寫。FlashAttention 改成在鍵與值的區塊上串流，把線上 softmax 的統計量保存在快速的晶片上 SRAM，於是用遠少的記憶體流量產生完全相同的結果。FlashAttention-2 重新工程該核心，把與 GPU 峰值吞吐的剩餘差距幾乎補平。

第二版做了三項具體改動。它減少非矩陣乘的浮點工作，也就是線上 softmax 的重新縮放記帳，因為張量核心做矩陣乘遠快於其他運算。它在序列長度維度上平行化，而非只在批次與頭上，使長序列能讓每個串流多處理器保持忙碌。它也改良執行緒區塊內各 warp 之間的工作切分，以減少共享記憶體的同步。三者合計約使原版吞吐翻倍，在 A100 等級硬體上達到理論浮點峰值的 50 至 70 個百分點。

結果是精確而非近似，與標準注意力相同（至浮點運算順序），因此是可直接替換的速度與記憶體雙贏，能在固定記憶體下支援更長上下文。其後繼 FlashAttention-3 進一步利用 Hopper 的非同步與 FP8；這一系列如今是多數訓練堆疊的預設注意力後端。