← 返回 2026-05-29

CONF-KV:用于长视界 LLM 推理的置信度感知 KV 缓存驱逐与混合精度存储 CONF-KV: Confidence-Aware KV Cache Eviction with Mixed-Precision Storage for Long-Horizon LLM

Yubo Li, Yidi Miao 📅 2026-05-24 👍 9 2026-07-13 08:36
KV 缓存管理 混合精度量化 置信度估计 自适应缓存策略 长上下文推理

利用模型置信度动态调整 KV 缓存预算,在保持质量的同时大幅降低内存占用

前置知识

KV 缓存(Key-Value Cache)

在 Transformer 自回归生成中,为了避免重复计算之前令牌的注意力,将每一层的 Key 和 Value 张量缓存起来。这些缓存会随着序列长度线性增长,在长上下文场景下(如 4K 生成令牌)可能占据数十 GB 显存,成为推理瓶颈。缓存管理需要在保留质量的前提下,通过驱逐不重要的令牌来控制内存使用。

本文的核心问题就是如何智能地管理这个缓存,理解 KV 缓存的工作原理是理解本文提出的驱逐策略的基础。

块式在线 softmax 注意力(Blockwise Online Softmax Attention)

FLASHATTENTION 引入的计算优化技术。它将注意力计算分块进行,避免存储完整的注意力矩阵,而是在遍历块时维护运行中的最大值和归一化因子,从而计算出精确的 softmax 结果。这种方法显著降低了激活内存,同时不改变计算结果。CONF-KV 利用这种紧凑布局来简化注意力内核。

本文的缓存管理器使用块式注意力读取,理解这种计算模式有助于理解为什么 CONF-KV 选择紧凑存储而不是页式间接表。

Needle-in-a-Haystack 测试

评估长上下文模型检索能力的标准基准测试。在海量无关文本(干草堆)中插入一个目标事实(针),然后查询该事实。测试通过变化海草堆长度(1K 到 32K 令牌)和针的深度(距离查询的位置 10% 到 90%)来评估模型能否保留关键的长程信息。这是检验 KV 缓存策略能否保留重要历史信息的严格测试。

本文的核心贡献之一就是在 NIAH 任务上达到 91.4% 的检索准确率,远超滑动窗口的 53.8%,理解这个测试能帮助理解为什么这个结果意义重大。

指数移动平均(Exponential Moving Average, EMA)

一种平滑数据的统计方法,对近期数据赋予更高权重。公式为 a_i ← λ a_i + (1-λ)A̅_i,其中 λ = 0.9 是衰减因子,A̅_i 是当前注意力质量。这种方式既能积累历史重要性,又能对新的注意力模式快速响应。CONF-KV 用它来跟踪每个令牌的历史注意力质量。

本文的排名器使用 EMA 来计算令牌的重要性,理解这个机制有助于理解为什么 CONF-KV 能比纯最近度策略更好地保留有价值的历史信息。

研究动机

长视界 LLM 应用(如 Web 代理、长文档分析、多轮对话、工具使用系统)在多轮交互中累积上下文,KV 缓存随序列长度和深度线性增长。在作者的测量中,Qwen-32B 在 4K 生成令牌时的 KV 内存分配器足迹达到 15.8 GB,在单块 80 GB H100 上显著限制了可行批量大小。现有的缓存策略要么使用静态窗口(如 512 令牌滑动窗口)盲目丢弃旧令牌,要么基于历史注意力信号(如 H2O、SCISSORHANDS)排名,但这些方法都没有利用当前解码步骤时模型的自身状态信息。滑动窗口在 Needle-in-a-Haystack 上只能达到 53.8% 的检索准确率,因为它会结构性丢失窗口之外的所有信息,无论这些信息是否语义上重要。

本文的目标是本文的核心目标是探索一个未被充分利用的信号——当前 next-token 分布——能否改善 KV 缓存驱逐的内存-质量权衡。作者希望回答:这个在采样前就已经可用的自由信号,是否能比静态策略更有效地控制何时驱逐令牌?具体来说,当模型在某个解码步骤不确定时,应该保留更多上下文;当模型自信时,可以积极修剪。这种自适应预算选择的目标是在匹配固定内存预算的情况下,尽可能接近完整 KV 缓存的质量。

