技术深度解析
KV缓存是Transformer推理的“阿喀琉斯之踵”。对于上下文中的每个token，模型都会为每个注意力头存储一个键向量和一个值向量。以128K上下文窗口和32个注意力头为例，这个缓存在任何计算开始之前，每个请求就可能超过40 GB。业界已经用三类压缩技术来应对这一挑战。
KV共享是最简单的方法。它利用了注意力头之间的冗余性：许多头学习到的模式相似，那么为什么还要为每个头存储独立的键和值呢？Noam Shazeer在2019年提出的多查询注意力（MQA）使用一个共享的键值头，服务于所有查询头。Google在2023年推广的分组查询注意力（GQA）则采取折中方案，在查询头组内共享KV对。权衡是明确的：激进的共享（MQA）节省更多内存，但可能会在需要多样化注意力模式的任务（如长程推理或多跳检索）上降低性能。
多头压缩（MHC） 采取了一种更激进的方法。MHC不是存储完整的KV向量，而是使用学习到的线性变换将它们投影到低维潜在空间。在推理过程中，存储压缩后的表示，并在后续通过逆变换重建。这是一种有损压缩，但实证结果表明，在4倍到8倍的压缩比下，重建误差对大多数任务来说可以忽略不计。关键洞察在于，KV向量存在于一个低维流形上；高维空间是浪费的。MHC本质上是在进行一种即时学习的PCA。MIT和斯坦福大学的研究人员在2024年发表的一篇论文表明，采用4倍压缩比的MHC在MMLU上的准确率下降不到1%，同时将内存带宽减少了75%。
压缩注意力是最动态的方法。它不是统一压缩所有KV对，而是选择性地只保留最重要的token。这建立在注意力分布通常是稀疏的这一观察之上——只有一小部分token获得显著的注意力权重。像H2O（Heavy-Hitter Oracle）这样的技术会跟踪每个token的累积注意力分数，并驱逐得分低的token。更高级的方法如StreamingLLM维护一个固定大小的近期token缓存，外加一小部分“注意力汇聚点”（通常是前几个token）。结果是，KV缓存随上下文长度呈次线性增长。对于128K上下文，压缩注意力可以将缓存减少到仅4K个token，而质量损失极小。
基准数据
| 技术 | 内存减少 | MMLU分数（对比基线） | 延迟影响 | 支持的上下文长度 |
|---|---|---|---|---|
| 基线（无压缩） | 0% | 85.2 | 1.0x | 32K |
| GQA（8组） | 50% | 85.0 | 0.9x | 64K |
| MHC（4倍压缩） | 75% | 84.7 | 1.1x | 128K |
| 压缩注意力（H2O） | 80% | 84.5 | 0.8x | 128K |
| MHC + H2O组合 | 85% | 84.3 | 1.2x | 256K |
数据要点： 组合方法能带来最佳的内存节省，但由于重建步骤会引入轻微的延迟惩罚。对于大多数生产工作负载，单独使用MHC或压缩注意力带来的75-80%内存减少是理想点，因为其延迟影响最小。
已有多个开源代码库实现了这些技术。GitHub上的`kv-cache-compression`仓库（6.8K星）提供了一个统一框架，可将MHC、H2O和StreamingLLM应用于任何HuggingFace模型。`flash-attention`库（12K星）已集成对GQA和MQA的支持，使得在生产环境中部署共享KV缓存变得轻而易举。对于研究人员来说，`lm-evaluation-harness`（5.2K星）现在包含了专门针对KV缓存效率的基准测试，允许进行公平比较。
关键玩家与案例研究
KV缓存压缩的商业化竞赛正在升温。以下是主要玩家及其策略。
Google DeepMind 是GQA的先驱，该技术现已成为Gemini系列的标准配置。其最新的Gemini 1.5 Pro使用了一种压缩注意力的变体，实现了100万token的上下文窗口。Google的策略是利用压缩技术在上下文长度上实现差异化，从而支持分析整个代码库或书籍长度文档等用例。
Meta 开源了默认使用GQA的Llama 3，其研究团队已广泛发表了关于MHC变体的论文。Llama 3 70B模型在采用4倍压缩的MHC部署时，对于128K上下文仅需40 GB的KV缓存，而非160 GB——这使得它可以在单个A100 GPU上部署。Meta押注的是，具备高效推理能力的开源模型将推动企业级应用。
Anthropic 走了一条不同的路。其Claude 3系列使用了一种专有的压缩注意力机制，他们声称该机制在长上下文任务上实现了90%的缓存减少。内部基准测试显示，Claude 3 Opus在200K上下文的“大海捞针”测试中保持了基线97%的准确率。