DynamicKV: Task-Aware Adaptive KV Cache Compression for Long Context LLMs
DynamicKV: Task-Aware Adaptive KV Cache Compression for Long Context LLMs
发表时间: 2024-12 · arXiv:2412.14838 (EMNLP 2025 Findings)
原文: https://arxiv.org/abs/2412.14838
作者/机构:Xiabin Zhou, Wenbin Wang, Minyan Zeng, Jiaxian Guo, Xuebo Liu, Li Shen, Min Zhang, Liang Ding (xiabinzhou0625, [email protected])
速读
一句话结论 本文提出了一种名为 DynamicKV 的任务感知自适应 KV Cache 压缩方法,通过在推理的预填充阶段动态调整每一层保留的 token 数量,在仅保留极低比例缓存的情况下依然维持了极高的长上下文任务性能。
要解决什么问题 大语言模型在处理长上下文(如 32K 到 1M 个 token)时,注意力机制的二次复杂度会带来显著的延迟。常见的解决方案是缓存键值状态(KV Cache)以减少重复计算,但这会消耗巨大的显存,例如在 LLaMA2-7B 中缓存 100K 个 token 需要超过 50GB 的显存。为了缓解显存压力,现有的 KV Cache 压缩方法(如丢弃 token、分层金字塔结构或固定长度保留)通常强制所有任务遵循一种固定的跨层缓存保留模式。然而,这种固定模式的卡点在于,它忽略了不同任务在模型各层中截然不同的激活特征。作者通过观察发现,不同任务的信息流聚合方式差异巨大:例如在摘要任务中,模型的高层只需要很小的缓存规模;而在代码补全等复杂的推理任务中,模型的高层反而需要维持更大的缓存规模。如果强行使用固定的跨层缓存分配策略,会导致关键信息的丢失,进而造成模型在特定任务上的精度下降。因此,如何打破固定模式的限制,根据具体任务的独特需求动态定制每一层的缓存保留策略,成为了提升长上下文模型推理效率与效果的核心卡点。
怎么做的 核心思路是摒弃跨层固定的缓存分配,转而根据当前任务的注意力分布特征,动态且自适应地计算并分配每一层所需的 KV Cache 预算。这种方法之所以能绕开固定模式的卡点,是因为它让显存预算严格对齐了特定任务的实际信息流需求。关键设计由动态预算分配和渐进式缓存更新两个部件构成。首先是动态预算分配部件,它负责初步筛选重要信息。传统方法通常优先保留包含最新上下文的当前窗口 token(窗口大小设为 $ws$)。对于窗口之外的 token,该部件会设定一个单层最大缓存保留预算 $B^l = (wt - ws) \times r_{max}$,其中 $wt$ 是考虑的总 token 数,$r_{max}$ 是缩放比例。接着,利用多头池化操作计算当前窗口内的局部注意力分数 $A^l = Pooling(A[:, ws])$,并据此选出得分最高的 $B^l$ 个 token 进行保留,形成初步压缩的缓存:
$$KV_{retained}^l = KV^l[argtopK(A^l, B^l)]$$其次是渐进式缓存更新部件,它负责进一步压缩中间层的缓存使用量。该部件将模型划分为包含 $m$ 个连续层的块。在处理每个块时,利用计数函数 $\Phi$ 统计分配给每一层的前 K 个注意力条目,计算出每层的保留计数 $C^l$:
效果如何 实验基于 LlaMA-3-8B-Instruct、Qwen-2-7B-Instruct、Mistral-7B-Instruct-v0.2 和 InternLM-2.5-7B-Chat-1M 四个模型搭建,上下文长度覆盖 8K 到 1M。评估基准包含 LongBench 的 16 个数据集(涵盖单/多文档问答、摘要、少样本学习、合成任务和代码补全)以及大海捞针测试。对比基线点名了四个代表性固定模式方法:代表注意力漏槽与滚动缓存路线的 StreamingLLM、代表重击者预言机驱逐路线的 H2O、代表按注意力头聚类选择路线的 SnapKV,以及代表单调递减金字塔模式路线的 PyramidKV。量化结果显示,在 LongBench 上将总 KV Cache 大小严格限制为 512(仅保留 6.9% 的 token)时,DynamicKV 击败了所有基线,在不同模型上分别维持了全量缓存 97%、96%、96% 和 89% 的性能。在 Mistral-7B 设定的 32K 上下文大海捞针测试中,即使在极端的 64 缓存大小(0.9% 压缩率)下,DynamicKV 依然保持了 90% 的原始性能,准确率比上述基线方法高出 11% 到 57%。在计算开销方面,面对 32K 输入和 8K 输出的设置,DynamicKV 相比全量缓存将每个输出 token 的时间提升了 129%,端到端延迟降低了 56%,同时节省了 11.2% 的显存。作者也承认了该方法的局限性:受限于计算预算,目前仅在 8B 规模及以下的模型上进行了验证,尚未在 70B 等超大模型或更前沿的架构上测试;此外,尽管在预填充阶段修改注意力计算带来的额外开销很小(因为计算仅在极小的窗口内进行),但这确实限制了该阶段对 FlashAttention 的兼容性。
主要贡献
在大型语言模型(LLMs)中,高效的键值(KV)缓存管理对于检索增强生成(RAG)和文本摘要等长上下文任务至关重要。随着LLMs扩展以处理超长上下文(如LlaMA3处理32K Tokens,InternLM处理1M Tokens),注意力机制的二次复杂度带来了显著的延迟。尽管缓存KV状态可以减少计算量,但这带来了高昂的内存成本(例如,在LLaMA2-7B中缓存100K Tokens仍需要超过50GB的内存)。现有的KV缓存压缩方法通常强制执行固定的模式(如分层金字塔结构或跨层固定长度),忽略了特定任务的特征,从而减少了关键信息的保留。
本文的核心目标是系统地研究不同任务下KV缓存压缩的设计原则,并提出一种自适应策略。作者观察到,在不同的任务中,不同层之间的激活模式存在显著差异,这突出了为每个任务的独特需求量身定制自适应策略的必要性。
基于上述观察,本文提出了以下主要贡献:
1. 揭示了任务类型对Token保留的影响:探讨了不同任务类型对LLM各层Token保留的影响,发现对于不同的任务,每层所需的Token保留量各不相同,因此有必要针对不同任务动态选择每层的Token保留量。
2. 提出DynamicKV方法:基于上述发现,提出了一种新颖的任务感知自适应KV缓存压缩方法——DynamicKV。该方法通过在推理的预填充(prefill)阶段动态调整每层保留的Token数量来优化Token保留。DynamicKV建立了全局和单层的最大KV缓存预算,临时保留当前层的最大预算,并定期更新所有前面层的KV缓存大小。
3. 卓越的压缩性能与实验结果:在广泛使用的长上下文理解基准测试LongBench上,DynamicKV仅保留了1.7%的KV缓存大小,就实现了完整KV缓存(Full KV)约90%的性能。值得注意的是,即使在极端的压缩率(0.9%)下,使用Mistral7B-Instruct-v0.2的DynamicKV在“大海捞针”(Needle-in-a-Haystack)测试中也比最先进的(SOTA)方法高出11%。
背景知识与设计原则
LLMs中注意力的潜在模式:Transformer架构通过对输入的逐层细化推动了NLP的进步。研究表明,中间层对于简单任务已经足够,这挑战了全深度推理的必要性。同时,交叉层注意力(CLA)和Ada-KV等工作通过分析跨层注意力模式来优化缓存大小和驱逐策略,突出了注意力动态与内存高效计算之间的相互作用。
KV缓存压缩中的Token丢弃策略:Token丢弃策略侧重于识别有影响力的Token。基于注意力的方法(如FastGen和Scissorhands)使用注意力模式进行修剪。内存感知方法(如StreamingLLM)通过注意力汇聚(attention sinks)优先考虑流式处理,而H2O则使用累积注意力评分进行贪婪驱逐。分层方法(如PyramidKV)按层自适应但缺乏泛化性。SnapKV通过为每个头选择关键位置提供与任务无关的压缩。
不同任务对跨层注意力机制的细粒度分析:为了系统地研究LLMs在长上下文输入下跨层的注意力机制,我们对四种不同类型的任务进行了细粒度分析:单文档和多文档问答(QA)、摘要、合成任务和代码完成。我们特别关注LlaMA模型,计算最近Token与所有其他Token之间的平均注意力分数,并识别出所有层中注意力最高的前 $128 \times k$(层数)个Token。现有方法在不同任务中使用固定模式,但LLMs根据任务使用不同的层,这表明KV缓存压缩过程中的Token保留也可能因任务而异。
跨层Token保留的观察与洞察:我们使用箱线图直观地呈现了四种不同类型的任务在不同层上的分布。我们发现不同的任务在固定层上显示出显著不同的Token保留率。例如,在早期层,分布范围很广,表明特定于任务的变异很大。为了进一步理解分布,我们可视化了每个任务在所有层上的Token保留率。我们发现:合成任务在早期层显示出较高的保留率;代码完成任务在早期层以及最后三层显示出较高的保留率;而文档QA和摘要表现出与其他任务不同的保留动态。由此得出的洞察是:每层要保留的Token应根据任务类型进行动态自适应。
方法细节
动态层自适应KV缓存压缩框架:以前关于KV缓存压缩的工作通常在LLM的各层中分配固定的KV缓存大小。然而,正如我们的分析所证明的,不同类型任务在不同层上的注意力模式并不相同。因此,在不同任务上跨层使用固定的KV缓存大小可能会导致次优性能。为此,我们提出了DynamicKV——一种动态的、层自适应的KV缓存压缩方法。DynamicKV包含两个步骤:动态预算分配和渐进式缓存更新。
当前窗口优先保留策略与最大预算定义:传统的Token丢弃方法通常优先考虑最近的Token,因为它们通常携带生成下一个输出最相关的上下文。我们将这组Token称为当前窗口,其大小记为 $ws$。窗口内的Token被赋予最高的保留优先级。为了有效地管理内存,我们首先定义每层的最大KV缓存保留预算,记为 $B^l$,计算公式为 $B^l = (wt - ws) \times r_{max}$,其中 $r_{max}$ 是缩放比例,$wt$ 是考虑的Token总数。
基于注意力的剩余Token选择:遵循Li等人【28,Snapkv: Llm knows what you are looking for before generation+2024+arXiv】的方法,我们指导选择当前窗口之外的剩余Token,这是基于它们相对于指令Token的注意力分数。具有较高注意力分数的Token被认为更相关,因此被优先保留在GPU缓存中。在一个标准的LLM中,注意力的计算方式为:
$$ A = softmax(\frac{Q \cdot K^T}{\sqrt{d_k}}) $$其中 $Q \in \mathbb{R}^{M \times d_k}$ 和 $K \in \mathbb{R}^{M \times d_k}$ 分别是查询和键矩阵,$d_k$ 是键/查询的维度,$M$ 是序列长度。
每层注意力分数池化与KV状态压缩:受Li等人【28,Snapkv: Llm knows what you are looking for before generation+2024+arXiv】和Zhang等人【43,Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling+2024+arXiv】的启发,我们使用多头池化操作在当前窗口上计算每层注意力分数 $A^l$:
$$ A^l = Pooling(A[:, ws]) $$然后,我们根据 $A^l$ 中的最高值选择前 $B^l$ 个Token。保留这些位置对应的KV状态以形成压缩缓存:
$$ KV_{retained}^l = KV^l[argtopK(A^l, B^l)] $$分块与初始保留阈值确定:为了进一步减少中间层的KV缓存使用,我们将模型划分为包含 $m$ 个连续层的块。对于每个这样的块,我们动态确定满足累积保留需求所需的最小初始保留阈值,同时刷新历史KV缓存。在每个 $m$ 层块的末尾,我们对保留分数进行归一化,以优先考虑操作上关键的Token。这个过程产生了一个特定层的预算分配 $Z^\prime$,它促进了缓存在各层之间高效且自适应的分配。
特定层预算分配与保留计数计算:具体来说,我们应用 top-K 选择来保留这些层中最相关的Token,并使用计数函数 $\Phi$ 计算每层的保留计数:
$$ C^l = Norm(\frac{1}{n} \cdot \Phi(TopK(A, (wt - ws) \times n)) $$其中 $n$ 是迄今为止处理的渐进更新层数,$(wt - ws)$ 表示当前窗口外的Token数量。
临时预算计算与全局归一化:接下来,我们通过相对于最大值缩放每层的保留分数来计算临时预算 $Z$:
$$ Z = \left[ \frac{B^l \times t}{max(C^l)} | t \in C^l \right] $$其中 $B^l$ 是每层保留预算。然后,为了确保遵守总预算 $B = (wt - ws) \times L$,跨层对其进行归一化:
$$ Z^\prime = \left[ k \cdot \frac{B}{\sum Z} | k \in Z \right] $$渐进更新的实际执行机制:在实践中,在前 $m$ 层的渐进更新期间,该机制使用注意力分数 $A$ 来估计每层需要保留的最佳Token数量。函数 $\Phi$ 计算分配给每层的 top-K 注意力条目,形成 $C^l$,然后将其归一化为 $Z$。最后,预算 $Z^\prime$ 控制着如何为每一层细化KV缓存,从而在不同的层之间实现自适应且有效的压缩策略。上述过程可以表示为算法1。
1: Input: 初始预算K/V缓存列表 K^b, V^b,最大比例 r_{max},更新间隔 m,平均Token长度 wt,窗口大小 ws,序列长度 S,头部维度 d_k,窗口大小的输入嵌入 X^{ws} \in \mathbb{R}^{ws * d_k},由窗口Token和其他Token计算的初始预算注意力列表 A^b
2: Output: 压缩后的 K/V 缓存 K^c, V^c
3: B^l = (wt - ws) * r_{max}
4: def Update_Buffer_Length(A, l):
5: Agather = cat([A for l in (1, l)], 0).view(-1)
6: cnts = Count_Elements(topk(Agather, k = (wt - ws) * H * l).indices / (L * S)) / l
7: 计算 cnts 的归一化,范围在 (0, 1)
8: Z = [int((B^l * t / max(norm))) for t in norm]
9: r = sum(Z) / ((wt - ws) * L)
10: Z' = [int(k/r) for k in Z]
11: Return Z'
12: for l = 1 to L do
13: 计算完整的KV状态 K^s, V^s
14: for h = 1 to H do
15: /* 计算窗口大小Token与所有其他Token之间的注意力 */
16: A_h^l = softmax((X^{ws} W_h^Q) \cdot K_h^T).mean(dim=-2).pooling(dim=-1)
17: end for
18: 将 A^l 追加到 A^b /* 当前 A_l 的形状是 [H, S] */
19: /* 计算当前层缓冲区KV缓存 */
20: indices = A^l.topk(B^l, dim=-1).indices.unsqueeze(-1).expand(-1, -1, d_k)
21: K_l^b = cat((K^s[:, :-ws, :].gather(dim=-2, indices), K^s[:, -ws:, :]), dim=-2)
22: V_l^b = cat((V^s[:, :-ws, :].gather(dim=-2, indices), V^s[:, -ws:, :]), dim=-2)
23: /* 逐步压缩 */
24: if l % m == 0 then
25: Z' = Update_Buffer_Length(A^l, l)
26: /* 更新缓冲区K/V缓存 */
27: for i = 1 to l do
28: K_i^b = cat((K_l^b[:, :Z'_i, :], K_l^b[:, -ws:, :]), dim=-2)
29: V_i^b = cat((V_l^b[:, :Z'_i, :], V_l^b[:, -ws:, :]), dim=-2)
30: end for
31: end if
32: end for
33: 从 K^b, V^b 更新 K/V 缓存 K^c, V^c
推理阶段的阶段划分与操作:在推理期间,该过程被分为两个阶段:预填充阶段(prefilling phase)和解码阶段(decoding phase),这与现有的推理引擎保持一致。我们的DynamicKV虽然在确定每层最佳KV缓存大小时可能会遇到特定样本的注意力模式,但它是在预填充阶段执行此步骤的。在解码阶段,不应用任何修改。
批处理推理的兼容性处理:Q1:DynamicKV处理批处理推理吗?A1:是的。事实上,现代LLM推理和服务引擎(例如,vLLM)通常在预填充阶段单独处理样本(即,批大小 $\rho = 1$),而解码阶段允许在批次中进行高效的并行计算。由于我们的DynamicKV在解码期间不引入任何修改,因此我们的方法与现有的推理引擎无缝对齐,确保解码阶段与批处理执行完全兼容,以实现高吞吐量生成。
FlashAttention的兼容性分析:Q2:DynamicKV如何与FlashAttention兼容?A2:我们的DynamicKV在解码阶段可以与FlashAttention兼容。虽然我们的DynamicKV在预填充阶段修改了注意力分数的计算,这限制了与FlashAttention的兼容性,但它仍然非常高效。这是因为注意力仅在一个非常小的窗口大小 $ws$ 内计算,其中 $ws \ll M$,即使没有FlashAttention也能将开销保持在最低水平。相反,在解码阶段不应用任何修改,我们利用FlashAttention显著提高计算效率。
实验环境
-
数据集:
- LongBench:用于评估LLMs的上下文理解能力,包含16个英语数据集(如NarrativeQA, Qasper, HotpotQA, GovReport, QMSum, MultiNews, TREC, TriviaQA, SAMSum, LCC, RepoBench-P等),涵盖单/多文档QA、摘要、少样本学习、合成任务和代码完成。
- Needle In A Haystack:用于评估跨长上下文的事实检索能力,上下文长度扩展至32K Tokens。
-
模型架构关键参数:
- LlaMA-3-8B-Instruct:隐层大小4096,32层,32个查询头,8个KV头,词表大小128,256,支持8K上下文。
- Mistral-7B-Instruct-v0.2:隐层大小4096,32层,32个查询头,8个KV头,词表大小32,000,支持32K上下文。
- Qwen2-7B-Instruct:隐层大小3584,28层,28个查询头,4个KV头,词表大小151,646,支持32K上下文。
- InternLM2.5-7B-Chat-1M:隐层大小4096,32层,32个查询头,8个KV头,词表大小92,544,支持1M上下文。
-
基线方法:StreamingLLM, H2O, SnapKV, PyramidKV。
实验结果
-
LongBench上的比较实验 (KV缓存大小=512):
- 实验内容:在KV缓存总大小限制为512的情况下,评估StreamingLLM、H2O、SnapKV、PyramidKV和DynamicKV相对于FullKV的性能保留情况。
- 实验结果:DynamicKV始终优于现有方法,即使在极低的缓存上下文比率(仅6.9%)下运行也是如此。在可比较的模型中,DynamicKV分别比表现最好的基线高出0.43%、0.19%、0.69%和0.53%,分别保留了FullKV性能的97%、96%、96%和89%(详见表1)。
- 分析结论:这些结果强调了DynamicKV在极端内存限制下保持接近FullKV级别性能的卓越能力。此外,DynamicKV不仅在代码完成等复杂任务上匹配并增强了PyramidKV的能力,而且显著提高了较低缓存容量下的性能上限。
-
消融研究 (不同KV缓存大小的影响):
- 实验内容:研究DynamicKV机制在不同键值缓存大小(64到1024)下的性能。
- 实验结果:随着所有评估模型的缓存大小增加,性能持续改善。对于LlaMA-3-8B-Instruct,随着KV缓存大小从64增加到1024,性能指标从34.93提高到41.22(详见图4)。
- 分析结论:这些发现强调了DynamicKV缓存在利用KV缓存压缩以维持长上下文功能方面的有效性。较大的缓存容量通常与卓越的性能相关,但在实际应用中需要平衡存储和计算资源。
-
计算开销分析:
- 实验内容:在LongBench上使用Llama模型比较DynamicKV与FullKV的计算开销,评估指标包括首次Token时间(TTFT)、每个输出Token时间(TPOT)、端到端延迟和GPU内存使用量。
- 实验结果:与FullKV相比,DynamicKV的TPOT提高了129%,延迟降低了56%,内存使用量从30.48GB降至27.06GB(详见表2)。
- 分析结论:实验结果表明,我们的DynamicKV在计算效率和内存使用方面都提供了显著的优势。
-
大海捞针 (Needle-in-Haystack) 任务上的可视化:
- 实验内容:使用Mistral-7B-Instruct-v0.2模型,上下文长度扩展至32K Tokens,在64的KV缓存大小下比较多种压缩策略。
- 实验结果:即使在极端的压缩下,DynamicKV也保留了模型原始性能的90%,比竞争方法(StreamingLLM, PyramidKV等)实现了57%、37%、41%和11%的准确率提升。在长达7K Tokens的上下文长度下,DynamicKV的极端压缩几乎实现了完全的准确性,超出此范围后继续显著优于所有基线(详见图5)。
- 分析结论:这些结果强调了DynamicKV在分层Token选择方面的卓越能力,并验证了我们的假设,即关键Token跨层的分布本质上是动态的。
结论
本文研究了处理长上下文输入的LLMs中特定任务的注意力模式,并发现了跨任务的不同注意力分布。为了解决这个问题,提出了一种层自适应KV缓存压缩框架DynamicKV,该框架动态优化每层的KV缓存分配。通过在LongBench基准的16个数据集上的实验评估,证明了其广泛的适用性和性能优势。主要结论如下:(1) 在复杂推理任务(如代码完成)中遵循波浪形模式;(2) 在合成和摘要任务中遵循金字塔形模式;(3) 动态分层自适应的DynamicKV方法能够根据不同任务制定相对合适的KV缓存保留策略。特别是在保持极小KV缓存大小的情况下,效果显著增强。未来的工作希望能找到一种更合适的方法,在不增加计算量的情况下执行KV缓存压缩。
局限性:受限于计算预算,仅在规模有限的模型(如7B/8B)上验证了DynamicKV。将其应用于超大模型(如70B)和更前沿的模型架构将更具说服力。此外,虽然在多种任务上进行了实验,但DynamicKV对其他任务或数据集的泛化能力尚未完全探索。
附录细节补充
大海捞针任务上的性能对比:如表5所示,我们比较了各种KV缓存压缩方法(StreamingLLM、H2O、SnapKV、PyramidKV和DynamicKV)在LlaMA-3-8B-Instruct和Qwen2-7B-Instruct两个模型上的Needle in a Haystack任务性能。在这两个模型中,我们的DynamicKV都取得了最高性能,LlaMA-3-8B-Instruct得分为0.9,Qwen-2-7B-Instruct得分为0.87。这些结果突出了DynamicKV在长上下文场景中保留任务关键信息的卓越能力。
不同输入输出长度下的效率评估:我们在不同的输入/输出长度下,针对标准方法(FullKV)评估了DynamicKV的效率。所有实验均在固定上下文窗口($m = 128$)下进行,测量了首次Token时间(TTFT)、每个输出Token时间(TPOT)、端到端延迟和GPU内存使用情况。
短序列场景的效率表现:在短序列($8k / 2k$)下:DynamicKV将TPOT提高了22.5%(27.63到33.85 tok/s),同时TTFT略微增加了6%(0.66s到0.70s),实现了总延迟降低18.2%(74.79s到61.21s),并减少了638MB的内存。结果表明,我们的DynamicKV始终实现与以前方法具有竞争力或更优越的性能。虽然FullKV产生了最高的平均分数,但DynamicKV在几个模型中实现了最佳或接近最佳的性能,特别是在Mistral-7B-Instruct-v0.2和InternLM-2.5-Chat-1M上表现出色,证明了有效的内存压缩且精度损失最小。
长序列场景的效率表现:在长序列($32k / 8k$)下:优势显著放大,DynamicKV提供了129%的更高TPOT(11.65到26.69 tok/s),降低了56%的延迟(706.56s到310.56s),并节省了11.2%的内存(31213MB到27713MB)。
可扩展性分析:FullKV表现出超线性的TPOT退化(在32k输入时为11.65 tok/s),而DynamicKV通过按需计算保持了稳定的吞吐量,证明了对长上下文生成具有更好的适应性。实验表明,动态KV缓存以微小的初始延迟为代价,换取了显著更好的持续生成速度和内存效率,这对长文本生成任务(输出Token $> 2k$)特别有利。
KV缓存大小为128时的LongBench性能比较:表7展示了在固定缓存大小为128的情况下,不同KV缓存压缩方法(StreamingLLM、H2O、SnapKV、PyramidKV和我们的DynamicKV)在LongBench上的性能比较。我们跨越了单文档QA、多文档QA、摘要、少样本学习、合成任务和代码完成等各种任务进行了实验,结果再次印证了DynamicKV在极低缓存预算下的优越性。
💬 评论讨论
欢迎在这里分享您的想法和见解!