BeaconKV: 基于Beacon查询的键值缓存压缩,用于高效大型推理模型推理

TL;DR BeaconKV 使用 beacon queries 预测推理中将被重访的 KV 对，解决 TRT 远程重关注导致近期 query 假设失效的问题，内存降 5.8倍、吞吐升 4.3倍且近无损。

问题

问题背景：大型推理模型（LRM）凭借长链思维（CoT）在复杂任务上表现优异，但 KV cache 随序列长度线性增长，成为长推理轨迹部署的主要内存瓶颈。

现有方法局限：主流 KV cache 压缩方法（如 H2O、SnapKV）依赖 recent queries 的注意力分数来评估历史 KV 的重要性，隐含假设近期查询模式能代表未来访问模式。然而，在 LRM 的长程推理中，某些解码步骤会生成 Thought Revisiting Tokens (TRT)，重新关注远距离的前期上下文（如任务规划），导致 recent queries 无法预测这些远期依赖。压缩时若丢弃早期关键 KV，会显著损害推理准确率。

为什么难/重要：推理过程中的注意力分布非平稳，全局 query 呈现聚类结构，但现有方法未建模全局 query 空间，且要求压缩算法训练自由、开销低，以适配不同模型与任务。高效压缩 KV cache 对降低推理成本、提升吞吐至关重要，业界高度关注在保持准确率前提下的内存优化。

行业类比：类似长文档问答或代码生成中，模型需要回溯先前段落或函数定义，若缓存压缩只基于最近 token 的局部信息，会丢失关键引用，导致回答不完整或代码错误。

核心洞察

• **Thought Revisiting Tokens (TRT)** 是长程推理中重新关注早期上下文的查询，它们在嵌入空间中聚成少数簇。BeaconKV 用**信标查询 (beacon queries)** 作为每个全局查询簇的紧凑代表，来预测哪些 KV 对将被再次访问，从而避免存储完整查询历史。这一视角不同于 SnapKV 等依赖最近查询作为未来注意力代理的方法，它直接建模了 LRM 中固有的长程重访模式，使缓存压缩与推理行为对齐，而非依赖局部性假设。
• BeaconKV 通过**持续最远点采样 (Continual FPS)** 在线选取代表性查询，无需训练即可维护全局重要性打分器。该方法的核心创新在于：用极小的内存开销（仅保存少量 beacon queries）实现了对远距离 KV 对的精确评分，在 5.8 倍压缩比下几乎保持全精度，并提升吞吐 4.3 倍以上。与现有训练无关方法相比，它不牺牲推理质量即可大幅降低缓存，为长 CoT 推理提供了工程上可落地的内存优化路径。

方法

输入

BeaconKV 的输入是 LRM 在 CoT 解码过程中逐步产生的 KV 缓存 与对应的 query 向量，无需任何训练或梯度更新。

关键模块

1. 周期性 KV 缓存驱逐
   当缓存达到容量上限时，不是逐 token 实时剪枝，而是按固定间隔（如每 128 步）执行一次全局驱逐。每次驱逐基于当前维护的 beacon queries 对历史 KV 对进行重要性评分，保留高分 KV，删除低分 KV。

2. 观察 query 选择 via Continual FPS
   在 query embedding 空间上持续运行 最远点采样 (Farthest Point Sampling, FPS)，从整个解码历史中选出少量代表性 query 作为 beacon queries。这些 beacon 是各个 query 簇的中心代表，能覆盖从普通 token 到 Thought Revisiting Tokens (TRT) 的不同 query 分布，FPS 的增量特性保证随解码推进不断更新代表集合。

3. 基于 beacon queries 的注意力 KV 评分
   用 beacon queries（而非仅最近窗口的 queries）作为代理，计算它们对所有历史 KV 的注意力权重（或聚合重要性分数），以此判断每个 KV 对在未来是否可能被 TRT 重新访问。重要性高的 KV 被保留，低分 KV 被驱逐。

输出

压缩后的 KV 缓存，显存占用大幅降低（最高 5.8×），同时几乎保持完整缓存的推理准确率，吞吐提升超过 4.3×。

与同类方法的差异

传统方法（如 SnapKV、H2O）只依赖最近 queries 估计长期重要性，无法应对 LRM 中 TRT 重新关注远期上下文的现象；BeaconKV 通过全局 beacon query 聚类预测远期重访，且全程 训练无关，更贴合长推理链的真实注意力模式。

实验

实验设计

