基于局部分布还原的高精度低位KV缓存量化

# 通过局部分布恢复实现高精度低位KV缓存量化
来源：https://arxiv.org/html/2607.16248
###### 摘要

长上下文大语言模型推理依赖 KV 缓存来避免冗余的注意力计算，但这会带来高昂的内存和带宽开销。低位 KV 缓存量化降低了这一成本，但严重损害了生成质量；特别地，在 Llama-3.1-8B 上使用 RULER 评估时，1 位量化将准确率从 84.2% 降至 47.8%。

与普遍认为的 logits 绝对误差是主要原因相反，我们发现根本原因是*结构化局部排序错误*，即 top-K 区域中 logits 的分布发生了偏移。因此，我们提出了*局部分布恢复*，这是一种新技术，它从量化 logit 特征中检测具有高局部分布风险的时间步，并在 token 选择之前仅恢复选中的 top-K 候选分布。我们实现了 DGAP 来完成局部分布恢复，并配备了高效的风险检测器和校正器。实验表明，在 Llama-3.1-8B 上，DGAP 将 K1V1 RULER 准确率从 47.8% 恢复至 83.2%，并将分布漂移从 0.38 降至 0.14；在 Llama、Mistral 和 Qwen 模型上，它在保持较低解码开销的同时，保留了持续的低位 KV 缓存占用。

## 1 引言

