💻 计算机科学

Joint Optimization of Retrieval-Augmented KV Caching and Adaptive Speculative Decoding for Autoregressive Generation

本文介绍了 RAKV-ASD，这是一个通过学习控制器来联合优化检索增强型 KV 缓存与自适应投机采样，从而动态选择草稿长度，在保持与标准解码完全等价的输出结果的同时，相比固定策略实现了显著的延迟和吞吐量提升的系统。

原作者： Chenyi Huang, Zitan Chen, Zhenning Guo, Wentao Zhang

发布于 2026-10-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Chenyi Huang, Zitan Chen, Zhenning Guo, Wentao Zhang

原始论文采用 CC BY 4.0 许可（https://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

当大型语言模型逐词编写句子时，它面临着一个反复出现的瓶颈：为了预测下一个词，它必须记住之前出现过的每一个词。这种记忆以被称为键值缓存（key-value cache）的海量数字形式存储，并随着每一个新标记（token）的产生而增长。在现代计算中，检索这种记忆往往比实际计算下一个词的速度还要慢。为了提高速度，工程师们开发了两种主要技巧。第一种是如果起始词相同，则重复使用之前请求的记忆，就像跳过一本你已经读过的书的前几页一样。第二种是投机采样（speculative decoding），即由一个更小、更快的模型预先猜测后续的几个词，然后由主模型快速检查这些猜测是否正确。如果猜测正确，系统就可以跳过逐一生成这些词的缓慢过程。

研究人员提出的问题是，将这两种技巧结合起来是否会创造出一个超快速的系统。直觉上，这看起来很合逻辑：复用你已有的记忆，然后通过猜测来节省时间。然而，计算机硬件的现实往往不像直觉暗示的那样宽容。由约翰斯·霍普金斯大学、卡内基梅隆大学、纽约大学和南加州大学的 Chenyi Huang 及其同事开展的一项新研究，调查了这种精确的结合。他们构建了一个并不盲目应用两种技巧的系统。相反，它扮演着一名谨慎管理者的角色，不断测量每一种可能行动的成本，并针对每一次特定的请求，决定是复用记忆、进行猜测，还是两者都不做。

研究人员在配备标准笔记本电脑显卡的设备上，使用两对语言模型测试了他们的系统，称之为 RAKV-ASD。他们进行了数千次计时实验，向系统输入各种长度的提示词，并要求其生成不同大小的文本。其工作的核心是一个控制器，该控制器观察当前情况并从一系列选项中选择最佳路径。这些选项包括：在不进行猜测的情况下复用记忆；在复用记忆的同时，向前猜测一个、两个或最多八个词；或者干脆完全忽略记忆缓存，从头开始。该系统测量了每一步在硬件上实际花费的时间，而不是依赖于理论估算。

对于那些期望获得普遍加速的人来说，结果令人惊讶。研究发现，简单地将记忆复用与固定量的猜测相结合，往往会让系统变慢而非变快。在许多情况下，生成猜测并验证它们所花费的时间，比跳过原始工作所节省的时间还要长。研究人员发现，“猜测”这一过程只有在系统几乎确定猜测会正确时才值得去做。在实验中，控制器经常选择不对猜测做任何处理，而是选择一种复用记忆但逐一生成文本的设置。这种决策并非系统的失败，而是其逻辑的成功：它识别出猜测的开销对于特定的测试条件而言过高。

当研究人员将他们的智能自适应系统与一个强制系统始终向前猜测四个词的版本进行比较时，差异非常显著。与强制猜测的方法相比，自适应系统减少了 64.0% 的文本生成时间，并将每秒生成的单词数增加了 170.8%。这是因为自适应系统避免了失败猜测带来的沉重惩罚。然而，与仅复用记忆而不进行任何猜测的系统相比，自适应系统仅慢了 0.45%。这微小的差异表明，对于测试中所使用的特定硬件和模型，向前猜测并没有带来实质性的收益。该系统非常擅长识别何时“不”应该进行猜测，以至于它表现得就像一个从未进行过猜测的系统。

