SCOPE:面向大語言模型長序列生成的雙階段KV緩存優(yōu)化框架
Key-Value (KV)緩存已成為大語言模型(LLM)長文本處理的關(guān)鍵性能瓶頸。當(dāng)前研究尚未充分關(guān)注解碼階段的優(yōu)化,這一階段具有同等重要性,因?yàn)椋?/span>
1、對需要完整上下文的場景,預(yù)填充階段的過度壓縮會(huì)顯著降低模型的推理理解能力
2、在長輸出推理任務(wù)中存在重要特征的顯著偏移現(xiàn)象
這篇論文提出SCOPE框架,通過分離預(yù)填充與解碼階段的KV緩存優(yōu)化策略,實(shí)現(xiàn)高效的緩存管理。該框架保留預(yù)填充階段的關(guān)鍵KV緩存信息,同時(shí)引入基于滑動(dòng)窗口的新型策略,用于解碼階段重要特征的高效選取。
關(guān)鍵發(fā)現(xiàn)
KV緩存的推理過程分析
LLM的每次請求包含兩個(gè)獨(dú)立階段:預(yù)填充階段處理完整輸入提示以生成初始輸出標(biāo)記;解碼階段通過迭代方式逐個(gè)生成后續(xù)輸出標(biāo)記。
預(yù)填充階段分析
下圖展示了在完整解碼緩存條件下,三個(gè)任務(wù)在不同預(yù)填充階段壓縮比率下的性能表現(xiàn):

Long-Bench中的PassageRetrieval-en和HotpotQA任務(wù)表現(xiàn)出顯著的壓縮容忍度:即使在20%的預(yù)填充階段壓縮比率下,模型仍能保持與完整緩存相近的性能水平,證實(shí)了模型在高壓縮率下維持上下文理解能力的魯棒性。然而在LONGGENBENCH的GSM8k+任務(wù)中,同樣的20%壓縮率導(dǎo)致準(zhǔn)確率下降約95%
實(shí)驗(yàn)結(jié)果表明:對于依賴完整上下文的推理任務(wù),預(yù)填充階段的過度壓縮會(huì)導(dǎo)致顯著的性能損失
解碼階段分析
下圖描述了在解碼步驟1、300和500時(shí),基于前15%注意力分?jǐn)?shù)選取的重要特征在0、13和31層的位置分布:

觀察發(fā)現(xiàn),各層中保留的重要特征主要源自解碼階段生成的KV緩存,在長輸出任務(wù)中,隨著輸出序列延長,特征偏移現(xiàn)象愈發(fā)顯著。這一現(xiàn)象要求在保持預(yù)填充階段識別的重要特征的同時(shí),對解碼階段新出現(xiàn)的重要特征實(shí)施有效管理。數(shù)據(jù)表明,長文本生成的解碼階段采用貪婪算法可能引發(fā)重要特征的顯著偏移
KV緩存預(yù)算分配優(yōu)化
下圖呈現(xiàn)了LONGGENBENCH中GSM8k+樣本第13層的注意力熱力圖,以及注意力分?jǐn)?shù)與生成標(biāo)記位置的對應(yīng)關(guān)系:

