使用 KVPress 掌握 LLM 中的长上下文

社区文章发布于 2025 年 1 月 23 日

什么是 KV 缓存，它为何重要？

问题：KV 缓存及其线性扩展的负担
KV 缓存的大小

KVPress：KV 缓存压缩工具包

KVPress 实战

简而言之：KVPress 包含了最新的 KV 缓存压缩技术，实现了内存高效的长上下文 LLM。🚀

大型语言模型 (LLM) 的关键特性之一是其上下文窗口——它们在单个请求中可以处理的最大标记数量。随着 LLM 的发展，其上下文窗口变得越来越大。

更大的上下文窗口解锁了惊人的可能性

上下文内检索：在单个查询中无缝引用大量文本。
上下文内学习：在同一会话中根据特定示例调整行为。
扩展推理：处理非常长的思维链而不会中断上下文。

然而，这些扩展窗口也带来了成本——KV 缓存中长上下文占用的内存变得难以管理。例如，使用 Llama 3-70B 在 float16 精度下处理 1M 标记需要 330GB 的 KV 缓存，这使得许多应用程序无法实现。

在这篇博客文章中，我们将讨论解决此问题的一种方案：压缩 KV 缓存以实现更高效的生成。为此，我们将探讨

KV 缓存是什么以及它为何重要。
KVPress，NVIDIA 设计的一个功能强大的工具包，旨在有效压缩 KV 缓存。
KVPress 的内部工作原理以及它如何实现压缩。

在开始之前，请在这个空间中探索 KVPress（如果需要，末尾会找到示例）

什么是 KV 缓存，它为何重要？

图 1：注意力模块中的键值缓存（来源：NVIDIA）

在自回归模型中，文本生成是逐标记进行的，每个预测都依赖于所有先前的标记作为上下文。例如

为了生成标记 1000，模型必须考虑标记 1 到 999 的表示。
为了生成标记 1001，必须再次处理相同的信息（标记 1 到 999），以及标记 1000。

随着序列的增长，这种重复计算变得低效，特别是对于大型模型。KV 缓存通过存储注意力层中的中间结果——键 (K) 和值 (V)——来优化此过程，因此模型可以重复使用它们用于未来的标记，而不是重新计算它们。

问题：KV 缓存及其线性扩展的负担

尽管 KV 缓存功能强大，但它有一个主要缺点——它随上下文窗口的大小线性扩展。虽然这听起来一开始并不令人担忧，但让我们分解一下，看看为什么这会成为一个严重的瓶颈。

KV 缓存的大小

KV 缓存中存储的值来自模型使用的所有注意力块。因此，其大小取决于模型架构，这决定了注意力头的数量。更具体地说，KV 缓存消耗的内存由以下方程确定

$\text{Size}(\text{KV}) = 2 \times \text{precision} \times n_{layers} \times n_{heads} \times d \times n_{tokens}$

这些因素中的每一个都导致内存使用量爆炸式增长。为了更具体地说明这一点，让我们考虑一个具体示例——Llama 3-70B 在 bfloat16 精度下运行（由模型作者推荐），上下文大小为 1M 标记

$\text{Size}(\text{KV}) = 2 \times 2 \times 80 \times 8 \times 128 \times 1M = 327.6 \text{GB}$

由于 bfloat16 每个参数使用 2 字节，因此模型权重本身就需要 140 GB（70B x 2 字节）。这意味着，以 1M 标记上下文大小运行模型需要大约 470 GB 的内存，其中 KV 缓存就占了惊人的 70%。

KVPress：KV 缓存压缩工具包

正如我们所见，KV 缓存既是关键的推动者，也是部署具有长上下文窗口的大型语言模型 (LLM) 的重要瓶颈。解决线性扩展内存问题需要创新的压缩技术，而这正是 KVPress 的用武之地。

KVPress 由 NVIDIA 开发，是一个 Python 工具包，旨在通过提供一套最先进的压缩技术来解决大型 KV 缓存的内存挑战。它还与其他方法集成，例如 KV 缓存量化，这是一种内置于 transformers 库中的方法，用于减少内存使用（上述方程中的精度项），进一步扩展其用途（此处有详细信息）。

对于专注于压缩的研究人员，KVPress 提供了一个灵活的模块化框架，使其易于理解和扩展新方法。对于开发人员，KVPress 简化了部署这些尖端技术的过程，从而能够快速高效地集成到实际应用程序中。

KVPress 实战

KVPress 的核心是利用 压榨器，这是一种专门设计用于减少 KV 缓存内存占用的高级压缩算法。

许多压榨器依赖于在每个头中使用的分数来修剪重要性最低的 KV 对。例如，KnormPress 修剪键值范数最低的 KV 对（论文），SnapKVPress 修剪与最新查询的低注意力权重相关的 KV 对（论文）。

这些压榨器通过前向钩子无缝集成到模型的注意力层中。

图 2：KV 压缩可视化（来源：NVIDIA）

在文本生成过程中，它们动态压缩 KV 缓存，在不影响模型生成连贯准确输出的能力的情况下减少内存使用。每个压榨器都具有 compression_ratio 属性，该属性决定了应用于 KV 缓存的压缩程度。

这些压榨器与自定义的 transformers 管道无缝集成，从而便于应用和实验。

以下是如何使用 KVPress 中的一种压榨器 ExpectedAttentionPress。此压榨器会修剪与未来查询的最低预期注意力权重相关的 KV 对。

from transformers import pipeline
from kvpress import ExpectedAttentionPress

pipe = pipeline(
"kv-press-text-generation",
model="meta-llama/Llama-3.1-8B-Instruct",
device="cuda",
model_kwargs={"attn_implementation": "sdpa"}
)

context = "A very long text you want to compress once and for all"
question = "\nA question about the compressed context"  # optional

press = ExpectedAttentionPress(compression_ratio=0.5)
answer = pipe(context, question=question, press=press)["answer"]

尝试直接在此 Hugging Face space 或此 Google colab notebook 中使用它！

通过针对预填充阶段，KVPress 确保在缓存最大时对其进行压缩，从而有助于减少数万甚至数百万个标记序列的内存开销。

下图展示了随着提示长度的增加，KVPress 压缩实现的 GPU 内存节省。对于较短的提示，大部分内存分配给模型权重——Llama 3.1 8B 在 bfloat16 中大约需要 15GB。然而，随着提示长度的增长，KV 缓存成为内存消耗的主要因素。对于 128k 上下文长度，应用 50% 压缩比的 KVPress 可将峰值内存使用量从 45GB 减少到 37GB。较小的 KV 缓存还提高了解码速度，在 A100 GPU 上从每秒 11 个标记提高到每秒 17 个标记（来源）。

图 3：内存使用量与上下文长度（来源：NVIDIA）

基准测试

研究社区一直积极开发各种 KV 缓存压缩技术。KVPress 鼓励研究人员贡献他们的方法，并且已经提供了十多种压榨器。

为了评估这些压榨器的性能，KVPress 包含一个简单的 CLI，用于在 RULER、InfiniteBench 和 Loogle 等标准长上下文数据集上对其进行基准测试。下图展示了在 RULER 数据集上，针对 4k 上下文长度和不同压缩比的 9 种不同压榨器的基准测试结果。在此数据集上表现最佳的压榨器是 AdaKVPress（论文）和 ExpectedAttentionPress 的组合，后者是 KVPress 作者创建的一种新的未发表的剪枝技术（更多信息在此处）。

图 4：平均分数与压缩比（来源：NVIDIA）