与已有工作不同的是,本文的独特切入角度在于从时间维度而非空间维度思考缓存管理。现有方法(H2O、SCISSORHANDS、SNAPKV、PYRAMIDKV)关注如何在不同头、层或精度之间分配预算,但这些预算决策通常是全局静态的或在提示阶段确定的。CONF-KV 则关注解码时的时间维度:当前的输出分布形状直接反映了模型在当前步骤的不确定性程度,这是一个前向信号——它不能恢复已被驱逐的令牌,但可以在模型处理困难段落时防止过早驱逐。这种时间自适应与空间自适应是正交的,可以组合使用。作者通过 KL 散度消融实验验证了这个假设:高置信度与上下文移除引起的分布变化呈现 r = -0.77 的强负相关,支持了置信度与上下文需求之间的单调关系。

核心方法

CONF-KV 的整体思路很直观:在每一步解码时,利用模型的 logits 计算置信度分数,根据这个分数选择宽松或严格的缓存预算,然后在预算内驱逐排名最低的令牌。整个流程可以分成三个阶段:置信度估计、预算选择、令牌排名与驱逐。置信度估计器将 logits 转换为三个信号——归一化熵 Ĥ = H(p) / log V、对数概率边缘 m = log p^(1) - log p^(2) 和 top-token 概率 p^(1),通过加权组合得到标量置信度分数 c = w_H(1-Ĥ) + w_m σ(m) + w_p p^(1),权重为 (0.4, 0.3, 0.3)。如果 c ≥ τ(默认 τ = 0.7),选择严格预算 N_high;否则选择宽松预算 N_low。在预算内,令牌按综合分数 s_i = α â_i + (1-α) r̂_i 排名,其中 â_i 是 EMA 注意力质量,r̂_i 是归一化最近度,α = 0.65。最后驱逐最低排名的令牌,但始终保护最近 P 个令牌(默认 32 或 64)。

CONF-KV 的核心创新在于使用模型当前的 next-token 分布作为缓存管理的控制信号。这与现有方法的本质区别在于:H2O 和 SCISSORHANDS 使用历史注意力质量,SNAPKV 使用提示阶段观察,PYRAMIDKV 使用层级预算,而 CONF-KV 使用的是解码时的实时不确定性。这种信号有三个优势:它是免费的(不增加额外计算开销)、它是前向的(预测当前步骤需要多少上下文)、它是自适应的(随模型状态动态调整)。结合混合精度存储(FP16 窗口 + INT8 历史缓存)和可选的金字塔式每层预算,CONF-KV 在 GPT-2 上关闭了 74% 的困惑度差距,而最好的静态基线 PYRAMIDKV 只关闭了 63%。作者通过匹配速率的隔离研究证明了这个信号的重要性:随机驱逐虽然使用相同的驱逐概率,但困惑度高达 36.54,而完整 CONF-KV 只有 30.92,证明置信度驱动的调度本身就有价值。

方法步骤详情

CONF-KV 的单步生成算法(Algorithm 1)包含以下步骤。输入是当前缓存 C、模型 f、输入令牌 x_t 以及超参数 τ, N_high, N_low, P, α, W。步骤 1-2:前向传播得到 logits 和注意力矩阵 A,计算 p = softmax(logits)。步骤 3:计算置信度分数 c = w_H(1-Ĥ(p)) + w_m σ(log p^(1) - log p^(2)) + w_p p^(1)。步骤 4:选择预算 N = N_high 若 c ≥ τ,否则 N = N_low。步骤 5-11:对每一层 ℓ,使用 A^(ℓ) 更新 EMA 注意力元数据;如果 |C_ℓ| > N,则对非保护令牌按 s_i = α â_i + (1-α) r̂_i 排名,驱逐最低排名的令牌并压缩幸存者;将保留位置中最新的 W 个生成令牌之外的保留位置量化到 INT8。步骤 12:结束循环。步骤 13:将新的 K/V 以 FP16 格式追加到缓存,从 p 中采样得到 x_{t+1}。整个过程的时间复杂度是线性的,驱逐事件率约为 42%,压缩成本为每步 0.22ms,元数据更新为 0.11ms,都远小于注意力节省。

技术新颖性

CONF-KV 的技术新颖性体现在三个方面。第一,它首次将 next-token 置信度信号用于控制 KV 缓存的令牌保留预算,这与早期退出和投机解码中使用置信度控制计算或接受决策是类似的思路,但应用到了内存管理这个新领域。第二,它将自适应驱逐与块式注意力、缓存压缩、混合 FP16/INT8 存储和金字塔式每层预算进行了系统性的设计组合,形成一个完整的工程方案。第三,它通过机制性测试验证了置信度与上下文需求之间的因果关系:消融最近 256 个令牌后测量 KL 散度变化,发现高置信度步骤的分布变化显著低于低置信度步骤(Pearson r = -0.77, p < 10^-30),这支持了策略的核心假设。虽然不确定性信号本身不是新概念,但将其用于解码时的自适应缓存预算选择是新颖的,而且这个信号与头级分配、层级分配、精度选择是正交的,可以组合使用而不是替代它们。

