LaCache: Ladder-Shaped KV Caching for Efficient Long-Context Modeling of Large Language Models

发表时间: 2025-07 · arXiv:2507.14204 (ICML 2025)

原文: https://arxiv.org/abs/2507.14204

Dachuan Shi, Yonggan Fu, Xiangchi Yuan, Zhongzhi Vu, Haoran You, Sixu Li, Xin Dong, Jan Kautz, Pavlo Molchanov, Yingyan (Celine) Lin

速读

一句话结论 本文提出了一种免训练的键值缓存优化方法 LaCache，通过阶梯状的缓存存储模式和迭代压缩机制，在固定显存预算下实现了大语言模型高精度、无内存溢出的连续长上下文生成。

要解决什么问题 大语言模型在自回归解码时通常使用键值缓存来避免重复计算注意力，但其显存开销会随序列长度呈 $\mathcal{O}(T)$ 线性增长，导致处理长序列时极易发生内存溢出。现有的缓存驱逐策略难以兼顾长依赖建模能力和持续生成能力：以 StreamingLLM 为代表的基于近期词的策略仅保留滑动窗口内的最新缓存，虽然显存复杂度降为 $\mathcal{O}(1)$ 从而避免了内存溢出，但会严重损失长上下文任务的精度；以 Quest 为代表的检索路线需要缓存所有历史状态并动态检索，依然面临 $\mathcal{O}(T)$ 的显存占用，最终仍会内存溢出；而以 H2O 为代表的基于重要性的动态驱逐方法，由于高度依赖完整的注意力权重图，无法与 FlashAttention 等不显式计算注意力图的系统级高效推理框架兼容，导致实际设备上的计算吞吐量大幅下降。这个卡点意味着，在真实业务中，工程师往往只能在长文本精度、无限生成不爆显存和推理速度这三者之间做妥协，缺乏一个统一的轻量级解决方案。

怎么做的 LaCache 的核心思路是打破所有网络层都保留同一批词的缓存这一传统假设，让不同层负责维护不同时间段的上下文信息。它由两个关键部件构成。第一个部件是阶梯状键值缓存模式，它在浅层网络中保留较早生成的词的缓存，随着网络层数加深，逐渐将保留重点转移到较新的词上，在二维的缓存矩阵中形成一个阶梯状的保留区域，区域外的缓存则被直接丢弃。这种设计在固定的存储预算下大幅拉长了模型能够捕捉的历史跨度。由于不断扩展的重复模式将覆盖率尽可能均匀地分配给每一层，它提升了整体信息保留的下界，并且因为完全不依赖注意力图，天然兼容 FlashAttention。该阶梯结构的形态由两个超参数定义：跨度 $S$ 代表同一个词的缓存被连续保留的层数，重叠度 $O$ 代表每一层保留的词的数量。在语言建模任务中，经验上将 $S$ 设为模型层数的四分之一，将 $O$ 设为 $S$ 的一半以保证语义连续性。第二个部件是面向无限长生成的迭代压缩机制。当键值缓存达到预设的容量上限时，LaCache 会对已经被压缩过的缓存再次应用阶梯状模式进行二次甚至多次压缩。在这一机制下，越老的词被压缩得越激进，越新的词被压缩得越少，从而在维持固定缓存大小的同时，动态腾出空间给新输入的词，并优先保障了近期信息的完整性。

效果如何 实验在 Llama2、Llama3、SmolLM2 和 LongChat 等多个模型上展开，覆盖了从 1.7B 到 13B 的参数规模，对比基线包括全量缓存、代表近期路线的 StreamingLLM，以及代表重要性路线的 H2O、TOVA、PyramidInfer 和 SnapKV。在长文本建模任务 Wikitext-2 中，使用 512 的缓存预算和 1K 输入长度时，LaCache 在 Llama2-7B-Chat 上的困惑度相比全量缓存仅退化了 5%，而 StreamingLLM 退化了 35%。在极端长文本 PG19 测试中，单张 A100 显卡下全量缓存路线在 160K 长度时就会内存溢出，而 LaCache 能够支持高达 600K 长度的连续生成并保持合理的困惑度。在极端苛刻的 80 个词缓存预算下，LaCache 依然在 Llama3-8B 上跑通并优于基线。在长文本理解任务大海捞针测试中，在 50% 的缓存预算下，LaCache 将 Llama3.2-3B-Instruct-128k 的检索准确率从 StreamingLLM 的 54.54% 提升到了 99.16%。在 RULER 评测的 13 个子任务中，LaCache 的平均准确率比基线高出 5.06%。在 LongBench 评测中，由于兼容 FlashAttention，LaCache 展现出了比 H2O 等基于重要性的方法高得多的吞吐量，实现了精度与速度的最佳折中。作者也指出了该方法的局限性：固定的阶梯状模式可能并非在所有场景下都是理论最优解，且作为一种免训练方法，它缺乏针对特定下游任务进行微调适配的能力。

