4比特量化三倍加速不掉點(diǎn)！清華即插即用的SageAttention迎來(lái)升級(jí)

AIGC動(dòng)態(tài)11個(gè)月前發(fā)布機(jī)器之心

從8-Bit到4-Bit。

原標(biāo)題：4比特量化三倍加速不掉點(diǎn)！清華即插即用的SageAttention迎來(lái)升級(jí)
文章來(lái)源：機(jī)器之心
內(nèi)容字?jǐn)?shù)：6979字

清華大學(xué)陳鍵飛團(tuán)隊(duì)提出SageAttention2：4-Bit即插即用注意力機(jī)制，實(shí)現(xiàn)3-4.5倍推理加速

本文總結(jié)了清華大學(xué)陳鍵飛團(tuán)隊(duì)最新提出的SageAttention2論文要點(diǎn)。該工作在之前的SageAttention基礎(chǔ)上，進(jìn)一步將注意力機(jī)制中的Q、K矩陣量化到INT4，實(shí)現(xiàn)了更高速的推理速度，并在多個(gè)大型模型上保持了端到端的精度。

1. 背景與挑戰(zhàn)

隨著大型模型序列長(zhǎng)度的增加，注意力機(jī)制(Attention)的計(jì)算開(kāi)銷(xiāo)成為瓶頸。雖然線性層的低比特量化已較為成熟，但注意力模塊大多仍使用高精度(FP16或FP32)運(yùn)算。SageAttention率先將Attention中的QK^T量化到INT8，取得了顯著加速效果。然而，INT8的矩陣乘法速度仍不及INT4，且FP16的乘法累加器加速僅在特定顯卡上有效。因此，將注意力機(jī)制量化到INT4成為進(jìn)一步提升效率的關(guān)鍵，但也面臨巨大挑戰(zhàn)：直接量化到INT4會(huì)導(dǎo)致精度嚴(yán)重下降。

2. SageAttention2的技術(shù)方案

為了解決INT4量化帶來(lái)的精度損失問(wèn)題，SageAttention2提出了以下技術(shù)方案：

對(duì)Q和K進(jìn)行平滑處理： 在對(duì)K進(jìn)行平滑處理的基礎(chǔ)上，新增對(duì)Q進(jìn)行平滑處理(Q – mean(Q))，并補(bǔ)償?shù)阶罱K結(jié)果中，有效降低了量化誤差。
Per-thread量化： 將Q、K矩陣的量化粒度細(xì)化到per-thread級(jí)別，提高了4-Bit QK^T乘法的精度，且不增加額外開(kāi)銷(xiāo)。
FP32寄存器累加FP8 PV乘法結(jié)果： 使用FP32寄存器累加FlashAttention分塊粒度的PV的FP22乘法結(jié)果，避免了累積誤差。
使用E4M3格式的FP8： 實(shí)驗(yàn)表明，E4M3格式的FP8精度最高，接近FP16。
可選的V矩陣平滑： 對(duì)V矩陣進(jìn)行平滑處理，進(jìn)一步提升PV矩陣乘法的精度。

3. 實(shí)驗(yàn)結(jié)果

SageAttention2在速度和精度上都取得了顯著提升：

速度： 相比FlashAttention2和xformers分別實(shí)現(xiàn)了3倍和4.5倍的推理加速，在不同顯卡上均有不同程度的加速。
精度： 在視頻、圖像、文本生成等多種大型模型上保持了端到端的精度，例如在CogvideoX-1.5-5B模型上實(shí)現(xiàn)了1.8倍的端到端加速，且視頻質(zhì)量無(wú)損。

4. 總結(jié)

SageAttention2通過(guò)一系列創(chuàng)新技術(shù)，成功地將注意力機(jī)制量化到INT4，實(shí)現(xiàn)了顯著的推理加速，并在多個(gè)大型模型上保持了精度。其即插即用的特性也方便了在實(shí)際應(yīng)用中的部署，為大型模型的效率提升提供了有力支持。該工作已開(kāi)源，方便開(kāi)發(fā)者使用。