CONF-KV pipeline. The logits at each decoding step are converted into confidence c, which selects a tight or loose cache budget. The cache manager ranks candidates by attention mass and recency, protects the most recent tokens, compacts the survivors, and serves the next attention call from a mixed-precision cache.
Figure 1: CONF-KV pipeline. The logits at each decoding step are converted into confidence c, which selects a tight or loose cache budget. The cache manager ranks candidates by attention mass and recency, protects the most recent tokens, compacts the survivors, and serves the next attention call from a mixed-precision cache.

实验结果

论文的核心发现是置信度信号确实能比静态策略更有效地控制驱逐时机。在 WikiText-2 持续性困惑度测试(GPT-2, 2048 生成令牌)中,匹配 CONF-KV+INT8 的平均峰值 KV 内存(38.7 MB ± 2 MB)时,CONF-KV+INT8 的困惑度为 31.26,比滑动窗口的 34.37 降低 3.11 点;CONF-KV-L(金字塔式变体)达到 30.48,降低 3.89 点,关闭了完整 KV(29.14)和滑动窗口之间 74% 的困惑度差距,而最好的静态基线 PYRAMIDKV 只关闭了 63%。在 Needle-in-a-Haystack 测试(gpt-oss-20b,海草堆 1K-32K,五种针深度)中,CONF-KV 达到 91.4% 的平均检索准确率,远超滑动窗口的 53.8% 和 H2O 的 80.6%。失败模式主要发生在模型在罕见实体查找前过度自信时,此时策略选择了严格预算。在 VisualWebArena 的 75 个任务中,CONF-KV 保留了 95.3% 的完整 KV 成功率(完整 KV 40.2%,CONF-KV 38.3%),峰值内存降低 2.8×,滑动窗口损失 11.1 个绝对点。在系统性能方面,2048 令牌生成时 CONF-KV 在 GPT-2 上降低 p50 延迟 1.8×(相对完整 KV),在 Qwen-32B 上也降低 1.8×。吞吐量随批量大小提升,在批量 8 时 CONF-KV 达到完整 KV 的 2.06× 吞吐量,与滑动窗口基线匹配在 1% 以内。这些结果表明自适应预算选择在匹配内存的同时显著提升了质量,而且性能提升不是仅仅来自驱逐更少——匹配速率的随机驱逐证明调度本身就有价值。