主要贡献

大型语言模型（LLMs）在处理长上下文和持续生成任务时，由于键值（KV）缓存随序列长度线性增长，面临着严重的内存不足（OOM）瓶颈。现有的KV缓存驱逐策略难以同时兼顾长程建模能力和无OOM的连续生成，例如StreamingLLM为了连续生成牺牲了长上下文准确性，而Quest和H2O为了保持准确性则面临高内存开销或与FlashAttention不兼容导致的计算缓慢问题。

为了应对这些挑战，本文提出了一种名为LaCache的免训练KV缓存优化范式，旨在实现LLMs高效且准确的生成推理。LaCache的核心创新点包括：
1. 阶梯状（Ladder-Shaped）KV缓存模式：不仅在每一层内按顺序（从左到右）存储KV对，还跨层（从浅层到深层）存储KV对。这种结构在固定的存储预算下扩展了捕获长程依赖的跨度，从而增强了模型的长程能力。
2. 迭代压缩（Iterative Compaction）机制：根据Token距离进行动态压缩，逐步对旧缓存进行更深度的压缩，在固定的缓存大小内为新Token腾出空间，从而在受限的缓存预算下实现更有效的连续生成。

背景知识与设计动机

现有方法的局限性。现有高效LLM生成方法在平衡生成准确性和内存效率方面存在局限，这对于避免OOM的连续长上下文生成至关重要。如图1(a)所示，基于最近性（recency-based）的方法（如StreamingLLM【1，Efficient streaming language models with attention sinks，2023，arXiv】）仅在固定长度滑动窗口内保留最新token的KV cache，其内存复杂度为 $\mathcal{O}(1)$，能够支持无限长度生成但会降低生成准确性。如图1(b)所示，基于检索（retrieval-based）的方法（如Quest【2，Quest: Query-aware sparsity for efficient long-context llm inference，2024，arXiv】）存储所有token的完整KV cache，并在生成新token时动态检索最相关的cache以提高计算效率，这种策略能保持高准确率，但由于需要存储整个KV cache，内存复杂度高达 $\mathcal{O}(T)$，在处理长上下文时易导致OOM问题。

LaCache的总体设计动机。鉴于上述两种方法的限制，如图1(c)所示，作者提出了LaCache。为了在保持重要历史信息的同时实现有效的KV压缩，LaCache不采用StreamingLLM在所有层保留相同token集合的策略，而是在较浅层保留早期token的KV状态，并在后续层逐渐将焦点转移到较晚的token上，形成阶梯状（ladder-like）结构。此外，为了支持无限长连续生成且不发生OOM，LaCache结合了迭代压缩策略。具体而言，当KV cache达到容量上限时，对已压缩的KV cache再次应用带有阶梯模式的LaCache。该策略确保旧token信息被进一步渐进压缩，而新输入的token被压缩得较少。

图1：(a) 基于最近性的KV缓存，(b) 基于检索的KV缓存，与 (c) 本文提出的具有阶梯状模式的LaCache的说明性比较。
图1：(a) 基于最近性的KV缓存，(b) 基于检索的KV缓存，与 (c) 本文提出的具有阶梯状模式的LaCache的说明性比较。

方法细节

核心洞察。与StreamingLLM跨所有层维护相同最近token集合的做法不同，作者的核心洞察是：虽然最近token的信息对生成准确性至关重要，但它们的KV cache可以通过较少的层来维护和处理。换言之，不同的层可以维护对应于不同token集合的KV cache。该方法的关键优势在于，在相同的KV cache预算下，可以在KV cache中保留更多的token，从而有效扩大上下文长度并保留更多的过去信息。

