Key-Value (KV)緩存已成為大語(yǔ)言模型(LLM)長(zhǎng)文本處理的關(guān)鍵性能瓶頸。當(dāng)前研究尚未充分關(guān)注解碼階段的優(yōu)化，這一階段具有同等重要性，因?yàn)椋?/span>

1、對(duì)需要完整上下文的場(chǎng)景，預(yù)填充階段的過(guò)度壓縮會(huì)顯著降低模型的推理理解能力

2、在長(zhǎng)輸出推理任務(wù)中存在重要特征的顯著偏移現(xiàn)象

這篇論文提出SCOPE框架，通過(guò)分離預(yù)填充與解碼階段的KV緩存優(yōu)化策略，實(shí)現(xiàn)高效的緩存管理。該框架保留預(yù)填充階段的關(guān)鍵KV緩存信息，同時(shí)引入基于滑動(dòng)窗口的新型策略，用于解碼階段重要特征的高效選取。

關(guān)鍵發(fā)現(xiàn)

KV緩存的推理過(guò)程分析

LLM的每次請(qǐng)求包含兩個(gè)獨(dú)立階段：預(yù)填充階段處理完整輸入提示以生成初始輸出標(biāo)記；解碼階段通過(guò)迭代方式逐個(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é)果表明：對(duì)于依賴完整上下文的推理任務(wù)，預(yù)填充階段的過(guò)度壓縮會(huì)導(dǎo)致顯著的性能損失

解碼階段分析

下圖描述了在解碼步驟1、300和500時(shí)，基于前15%注意力分?jǐn)?shù)選取的重要特征在0、13和31層的位置分布：

觀察發(fā)現(xiàn)，各層中保留的重要特征主要源自解碼階段生成的KV緩存，在長(zhǎng)輸出任務(wù)中，隨著輸出序列延長(zhǎng)，特征偏移現(xiàn)象愈發(fā)顯著。這一現(xiàn)象要求在保持預(yù)填充階段識(shí)別的重要特征的同時(shí)，對(duì)解碼階段新出現(xiàn)的重要特征實(shí)施有效管理。數(shù)據(jù)表明，長(zhǎng)文本生成的解碼階段采用貪婪算法可能引發(fā)重要特征的顯著偏移

KV緩存預(yù)算分配優(yōu)化

下圖呈現(xiàn)了LONGGENBENCH中GSM8k+樣本第13層的注意力熱力圖，以及注意力分?jǐn)?shù)與生成標(biāo)記位置的對(duì)應(yīng)關(guān)系：

圖中最左側(cè)和最右側(cè)分別對(duì)應(yīng)預(yù)填充和解碼階段。對(duì)于需要并行處理多個(gè)問(wèn)題的推理任務(wù)，準(zhǔn)確定位當(dāng)前預(yù)測(cè)位置至關(guān)重要。如圖所示，這些關(guān)鍵信息可通過(guò)貪婪算法識(shí)別的重要特征有效捕獲。預(yù)填充和解碼階段的KV緩存預(yù)算需要獨(dú)立分配以實(shí)現(xiàn)最優(yōu)性能。這一發(fā)現(xiàn)啟發(fā)了SCOPE框架的設(shè)計(jì)理念：通過(guò)解耦預(yù)填充和解碼階段的壓縮過(guò)程，實(shí)現(xiàn)KV緩存預(yù)算的精確分配，在保留預(yù)填充階段全部KV緩存的同時(shí)，優(yōu)化緩存預(yù)算的重分配效率

方法論

KV緩存壓縮機(jī)制重構(gòu)

初始化過(guò)程

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對(duì)表示為KPVP
• 注意力權(quán)重AttP通過(guò)P和KPVP計(jì)算得出，預(yù)填充階段的壓縮表達(dá)式：

其中：

·表示張量連接操作

函數(shù)Ψα1(AttP)從AttP[: ?α2]中選擇具有最高Top-α1注意力權(quán)重的KV緩存

解碼階段實(shí)現(xiàn)

解碼階段復(fù)用預(yù)填充階段的KV緩存，并通過(guò)序列生成過(guò)程持續(xù)更新，對(duì)于時(shí)間步t處的新標(biāo)記張量Xt,t∈{1,T }，鍵值計(jì)算如下：