大语言模型 (LLM) 越来越多地依赖长上下文推理来进行文档理解、检索增强生成、推理和代码分析 (Guo et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib40); OpenAI 2024 (https://arxiv.org/html/2607.16248#bib.bib41))。然而，KV 缓存随着上下文长度和批次大小线性增长；一个 Llama 风格的 8B 模型在 128K token 序列下已经需要约 8 GB 的全精度 (FP16) KV 缓存，使得 KV 缓存的内存和带宽成为关键的服务瓶颈 (Kwon et al. 2023 (https://arxiv.org/html/2607.16248#bib.bib15); Yang et al. 2025b (https://arxiv.org/html/2607.16248#bib.bib20))。KV 缓存量化通过以低精度存储缓存的键和值来降低这一成本，最近的方法通过非对称格式 (Liu et al. 2024 (https://arxiv.org/html/2607.16248#bib.bib4))、异常值处理 (Hooper et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib6)) 以及混合精度或误差补偿设计 (Liu et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib16); Zandieh et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib17)) 改进了量化效果。

尽管取得了这些进展，我们的实验表明，激进的低比特 KV 量化仍然会严重损害生成质量。在 K1V1 (键和值均以 1 位存储) 设置下，Llama-3.1-8B 在 RULER 上的准确率从 84.2% 下降到 47.8%，Mistral-7B 和 Qwen2.5-14B 也表现出类似的趋势。这促使我们更深入地分析低位 KV 量化如何在 token 选择之前改变局部候选分布，而不只是关注内存使用、吞吐量和最终任务分数所揭示的问题。

我们发现，这种退化并不能仅用 logit 幅值误差来解释。在缩放点积注意力中，注意力权重计算为 softmax(QK⊤/dk)V (Vaswani et al. 2017 (https://arxiv.org/html/2607.16248#bib.bib37); Xu et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib32); Tao et al. 2024 (https://arxiv.org/html/2607.16248#bib.bib42))。因此，键的量化可以直接改变查询-键分数，重新排序被关注的缓存位置，并改变用于形成下一 token logits 的上下文混合。一项将分数幅度与候选排序分开的控制分析表明，将量化后的排序强加于 FP16 分数会导致 31.8 个百分点的准确率下降，而保留排序的分数变化与 FP16 的差距保持在 1.2 个百分点以内。这些结果表明，激进的低位 KV 量化主要通过重新排序局部高概率候选分布来损害解码，而不是通过均匀地偏移 logit 值。

重要的是，这种分布偏移通常是局部可恢复的。尽管量化扭曲了概率质量并重新排序了高概率候选，但这些候选通常不会从考虑中移除；相反，它们仍然保持在紧凑的量化 top-K 候选区域内。我们将这种现象称为*结构化局部排序错误*。这表明了一种实用的替代方案，无需增加持续的 KV 精度或纠正整个词汇表：在 token 选择之前仅恢复受影响的局部候选分布。

我们提出了 DGAP (Disagreement-Guided Adaptive Precision，分歧引导的自适应精度)，这是一种轻量级层，用于在低位 KV 缓存解码中恢复局部分布保真度。DGAP 包含两个可训练模块。一个*分布风险检测器*根据量化 logit 特征估计局部分布风险，而一个*选择性 top-K logit 校正器*仅在风险超过校准阈值时恢复选中的候选分布。这种分离保持了较低的运行时开销：检测器避免了在稳定步骤进行不必要的校正，校正器只更新紧凑的候选区域而非整个词汇表。剩余的 logits 和持续的低位 KV 缓存保持不变。

DGAP 是现有低位 KV 缓存解码器的附加组件，不需要基础模型重新训练或注意力内核修改。在 Llama-3.1-8B、Mistral-7B 和 Qwen2.5-14B 上使用 LongBench、RULER、MMLU 和 WikiText-2 进行的实验表明，DGAP 在 K1V1–K8V8 设置下提高了生成质量和局部分布保真度。在 Llama-3.1-8B 的 K1V1 设置下，DGAP 将 RULER 准确率从 47.8% 提升至 83.2%，接近 FP16 的 84.2% 结果，将分布漂移从 0.38 降至 0.14，并且仅以 (1.06×) 的相对解码延迟代价保持了持续的 1 位 KV 缓存占用。

这项工作从局部分布层面提供了对激进低位 KV 缓存量化的机制分析。DGAP 将这一分析转化为实用的恢复机制，表明可恢复的候选分布重新排序可以在不重建 FP16 KV 状态或增加持续 KV 精度的情况下被本地检测和纠正。

总之，本文做出了以下贡献：

- • 我们提供了对激进低位 KV 缓存量化的机制分析，表明质量退化与局部候选分布重新排序密切相关，而不仅仅是 logit 幅值误差。
- • 我们识别了*结构化局部排序错误*，其中高概率候选被量化重新排序，但通常仍保持在紧凑的量化 top-K 区域内，使得局部分布恢复成为可能。
- • 我们提出了 DGAP，一个轻量级分布恢复层，配备风险检测器和选择性 top-K 校正器，可在不改变持续低位 KV 缓存的情况下恢复局部候选分布。
- • 我们在 Llama、Mistral 和 Qwen 模型上，在长上下文和语言基准测试中评估了 DGAP，展示了在 K1V1–K8V8 设置下以适度的解码开销实现的质量和分布保真度提升。

## 2 相关工作

#### 长上下文 KV 缓存效率。

长上下文 LLM 推理通常受限于 KV 缓存的内存和带宽，因为缓存的键和值随序列长度和批次大小线性增长 (Sheng et al. 2023 (https://arxiv.org/html/2607.16248#bib.bib43))。先前的工作通过分页缓存管理 (Kwon et al. 2023 (https://arxiv.org/html/2607.16248#bib.bib15); Yang et al. 2025b (https://arxiv.org/html/2607.16248#bib.bib20))、I/O 感知注意力内核 (Dao et al. 2022 (https://arxiv.org/html/2607.16248#bib.bib13)) 以及基于注意力汇聚点、重击者或提示级重要性进行稀疏 KV 保留或驱逐 (Xiao et al. 2024 (https://arxiv.org/html/2607.16248#bib.bib7); Zhang et al. 2023 (https://arxiv.org/html/2607.16248#bib.bib1); Li et al. 2024 (https://arxiv.org/html/2607.16248#bib.bib9); Cai et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib14)) 来改善服务效率。这些方法减少了内存碎片、内存流量或保留的 KV 状态数量。然而，它们主要优化 KV 状态的存储、访问或选择方式，并未直接分析激进的低位 KV 表示如何在解码过程中改变局部下一 token 候选分布。

#### KV 缓存量化与分布保持。

KV 缓存量化通过以低精度存储缓存的键和值来减少解码内存和带宽。现有方法通过非对称键/值量化 (Liu et al. 2024 (https://arxiv.org/html/2607.16248#bib.bib4))、异常值感知和非均匀量化 (Hooper et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib6))、稀疏或低秩误差补偿 (Kang et al. 2024 (https://arxiv.org/html/2607.16248#bib.bib5))、渐进式混合精度 (Liu et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib16))、在线向量量化 (Zandieh et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib17)) 以及用于推测性解码的分层量化缓存 (Tiwari et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib18)) 来改善低位质量。另一类相关的量化感知或分布感知方法使用校准、损失感知目标、旋转或蒸馏来减少低精度模型与全精度模型之间的不匹配 (Liu et al. 2023 (https://arxiv.org/html/2607.16248#bib.bib10); Zhang and Shrivastava 2025 (https://arxiv.org/html/2607.16248#bib.bib28); Shao et al. 2026 (https://arxiv.org/html/2607.16248#bib.bib29))。这些方法主要在量化器设计、校准或训练期间改进量化表示，而运行时在激进的低位 KV 缓存解码下的局部分布漂移和候选重新排序仍较少被探索。

参见图注 图 1：左：K1V1 急剧降低 RULER 准确率。右：控制变量实验识别出候选重新排序是主要原因。

## 3 动机

### 3.1 激进 KV 量化下的准确率损失

尽管 KV 缓存量化减少了内存占用 (Liu et al. 2024 (https://arxiv.org/html/2607.16248#bib.bib4); Sun et al. 2025 (https://arxiv.org/html/2607.16248#bib.bib31))，激进的低位设置会严重降低模型质量。

#### 方法。

我们评估 Llama-3.1-8B 在 RULER 上的 FP16 和四种 KV 精度设置：K8V8、K4V4、K2V2 和 K1V1，其中 KbVb 表示 b 位键和 b 位值。所有实验均使用 NVIDIA A100-80GB GPU 以及第 5 节 (https://arxiv.org/html/2607.16248#S5) 中的设置。

#### 观察 #1：激进的低位 KV 量化大幅降低任务准确率。

如图 1 (https://arxiv.org/html/2607.16248#S2.F1) (左) 所示，K1V1 将 RULER 准确率从 FP16 下的 84.2% 降低到 47.8%。较高位设置保持接近 FP16，K8V8、K4V4 和 K2V2 分别达到 84.0%、83.7% 和 81.4%。因此，在激进的 1 位 KV 缓存机制下，任务退化变得尤为严重。

参见图注 图 2：Top-K 恢复。量化扭曲了局部 top-K 分布，而恢复则恢复了选中的候选。

### 3.2 解码时准确率损失分析

任务分数的急剧下降引出了一个局部的解码时问题：低位 KV 量化如何改变 token 选择前的高概率候选分布？

#### 分析。

我们考虑两个可能的原因。首先，量化可能改变 logit 幅度，同时保持高概率候选的相对顺序。其次，量化可能扭曲概率质量并重新排序局部候选分布。由于注意力权重是从查询-键分数计算得出的，键的量化可以直接改变被关注缓存位置的排名，这反过来又会改变用于形成下一 token logits 的上下文混合。因此，局部候选重新排序可能比单独的 logit 幅度误差更具破坏性。

#### 方法。

为了将 logit 幅度变化与候选排名变化分开，我们根据 FP16 logits 构建了两个控制变量。第一个保留 FP16 logit 值，但根据量化排名对其进行排列，从而分离出候选重新排序的影响。第二个对 FP16 logits 应用保留排序的分数变换。在我们的控制中，这实现为恒定偏移，它会改变原始 logit 值，但保持 softmax 分布和 token 排名不变。

#### 观察 #2：局部分布重新排序主导 logit 值变化。

如图 1 (https://arxiv.org/html/2607.16248#S2.F1) 所示，保留分布的控制变体与 FP16 的差距保持在 1.2 个百分点以内，而用量化排名替换 FP16 排名则导致 31.8 个百分点的准确率下降。这表明，激进的低位 KV 量化下的质量损失主要由局部候选分布重新排序驱动，而不仅仅是 logit 幅度变化。图 2 (https://arxiv.org/html/2607.16248#S3.F2) 说明了这种结构化的局部排序错误：量化重新排序了高概率的 FP16 候选，但这些候选中的许多仍保留在紧凑的量化 top-K 区域内。因此，保留局部候选分布保真度和相对排名结构对于激进的低位 KV 解码至关重要。

一个简单的成对分析解释了为什么这种重新排序会改变局部分布。对于两个候选 i 和 j，且 zifp > zjfp，FP16 分布为 i 分配的概率大于 j，因为

pifp / pjfp = exp(zifp - zjfp) > 1。

在量化 logits ẑi = zifp + ei 下，相应的概率比变为

p̂i / p̂j = exp( (zifp - zjfp) + (ei - ej) )。

当 ẑj > ẑi ⇔ ej - ei > zifp - zjfp 时，局部顺序发生反转。因此，当 FP16 的差距较小时，即使是中等程度的非均匀误差也能重新分配概率质量并重新排序候选。相比之下，添加一个恒定偏移 c 得到

softmax(zfp + c1) = softmax(zfp)，

因此 token 概率和排名保持不变。

### 3.3 关键思想：局部分布恢复

上述分析表明，通过恢复 token 选择前受影响的局部 top-K 候选分布，可以缓解激进的低位 KV 缓存量化造成的质量损失。一个朴素的分布恢复流程包括三个步骤。首先，将量化解码分布与 FP16 参考进行比较，以确定当前步骤是否发生了漂移。其次，搜索整个词汇表以定位概率质量或相对顺序发生变化的 token。第三，调整受影响的 logits，使得结果分布更好地匹配 FP16 分布。

#### 挑战。

尽管这个流程在概念上很直接，但在自回归解码内部进行分布恢复的成本很高。首先，漂移检测不能在运行时依赖 FP16 logits 或额外的参考前向传递，因为这将违背低位解码的目的。其次，全词汇表恢复是不切实际的，因为在每个解码步骤搜索和纠正数万个 logits 将增加大量的计算和内存流量。第三，纠正必须保持局部性和可预测性；失败########################