打開網(wǎng)易新聞 查看精彩圖片

作者|CalibQuant 團隊

在 InternVL-2.5 上實現(xiàn)10 倍吞吐量提升,模型性能幾乎無損失。

最新 1-bit 多模態(tài)大模型 KV cache 量化方案CalibQuant來了。

通過結(jié)合后縮放和校準方法,可顯著降低顯存與計算成本,無需改動原模型即可直接使用。

打開網(wǎng)易新聞 查看精彩圖片

即插即用、無縫集成

模態(tài)大語言模型在各種應(yīng)用中展現(xiàn)出了卓越的性能。然而,它們在部署過程中的計算開銷仍然是一個關(guān)鍵瓶頸。

雖然 KV cache 通過用顯存換計算在一定程度上提高了推理效率,但隨著 KV cache 的增大,顯存占用不斷增加,吞吐量受到了極大限制。

為了解決這一挑戰(zhàn),作者提出了 CalibQuant,一種簡單卻高效的視覺 KV cache 量化策略,能夠大幅降低顯存和計算開銷。具體來說,CalibQuant 引入了一種極端的 1 比特量化方案,采用了針對視覺 KV cache 內(nèi)在模式設(shè)計的后縮放和校準技術(shù),在保證高效性的同時,不犧牲模型性能。

作者通過利用 Triton 進行 runtime 優(yōu)化,在 InternVL-2.5 模型上實現(xiàn)了 10 倍的吞吐量提升。這一方法具有即插即用的特性,能夠無縫集成到各種現(xiàn)有的多模態(tài)大語言模型中。

動 機

當前的多模態(tài)大語言模型在實際應(yīng)用中常常需要處理大尺寸、高分辨率的圖像或視頻數(shù)據(jù),KV cache 機制雖然能提升效率,但其顯存占用與輸入長度(如視覺幀數(shù)、圖像尺寸等)成正比。

當輸入數(shù)據(jù)的規(guī)模增大(例如更多的視覺幀、更高的圖像分辨率)時,KV 緩存的顯存使用量迅速增加,成為限制吞吐量的瓶頸。盡管當前有些針對 LLM KV cache 量化的方法可以將其壓縮至 2 比特,但這些方法沒有針對多模態(tài)問題中特有的視覺冗余做分析優(yōu)化,導(dǎo)致其無法在極限情況 1 比特下被使用

本文通過分析多模態(tài)大語言模型中的視覺 KV cache 的冗余,設(shè)計了適合多模態(tài)模型特有的 KV cache 量化方案。

方 法

本文在通道維度量化的基礎(chǔ)上提出了針對反量化計算順序的后縮放優(yōu)化方案和針對注意力權(quán)重優(yōu)化的校準策略。

1. 通道維度 KV cache 量化:

一種廣泛使用的方法是均勻整數(shù)量化。給定一個比特寬度 b>0 和一個輸入值 x,它位于某個范圍 [α,β] 內(nèi),則將其映射到一個離散整數(shù),計算過程為:

打開網(wǎng)易新聞 查看精彩圖片

這里的???表示取整運算符。最樸素的方法是使用全局統(tǒng)計量來計算這些極值,但是模型性能會受較大影響,作者選擇在通道維度上細化統(tǒng)計范圍。具體來說,令表示一個 K cache,其中 n 和 d 分別表示 token 的數(shù)量和 head 的維度。 定義兩個向量如下:

然后,通過上述過程對 K 中的每一行向量進行量化,其中乘法操作是逐元素進行的。作者同樣將這種按通道的量化方法應(yīng)用于 V cache。

2. 后縮放 KV cache 管理策略:

量化后的 K cache 可以用離散化的整數(shù)值、一個縮放因子(scale factor)和一個偏置項(bias term)來表示。在解碼階段,這些值被用于對 K cache 進行反量化,并隨后與 Q 相乘。然而,通道維度的量化需要為每個通道分別指定不同的縮放因子和偏置向量,這將導(dǎo)致產(chǎn)生大量不同的數(shù)值,增加了反量化過程中的計算開銷。此外,這種方式也使得 CUDA 內(nèi)核中的計算效率降低。作者觀察到量化后的 K 僅具有有限數(shù)量的離散取值(例如,對于 2 比特量化,其取值僅為 0、1、2、3),于是提出利用簡單的計算順序重排來減少存儲需求,并提高計算效率。具體過程如下:

設(shè)是 K cache 矩陣中的任意一行向量,

為其進行 b 比特整數(shù)量化后的結(jié)果,并伴隨有逐通道的縮放因子α,β。給定一個查詢向量,在生成 token 過程中注意力計算如下:

打開網(wǎng)易新聞 查看精彩圖片

其中,符號?和⊙分別表示向量之間的內(nèi)積和逐元素乘積。通道維度上的反量化操作被延遲執(zhí)行,并高效地集成到后續(xù)的向量乘法運算中。因此,這種方法僅存儲經(jīng)過 b 比特整數(shù)量化后的數(shù)值,并且避免了全精度反量化計算過程。這種方法確保了低比特反量化執(zhí)行的高效性。這種后縮放方法也可以自然地應(yīng)用到 V cache 的反量化過程中。