BeaconKV 在 四个开源 LRM 和多个推理基准上评估，涉及长 CoT 生成场景。方法为 training-free，无需额外训练或微调。对比基线包括现有 KV 缓存压缩方法。

关键发现

• BeaconKV 在多个模型中普遍优于现有压缩方法。
• 最高实现 5.8× 内存减少，且接近保住全缓存精度。
• 吞吐量提升超过 4.3×，有效缓解长推理的显存瓶颈。
• 关键机制在于利用 beacon queries 预测将被重新访问的 KV 对，有效捕捉 Thought Revisiting Tokens (TRT) 的注意力模式。

与基线对比解读

现有方法依赖最近查询估计未来重要性，在长程推理中失准。BeaconKV 通过维护全局查询聚类的紧凑代表，避免了存储全部查询历史，更准确地预判远距离重访问。相比 SnapKV 这类依赖最近窗口的方法，BeaconKV 在保持精度的同时压缩比更高，吞吐提升显著。该差异归因于对查询空间几何结构的利用，而非仅时间局部性。

行业影响

落地场景

BeaconKV 主要适用于长推理链 场景，如数学解题助手、代码调试 Agent、复杂决策支持系统。这些服务常因 KV cache 内存膨胀导致 OOM 或被迫截断推理，BeaconKV 可在不牺牲准确率前提下压缩缓存。例如，在线教育平台的逐步解题工具可支持更长的推理步骤而不增加 GPU 成本；企业级金融分析报告生成中，模型需要回顾早期规划，BeaconKV 可避免重算并保持完整性。

商业价值

价值主要体现在降本 与 体验提升：

• 降本：5.8× 内存缩减与 4.3× 吞吐提升意味着相同硬件可服务更多并发请求，直接降低 per-token 成本。
• 体验提升：支持更长推理而不 OOM，减少用户遭遇截断或失败，尤其付费 API 场景可提高客户留存。

与现有工作流集成

BeaconKV 是训练无关 的动态淘汰策略，可无缝嵌入主流推理框架（如 vLLM、TensorRT-LLM）的 KV cache 管理层。集成方式：

1. 在 attention 计算前，用 continual FPS 选取少量 beacon queries 代表全局 query 分布。
2. 用 beacon queries 计算注意力分数，替代传统 recent-query 打分，决策哪些 KV 对保留。
3. 周期性触发淘汰，保持内存上限稳定。

这种方式不需要重新训练模型，可作为现有 SnapKV、H2O 的替代插件。对于已经部署长推理模型的服务商，只需替换 cache 淘汰模块即可获得收益。

原文指出“certain decoding steps generate Thought Revisiting Tokens (TRT) that re-attend to distant previous context”，这解释了为何仅依赖 recent queries 的方法在 LRMs 上失效。BeaconKV 的 beacon queries 设计直接针对该模式。

局限

• **泛化性验证有限**：论文在四个开源 LRM（可能包括 LLaMA、Qwen 等，但未具体）和几个推理基准上评估，但未覆盖闭源商业模型（如 OpenAI o1/o3、Claude 系列）或不同架构（如非 Transformer 或状态空间模型）。实验主要关注长 CoT 推理任务，对短文本生成、代码补全、多轮对话等非推理场景的效果未充分验证。若任务中没有明显的 Thought Revisiting Tokens 聚类现象，`beacon queries` 可能难以准确预测需要保留的 KV pairs，导致压缩后性能下降或 cache 预算浪费。
• **在线聚类与采样开销**：BeaconKV 需要在线维护 beacon queries，采用 Continual FPS 等策略选择代表性查询。虽然避免了存储完整 query 历史，但每次解码步骤仍需计算新 query 与 beacon 的相似度、可能更新 beacon 集合，这引入了额外的计算和逻辑复杂度。论文未详细报告这部分延迟占比，在某些低延迟批处理或极长序列场景下，该开销可能抵消部分吞吐收益。与 SnapKV、H2O 等仅使用最近窗口的直接方法相比，实现更复杂，对部署的友好度有影响。
• **超参数敏感性与压缩率限制**：方法依赖多个超参数（如 beacon 数量、聚类更新频率、压缩预算），论文虽做了消融，但未给出普适的默认配置或自动调参方案。在实际推理服务中，不同模型、任务和序列长度可能需要重新调参。另外，5.8 倍内存缩减虽显著，但在极端压缩率（如 10 倍以上）下是否仍能保持精度尚未探索，可能不适合移动端或严格内存约束的场景。