北大華為聯(lián)手推出KV cache管理新方式，推理速度比前SOTA提升4.7倍！

大模型處理長(zhǎng)序列時(shí)，KV cache的內(nèi)存占用隨序列長(zhǎng)度線性增長(zhǎng)，已成為制約模型部署的嚴(yán)峻瓶頸。

為此，來自北京大學(xué)與華為的研究團(tuán)隊(duì)聯(lián)合提出了LouisKV——一個(gè)專為長(zhǎng)輸入、長(zhǎng)輸出等各類長(zhǎng)序列場(chǎng)景設(shè)計(jì)的高效KV cache 檢索框架。

它通過創(chuàng)新的語(yǔ)義感知檢索策略與解耦的精細(xì)化管理機(jī)制，在幾乎不損失模型精度的前提下，實(shí)現(xiàn)了高達(dá)4.7倍的推理加速，為突破LLM長(zhǎng)序列推理瓶頸提供了全新的解決方案。

關(guān)鍵洞察

傳統(tǒng)上，學(xué)術(shù)界與工業(yè)界提出了多種KV cache優(yōu)化方案，其中KV Cache Retrieval是極具前景的方向之一。

該類方法將完整的KV cache卸載至容量更大的CPU內(nèi)存中，并在推理時(shí)僅將最關(guān)鍵的KV子集檢索回GPU進(jìn)行計(jì)算，從而有效緩解GPU 顯存壓力。

然而，現(xiàn)有的KV retrieval方法仍面臨著效率和精度的雙重瓶頸：

• 現(xiàn)有方法通常在生成每個(gè) token 時(shí)都觸發(fā)一次檢索操作，這引入了重要性評(píng)估的計(jì)算開銷與 CPU-GPU 間的數(shù)據(jù)傳輸開銷。在需要生成數(shù)千甚至數(shù)萬(wàn) token 的長(zhǎng)輸出任務(wù)中，檢索操作帶來的累積開銷尤為突出，導(dǎo)致模型推理效率不高；
• 現(xiàn)有方法普遍采用固定大小的頁(yè)（page）作為檢索的基本單元。這種粗粒度的劃分方式，常常導(dǎo)致被檢索的頁(yè)中僅包含少量真正關(guān)鍵的 KV 條目，而大量無(wú)關(guān)條目占用了寶貴的 CPU-GPU 帶寬和 GPU 上的緩存預(yù)算。這不僅造成了數(shù)據(jù)傳輸?shù)睦速M(fèi)，更重要的是，在有限的預(yù)算下，它擠占了本可以留給其他更關(guān)鍵信息的位置，導(dǎo)致模型推理精度的下降。

為了設(shè)計(jì)更高效的檢索策略，研究團(tuán)隊(duì)首先對(duì)不同長(zhǎng)序列任務(wù)中關(guān)鍵 KV 的訪問模式進(jìn)行實(shí)驗(yàn)分析，得到了兩個(gè)關(guān)鍵洞察。

一是訪問模式的時(shí)序局部性。

該特性表現(xiàn)為，在解碼過程中生成一個(gè)語(yǔ)義連貫的segment時(shí)，segment內(nèi)相鄰token所關(guān)注的關(guān)鍵KV集合高度重疊。

如下圖(a)和(b)左下角的相似度曲線所示，在生成當(dāng)前segment的過程中，相鄰token關(guān)鍵KV集合的Jaccard相似度始終維持在0.8以上。

該現(xiàn)象符合直覺，在數(shù)學(xué)推導(dǎo)的某一步驟中，其內(nèi)部的各個(gè)token會(huì)持續(xù)關(guān)注相同的上文引理或條件。

這一洞察揭示了逐 token 檢索策略的內(nèi)在冗余性——既然模型在同一語(yǔ)義段內(nèi)的關(guān)注點(diǎn)保持穩(wěn)定，頻繁的檢索便非必要。

二是關(guān)鍵KV的分布模式差異性。

該特性指關(guān)鍵KV在長(zhǎng)輸入序列和長(zhǎng)輸出序列中通常表現(xiàn)出差異的分布模式：

• 長(zhǎng)輸入序列中的稀疏分布：在長(zhǎng)文檔問答（圖a）等任務(wù)中，生成答案所需的關(guān)鍵信息在長(zhǎng)篇輸入中呈稀疏、離散的分布狀態(tài)；
• 長(zhǎng)輸出序列中的密集分布：在數(shù)學(xué)推理（圖b）等任務(wù)中，模型的注意力會(huì)高度集中于先前生成的中間步驟，使得關(guān)鍵 KV 在局部區(qū)域內(nèi)呈現(xiàn)密集分布。

這一洞察啟發(fā)我們，傳統(tǒng)的、粗粒度的頁(yè)式KV檢索策略過于粗糙，無(wú)法高效應(yīng)對(duì)輸入輸出序列不同的注意力分布模式。

核心設(shè)計(jì)

基于上述洞察，研究團(tuán)隊(duì)提出了一個(gè)高效的KV cache檢索框架 LouisKV。該框架通過算法與系統(tǒng)的協(xié)同設(shè)計(jì)，解決了現(xiàn)有方法的瓶頸。

其核心包含三大創(chuàng)新。

首先是語(yǔ)義感知的KV檢索策略（Semantic-Aware KV Retrieval），為利用時(shí)序局部性，LouisKV摒棄了“逐token檢索”的低效模式，引入了一種自適應(yīng)的檢索策略。