圖中最左側(cè)和最右側(cè)分別對應(yīng)預(yù)填充和解碼階段。對于需要并行處理多個(gè)問題的推理任務(wù),準(zhǔn)確定位當(dāng)前預(yù)測位置至關(guān)重要。如圖所示,這些關(guān)鍵信息可通過貪婪算法識別的重要特征有效捕獲。預(yù)填充和解碼階段的KV緩存預(yù)算需要獨(dú)立分配以實(shí)現(xiàn)最優(yōu)性能。這一發(fā)現(xiàn)啟發(fā)了SCOPE框架的設(shè)計(jì)理念:通過解耦預(yù)填充和解碼階段的壓縮過程,實(shí)現(xiàn)KV緩存預(yù)算的精確分配,在保留預(yù)填充階段全部KV緩存的同時(shí),優(yōu)化緩存預(yù)算的重分配效率
方法論
KV緩存壓縮機(jī)制重構(gòu)
初始化過程
KV緩存壓縮的核心在于基于預(yù)設(shè)緩存預(yù)算進(jìn)行動(dòng)態(tài)調(diào)整。本文構(gòu)建了緩存池Φ,其包含兩個(gè)子集:存儲(chǔ)預(yù)填充階段KV緩存的Φp和存儲(chǔ)解碼階段KV緩存的Φd。緩存池在時(shí)間步t實(shí)時(shí)更新,記為Φt。采用廣泛驗(yàn)證的貪婪算法函數(shù)ΨK(Att)從給定注意力權(quán)重Att中篩選Top-K KV緩存
預(yù)填充階段設(shè)計(jì)
定義輸入提示張量P ∈ RM×D,其結(jié)構(gòu)為P = {P1,P2, . . . ,PM},其中:
- Pi表示第i個(gè)標(biāo)記的嵌入向量
 - M為輸入標(biāo)記數(shù)量
 - D為模型隱藏維度
 
鍵值張量的計(jì)算公式如下:

其中:
- WK,WV ∈ RD×D分別為鍵和值的投影權(quán)重矩陣
 - KV對表示為KPVP
 - 注意力權(quán)重AttP通過P和KPVP計(jì)算得出,預(yù)填充階段的壓縮表達(dá)式:
 

其中:
·表示張量連接操作
函數(shù)Ψα1(AttP)從AttP[: ?α2]中選擇具有最高Top-α1注意力權(quán)重的KV緩存
解碼階段實(shí)現(xiàn)
解碼階段復(fù)用預(yù)填充階段的KV緩存,并通過序列生成過程持續(xù)更新,對于時(shí)間步t處的新標(biāo)記張量Xt,t∈{1,T },鍵值計(jì)算如下:

過程中,KtVt與緩存池Φ中的歷史KV緩存進(jìn)行連接,形成當(dāng)前時(shí)刻的KV對。隨后與查詢張量Xt計(jì)算得到注意力權(quán)重Attt
- 本方法與現(xiàn)有KV壓縮方案的核心區(qū)別在于Φp和Φd在緩存池Φ中的動(dòng)態(tài)分配機(jī)制
 
SCOPE框架詳解
針對解碼階段設(shè)計(jì)了三種優(yōu)化策略:滑動(dòng)窗口、自適應(yīng)調(diào)整和不連續(xù)更新,這些策略均專注于Φd的優(yōu)化

滑動(dòng)窗口機(jī)制
通過動(dòng)態(tài)調(diào)整解碼必要?dú)v史窗口β1和解碼局部窗口β2實(shí)現(xiàn)解碼階段的KV緩存壓縮:
- β1負(fù)責(zé)捕獲當(dāng)前預(yù)測位置的上下文信息
 - β2保存與歷史標(biāo)記高度相關(guān)的全局特征
 
滑動(dòng)策略在t > M + β1 + β2時(shí)觸發(fā),通過函數(shù)Ψβ1 (Attt[α1 + α2 : ?β2])實(shí)現(xiàn)對Φd的定向更新,同時(shí)保持Φp不變。該機(jī)制通過限制選擇函數(shù)Ψ的操作范圍(從α1 +α2開始的Attt),有效規(guī)避了預(yù)填充階段注意力權(quán)重的干擾
自適應(yīng)調(diào)整策略
針對解碼生成標(biāo)記長度較短的情況,引入動(dòng)態(tài)調(diào)整機(jī)制,避免過長的解碼必要?dú)v史窗口β1導(dǎo)致的資源浪費(fèi)。設(shè)計(jì)了基于時(shí)間步驟t和最大長度T的自適應(yīng)函數(shù),用于動(dòng)態(tài)調(diào)整β1長度,其中T ? β1 + β2。窗口大小從基準(zhǔn)值β2起始,隨時(shí)間步t線性增長:

特征:
- 當(dāng)t < T時(shí),Φdt的預(yù)算規(guī)模為β2+(t?β2)·β1/ T?β2
 - 該設(shè)計(jì)有效優(yōu)化內(nèi)存利用率,因比率(t?β2 (T?β2)恒小于1
 - t達(dá)到T時(shí),Φdt規(guī)模收斂于β1 + β2
 - 此調(diào)整機(jī)制與LLM的自回歸編碼特性高度契合,在不引入額外超參數(shù)的前提下提升資源利用效率
 
不連續(xù)更新機(jī)制
傳統(tǒng)策略中,Top-K選擇操作ΨK(Att)需執(zhí)行T ? β2次,這種頻繁的GPU I/O操作帶來顯著性能開銷?;谶B續(xù)查詢傾向于選擇相似鍵的特性,提出不連續(xù)更新策略:
1、將Top-K選擇操作ΨK(Att)的執(zhí)行頻率優(yōu)化為每T?β2/β1間隔執(zhí)行一次ζ
2、相比于傳統(tǒng)的逐步執(zhí)行方式,顯著降低了計(jì)算開銷
實(shí)驗(yàn)評估
實(shí)驗(yàn)環(huán)境配置
基于兩個(gè)主流開源大語言模型構(gòu)建實(shí)驗(yàn)平臺:
- LLaMA-3.1–8B-Instruct
 - Mistral-7B-Instruct-v0.3
 
Φp參數(shù)配置:
- LongGenBench-4K場景:α1 + α2 = 2048(約占輸入長度60%)
 - LongGenBench-8K場景:α1 + α2 = 4096(約占輸入長度60%)
 - α2統(tǒng)一設(shè)置為8
 - β1 + β2參數(shù)設(shè)置:
 
- 4K輸出配置:512
 - 8K輸出配置:1024
 - β2固定為256,用于適配答案中的推理鏈(Chain-of-Thought)長度,避免序列過短引發(fā)的性能退化
 
性能評估結(jié)果
基準(zhǔn)系統(tǒng)對比分析
下表展示了在LONGGENBENCH基準(zhǔn)測試中,基于LLaMA-3.1–8B-Instruct的SCOPE三種策略與現(xiàn)有方法的性能對比:

實(shí)驗(yàn)結(jié)果分析:
SCOPE框架的三種策略在所有解碼壓縮方案中均實(shí)現(xiàn)最優(yōu)性能,其中針對內(nèi)存使用和傳輸優(yōu)化的不連續(xù)策略表現(xiàn)尤為突出。在高難度GSM8K+/GSM8K++任務(wù)上,SCOPE通過保留預(yù)填充階段KV緩存的策略展現(xiàn)顯著優(yōu)勢,而其他壓縮方法出現(xiàn)明顯的性能衰減。PyramidInfer與H2O的性能差異不顯著,表明在長序列輸出任務(wù)中,層間稀疏性特征的影響相對有限
預(yù)填充方法集成驗(yàn)證
下表呈現(xiàn)了LLaMA3.1-8B在LONGGENBENCH-4K的GSM8K+任務(wù)上的方法集成實(shí)驗(yàn)數(shù)據(jù):

關(guān)鍵發(fā)現(xiàn):
- 部分優(yōu)化策略在僅使用65%原始KV緩存的情況下,實(shí)現(xiàn)了超越完整緩存的性能表現(xiàn)
 - PyramidKV(SnapKV的變體)通過跨層預(yù)算重分配的嘗試未能帶來顯著改善,驗(yàn)證了前期實(shí)證研究結(jié)果
 - 預(yù)填充階段保留的KV緩存表現(xiàn)出類似StreamingLLM中"注意力匯聚點(diǎn)"的特性
 
深入分析
必要特征損失的緩解效果
H2O等統(tǒng)一壓縮方法由于重要特征偏移導(dǎo)致關(guān)鍵KV緩存損失,影響了上下文理解能力,下圖描述了預(yù)測位置與模型性能的關(guān)聯(lián)關(guān)系:

數(shù)據(jù)顯示H2O在后期預(yù)測中出現(xiàn)顯著性能下降,而SCOPE的三種策略有效緩解了這一問題,證實(shí)了預(yù)填充KV緩存保留策略的有效性
關(guān)鍵參數(shù)影響分析
解碼階段的兩個(gè)核心參數(shù):
- KV緩存預(yù)算 β1+β2
 - 特征選擇算法 ΨK(Att)
 
下圖展示了基于兩種主流top-K選擇算法對預(yù)算參數(shù)進(jìn)行縮放的實(shí)驗(yàn)結(jié)果:

與預(yù)填充階段相比,解碼階段展現(xiàn)出更強(qiáng)的壓縮容忍度:25%的壓縮率僅導(dǎo)致15%的性能降低,而類似壓縮率在預(yù)填充階段會(huì)導(dǎo)致GSM8k+任務(wù)性能的顯著下降
系統(tǒng)效率評估
基于滑動(dòng)窗口策略的自適應(yīng)和不連續(xù)優(yōu)化顯著提升了內(nèi)存效率,具體數(shù)據(jù)如下表:

相較于完整緩存和單一預(yù)填充壓縮方案,本文提出的方法和統(tǒng)一壓縮策略通過降低KV緩存存儲(chǔ)量,有效緩解了內(nèi)存壓力,自適應(yīng)策略通過動(dòng)態(tài)預(yù)算調(diào)整機(jī)制進(jìn)一步優(yōu)化了系統(tǒng)性能。
泛化能力驗(yàn)證
下圖展示了在β1 + β2 = 512配置下,∞BENCH中En.Sum任務(wù)的實(shí)驗(yàn)結(jié)果:

實(shí)驗(yàn)表明SCOPE的三種策略均優(yōu)于完整緩存配置,充分驗(yàn)證了該框架的泛化能力
局限性分析
SCOPE框架通過分離預(yù)填充和解碼階段實(shí)現(xiàn)長文本生成任務(wù)的優(yōu)化,在兩個(gè)階段均采用Top-K算法進(jìn)行特征選擇。當(dāng)前框架存在以下待改進(jìn)方向:
1、預(yù)填充階段仍沿用傳統(tǒng)top-K算法,未來研究可探索分塊處理或其他新型技術(shù),以提升歷史標(biāo)記的估計(jì)精度
2、解碼階段每步執(zhí)行Top-K操作導(dǎo)致頻繁的GPU I/O開銷。雖然不連續(xù)策略在一定程度上優(yōu)化了操作頻率,但仍可通過減少I/O數(shù)據(jù)量進(jìn)一步降低系統(tǒng)延遲
3、當(dāng)前SCOPE主要驗(yàn)證了文本模態(tài)的長輸出任務(wù)優(yōu)化效果??蚣芫哂袛U(kuò)展到視覺領(lǐng)域的潛力,特別是在多圖像生成等需要大量KV緩存的任務(wù)場景
總結(jié)
SCOPE框架針對LLM長文本生成中的KV緩存優(yōu)化問題提供了系統(tǒng)性解決方案。
通過實(shí)驗(yàn)觀察發(fā)現(xiàn)兩個(gè)關(guān)鍵問題:
預(yù)填充階段的過度壓縮對推理能力造成顯著影響,解碼過程中存在重要特征的偏移現(xiàn)象。
SCOPE通過以下機(jī)制解決上述問題:
保持預(yù)填充階段必要的KV緩存完整性;引入滑動(dòng)窗口策略實(shí)現(xiàn)解碼階段KV緩存的高效管理。
大規(guī)模實(shí)驗(yàn)驗(yàn)證表明:
SCOPE僅使用35%原始內(nèi)存即可達(dá)到接近完整KV緩存的性能水平,并且保持了與現(xiàn)有預(yù)填充壓縮方法的良好兼容性。















 
 
 



















 
 
 
 