该研究还确定了一个明确的界限，即猜测何时才会变得有用。研究人员计算出，只有当猜测正确的概率达到一个极高的阈值——在 89.4% 到 97.3% 之间时，系统才会从猜测中获益。在他们的测试中，猜测的实际接受率低于此，这意味着条件尚不具备进行额外加速的优势。系统的控制器在做出这种区分方面表现得异常精准。尽管它无法完美预测任务执行的确切时间，但在决定是否进行猜测方面达到了 100% 的准确率。它正确地识别出，在几乎所有测试场景中，猜测的成本都超过了其收益。

这项工作为未来人工智能系统的发展提供了一个至关重要的教训：并非越复杂的策略就越好。将多种优化技术层叠在一起的想法听起来很高效，但如果这些部分不能与硬件能力完美匹配，它们就会相互抵消。研究人员表明，一个知道何时停止的系统通常比一个试图做所有事情的系统更快。通过将“猜测”这一决策视为一个变量选择而非固定规则，他们创造了一种能够适应机器现实的方法。研究结果表明，对于当前的硬件，最高效的路径通常是复用已知信息并谨慎地生成剩余内容，而不是带着可能需要被修正的猜测匆忙前进。

该研究在细节处理上保持了严谨性，运行了 3,600 次计时执行，以确保结果并非偶然。他们验证了自适应系统生成的文本与标准的、较慢的方法生成的文本完全一致，证明了速度的提升并非以牺牲准确性为代价。虽然研究结果针对的是实验中所使用的模型和硬件，但其原则具有更广泛的真理：在大型语言模型的复杂世界中，最明智的举动往往是识别出什么时候捷径实际上是绕路。研究人员并没有找到一个能加速所有 AI 生成过程的“万灵药”，但他们确实找到了一种可靠的方法，可以避免因试图表现得过于聪明而导致的减速。他们的工作为何时进行推测是有益的、何时是有害的，提供了一个基于数据的衡量边界，为未来构建更高效的系统提供了清晰的指南。

技术摘要：检索增强型 KV 缓存与自适应投机采样（Speculative Decoding）的联合优化

问题陈述
大语言模型（LLM）中的自回归生成依赖于重用键值（KV）状态，以避免重新计算因果历史。虽然检索辅助的 KV 缓存（在请求间重用状态）和投机采样（使用较小的模型进行草拟）解决了不同的瓶 quite 瓶颈，但它们的直接组合会产生效率低下。具体而言，当前缀命中（prefix hit）缩短了目标模型的预填充（prefill）阶段时，投机采样的固定开销（草拟、验证及潜在的拒绝）可能会在不再具有收益时依然存在。核心挑战在于：如何为每个请求确定是仅依赖于 KV 重用（不进行投机），还是启用投机采样，以及如果启用，应该草拟多少个 token（ $k$ ）。

方法论：RAKV-ASD
作者提出了 RAKV-ASD（检索增强型 KV 缓存与自适应投机采样），该框架联合优化了“重用 KV 状态的决策”与“投机草拟长度（ $k \in \{0, 1, 2, 4, 6, 8\}$ ）的选择”。

安全的跨请求 KV 重用：
- 系统存储完整提示词的 KV 张量、Token ID 和元数据。
- 通过精确的因果前缀验证而非语义相似性来确保安全性。使用前 16 个 token 的 12 字节 BLAKE2b 摘要选择候选桶，随后进行精确的最长公共前缀（LCP）验证。
- 对于 $q$ 个 token 的部分命中，克隆 $q-1$ 个位置，并重新计算边界处的 token。精确重复的部分可以重用存储的下一个 token 的 logits。
- 系统在固定内存预算（512 MiB）下运行，并采用 LRU 驱逐策略。
自适应控制器与成本模型：
- RAKV-ASD 不采用顺序优化机制的方法，而是使用一个实测控制器来选择最优动作 $(r, k)$ ，其中 $r \in \{0, 1\}$ 表示是否重用， $k$ 为草拟长度。
- 成本模型： 控制器利用实测的运行时统计数据而非理论预测来估计延迟。它考虑了目标模型/草拟模型的预填充时间、单 token 步进、针对 $k$ 个 token 的验证成本以及内存惩罚项。
- 决策逻辑： 控制器计算所有可行动作的预期成本 $C(k, r)$ 。它明确将 $k=0$ （不进行投机）作为一种有效的选项。超过内存预算的动作将被拒绝。系统选择最小化估计成本的动作。
- 在线统计： 系统追踪指数加权接受率（ $\alpha_t$ ）以更新预期接受进度 $E_k(\alpha)$ 。