如下圖(a)所示，該策略通過輕量級(jí)機(jī)制監(jiān)控語(yǔ)義變化。在每個(gè)解碼步，它會(huì)計(jì)算當(dāng)前token與前一token的query向量之間的余弦相似度r。

• 若r高于閾值τ，表明模型關(guān)注點(diǎn)未發(fā)生顯著偏移，此時(shí)不觸發(fā)檢索，直接復(fù)用上一個(gè)token檢索得到的關(guān)鍵KV cache；
• 僅當(dāng)r低于閾值τ，表明出現(xiàn)語(yǔ)義邊界，才觸發(fā)一次檢索操作，從CPU的KV cache pool中加載新的關(guān)鍵KV cache。

該策略的核心優(yōu)勢(shì)是將昂貴的檢索開銷均攤到多個(gè)token的生成過程中，極大地降低計(jì)算與數(shù)據(jù)傳輸帶來的總開銷，顯著提升推理效率。

第二是解耦的細(xì)粒度KV管理方案（Decoupled Fine-grained KV Management），為應(yīng)對(duì)分布差異性，LouisKV為輸入和輸出序列定制了不同的KV管理方式，以實(shí)現(xiàn)更精確的檢索。

• 輸入序列（Prefill Stage）：針對(duì)關(guān)鍵KV稀疏分布的特點(diǎn)，LouisKV采用K-Means聚類。如圖(b)所示，它將語(yǔ)義上相似但物理位置上分散的KV聚合為語(yǔ)義簇（Semantic Clusters）；
• 輸出序列（Decode Stage）：針對(duì)關(guān)鍵KV局部密集的特點(diǎn)，LouisKV將連續(xù)生成的token組織成時(shí)序（Temporal Segments）。這與模型生成連貫推理步驟的行為天然對(duì)齊。

通過這種細(xì)粒度的管理，LouisKV創(chuàng)建的檢索單元（語(yǔ)義簇/時(shí)序段）與模型的實(shí)際注意力模式高度匹配，避免了傳統(tǒng)頁(yè)式管理中大量無(wú)關(guān)KV的冗余傳輸，顯著提升了檢索精度。

最后，為了將算法的理論優(yōu)勢(shì)完全轉(zhuǎn)化為運(yùn)行效率，LouisKV在底層進(jìn)行了內(nèi)核級(jí)系統(tǒng)優(yōu)化（Kernel-Level System Optimization）。

具體實(shí)現(xiàn)上，團(tuán)隊(duì)開發(fā)了定制化的Triton和CUDA內(nèi)核。內(nèi)核專門用于加速框架中的關(guān)鍵計(jì)算密集型操作，包括KV聚類和檢索過程。

通過軟硬件協(xié)同優(yōu)化，LouisKV確保了創(chuàng)新算法能夠高效地在硬件上運(yùn)行，實(shí)現(xiàn)了高吞吐率與低延遲的卓越性能。

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

為了全面驗(yàn)證LouisKV的高效性，研究團(tuán)隊(duì)在多個(gè)主流的長(zhǎng)序列任務(wù)上進(jìn)行了詳盡測(cè)試。

這些任務(wù)涵蓋了長(zhǎng)輸入-短輸出（如文檔問答）、短輸入-長(zhǎng)輸出（如數(shù)學(xué)推理）和長(zhǎng)輸入-長(zhǎng)輸出（如長(zhǎng)文推理）等多種應(yīng)用場(chǎng)景。

實(shí)驗(yàn)結(jié)果表明，LouisKV成功地在推理精度和推理效率之間取得了當(dāng)前最佳的平衡。

推理精度層面，在所有12個(gè)基準(zhǔn)測(cè)試中，LouisKV的性能表現(xiàn)都極其接近將全部KV cache保留在GPU中的FullCache方案（灰色虛線），后者代表了理論上的精度上限。

同時(shí)，無(wú)論是與KV cache dropping方法（如H2O、RaaS），還是與KV cache retrieval方法（如Arkvale、Quest）相比，LouisKV在同等KV cache預(yù)算下均展現(xiàn)出更優(yōu)的推理精度。

這證明了LouisKV的語(yǔ)義感知檢索和細(xì)粒度管理策略能夠精準(zhǔn)地識(shí)別并保留對(duì)模型推理最關(guān)鍵的信息，有效避免了精度損失。

推理效率上，LouisKV在三種典型的輸入輸出場(chǎng)景下表現(xiàn)出卓越的性能。

• 大幅降低延遲：與先進(jìn)的KV檢索方法Arkvale相比，LouisKV實(shí)現(xiàn)了高達(dá)1.4倍至4.7倍的端到端推理加速；
• 支持更大批量：當(dāng)處理大批量任務(wù)時(shí)，F(xiàn)ullCache會(huì)因顯存不足而失效。相比之下，LouisKV能夠在此類高負(fù)載場(chǎng)景下穩(wěn)定運(yùn)行，從而顯著提升了系統(tǒng)的有效吞吐量。

這種顯著的效率提升主要得益于LouisKV對(duì)系統(tǒng)核心開銷的精準(zhǔn)優(yōu)化。相較于Arkvale ，LouisKV大幅降低了數(shù)據(jù)傳輸（Transfer）和重要性評(píng)估（Estimation）帶來的開銷。

論文地址：https://arxiv.org/abs/2510.11292