Speculative KV Coding：利用预测模型无损压缩 KV 缓存高达 ~4 倍

背景

随着大型语言模型（LLM）上下文窗口（Context Window）的不断扩展，KV 缓存（KV Cache）的大小也在与日俱增。KV 缓存机制通过“以计算换内存”的方式，避免了模型对已完成的工作进行重复预填充（re-prefill），从而使得运行长上下文变得经济可行。然而，随着智能体工作流（Agentic workflows）推动上下文越来越长，存储和传输这些缓存的数据量开始主导整个系统的开销。

为了突破 LLM 能力的下一个数量级瓶颈，必须解决缓存体积过大的问题。目前主要有两种压缩路径：

1. 有损压缩（Lossy Compression）：例如 TurboQuant 等方法，通过降低 K 和 V 的位宽来减小体积。这种方法存在争议，且其质量损失难以预先精确控制，通常需要通过运行评估（evals）来发现模型退化情况，具有不确定性。
2. 无损压缩（Lossless Compression）：通过精确重建缓存来绕过质量损失的问题。

研究表明，LLM 权重的无损压缩潜力有限。经验数据显示，bf16 格式缓存的字节级熵约为每标量 11 比特，比原始表示小约 30%。更糟糕的是，随着位宽降低，这种冗余空间会迅速收缩。例如，FP4 权重的熵值已接近其格式的理论极限（饱和状态），低精度缓存也存在类似情况。鉴于低位宽对性能带来的显著优势，将其作为基线处理是必要的。

在此背景下，研究人员提出了一种名为 Speculative KV Coding 的新方法，旨在利用一个更便宜的预测模型，对大型目标模型的 KV 缓存进行无损压缩，压缩率最高可达约 4 倍（这是在 FP8 有损压缩基础上的额外增益，总增益更大）。

核心内容

Speculative KV Coding 的核心思想类比于推测解码（Speculative Decoding）。该方法使用一个更快的预测模型，在相同的提示词（Prompt）上运行前向传播，以预测目标模型的 KV 缓存内容。随后，使用算术编码器（Arithmetic Coder）对目标模型的真实缓存进行编码，编码的比特率取决于预测模型与目标模型之间的拟合程度。

KV 缓存并非真正的随机数据

要理解该方法的原理，首先需要回顾熵编码的基本概念。根据香农信源编码定理，任何无损编码器在平均情况下能达到的最佳效果是每符号 $H(X)$ 比特。然而，现实中的编码器使用的是模型分布 $Q$ 而非真实分布 $P$，因此实际比特率为 $H(P) + D_{KL}(P || Q)$。其中，KL 散度项代表了因模型不完美而支付的额外代价。

对于 KV 缓存而言，它并非来自随机源的一组样本，而是已知权重在已知提示词上经过确定性前向传播的输出。这意味着“真实”分布是一个在该特定张量上的狄拉克 $\delta$ 函数，其熵为零。因此，编码器花费的每一比特都是纯粹的 KL 散度代价。比特率直接衡量了模型对正确 KV 缓存赋予的权重大小。

我们需要的是一个能够校准的模型，用于描述产生该缓存的具体前向传播过程，并将其转化为算术编码器可消费的分布形式。

预测模型的设计与优化

假设我们有一个预测模型，给定提示词后，能产生合理的每标量 KV 缓存预测值 $\hat{k}$，并校准其误差范围。最自然的方式是使用以 $\hat{k}$ 为中心、方差为 $\sigma^2$ 的高斯分布。在这种情况下，编码真实值的成本（每标量纳特）可以分解为两项：

1. 扩散成本（Spread Cost）：与 $\log \sigma$ 相关，希望 $\sigma$ 越小越好。
2. 失配成本（Miss Cost）：当真实值 $k$ 远离 $\hat{k}$ 相对于 $\sigma$ 时产生惩罚，如果 $\sigma$ 设置得比实际误差小太多，该项会急剧上升。