阶梯状模式的设计与实现。上述洞察启发了阶梯状KV cache模式的设计。如图1(c)所示，LaCache采用了一种简单而有效的策略来跨不同层缓存不同token的KV状态：首先在较早的层中保留早期token的KV状态，接着在后续层中逐渐将焦点转移到较晚的token上，这与时间和顺序处理的本质相一致。这种方法形成了阶梯状模式，既保证了存储效率，又保留了过去token中的基本信息。具体而言，如图2所示，为了实现由绿色框表示的阶梯状模式，首先丢弃落入该模式之外的KV状态，然后将原始的2D KV cache压缩成具有较小缓存大小的更紧凑结构。

图2：LaCache的KV缓存存储模式图解。LaCache用于将原始的完整KV缓存压缩成紧凑的阶梯状模式，在相同的KV缓存预算下，与StreamingLLM相比，能够存储更长程token的信息。
图2：LaCache的KV缓存存储模式图解。LaCache用于将原始的完整KV缓存压缩成紧凑的阶梯状模式，在相同的KV缓存预算下，与StreamingLLM相比，能够存储更长程token的信息。

阶梯状模式的理论与实证分析。阶梯状模式通过提高信息保留的下界，有效覆盖了潜在的重要token。该方法故意不依赖注意力图来识别重要token，从而避免与现有的高效注意力计算优化（如FlashAttention【3，Flashattention-2: Faster attention with better parallelism and work partitioning，2023，arXiv】）产生冲突。其背后的两个基本原理是：
* 第一，不断扩展重复模式并尽可能均匀地为每一层分配覆盖范围（由阶梯状模式保证），提高了信息保留的下界。因为在最坏情况下，重要token集可能出现在覆盖范围最小的层，不均匀的覆盖策略会导致准确率下降。
* 第二，由于自然语言中相邻token通常具有更高的语义相关性，阶梯状模式为每个保留的缓存段结合了平滑过渡。因此，带有部分重叠的不断扩展的阶梯模式使得旧token能够更平滑地淡出，保持稳定的信息保留。
为了实证验证这些原理的优势，作者在不同的KV cache大小下随机生成了1500多种模式进行探索，并在图3中可视化了实现的困惑度（PPL）与缓存大小之间的权衡。如图所示，LaCache的阶梯状模式位于帕累托最优边界上。

图3：LaCache与超过1500个随机采样的KV缓存模式在PPL和缓存大小之间的权衡可视化。
图3：LaCache与超过1500个随机采样的KV缓存模式在PPL和缓存大小之间的权衡可视化。

实现细节与关键超参数。为了平衡存储效率和生成准确性，需要确保该方法在消除存储的KV状态中的冗余的同时，通过保留足够的历史KV状态来准确保存过去的上下文信息。为了满足这些原则，设计了两个关键因素来实现最佳权衡：
* 跨度 $S$：即跨连续层的跨度，表示用于保存同一token对应KV状态的层数。$S$ 越大，记录同一token上下文的层数越多，上下文保存越好，但存储成本增加。
* 重叠 $O$：即每一层的重叠，表示每一层保留其KV状态的token数量。$O$ 越大，每一层保留的KV状态越多，上下文记录越准确，但存储效率降低。

迭代压缩策略的动机与优势。为了在LLMs中实现无OOM问题的连续生成，即使是无限生成长度，也非常需要保持恒定的KV cache大小。为了实现这一点，需要为LaCache增加一种驱逐机制，当预定义的缓存大小被完全利用时，移除过去token的KV状态。为此，提出了一种迭代压缩策略。该方法的基本原理是：一旦使用LaCache压缩过的KV cache达到满载，再次应用LaCache对其进行进一步压缩。这种方法的优势包括：
* 得益于阶梯状模式设计，在对已压缩的KV cache应用LaCache时，早期的KV cache会被首先丢弃（如图4所示）。这种对早期KV cache使用更大压缩比、对晚期KV cache使用更小压缩比的做法，符合基于最近性方法的原理。
* 从部署角度来看，使用LaCache进行迭代压缩提供了一个统一的解决方案和干净的接口，便于更广泛的使用。