KV-cache management methods differ in the signal used to decide token utility and in whether the budget adapts. CONF-KV is distinctive because the budget reacts to the current next-token distribution.
Table 1: KV-cache management methods differ in the signal used to decide token utility and in whether the budget adapts. CONF-KV is distinctive because the budget reacts to the current next-token distribution.
Matched-memory comparison on GPT-2 at 2048 generated tokens. PPL is WikiText-2 continuation; latency is per-step. Baselines are calibrated to match CONF-KV+INT8's average peak KV memory (38.7 MB ± 2 MB).
Table 2: Matched-memory comparison on GPT-2 at 2048 generated tokens. PPL is WikiText-2 continuation; latency is per-step. Baselines are calibrated to match CONF-KV+INT8's average peak KV memory (38.7 MB ± 2 MB).
Memory–quality and scaling behavior. Left: GPT-2 memory–PPL Pareto at 2048 tokens. Right: peak KV memory across generated-token lengths for GPT-2 and Qwen-32B.
Figure 2: Memory–quality and scaling behavior. Left: GPT-2 memory–PPL Pareto at 2048 tokens. Right: peak KV memory across generated-token lengths for GPT-2 and Qwen-32B.
Signal validation. Left: matched-rate isolation shows that random eviction with the same schedule fails, while attention and recency each help. Right: high confidence anticorrelates with the KL shift caused by ablating recent context.
Figure 3: Signal validation. Left: matched-rate isolation shows that random eviction with the same schedule fails, while attention and recency each help. Right: high confidence anticorrelates with the KL shift caused by ablating recent context.
Needle-in-a-Haystack retrieval accuracy on gpt-oss-20b. Columns are haystack lengths; rows use distance-from-query depth, so larger percentages place the needle farther from the retained recent window.
Figure 4: Needle-in-a-Haystack retrieval accuracy on gpt-oss-20b. Columns are haystack lengths; rows use distance-from-query depth, so larger percentages place the needle farther from the retained recent window.
Task quality and latency. Left: VWA success rate, mean ± s.d. over three seeds. Right: p50 per-step latency at 2048 generated tokens across four models.
Figure 5: Task quality and latency. Left: VWA success rate, mean ± s.d. over three seeds. Right: p50 per-step latency at 2048 generated tokens across four models.
Systems behavior. Left: per-step latency decomposition on GPT-2. Right: GPT-2 throughput versus batch size at 2048 generated tokens; full-KV points beyond batch 8 were unavailable in our prototype allocator run and are not used for speedup claims.
Figure 6: Systems behavior. Left: per-step latency decomposition on GPT-2. Right: GPT-2 throughput versus batch size at 2048 generated tokens; full-KV points beyond batch 8 were unavailable in our prototype allocator run and are not used for speedup claims.
Ablations. Left: τ controls the eviction-rate/quality trade-off. Right: budget and FP16-window sweeps place the default Nhigh = 128 and W = 128 near the knees of the curves.
Figure 7: Ablations. Left: τ controls the eviction-rate/quality trade-off. Right: budget and FP16-window sweeps place the default Nhigh = 128 and W = 128 near the knees of the curves.
Trace-level behavior. Left: cache size over 500 generation steps. Right: confidence trace and empirical histogram.
Figure 8: Trace-level behavior. Left: cache size over 500 generation steps. Right: confidence trace and empirical histogram.
查看结构化数据
任务指标本文基线提升
WikiText-2 持续性困惑度(GPT-2,2048 生成令牌) 困惑度(PPL) CONF-KV-L: 30.48 完整 KV: 29.14, 滑动窗口-512: 34.37, PYRAMIDKV: 31.09 相比滑动窗口降低 3.89 点,关闭完整 KV-滑动窗口差距的 74%,超过 PYRAMIDKV 的 63%
Needle-in-a-Haystack 检索准确率(gpt-oss-20b,海草堆 1K-32K) 平均准确率 91.4% 滑动窗口-512: 53.8%, H2O: 80.6% 相比滑动窗口提升 37.6 个百分点,相比 H2O 提升 10.8 个百分点
VisualWebArena 任务成功率(75 个任务) 成功率 38.3% 完整 KV: 40.2%, 滑动窗口-512: 29.1%, H2O: ~35% 保留完整 KV 成功率的 95.3%,峰值内存降低 2.8×,相比滑动窗口提升 9.2 个百分点
吞吐量(GPT-2,2048 生成令牌) 令牌/秒(相对完整 KV) 批量 8: 2.06× 滑动窗口-512: 2.06×, H2O: ~1.9× 在批量 8 时达到完整 KV 的 2.06× 吞吐量,与滑动窗口基线匹配
峰值 KV 内存(Qwen-32B,4K 生成令牌) 内存占用(GB) 2.6 GB 完整 KV: 15.8 GB 降低 13.2 GB(83.5%),在单块 80 GB H100 上显著增加可行批量大小

局限与改进

作者承认的局限性包括三个方面。第一,CONF-KV 在从未超过驱逐阈值的短上下文上是空操作,即如果上下文足够短,策略不会发挥作用,也就没有速度提升。第二,当 MLP 计算而非注意力或 KV 带宽主导延迟时,速度提升会变小,因为缓存管理节省的主要是注意力计算和内存带宽。第三,置信度信号在高温度采样下变得不那么有信息量,因为熵会饱和,可能需要重新调优阈值 τ。作者还提到了工程上的权衡:连续压缩策略在他们的环境中足够快,但将令牌级驱逐映射到页式服务布局不是免费的——细粒度驱逐会创建部分死块,除非实现粗粒度驱逐或添加稀疏掩码。我观察到的额外局限性包括:置信度阈值 τ 和预算 N_high, N_low 需要针对每个模型和工作负载调优,缺乏自适应机制;策略依赖于 logits 分布的形状,在不同架构和训练方法之间的泛化性需要更多验证;虽然理论上是模型无关的,但所有实验都在类 GPT 架构上,在其他架构(如 RWKV、Mamba)上的效果未知;INT8 量化虽然效果好,但不同硬件的 INT8 支持可能有差异,影响实际部署效果。

