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

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

通過(guò)結(jié)合后縮放和校準(zhǔn)方法，可顯著降低顯存與計(jì)算成本，無(wú)需改動(dòng)原模型即可直接使用。

即插即用、無(wú)縫集成

多模態(tài)大語(yǔ)言模型在各種應(yīng)用中展現(xiàn)出了卓越的性能。然而，它們?cè)诓渴疬^(guò)程中的計(jì)算開(kāi)銷(xiāo)仍然是一個(gè)關(guān)鍵瓶頸。

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

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

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

動(dòng)機(jī)

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

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

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

方法

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

1、通道維度KV cache量化

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

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

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

2、后縮放KV cache管理策略

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

設(shè)是K cache矩陣中的任意一行向量，為其進(jìn)行b比特整數(shù)量化后的結(jié)果，并伴隨有逐通道的縮放因子α,β。給定一個(gè)查詢(xún)向量，在生成token過(guò)程中注意力計(jì)算如下：

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

3、量化后的校準(zhǔn)

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

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

隨后對(duì)注意力分?jǐn)?shù)進(jìn)行如下調(diào)整：

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

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

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

在不同比特?cái)?shù)（8，4，2，1）下，本文提出的方法在大部分測(cè)試指標(biāo)上都優(yōu)于其他兩種方法。例如對(duì)于llava-1.5-7b，本文的方法在8比特下達(dá)到最高的CIDEr 分?jǐn)?shù) 1.105，與全精度持平，并在1比特下提升至1.109，超過(guò)了VLCache（1.053）。同樣地，對(duì)于InternVL-2.5-26B，本文的方法在4比特和2比特下分別取得了最高的CIDEr分?jǐn)?shù)1.32和1.313，均優(yōu)于VLCache和KIVI。

Runtime分析

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

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

總結(jié)

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

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

論文標(biāo)題：CalibQuant:1-Bit KV Cache Quantization for Multimodal LLMs
論文地址：https://arxiv.org/abs/2502.14882
代碼地址：https://github.com/insuhan/calibquant