4比特量化三倍加速不掉點！清華即插即用的SageAttention迎來升級

作者：機器之心 2024-12-27 09:30:00

近日，陳鍵飛團隊進一步提出了 4-Bit 的即插即用 Attention（SageAttention2），相較于 FlashAttention2 和 xformers 分別實現了 3 倍以及 4.5 倍的即插即用的推理加速，且在視頻、圖像、文本生成等大模型上均保持了端到端的精度表現。

論文共同第一作者張金濤、黃浩峰分別來自清華大學計算機系和交叉信息研究院，論文通訊作者陳鍵飛副教授及其他合作作者均來自清華大學計算機系。

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

此前，清華大學陳鍵飛團隊提出的 8-Bit 的即插即用 Attention（SageAttention），將 Attention 中的 QK^T 量化至 INT8，將 PV 保持為 FP16 精度并使用 FP16 精度的矩陣乘法累加器，同時提出 Smooth K 技術保持了量化 Attention 的精度，實現了 2 倍加速于 FlashAttention2，且在各類大模型上均保持了端到端的精度表現。

目前，SageAttention 已經被業界及社區廣泛地使用于各種開源及商業大模型中，比如 CogvideoX、Mochi、Flux、Llama3、Qwen 等。

近日，陳鍵飛團隊進一步提出了 4-Bit 的即插即用 Attention（SageAttention2），相較于 FlashAttention2 和 xformers 分別實現了 3 倍以及 4.5 倍的即插即用的推理加速，且在視頻、圖像、文本生成等大模型上均保持了端到端的精度表現。

論文標題：SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization
論文鏈接：https://arxiv.org/abs/2411.10958
開源代碼：https://github.com/thu-ml/SageAttention

即插即用舉例

SageAttention2 實現了高效的 Attention 算子，可以實現即插即用的推理加速。輸入任意 Q, K, V 矩陣，SageAttention2 可以快速返回 Attention Output (O)。

具體來說，SageAttention2 使用起來很方便，克隆倉庫（git clone https://github.com/thu-ml/SageAttention）并執行 python setup.py install 后，只需一行代碼便可以得到 Attention 的輸出，可以使用該接口方便地替換任意模型中的 Attention 函數：

效果上，以開源視頻生成模型 CogvideoX-1.5-5B 為例，使用 SageAttention2 可以端到端加速 1.8 倍，且生成的視頻無損：

使用全精度 Attention

使用 SageAttention2

更重要的是，SageAttention2 提供了比 SageAttention 更廣泛的硬件支持。除了在 RTX 4090 上可以 3 倍加速于 FlashAttention 外，在 L20、L40、L40S 可以實現 2 倍的加速，在 A100、A800、A6000 上可以實現 1.45-1.6 倍的加速（基于 SageAttention）。

接下來，研究團隊將從前言、挑戰、方法以及實驗效果四個方面介紹 SageAttention2（總體流程圖如下圖）。

前言

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

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

盡管 SageAttention 提出將 Q，K 量化至 INT8，將 P，V 保持 FP16 精度且采用 FP16 的矩陣乘法累加器來加快 Attention 的速度。然而，這樣做的缺點是：1）INT8 的矩陣乘法只達到了一半的 INT4 矩陣乘法的速度，2）使用 FP16 的乘法累加器的 FP16 的矩陣乘法的加速只在 RTX4090 和 RTX3090 顯卡上有效。

為了克服上述缺點，SageAttention2 提出將 Q, K 量化至 INT4，并將 P, V 量化至 FP8 來加速 Attention。然而，這樣做的挑戰是很大的。

4-Bit 注意力量化有什么問題？

研究團隊發現直接將注意力運算中的 Q, K 量化為 INT4 后將會導致在幾乎所有模型和任務上都會得到極差的結果，例如，在 CogVideoX 文生視頻模型中，會得到完全模糊的視頻；Llama2-7B 進行四選一選擇題任務上得到 25% 的準確率。

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