独立分析的弱点

CONF-KV 的几个关键弱点值得在后续工作中改进。第一,超参数敏感性。置信度阈值 τ、预算大小 N_high, N_low、保护窗口 P、排名权重 α 等都需要手动调优,而且不同模型和工作负载的最优值可能不同。改进方向可以是学习式调参或自适应阈值:例如根据历史置信度分布动态调整 τ,或者使用元学习找到特定任务的超参数配置。第二,模型依赖的置信度校准。置信度分数的计算权重 (0.4, 0.3, 0.3) 是在 GPT-2 上通过网格搜索找到的,虽然作者声称在评估模型中稳定,但这种稳定性缺乏系统性验证。改进方向可以是模型无关的置信度估计器,例如使用预训练的置信度分类器,或者通过校准数据集学习动态权重。第三,细粒度驱逐与页式服务的冲突。当前的连续压缩策略简单高效,但与 PAGEDATTENTION 等生产级服务系统不完全兼容。改进方向可以是块级驱逐策略(一次驱逐整个块而不是单个令牌),或者设计支持稀疏块的页式布局,在内存效率和计算效率之间取得平衡。第四,缺乏对极端情况的处理。当模型在罕见实体查找前过度自信时,策略会错误地选择严格预算,这是 NIAH 中的主要失败模式。改进方向可以是检测罕见实体模式(如命名实体识别),在这些情况下强制使用宽松预算,或者引入上下文紧急信号来覆盖置信度决策。

未来方向

基于 CONF-KV 的成果,可以延伸出多个研究方向。作者提出的未来方向包括将置信度信号与其他自适应机制组合:CONF-KV 与头级分配(ADA-KV)、层级分配(ZIGZAGKV)和精度选择(Adaptive Precision)是正交的,可以组合形成更强大的多轴自适应系统。可以探索的工作包括:第一,多信号融合缓存管理。除了 next-token 置信度,还可以利用其他信号如梯度信息(用于估计令牌对损失的影响)、局部模式检测(检测是否在列举或重复模式)、外部知识触发(检测到特定查询类型时扩大预算)。第二,自适应超参数学习。使用强化学习或贝叶斯优化自动学习最优的超参数配置,或者根据运行时统计(如平均驱逐率、困惑度变化)动态调整超参数。第三,多模态扩展。在视觉-语言模型中,图像令牌通常比文本令牌更重要,可以探索模态感知的驱逐策略,或者在置信度计算中加入多模态信号。第四,安全性和可控性研究。作者提到降低 KV 缓存内存可以降低长期自动化的成本,但也可能降低不良长期自动化的成本。可以研究缓存策略如何影响模型的可控性和安全性,例如限制记忆长度是否有助于防止不良行为模式的累积。第五,跨架构泛化。在非 Transformer 架构(如 RWKV、Mamba、线性注意力)上探索类似的置信度驱动策略,这些架构有不同的状态表示方式,但自适应的核心思想仍然适用。

复现评估

论文的复现性评估需要考虑几个方面。开源情况:作者声明代码将在接受后发布,但论文撰写时代码尚未公开,这限制了即时复现。数据:主要基准(WikiText-2、Needle-in-a-Haystack、VisualWebArena)都是公开数据集,WikiText-2 是确定性语言建模基准,NIAH 和 VWA 提供了标准协议。算力需求:实验在 NVIDIA H100 80 GB GPU 上进行,使用 CUDA 12.8、PyTorch 2.9 和 Transformers 4.51。GPT-2 等小模型可以在消费级 GPU 上运行,但 Qwen-32B 等大模型需要高内存 GPU。复现难度:中等。方法本身不需要训练,只是推理时的缓存管理策略,实现相对直接。主要挑战在于与现有框架(如 vLLM、TGI)集成,这些框架有自己的 KV 缓存管理机制。论文提供了详细的算法描述和超参数设置(τ = 0.7, N_high = 128 for WikiText, N_low = 256, P = 32, α = 0.65, W = 128, B = 128),这有助于独立复现。实验报告了三次随机种子的平均值和标准差(NIAH、VWA、吞吐量),以及 WikiText-2 的确定性结果,统计方法合理。性能分析提供了详细的分解(注意力从完整 KV 的 62% 降至 CONF-KV 的 47%,压缩添加 0.22ms,元数据更新添加 0.11ms),这有助于验证结果和瓶颈分析。总体而言,一旦代码发布,复现性应该良好,但需要与现有推理框架集成才能获得生产级性能。