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

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

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

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

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

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

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

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

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

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

一開(kāi)始，Ada-KV團(tuán)隊(duì)首先思考：

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

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

經(jīng)驗(yàn)性分析

Ada-KV團(tuán)隊(duì)發(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ù)算分配對(duì)于壓縮質(zhì)量的提升有著巨大潛力。

理論性分析

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

如果你對(duì)后續(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