量化后的校準:

1 比特量化的一個限制是經(jīng)過反量化之后的數(shù)值往往會包含大量的極端值。這是因為 1 比特量化的碼本總是包含了最小值和最大值,導(dǎo)致那些接近邊界的輸入值在反量化后直接映射到了極端值。

因此,重建后的 KV cache 通常包含過多的大絕對值,最終導(dǎo)致注意力分數(shù)產(chǎn)生明顯的失真。為了解決這個問題,作者提出了一種量化后校準方法,用于調(diào)整 softmax 之前注意力分數(shù)的峰值。具體來說,假設(shè)的所有元素都位于區(qū)間內(nèi)。給定,定義一個線性變換 g 將區(qū)間圖片映射到圖片,其表達式如下:

打開網(wǎng)易新聞 查看精彩圖片

隨后對注意力分數(shù)進行如下調(diào)整:

打開網(wǎng)易新聞 查看精彩圖片

圖片如下圖所示,校準方法(Quant-C,紅色)有效減輕了極端值的影響,使調(diào)整后的注意分數(shù)分布相較于未經(jīng)校準的量化方法(Quant,藍色)更接近全精度(Exact)分布。

打開網(wǎng)易新聞 查看精彩圖片

實驗結(jié)果

作者將提出的量化方法分別應(yīng)用在 LLaVA 和 InternVL model 上,測試了其在 captioning,VQA,Video QA 三個不同的任務(wù)上的性能。以 captioning 任務(wù)為例,下圖展示了本文所提出的方法在 cococaption benchmark 下和其他方法如 KIVI,VLCache 的對比。

在不同比特數(shù)(8,4,2,1)下,本文提出的方法在大部分測試指標上都優(yōu)于其他兩種方法。例如對于 llava-1.5-7b,本文的方法在 8 比特下達到最高的 CIDEr 分數(shù) 1.105,與全精度持平,并在 1 比特下提升至 1.109,超過了 VLCache(1.053)。同樣地,對于 InternVL-2.5-26B,本文的方法在 4 比特和 2 比特下分別取得了最高的 CIDEr 分數(shù) 1.32 和 1.313,均優(yōu)于 VLCache 和 KIVI。

打開網(wǎng)易新聞 查看精彩圖片

Runtime 分析

為了展示本文提出的量化方法對解碼效率的影響,作者使用 InternVL-2.5 系列模型,將所提出的 1 比特量化方法與 16 比特基線進行了吞吐量評估(即每秒生成的 token 數(shù))。作者考慮了兩種視覺 token 長度的情況:n=3328 和 8192。作者將 GPU 最大內(nèi)存從 5GB 變化到 30GB,并在每種內(nèi)存限制下,尋找能夠容納的最大 batch size,測量解碼階段的吞吐量。

如下圖展示,1 比特量化方法在所有顯存預(yù)算下始終優(yōu)于基線方法。例如,當 n=3329 且使用 80 億參數(shù)模型時,本文的方法在 5GB 顯存下實現(xiàn)了 126.582tokens/s 的吞吐量(基線為 11.628tokens/s),在 30GB 下提升至 459.016tokens/s(基線為 40.816tokens/s)。這意味著相比基線,本文方法的吞吐量提升約為 9.88×到 11.24×,充分展示了該方法在受限顯存條件下顯著提升解碼速率。

打開網(wǎng)易新聞 查看精彩圖片

總 結(jié)

本文探討了多模態(tài)大語言模型中視覺 KV cache 的壓縮方法。簡單地將量化應(yīng)用到極低比特數(shù)常常會引發(fā)分布偏移,導(dǎo)致模型性能下降。為了解決這一問題,本文提出了一種新穎的校準策略,作用于 softmax 之前的注意力分數(shù),有效緩解了量化帶來的失真。此外,本文還引入了一種高效的通道維度后縮放技術(shù)以提高計算和存儲效率。

作者在 InternVL 和 LLaVA 模型系列上,針對 COCO Caption、MMBench-Video 和 DocVQA 等基準任務(wù)進行了實驗,結(jié)果驗證了所提出方法的有效性。作者利用 Triton 實現(xiàn)了本文所提出的方法,runtime 分析表明本文提出的方法相較于全精度模型有大約 10 倍的吞吐量提升。

論文標題:CalibQuant:1-Bit KV Cache Quantization for Multimodal LLMs

論文地址:https://arxiv.org/abs/2502.14882

代碼地址:https://github.com/insuhan/calibquant

會議推薦

AICon 2025 強勢來襲,5 月上海站、6 月北京站,雙城聯(lián)動,全覽 AI 技術(shù)前沿和行業(yè)落地。大會聚焦技術(shù)與應(yīng)用深度融合,匯聚 AI Agent、多模態(tài)、場景應(yīng)用、大模型架構(gòu)創(chuàng)新、智能數(shù)據(jù)基建、AI 產(chǎn)品設(shè)計和出海策略等話題。即刻掃碼購票,一同探索 AI 應(yīng)用邊界!

今日薦文

你也「在看」嗎?