過(guò)程中，KtVt與緩存池Φ中的歷史KV緩存進(jìn)行連接，形成當(dāng)前時(shí)刻的KV對(duì)。隨后與查詢張量Xt計(jì)算得到注意力權(quán)重Attt

• 本方法與現(xiàn)有KV壓縮方案的核心區(qū)別在于Φp和Φd在緩存池Φ中的動(dòng)態(tài)分配機(jī)制

SCOPE框架詳解

針對(duì)解碼階段設(shè)計(jì)了三種優(yōu)化策略：滑動(dòng)窗口、自適應(yīng)調(diào)整和不連續(xù)更新，這些策略均專注于Φd的優(yōu)化

滑動(dòng)窗口機(jī)制

通過(guò)動(dòng)態(tài)調(diào)整解碼必要?dú)v史窗口β1和解碼局部窗口β2實(shí)現(xiàn)解碼階段的KV緩存壓縮：

• β1負(fù)責(zé)捕獲當(dāng)前預(yù)測(cè)位置的上下文信息
• β2保存與歷史標(biāo)記高度相關(guān)的全局特征

滑動(dòng)策略在t > M + β1 + β2時(shí)觸發(fā)，通過(guò)函數(shù)Ψβ1 (Attt[α1 + α2 : ?β2])實(shí)現(xiàn)對(duì)Φd的定向更新，同時(shí)保持Φp不變。該機(jī)制通過(guò)限制選擇函數(shù)Ψ的操作范圍（從α1 +α2開(kāi)始的Attt），有效規(guī)避了預(yù)填充階段注意力權(quán)重的干擾

自適應(yīng)調(diào)整策略

針對(duì)解碼生成標(biāo)記長(zhǎng)度較短的情況，引入動(dòng)態(tài)調(diào)整機(jī)制，避免過(guò)長(zhǎng)的解碼必要?dú)v史窗口β1導(dǎo)致的資源浪費(fèi)。設(shè)計(jì)了基于時(shí)間步驟t和最大長(zhǎng)度T的自適應(yīng)函數(shù)，用于動(dòng)態(tài)調(diào)整β1長(zhǎng)度，其中T ? β1 + β2。窗口大小從基準(zhǔn)值β2起始，隨時(shí)間步t線性增長(zhǎng)：

特征：

