又快又準，即插即用！清華8比特量化Attention，兩倍加速于FlashAttention2，各端到端任務均不掉點！

作者：機器之心 2024-10-21 12:10:00

大模型中，線性層的低比特量化（例如 INT8, INT4）已經逐步落地；對于注意力模塊，目前幾乎各個模型都還在用高精度（例如 FP16 或 FP32）的注意力運算進行訓練和推理。

論文第一作者張金濤來自清華大學計算機系，論文通訊作者陳鍵飛副教授及其他合作作者均來自清華大學計算機系。

大模型中，線性層的低比特量化（例如 INT8, INT4）已經逐步落地；對于注意力模塊，目前幾乎各個模型都還在用高精度（例如 FP16 或 FP32）的注意力運算進行訓練和推理。然而，隨著大型模型需要處理的序列長度不斷增加，Attention（注意力運算）的時間開銷逐漸成為網絡優化的主要瓶頸。

為了提高注意力運算的效率，清華大學陳鍵飛團隊提出了 8Bit 的 Attention（SageAttention）。實現了 2 倍以及 2.7 倍相比于 FlashAttention2 和 xformers 的即插即用的推理加速，且在視頻、圖像、文本生成等大模型上均沒有端到端的精度損失。

論文標題：SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration
論文鏈接：https://arxiv.org/abs/2410.02367
開源代碼：https://github.com/thu-ml/SageAttention

即插即用舉例

SageAttention 可以一行代碼輕松替換掉 torch 中當前最優的 Attention 接口（scaled_dot_product_attention），實現即插即用的推理加速。

具體來說，SageAttention 的使用非常方便，使用 pip install sageattention 后，

只需要在模型的推理腳本前加入以下三行代碼即可：

效果上，以開源視頻生成模型 CogvideoX 為例，使用 SageAttention 可以端到端加速 35%，且生成的視頻無損：

全精度 Attention

SageAttention

接下來，將從背景與挑戰，技術方案，以及實驗效果介紹 SageAttention。

背景

隨著大模型需要處理的序列長度越來越長（比如 Llama3.1 支持 128K 的序列長度），Attention 的速度優化變得越來越重要。下圖展示了一個標準的 Transformer 模型中各運算隨著序列長度變化的時間占比：

挑戰

為了方便指代注意力元算中包含的矩陣，我們先回顧一下注意力的計算公式：

將神經網絡中各運算的數值類型從高比特量化至低比特是一種有效提升計算和訪存效率的方法。然而，研究團隊發現直接將注意力運算中的 Q, K, P, V 從 FP16 量化為 INT8 或者 FP8 后將會導致在幾乎所有模型和任務上都會得到極差的結果，例如，在 Unidiffuser 文生圖模型中，會得到一張完全模糊的圖像；在 Llama2-7B 進行四選一選擇題任務上得到 25.5% 的準確率。

經過仔細分析后，研究團隊發現主要是兩個原因導致了量化注意力的不準確：

大多視頻、圖像生成模型中，矩陣 K 表現出了極強的通道維度的異常值分布，直接使用 INT8 或者 FP8 數據類型對其進行量化會導致巨大的誤差。

在所有模型中，對矩陣 P, V 進行量化不能保證一個模型中所有層的精度。下表展示了對 P, V 量化后，Llama2-7B 和 Unidiffuser 模型所有層中，最差情況的層對應的量化注意力的準確度，（該準確度為量化注意力相比全精度注意力的誤差），可以發現不管對 P, V 矩陣進行何種 8Bit （INT8，E4M3，E5M2）量化，總有些層的準確率非常差，導致了端到端效果的下降。

技術方案

為了解決上述的兩個關鍵問題，研究團隊提出了對應的解決辦法。

對 K 進行平滑處理。SageAttention 采用了一個簡單但非常實用的方法來消除矩陣 K 的異常值：K = K – mean (K) 其中 mean (K) 是沿著通道維度求平均值。這個簡單的做法不僅不會影響注意力計算的正確性 Softmax (QK^T) = Softmax (Q (K-mean (K))^T) ；且對整個 Attention 速度的影響只有 0.2%；同時還保證了量化后的注意力運算的精度：

對 Q, K 進行分塊 INT8 量化。對于矩陣 Q, K，SageAttention 采用了以 FlashAttention 的分塊大小為粒度的 INT8 量化。這是因為：1. 對 Q, K 矩陣進行 INT8 量化相比于進行 FP8 量化，注意力的精度更高。2. 在一些常用卡上，比如 RTX4090，INT8 矩陣乘法（INT32 為累加器）的速度是 FP8（FP32 為累加器）的兩倍。
對 P, V 采用 FP16 數據類型的矩陣乘法累加器。對于矩陣 P, V，SageAttention 采用了保留 P, V 為 FP16 的類型，但進行矩陣乘法時采用 FP16 數據類型的累加器。這是因為：1. PV 矩陣乘法的數值范圍始終在 FP16 的表示范圍內，且經過大量實驗驗證，FP16 作為累加器的數據類型不會帶來任何精度損失（見下表）。2. 在一些常用卡上，比如 RTX4090，以 FP16 為累加器數據類型的矩陣乘法的速度是 FP32 作為累加器的兩倍。