HyperAIHyperAI

Command Palette

Search for a command to run...

LLM
推理

随机注意力:重新思考 KV 缓存淘汰以实现高效推理

Heng Wang Jielin Qiu Wenting Zhao Cheng Qian Liangwei Yang Jiawei Han Heng Ji Silvio Savarese Shelby Heinecke Huan Wang

摘要

大语言模型在需要扩展推理的任务上取得了卓越性能,但长思维链使 KV 缓存成为严重的内存瓶颈。现有的 KV 缓存压缩方法共享一种范式:通过某种对未来重要性的估计为每个缓存 token 打分,并保留得分最高的 token。我们证明,这种选择信号几乎不起作用。随机注意力保留提示,并在每个注意力头内均匀随机地进行淘汰,完全不计算任何分数;在四个模型和六个推理任务上,它的表现与最强的先前淘汰器相当,而在 vLLM 部署中,其吞吐量比该淘汰器高出 32–43%。受控实验对此进行了解释,表明:1)提示是缓存中脆弱的部分,不同选择器之间的差距主要在于它们的选择信号是否恰好保留了提示;2)推理轨迹通过两个层面的冗余来保护自身免受淘汰,即在文本层面(模型在工作过程中重申其仍然需要的内容)和跨注意力头层面(每个头保留自己的轨迹副本),因此一旦提示安全,随机抽取就能保留足够多的模型仍需要的副本,无需任何分数来挑选它们。我们的代码已在 https://github.com/SalesforceAIResearch/Random-Attention 上公开。

一句话总结

Random Attention 由 Salesforce AI Research 和伊利诺伊大学厄巴纳-香槟分校提出,是一种 KV 缓存淘汰方法,它保留提示词(prompt)并在每个注意力头内随机淘汰 token,无需评分,在四个模型和六个推理任务上匹配最强的前置淘汰器,同时在 vLLM 部署中提供 32\textendash43%32\textendash43\%32\textendash43% 更高的吞吐量,并证明选择信号是不必要的,因为提示词是脆弱的,而推理轨迹是冗余的。

核心贡献

  • Random Attention 是一种 KV 缓存淘汰策略,它保留提示词并在每个注意力头内均匀随机地淘汰 token,无需计算重要性分数。它在四个模型和六个推理任务上匹配最强的前置淘汰器,同时在 vLLM 部署中提供 32–43% 更高的吞吐量。
  • 受控实验表明,提示词是缓存中脆弱的部分,一旦被保留,选择信号之间的性能差距几乎消失。这解释了为什么随机淘汰能匹配复杂的评分方法。
  • 推理轨迹通过两个层面的冗余保护自身:文本重述和跨头复制,因此当提示词安全后,随机抽取保留了足够多的所需信息副本,使得选择分数变得不必要。

引言

为推理任务设计的大型语言模型会生成极长的思维链,导致键值(KV)缓存线性增长并造成严重的内存瓶颈。先前的工作通过在固定预算下淘汰缓存的 token 来解决这一问题,使用启发式评分规则(累积注意力、值大小、键统计量),假设分数决定了压缩下的准确度。作者直接检验了这一前提,发现选择信号的贡献几乎为零。他们引入了 Random Attention,一种简单的策略,保留提示词并在每个注意力头内均匀随机淘汰 token,它在多个模型和推理任务上匹配或超越强基线,同时提供 32–43% 更高的吞吐量,因为它从不运行评分过程。

方法

作者提出了 Random Attention,一种简单且无信号的淘汰策略,用于在自回归生成过程中管理 KV 缓存。其核心思想是将不可替代的输入(提示词)与模型生成的推理轨迹分开,并对二者应用不同的保留策略。提示词仅被陈述一次,一旦被淘汰便无法恢复,因此必须完全保护。相反,推理轨迹在生成过程中不断被重新访问和重述,使其高度冗余,适合随机采样。

该方法包含两条结构规则。第一,所有属于预填充(prefill)的位置——包括系统提示词、聊天模板和用户问题——都被永久保留。这些位置 1,,p1,\dots,\ell_p1,,p 永远不会被淘汰。第二,每个其他缓存位置获得一个独立同分布的均匀随机分数,每个 KV 头独立抽取。然后每个头根据这些随机分数保留其 top-kkk 个位置,其中 kkk 由可用的缓存预算决定。这种逐头独立性将保留的预算均匀地分散在整个推理轨迹上,并确保不同头保留轨迹的不同子集。

形式上,每个缓存位置 iii 的分数 sis_isi 定义为