（1）INT4 的數值范圍相比 INT8 非常小，導致其量化誤差在 Q，K 矩陣中出現一些異常值時會變得十分明顯，恰好大多模型都在 Q, K 中表現出來了較大的通道維度的異常值。這極大削減了 QK^?矩陣乘法的精度。

（2）研究團隊發現 Nvidia 的顯卡上，FP8 的矩陣乘法指令 (mma.f32.f8.f8.f32) 的乘法累加器并不是官方宣稱的 FP32 精度，而是只有 FP22 精度，這導致了 PV 矩陣乘法出現較大的累加誤差。

技術方案

為了解決上述的兩個挑戰，研究團隊提出了對應的解決辦法。

（1）保留 SageAttention 中對 K 進行平滑處理的同時，提出對 Q 進行平滑處理：Q – mean (Q)。其中 mean (Q) 是沿著通道維度的平均值向量。完成該平滑操作后需要在 Attention 計算過程中將 mean (Q) 和 K^T 的向量與矩陣乘法的結果補償到 S 中。

這使得相比直接量化 Q, K 至 INT4 的準確度有質的改變，如下表展示了對比了該方法和直接量化 Q, K 至 INT4 在 Cogvideo 和 Llama3.1 上的端到端表現。

矩陣 Q 平滑前后的數據分布可視化的結果如下，可以發現平滑后的 Q 對 INT4 數據范圍的利用度更高：

（2）對 Q, K 進行 Per-thread 量化。對于矩陣 Q, K，SageAttention2 采用了根據 mma 指令對矩陣內存排布的要求，對 Q，K 中的 Token 按照 GPU 線程進行分組，使量化粒度比 SageAttention 中的 per-block 細化 16 倍，極大提高了 4Bit 的 QK^?乘法準確度的同時不引入任何額外開銷。

具體來說，在 SageAttention 中，每個 Q 的塊將被劃分為 c_w 個段，由 GPU 流處理器（SM）中的 c_w 個 GPU warp 處理。然后，每個包含 32 個線程的 warp 會使用 NVIDIA 的 mma.m16n8k64 PTX 指令來執行 QK^?運算。根據這一指令的布局要求，研究團隊發現一個 warp 內的 Q [8×(n%8)] 可以共用一個量化縮放參數，而一個 warp 內的 K [8×(n%8)] 和 K [8×(n%8+1)] 也可以共用一個量化縮放參數，其中 n 是 token 索引。

這種量化方法更為細致且不增加額外開銷。這是因為它根據 MMA 指令的布局將不同的 GPU 線程分配到不同的量化 Token 組，每個線程只對應一個量化縮放參數進行反量化。而非 Per-token 量化那樣，每個線程對應多個量化縮放參數。

如下表所示，可以發現 per-thread 量化的準確度比 SageAttention 中采用的 per-block 量化高得多，準確度和 per-token 量化幾乎沒有差別。

（3）對 FP8 的 PV 矩陣乘法采用 FP32 的寄存器將每次 FlashAttention 分塊粒度的 PV 的 FP22 的乘法結果累加起來。這種做法可以有效地避免 FP22 的乘法累加器沿著序列長度累積過多的誤差，將 FP22 累加器帶來的誤差控制在 FlashAttention 分塊的粒度中，提高了 FP8 的 PV 乘法的準確度。

（4）針對 P 和 V，研究團隊對比了多種量化的數據類型，對比發現使用 E4M3 數據格式的 FP8 精度最準確，基本接近了 FP16 的準確度。因此采用將 P 和 V 量化至 E4M3。

下圖展示了 SageAttention2 的算法流程：

SageAttention2 共實現了兩種 Kernel，區別在于對 Q, K 進行 INT4 量化還是 INT8 量化：

此外，SageAttention2 還提出一種可選的對矩陣 V 進行平滑處理的技術，可以進一步提高 PV 矩陣乘法的準確度。具體來說，當某些模型中 V 矩陣具有通道維度的偏移時，可以將 V 減去其通道維度的平均值 mean (V) 來去除偏移，之后進行正常的量化 Attention 運算。只需要對最終 Attention 的 Output 加上 mean (V) 即可保持計算的正確性。