迭代压缩的执行过程。如图4所示，高亮部分展示了部分存储的KV状态在迭代压缩后的变化，而未高亮部分也被其他KV状态占据。当KV cache达到其容量时，LaCache被应用于已存储的KV状态，这些状态在首次进入KV cache时已经被LaCache压缩过。接着，落入阶梯状模式之外的KV状态被丢弃，最后将释放的空间分配给新输入的token，从而实现连续的无限长生成。在图4的第二次迭代中，旧的KV状态被压缩得更多，而新的KV状态被压缩得更少，从而更好地保留了最近的信息。

图4：LaCache迭代压缩图解。引入迭代压缩是为了支持无内存不足的连续生成。一旦KV缓存达到预定义大小，LaCache的阶梯状模式将被应用于已压缩的KV缓存，为新token腾出空间。
图4：LaCache迭代压缩图解。引入迭代压缩是为了支持无内存不足的连续生成。一旦KV缓存达到预定义大小，LaCache的阶梯状模式将被应用于已压缩的KV缓存，为新token腾出空间。

实验环境

实验结果

图5：在拼接的PG19数据集的前十本书（对应600K tokens）上评估LaCache。
图6：在整个拼接的PG19数据集（对应1000万tokens）上评估LaCache。

图7：在单张H200 GPU上评估各方法在LongBench上的得分-吞吐量权衡。
图7：在单张H200 GPU上评估各方法在LongBench上的得分-吞吐量权衡。

图8：在50%缓存预算下，使用Llama3.2-3B-Instruct-128k评估LaCache和StreamingLLM。
图9：在25%缓存预算下，使用LongChat-7b-v1.5-32k评估LaCache和StreamingLLM。

补充细节

跨度 $S$ 的影响。在LongBench等长上下文理解任务中，跨度 $S$ 被设置为一个约等于模型层数乘以整体压缩比的整数，旨在实现均匀的压缩比分布。例如，在50%缓存预算下，将 $S$ 设为模型层数的一半，可以使不同位置的压缩比保持在约50%，从而避免某些位置过度压缩而其他位置压缩不足的情况。对于语言建模任务，根据消融实验的经验结果（如图10所示，使用Llama2-7B-Chat模型、256 KV缓存预算和Wikitext-2数据集），$S$ 被设置为模型层数的1/4。

图10：语言建模超参数的消融研究。图表中报告了困惑度（越低越好）。
图10：语言建模超参数的消融研究。图表中报告了困惑度（越低越好）。

重叠 $O$ 的影响。重叠 $O$ 的选择取决于任务类型。具体而言，较大的 $O$ 允许单个token的信息分布在更多的位置，这更适合需要复杂语义理解和更大全局上下文的任务。相反，较小的重叠将信息集中在较少的位置，更适合答案出现在极窄窗口中的任务。对于语言建模任务，$O$ 设置为 $S$ 的1/2以实现更好的语义连续性。对于长上下文理解任务，较大的重叠一致地提高了需要更多全局信息的合成任务（如PassageCount等）的性能，同时降低了更依赖局部信息的QA任务（如NarrativeQA等）的性能。

结论

本文提出了LaCache，一个新颖、免训练且易于部署的KV缓存优化框架，旨在提高LLMs在长上下文生成任务中的效率和有效性。通过阶梯状KV缓存存储模式和迭代压缩机制，LaCache解决了现有方法的局限性。这些创新使LLMs能够更好地捕获长程依赖，优化内存使用，并在固定的存储限制下维持连续生成。实验结果表明，LaCache在保持高生成质量的同时显著提高了内存效率，在各种基准测试中均优于基线方法。

未来工作展望：虽然阶梯状模式是有效的，但它可能并非在所有场景下都是最优的。未来的工作可以基于“最近token对准确性至关重要，但可以通过较少层处理”的核心洞察，探索多样化的KV存储配置。此外，虽然LaCache采用免训练设置以确保高效部署，但未来可以探索将其扩展以支持微调，使其适应特定任务，并与依赖训练的方法进行性能对比。

引用文献汇总