改進(jìn)KV緩存壓縮，大模型推理顯存瓶頸迎來新突破——

中科大研究團(tuán)隊提出Ada-KV，通過自適應(yīng)預(yù)算分配算法來優(yōu)化KV緩存的驅(qū)逐過程，以提高推理效率。

打破KV Cache壓縮將所有注意力頭分配相同壓縮預(yù)算的常規(guī)做法，針對不同的注意力頭進(jìn)行適配性壓縮預(yù)算分配

展開來說，由于大模型在自回歸生成過程中，每生成一個新token都需要將對應(yīng)的KV矩陣存儲下來，這導(dǎo)致緩存隨著生成序列長度的增加而急劇膨脹，引發(fā)內(nèi)存和I/O延遲問題，尤其在長序列推理中尤為突出。

因此，KV緩存壓縮成為了一項必要的優(yōu)化。

不過令人頭禿的是，現(xiàn)有壓縮方法往往在各個注意力頭之間平均分配預(yù)算，未能考慮其特性差異。

而中科大團(tuán)隊在注意到——不同注意力頭關(guān)注度存在差異后，對其進(jìn)行適配性壓縮預(yù)算分配，通過精細(xì)化運(yùn)作帶來更高的壓縮質(zhì)量。

相關(guān)研究不僅在學(xué)術(shù)界引起討論，更實現(xiàn)了工業(yè)界開源落地。

例如，Cloudflare workers AI團(tuán)隊進(jìn)一步將其改進(jìn)落地于工業(yè)部署常用的vLLM框架中，并發(fā)布技術(shù)報告，開源全部代碼。

KV緩存壓縮從均勻性預(yù)算分配→適配性預(yù)算分配

一開始，Ada-KV團(tuán)隊首先思考：

注意力頭間的適配性壓縮預(yù)算分配是必要的嗎？

通過從經(jīng)驗性和理論性兩個角度進(jìn)行分析后，團(tuán)隊的回答是：yes!

經(jīng)驗性分析

Ada-KV團(tuán)隊發(fā)現(xiàn)，在大模型中注意力頭之間存在著顯著不同的關(guān)注集中度差異：

大部分注意力頭關(guān)注度集中在少量KV cache上，只需很少的KV cache（例如，1%）就可以幾乎收集接近0.9的注意力權(quán)重；

而少數(shù)注意力頭則傾向于分散注意力，往往需要接近50%的KV Cache才能夠?qū)⒆⒁饬?quán)重聚集到0.9。

考慮到如此巨大的關(guān)注度集中度的差異，注意力頭間的適配性壓縮預(yù)算分配對于壓縮質(zhì)量的提升有著巨大潛力。

理論性分析

Ada-KV研究團(tuán)隊進(jìn)一步從壓縮輸出損失的角度出發(fā)，形式化了在不同分配策略下KV Cache壓縮對注意力輸出的損失影響：

他們基于這一理論提出了一種以注意力權(quán)重為基礎(chǔ)的自適應(yīng)分配方案，并發(fā)現(xiàn)這種跨注意力頭的預(yù)算分配策略始終能夠降低損失上界。

此外，這一理論上的更低損失上界在實際實驗中也展現(xiàn)出更低的注意力輸出損失：

理論與實際結(jié)果一致驗證了這一結(jié)論：注意力頭間的適配性預(yù)算分配能夠顯著提升KV緩存壓縮的效果。

通過適配性頭間預(yù)算分配增強(qiáng)KV Cache壓縮質(zhì)量

作者將Ada-KV這一適配性預(yù)算分配策略結(jié)合到現(xiàn)有的兩個領(lǐng)先的Cache壓縮方案：SnapKV和PyramidKV中，分別得到兩種適配性壓縮方案：Ada-SnapKV和Ada-Pyramid。

他們進(jìn)一步在廣泛使用的長序列開源大模型Mistral-7B-Instruct-32K和LWM-Text-Chat-1M和長文本任務(wù)評估基準(zhǔn)LongBench上的16個數(shù)據(jù)集上進(jìn)行了充分的評估。

實驗結(jié)果顯示，所有適配性預(yù)算分配增強(qiáng)的壓縮方法（Ada-SnapKV和Ada-Pyramid）全部優(yōu)于原有的均勻預(yù)算分配壓縮方法（SnapKV和Pyramid）。

Ada-KV團(tuán)隊在算法實現(xiàn)的同時，也考慮到了執(zhí)行效率的優(yōu)化。

他們開發(fā)了一種展平的KV Cache管理布局，并定制了CUDA kernel，以實現(xiàn)高效的Cache更新管理。

結(jié)合Flash Attention技術(shù)，該方案在適應(yīng)性預(yù)算分配的情況下，實現(xiàn)了高效推理，并在相同預(yù)算下保持了與先前Cache壓縮方案一致的計算效率。

目前，代碼已在GitHub上完全開源，助力推動注意力頭間適應(yīng)性壓縮預(yù)算分配的研究。

Cloudflare推動Ada-KV于工業(yè)界部署落地

Cloudflare公司旗下的Workers AI團(tuán)隊針對實際并發(fā)服務(wù)場景中存在的內(nèi)存碎片問題，基于Paged Attention重新實現(xiàn)了Ada-KV算法，并將其落地于實際部署使用的推理框架vLLM中。

他們發(fā)布了技術(shù)報告，對該方案進(jìn)行了詳細(xì)評估，同時開源了相關(guān)代碼，助力Ada-KV在工業(yè)界的快速應(yīng)用和落地。

如果你對后續(xù)進(jìn)展感興趣，歡迎持續(xù)關(guān)注~

Ada-KV Paper：https://arxiv.org/abs/2407.11550
Ada-KV Code：https://github.com/FFY0/AdaKV
Cloudflare Technical Report：https://arxiv.org/abs/2410.00161
Cloudflare Code：https://github.com/IsaacRe/vllm-kvcompress