si={+,ip(the prompt),uiUniform(0,1),otherwise,s_i = \begin{cases} +\infty, & i \leq \ell_p \quad (\text{the prompt}), \\ u_i \sim \text{Uniform}(0,1), & \text{otherwise}, \end{cases}si={+,uiUniform(0,1),ip(the prompt),otherwise,

随机抽取在每个淘汰步骤中对每个 KV 头独立进行。然后淘汰步骤简单地选择每个头中分数最高的位置,丢弃其余部分。每次淘汰的计算成本极低:每个头只需一次随机数生成和一次 top-kkk 操作。

Random Attention 既是一种实用的部署方法,也是 KV 缓存淘汰的零假设。因为它不使用来自注意力权重或隐藏状态的任何信号,任何基于信号的淘汰器在相同预算下若不能超越它,则说明其信号未能提取出有用信息。作者随后证明,这种随机策略出人意料地有效,因为提示词是脆弱的且必须保留,而工作状态在文本重述和多个注意力头中冗余存储。独立的逐头随机抽取最大化了至少一个头保留任何所需值副本的机会,利用了模型固有的跨头冗余,而无需任何选择启发式。

实验

评估在 Qwen3 和 Phi-4-reasoning 模型上,将随机 KV 缓存淘汰(Random Attention)与几种基于信号的淘汰器在数学、科学和代码推理任务上进行比较,压缩比约为 4 倍。主要结果表明,Random Attention 在数学和科学任务上匹配或优于学习到的淘汰器,而在代码推理中,保护提示词消除了大部分性能差距。进一步分析发现,提示词是脆弱的部分,而模型自身的推理轨迹在文本和注意力头之间高度冗余,使得随机淘汰具有鲁棒性。效率测试表明,跳过评分过程在分页服务下可带来 32–43% 更高的吞吐量,因为评分开销会随着并发请求而累积。

一种简单的随机淘汰策略(Random Attention)在高压缩比下,在数学、科学和代码推理任务上的准确度匹配或超过注意力评分方法。在数学和科学基准测试中,没有淘汰器显著优于 Random Attention,而在具有长提示词的代码任务中,SnapKV 和 VaSE 出现了大幅下降。基于选择的方法唯一显著的优势是 TriAttention 在一个模型的代码任务上获得了适度提升。Random Attention 在所有模型上在 MATH500 和 GPQA-D 上显著优于 SnapKV 和 VaSE,在 Qwen3-4B 上优于 R-KV。在 LiveCodeBench 上,SnapKV 在每个模型上比 Random Attention 低 20–35 分,VaSE 在 Phi-4-reasoning 上崩溃,落后 29 分。

保护提示词能显著提升所有方法的性能,而一旦提示词被保留,选择信号的选择就几乎无关紧要。默认丢失提示词的方法,如 SnapKV 和无信号策略,获得了最大的提升,带有提示词保护的 Random Attention 成为最佳整体策略。SnapKV 的默认分数丢弃了大部分提示词,当显式保留提示词时,其性能提升高达 22.5 分。带有提示词保护的 Random Attention 匹配或超过所有基于信号的淘汰器,证实了提示词保留而非评分方法是准确度的主要驱动因素。

当一个密码只被陈述一次,并在 57 轮压缩后需要使用时,只有 R-KV(在整个历史上累积注意力)能可靠地检索到它(84% 的轨迹,接近确定的对数概率)。Random Attention 和优先考虑最近窗口的方法几乎完全失败,而 VaSE 大约有三分之一的时间能检索到。Random Attention 从未复现密码,其平均对数概率为 -18.35,表明信息已从缓存中有效丢失。R-KV 在 83.6% 的轨迹中检索到密码,对数概率为 -0.71,接近确定。VaSE 在 34.4% 的时间检索到密码,而 SnapKV 和 TriAttention 的检索率低于 2%,与随机选择类似。检索成功率直接与每种策略所评分的注意力统计量相关:全历史累积保留了事实,而最近窗口信号则不能。

Random Attention 在四个模型上使用 vLLM 和 PagedAttention 时,吞吐量是全注意力的 1.6–2.7 倍,比 TriAttention 高 32–43%。这些提升来自于更小的 KV 缓存,允许更多并发请求,以及消除评分过程,这降低了每次压缩的延迟并避免了批处理服务中的同步停滞。Random Attention 在所有测试模型上实现了全注意力 1.6–2.7 倍的吞吐量,比 TriAttention 高 32–43%。跳过评分过程将每次淘汰时间从 1.47–1.64 毫秒减少到 0.30 毫秒,累积的节省转化为观察到的服务加速。

Random Attention 在数学、科学和代码推理任务上匹配或超越注意力评分淘汰方法,提示词保留成为准确度的主要驱动力,而非评分信号的选择。在长距离检索测试中,只有在整个历史上累积注意力的策略能可靠地保留远距离信息,而随机和最近窗口淘汰器几乎完全丢失了它。Random Attention 还通过消除评分过程并实现更高并发,相比全注意力和基于评分的淘汰器,提供了显著的吞吐量提升。


用 AI 构建 AI

从创意到上线——通过免费 AI 协同编码、开箱即用的环境和最优惠的 GPU 价格,加速您的 AI 开发。

AI 协同编码
开箱即用的 GPU
最优定价

HyperAI Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供