通过对数据取期望并最小化，最优解出现在 $\sigma$ 等于典型误差幅度时。这意味着更好的预测精度直接带来比特数的节省，而校准不良则会浪费比特（$\sigma$ 过大导致不必要的扩散成本，过小导致失配成本失控）。

预测器的选择与权衡

预测器的选择本质上是一个成本与比特率的权衡，存在两个极端：

1. 极端 A（目标模型本身）：预测器就是目标模型。此时条件熵为零，无需传输任何比特。但这需要额外执行一次完整的目标模型前向传播，为了“压缩”而付出巨大的计算代价，得不偿失。
2. 极端 B（纯噪声）：预测器输出纯噪声。此时预测器几乎免费，但算术编码器需要支付接近 16 比特/标量的成本（覆盖 bf16 范围）。

实际的预测器介于两者之间。关键在于如何以最低的成本降低条件熵。

具体实现：量化作为预测器

最直接且高效的预测器是对目标模型进行结构保持的优化版本，例如量化（Quantization）。

• 原理：预测器可以是目标模型的量化版本（如 FP8、INT4、MXFP4 等）。量化引入的舍入噪声构成了前向传播的微小扰动，使得量化模型的输出与原始模型输出在元素级别上非常接近。
• 优势：量化模型的残差（Residual）小且具有结构性，这正是高斯模型下高效编码的理想场景。
• 流程：
  1. 编码端：运行量化预测模型生成 $\hat{k}$，运行目标模型生成真实 $k$。计算残差，结合校准后的方差 $\sigma$，通过算术编码器生成比特流。
  2. 解码端：仅运行量化预测模型生成 $\hat{k}$，接收比特流，通过算术解码器恢复残差，最终重建出精确的 $k$。
  3. 共享信息：双方共享提示词和预测模型参数。残差的每通道统计信息（用于 $\sigma$）可以在小规模校准集上测量一次并冻结。

初步结果

最简单的具体实例是：预测器直接取自目标模型的 FP8 版本。方差 $\sigma$ 是在训练数据上针对每个 (KV, Head, Channel) 的经验残差方差拟合的，并在位置上池化。实验显示，这种方法能够在保持无损的前提下，显著降低 KV 缓存的传输和存储开销。

关键要点

• 无损压缩的必要性：相比有损压缩带来的不可控质量下降，无损压缩通过精确重建缓存，确保了模型输出的准确性，适合对精度要求极高的场景。
• 利用相关性而非随机性：KV 缓存是确定性计算的产物，其内部存在极强的结构相关性。Speculative KV Coding 利用这种确定性，将压缩问题转化为预测误差的编码问题。
• 预测模型即量化模型：使用量化后的模型（如 FP8）作为预测器是最佳实践之一。量化模型运行速度快、成本低，且其输出与全精度模型高度相关，产生的残差分布集中，易于高效编码。
• 算术编码的应用：该方法不直接存储压缩后的 KV 值，而是存储预测值与真实值之间的残差。算术编码器根据预测的不确定性（方差）动态调整编码效率，不确定性低时编码更短。
• 计算与通信的权衡：该方法在编码端需要运行一次目标模型（计算开销），但大幅减少了传输/存储的数据量（通信/存储开销）。在长上下文推理中，通信和存储往往成为瓶颈，因此这种权衡是划算的。
• 压缩率可观：在 FP8 有损压缩的基础上，Speculative KV Coding 还能提供额外的 ~4 倍无损压缩增益，总压缩效果显著。

意义与影响

Speculative KV Coding 为长上下文 LLM 推理中的内存和带宽瓶颈提供了一条新的解决路径。

1. 降低部署成本：通过大幅减少 KV 缓存的大小，可以降低对 GPU 显存的需求，或者在相同的显存容量下支持更长的上下文窗口。这对于边缘设备部署或大规模并发服务尤为重要。
2. 加速推理流水线：在分布式推理或模型服务中，KV