引言
本文将探讨KV-Cache如何通过在内存使用和计算时间之间进行巧妙的权衡，使像ChatGPT和DeepSeek这样的语言模型在生成文本时更快。
总结11篇最近的研究论文，归纳三大类： token选择、后处理压缩技术和架构重新设计 。包括DeepSeek的多头潜在注意力（MLA），这些论文在这一基本思想的基础上，进一步提高了大型语言模型（LLM）推理的时间效率。
一、思考
为什么文本生成如此缓慢
让我们从一个简单的类比开始。想象你在写一个故事，每写一个新词，你都需要重新阅读到目前为止的整个故事以保持一致性。故事越长，重新阅读的时间就越长。这正是大型语言模型在文本生成时所面临的问题。
自注意力的基本构建块
现代语言模型的核心是一种称为自注意力的机制。对于一个由
个标记（大致对应单词）组成的序列，每个标记都需要“查看”或“关注”所有其他标记以理解上下文。
这种查看一切的过程的计算成本随着序列长度的增长而增长：
- 对于
个标记，每个标记都需要查看所有nn个标记
- 这意味着成本与
成正比
- 用数学符号表示，我们将其写为
的复杂度
真正的问题：一次生成一个标记
当语言模型生成文本时，它一次生成一个标记，这就是事情变得计算密集的地方：
- 第一个标记：查看1个标记（成本：
）
- 第二个标记：查看2个标记（成本：
）
- 第三个标记：查看3个标记（成本：
）
- 以此类推，直到第
个标记：查看
个标记（成本：
）
如果我们将生成长度为
的序列的所有这些成本加起来，我们得到：
这种
的成本意味着随着文本的增长，生成时间会极其迅速地增长。例如，生成两倍长的序列大约需要八倍的时间！显然，我们需要一个更好的方法。
解决方案：键值（KV）缓存
KV 缓存背后的关键是，我们正在做大量冗余工作。在生成每个新标记时，我们会重新计算之前已经处理过的所有先前标记。让我们看看如何解决这个问题。
什么是键值缓存？
可以将 KV 缓存想象成一个智能记事本，我们会在第一次看到每个 token 时记下有关它的重要信息。对于每个 token，我们计算并存储两件事：
- 键（k）：可以将其视为一种寻址机制——它有助于确定此标记与未来标记的相关性
- 值（v）：可以将其视为当此标记被发现相关时实际使用的信息
从数学上，我们计算这些为：
- 键：
（其中
是标记，
是一个学习到的变换）
- 值：
（其中
是另一个学习到的变换）
在生成一个新标记时，我们使用它的查询（计算方式类似于键）通过将其与所有存储的键进行比较来在我们的缓存中找到相关信息。然后使用匹配的值来帮助生成标记。
KV缓存如何加速
有了KV缓存，处理过程变得更加高效：
- 当我们遇到一个新token时，只需要计算它的key和value一次
- 对于所有后续的token，我们可以直接从缓存中查找这些预计算的值
- 这意味着每个新token只需要做少量新的计算，而不是重新做所有之前的计算
显然有一个权衡：
- 我们需要更多的内存来存储所有的keys和values。对于一个具有：
- 层
- 注意力头
- 序列长度
- key/value维度
，总的内存开销为
值（这个2是因为需要存储keys和values）。这会随着序列长度
线性增长
，但对于大模型来说，常数因子可能非常大。
- 但作为回报，我们将计算成本从
降低到
。
要理解为什么是
，让我们看一下每一步的成本：
- 第一步：处理一个token
成
本 2. 第二步：处理一个新token + 查找1个缓存的token
成
本 3. 第三步：处理一个新token + 查找2个缓存的token
成
本 4. 依此类推...
将这些加起来：
这相比
是一个显著的改进！虽然我们仍然需要做查看所有前面的tokens的基础工作
，但我们避免了每一步都进行昂贵的重新计算。
内存挑战：为什么我们需要更好的解决方案
虽然KV缓存是一个强大的优化手段，但它伴随着显著的内存开销。让我们通过一个具体的例子来看看，使用像Llama3 70B这样的现代大语言模型：
- 层
- 注意力头
- 批量大小为8个序列
- key/value维度
- 16位精度
处理一个批量（8个序列，每个序列1000个token）所需的内存为：
字节
字节
GB
这种巨大的内存使用带来了几个挑战：
- 随着序列长度线性增长
- 与批量大小成倍增长，支持并行处理
- 限制了我们可以处理的最大上下文长度
- 限制了在内存受限设备上的部署
这些挑战激发了研究界的一波创新，导致了各种优化KV缓存使用的技术。接下来，将探讨这些前沿的解决方案。
二、如何改善传统的KV缓存？
以下论文代表了KV缓存优化的关键创新。我们将通过三大主要方法来探索它们：token选择、后处理压缩技术和架构重设计。
2.1 Token 选择和修剪方法（Token Selection and Pruning Approaches）
- Heavy-Hitter Oracle (H2O)
H2O 引入了在KV缓存中识别和保留重要token的概念：
- 重型Token（Heavy-Hitter Tokens） ：H2O 识别在生成过程中具有最高累计注意力分数的token，这些token遵循幂律分布。这些token对于模型的功能至关重要，因此在缓存中优先处理。
- 动态次模撤销（Dynamic Submodular Eviction） ：该方法将缓存管理问题框架化为一个优化问题，目标函数为次模函数
，用于量化token集合
的重要性：
其中
是token
的累计注意力分数。缓存
通过以下方式更新：
确保每次最多只移除一个token。这个贪心算法在计算上高效，并在次模约束下保证接近最优的性能。
- 结果 ：通过该方法，KV缓存大小减少了5倍，几乎没有精度损失，并且吞吐量提升了高达29倍。
- StreamLLM
- 作者观察到 注意力汇聚（Attention Sinks） 现象：解码过程中，初始token充当自然的注意力锚点。
- 如果没有这些注意力汇聚的token，传统窗口注意力方法的性能会下降。
- 基于这一观察，他们引入了 滚动缓存（Rolling Cache） ，它保留了初始token，并处理最近的上下文，从而实现了无限长度序列的处理。
- 他们还展示了这些汇聚token可以通过训练获得，作为专用的注意力锚点，从而减少对多个初始token的依赖。
- Value-Aware Token Pruning (VATP)
VATP 扩展了 H2O 的 token 重要性概念，考虑了注意力模式和价值向量的属性：
- 重要性评分 ：结合了注意力分数和价值向量的信息：
其中
是累计注意力分数，
是价值向量的 L1 范数。
- Token修剪 ：根据
排名token，最低分数的token被修剪，而注意力汇聚token（例如，开始或换行token）被保留，以防止性能下降。
性能与效率：
- 在16个 LongBench 任务中，VATP 在12-14个任务中超越了 H2O 和 Scissorhands 等基准。
- 在保持最小性能损失的情况下，实现了50%的有效压缩。
- 引入的计算开销几乎可以忽略不计，并且与 Scissorhands 集成时兼容 FlashAttention。
2.2 后处理压缩技术（Post-hoc Compression Techniques）
这些方法压缩或优化KV缓存，同时保持标准的Transformer架构。
- Adaptive KV Compression (FastGen)
FastGen 通过观察运行时的注意力模式引入了自适应压缩：
- 注意力分析 ：在提示编码过程中，FastGen 识别注意力模式，并选择压缩策略
，以最小化内存开销，同时保留注意力恢复：
自适应压缩策略：
- 压缩策略包括：
- 特殊 token (
)：仅保留特殊 token。
- 局部性 (
)：逐出超过相对距离
的 token。
- 频率 (
)：保留具有高累计注意力分数的 token (
)。
- 混合策略 结合这些策略，首先采用
，并根据每个头的需要适应性地应用：
Token 生成：
- 在解码过程中，预先选择的压缩策略有效地管理 KV 缓存：
- 动态内存压缩（DMC）
DMC 引入了自适应的 token 合并：
- 决策机制 ：在时刻
，预测合并决策
和权重
：
- 加权合并 ：当
时，合并当前和先前的条目：
其中
累积重要性权重。
训练：
- 使用 Gumbel-Sigmoid 放松来训练
，支持端到端的梯度下降训练：
其中
是温度参数。
- 优化组合目标：
其中
是语言建模损失，第二项鼓励模型匹配目标压缩比（CR）。
- 结果 ：达到 8 倍的压缩率，保持性能。
范数基础的压缩
本文提出了一个令人惊讶的观察：缓存 KV 对的
范数与注意力分数之间存在明确的相关性，低
范数的键嵌入通常会导致解码时的高注意力分数。因此，提出了一个简单但有效的压缩目标：
- 基于范数的选择 ：对于一组缓存键
，计算并排序键的范数：
- 排序和选择 ：为了压缩 KV 缓存，按
范数值对所有键进行排序：
保留范数最小的前
个键，其中
，
为压缩比。
-
压缩缓存 ：压缩后的键值缓存为：
-
由于其简洁性，该方法与 FlashAttention 保持兼容。
2.3 体系结构重设计
这些方法改变了 Transformer 架构，以更高效地处理 KV 缓存，通常将压缩直接集成到架构中。
- 多查询注意力（MQA）
- https://arxiv.org/pdf/2305.13245
- 核心思想 ：MQA 通过共享单个键值头跨所有查询头来减少 KV 缓存大小，替代传统的多头注意力（MHA）：
其中
和
是共享的键和值投影。
- 优点 ：将 KV 缓存大小减少了
（注意力头的数量），显著降低了内存带宽开销。
- 权衡 ：虽然 MQA 更快，但在需要多样化注意力模式的任务中，通常会遭遇质量下降。
- 分组查询注意力（GQA）
- https://arxiv.org/abs/2305.13245
- 核心思想 ：GQA 在完全多头注意力和 MQA 之间进行插值，提供了推理速度和模型质量之间的可扩展权衡。它将查询头分为
组，每组共享一个单独的键值头：
- GQA-1 ：等价于 MQA
。
- **GQA-
**：等价于 MHA
。
- 训练 ：通过微调将 GQA 引入现有的预训练模型：
- 首先，将 MHA 权重通过均值池化转换为 GQA。
- 然后进行微调（“上训练”）以适应新的注意力模式。
- 该适应过程仅需原始预训练计算的 5%，使其非常高效。
- 结果模型保持质量，同时获得 GQA 的内存优势。
- 多头潜在注意力（MLA）
DeepSeek的多头潜在注意力（MLA）采用了一种新颖的方法来减少KV缓存开销。虽然MQA和GQA通过头共享来实现这一目标，MLA则采用低秩潜在压缩技术，在保持多头注意力的优点的同时，减少了KV缓存的大小。
- MLA通过将键（keys）和值（values）压缩成低维度的潜在向量，来减少KV缓存的大小。
- 它将键值嵌入（key-value embeddings）降投到一个压缩的潜在空间：
其中，
是降投矩阵，
、
是键和值的上投矩阵。
- 通过压缩表示，MLA保持了每个头的灵活性，不同于MQA的完全头共享。
- 它引入了旋转位置嵌入（RoPE）来解耦位置感知的键：
这进一步减少了KV缓存的存储，仅缓存压缩的潜在向量
和位置键
。
- SnapKV
- https://arxiv.org/pdf/2404.14469
- SnapKV引入了观察窗口（Observation Window）：使用提示结束的tokens来识别注意力模式：
其中，
表示注意力权重，
由压缩率决定。
- 压缩：使用池化层围绕选定位置聚类特征，以保持上下文完整性。
- 只缓存一次（YOCO）
YOCO修改了Transformer架构以优化缓存：
- 全局缓存：使用解码器-解码器设计，只有一个共享的KV缓存。
- 复杂度减少：将内存从
减少到
，其中
是序列长度，
是层数。
- 高效注意力：自解码器采用滑动窗口注意力或门控保留机制，使内存使用保持恒定（
，其中
是小窗口大小）。
结论
KV-Cache技术是将Transformer模型扩展和优化到实际应用中的核心。像动态逐出、压缩和结构化近似等创新，持续推动着在长上下文或资源受限的场景中实现更高效的技术。KV-Cache仍然是一个活跃的研究领域，既提供了理论上的见解，也带来了实际的改进。