• 當(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操作帶來(lái)顯著性能開(kāi)銷?；谶B續(xù)查詢傾向于選擇相似鍵的特性，提出不連續(xù)更新策略：

1、將Top-K選擇操作ΨK(Att)的執(zhí)行頻率優(yōu)化為每T?β2/β1間隔執(zhí)行一次ζ

2、相比于傳統(tǒng)的逐步執(zhí)行方式，顯著降低了計(jì)算開(kāi)銷

實(shí)驗(yàn)評(píng)估

實(shí)驗(yàn)環(huán)境配置

基于兩個(gè)主流開(kāi)源大語(yǔ)言模型構(gòu)建實(shí)驗(yàn)平臺(tái)：

• LLaMA-3.1–8B-Instruct
• Mistral-7B-Instruct-v0.3

Φp參數(shù)配置：

• LongGenBench-4K場(chǎng)景：α1 + α2 = 2048（約占輸入長(zhǎng)度60%）
• LongGenBench-8K場(chǎng)景：α1 + α2 = 4096（約占輸入長(zhǎng)度60%）
• α2統(tǒng)一設(shè)置為8
• β1 + β2參數(shù)設(shè)置：

• 4K輸出配置：512
• 8K輸出配置：1024
• β2固定為256，用于適配答案中的推理鏈（Chain-of-Thought）長(zhǎng)度，避免序列過(guò)短引發(fā)的性能退化

性能評(píng)估結(jié)果

基準(zhǔn)系統(tǒng)對(duì)比分析

下表展示了在LONGGENBENCH基準(zhǔn)測(cè)試中，基于LLaMA-3.1–8B-Instruct的SCOPE三種策略與現(xiàn)有方法的性能對(duì)比：

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

SCOPE框架的三種策略在所有解碼壓縮方案中均實(shí)現(xiàn)最優(yōu)性能，其中針對(duì)內(nèi)存使用和傳輸優(yōu)化的不連續(xù)策略表現(xiàn)尤為突出。在高難度GSM8K+/GSM8K++任務(wù)上，SCOPE通過(guò)保留預(yù)填充階段KV緩存的策略展現(xiàn)顯著優(yōu)勢(shì)，而其他壓縮方法出現(xiàn)明顯的性能衰減。PyramidInfer與H2O的性能差異不顯著，表明在長(zhǎng)序列輸出任務(wù)中，層間稀疏性特征的影響相對(duì)有限

預(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的變體）通過(guò)跨層預(yù)算重分配的嘗試未能帶來(lái)顯著改善，驗(yàn)證了前期實(shí)證研究結(jié)果
• 預(yù)填充階段保留的KV緩存表現(xiàn)出類似StreamingLLM中"注意力匯聚點(diǎn)"的特性

深入分析

必要特征損失的緩解效果

H2O等統(tǒng)一壓縮方法由于重要特征偏移導(dǎo)致關(guān)鍵KV緩存損失，影響了上下文理解能力，下圖描述了預(yù)測(cè)位置與模型性能的關(guān)聯(lián)關(guān)系：

數(shù)據(jù)顯示H2O在后期預(yù)測(cè)中出現(xiàn)顯著性能下降，而SCOPE的三種策略有效緩解了這一問(wèn)題，證實(shí)了預(yù)填充KV緩存保留策略的有效性

關(guān)鍵參數(shù)影響分析

解碼階段的兩個(gè)核心參數(shù)：

• KV緩存預(yù)算 β1+β2
• 特征選擇算法 ΨK(Att)

下圖展示了基于兩種主流top-K選擇算法對(duì)預(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)效率評(píng)估

基于滑動(dòng)窗口策略的自適應(yīng)和不連續(xù)優(yōu)化顯著提升了內(nèi)存效率，具體數(shù)據(jù)如下表：

相較于完整緩存和單一預(yù)填充壓縮方案，本文提出的方法和統(tǒng)一壓縮策略通過(guò)降低KV緩存存儲(chǔ)量，有效緩解了內(nèi)存壓力，自適應(yīng)策略通過(guò)動(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框架通過(guò)分離預(yù)填充和解碼階段實(shí)現(xiàn)長(zhǎng)文本生成任務(wù)的優(yōu)化，在兩個(gè)階段均采用Top-K算法進(jìn)行特征選擇。當(dāng)前框架存在以下待改進(jìn)方向：

1、預(yù)填充階段仍沿用傳統(tǒng)top-K算法，未來(lái)研究可探索分塊處理或其他新型技術(shù)，以提升歷史標(biāo)記的估計(jì)精度

2、解碼階段每步執(zhí)行Top-K操作導(dǎo)致頻繁的GPU I/O開(kāi)銷。雖然不連續(xù)策略在一定程度上優(yōu)化了操作頻率，但仍可通過(guò)減少I/O數(shù)據(jù)量進(jìn)一步降低系統(tǒng)延遲

3、當(dāng)前SCOPE主要驗(yàn)證了文本模態(tài)的長(zhǎng)輸出任務(wù)優(yōu)化效果?？蚣芫哂袛U(kuò)展到視覺(jué)領(lǐng)域的潛力，特別是在多圖像生成等需要大量KV緩存的任務(wù)場(chǎng)景

總結(jié)

SCOPE框架針對(duì)LLM長(zhǎng)文本生成中的KV緩存優(yōu)化問(wèn)題提供了系統(tǒng)性解決方案。

通過(guò)實(shí)驗(yàn)觀察發(fā)現(xiàn)兩個(gè)關(guān)鍵問(wèn)題：

預(yù)填充階段的過(guò)度壓縮對(duì)推理能力造成顯著影響，解碼過(guò)程中存在重要特征的偏移現(xiàn)象。

SCOPE通過(guò)以下機(jī)制解決上述問(wèn)題：

保持預(yù)填充階段必要的KV緩存完整性；引入滑動(dòng)窗口策略實(shí)現(xiàn)解碼階段KV緩存的高效管理。

大規(guī)模實(shí)驗(yàn)驗(yàn)證表明：

SCOPE僅使用35%原始內(nèi)存即可達(dá)到接近完整KV緩存的性能水平，并且保持了與現(xiàn)有預(yù)填充壓縮方法的良好兼容性。