核心贡献

精确前缀验证： 实现了一种带有边界内存限制的模型键值前缀验证，确保了安全性而无需语义近似。
实测成本模型： 一种基于观测到的运行时成本而非校准后的墙钟时间预测来对动作进行排序的实现匹配型成本模型。
联合选择机制： 一个将“无投机”（ $k=0$ ）视为一等公民动作的框架，根据当前工作负载特征动态地在仅重用路径和投机路径之间进行选择。
全面的评估： 一项涉及两个 GPT-2 系列模型对、3600 次定时执行的研究，包括敏感性分析和更大规模模型对的可行性研究。

实验结果
评估是在一台配备 8-GB 笔记本 GPU 的设备上进行的，使用了 GPT-2 和 DistilGPT-2（以及 GPT-2 Medium/GPT-2 对）在 WikiText-2 和 HumanEval 提示词上的贪婪解码任务。

性能对比基准：
- 与固定 $k=4$ 的联合方法（Joint-k4）相比，RAKV-ASD 实现了 64.0% 的延迟降低 和 170.8% 的吞吐量提升。这一改进归功于控制器频繁选择 $k=0$ ，从而避免了强制投机的开销。
- 与仅使用 KV 重用的方法（无投机）相比，RAKV-ASD 的延迟保持在 KV 仅重用基准的 0.45% 以内，吞吐量保持在 0.81% 以内。这表明 RAKV-ASD 并没有提供超越仅重用的普遍加速，但成功避免了次优投机带来的惩罚。
- 所有 3600 次执行生成的输出均与仅目标模型解码（AR）一致。
控制器行为：
- 控制器在选择最优 $k$ （在测试阶段主要为 $k=0$ ）方面达到了 100% 的准确率，在全动作选择方面达到了 83.3% 的准确率。
- 尽管在延迟的点预测方面平均绝对百分比误差（MAPE）高达 32.1%，但其对动作的排序非常准确（Spearman 相关系数为 0.780），导致平均遗憾值（mean regret）仅为 0.32%。
- 控制器正确识别出在测试场景中投机是不经济的，因为在这些场景下，最优的非零 $k$ 动作比 $k=0$ 慢 14–235%。
敏感性与边界：
- 研究确定了一个 预测接受率交叉点 $\alpha \ge 0.894$ 。只有当预测接受率超过此阈值时，投机才会变得有益。
- 在测试的 GPT-2 模型对中，接受率并未一致达到该边界，导致控制器选择了 $k=0$ 。
- 针对更大规模模型对（GPT-2 Large/GPT-2 Medium）的研究显示，虽然 $k=0$ 与非零 $k$ 动作之间的性能差距缩小了，但仍接近但未跨越理论边界。

意义与主张
本文主张 RAKV-ASD 提供的是一个 测量的边界和一种避免适得其反的投机的机制，而非普遍的加速。其主要意义在于证明了：

联合优化至关重要： 如果系统在仅靠重用即可满足需求时仍强制进行投机，那么简单的 KV 重用与投机采样组合可能会产生负面影响。
$k=0$ 是一个有效且往往是最优的策略： 在预填充阶段因重用而缩短的场景下，草拟和验证的串行成本往往超过其收益。
排序优于校准： 一个能够准确对动作进行排序（即确定哪个更快）的成本模型，在决策过程中比一个提供精确但可能失准的延迟预测的模型更有价值。

作者总结道，只有当测量的成本和预测的接受率支持时，才应启用投机采样，并且系统应在硬件或模型发生变化后进行重新校准。研究结果特定于所测试的 GPT-2 系列、贪婪解码以及单 GPU Batch-1 推理，并未声称其在多亿级参数模型、采样或连续批处理（continuous batching）中的普适性，除非经过进一步验证。

技术摘要：检索增强型 KV 缓存与自适应投机采样（Speculative Decoding）的联合优化

类似论文