论文

Random Attention: 重新思考高效推理的 KV Cache 驱逐

Random Attention: 重新思考高效推理的 KV Cache 驱逐

大型语言模型在需要扩展推理的任务上表现出色,但长思维链使 KV cache 成为严重的内存瓶颈。现有的 KV cache 压缩方法遵循同一范式:通过某种估计为每个缓存 token 打分以预测其未来重要性,并保留得分最高者。我们证明这种选择信号几乎毫无贡献。 Random Attention 保留 prompt 并在每个注意力头内均匀随机驱逐,无需计算任何分数;在四个模型和六个推理任务上,它匹配了最强的现有驱逐器,并在 vLLM 部署中吞吐量高出 32–43%。受控实验解释了这一现象: 1. prompt 是缓存中最脆弱的部分,选择器之间的差距大多只是其选择信号是否恰好保留了 prompt。 2. 推理轨迹通过两级冗余自我保护——文本层面(模型在推理中复述其仍需的内容)和跨注意力头层面(每个头保留轨迹的副本)——因此一旦 prompt 安全,随机抽取便足以保留模型所需的副本,无需打分。 我们的代码已开源:https://github.com/SalesforceAIResearch/Random-Attention。

论文精读

TL;DR Random Attention 发现 KV cache 逐出无需复杂评分:保留 prompt 并在每个注意力头内随机逐出 token,即可匹配最强选择器,vLLM 部署吞吐提升 32–43%。

问题

问题背景:长推理任务(如竞赛数学、代码生成)依赖长思维链提升性能,但 KV cache 内存随序列长度线性增长,成为显存与吞吐的硬瓶颈。业界正积极寻找低开销的 KV cache 压缩方案。

现有方法局限:主流方法对每个缓存 token 计算重要性评分(如注意力权重、梯度范数),保留高分 token。但这类方法在工程上存在明显缺陷:

  • 评分计算本身引入额外前向开销,抵消部分吞吐收益;
  • 选择信号在推理轨迹上区分度很低:模型通过文本复述和多头注意力的副本冗余保护关键信息,使得“未来重要性”难以提前准确预测;
  • 对 prompt 的保护不足,许多方法在长 prompt 场景下性能显著下降,因为 prompt 一旦被误驱逐,推理质量将不可逆受损。

为什么难/重要:长推理的 KV cache 可达数十万 token,显存压力极大;生产环境需要低延迟、高吞吐,eviction 策略必须既轻量又鲁棒。设计一个通用且开销低的评分器本身困难,且不同任务下最优信号不一致,增加了部署复杂度。这个问题直接关系到长上下文推理模型的部署成本与可扩展性。

行业类比:类似流式视频处理中的关键帧选择——与其费力预测哪些帧重要,不如随机抽样并始终保留起始段,因为后续内容自带冗余,过度优化往往得不偿失。

核心洞察

  • 选择信号在推理任务的 KV 缓存淘汰中几乎不贡献价值。随机保留 prompt 并在各注意力头内均匀随机淘汰,与最强选择器性能持平。这动摇了“评分函数决定缓存质量”的假设,揭示在长链推理中,正确保留 prompt 就够了,其余 token 的取舍并不依赖于复杂的重要性估计。
  • 推理轨迹通过文本冗余和跨头复制形成自保护:模型在推理中不断重申所需信息,且每个注意力头有自己的副本。因此,随机淘汰后仍有足够副本可用。这解释了为何无需评分挑选,也为 KV 缓存管理提供了新思路——利用模型自身冗余,而非设计更精细的选择器。
  • 在 paged serving 下跳过评分 pass 使 Random Attention 吞吐比 TriAttention 高 32-43%。因为评分计算成为瓶颈。这意味着对长推理负载,减少每次前向的额外计算比提升淘汰精度更能改善端到端性能。工程上应优先考虑轻量缓存策略。

方法

方法核心:随机驱逐与 Prompt 保留

Random Attention 重新审视了长链推理场景下的 KV cache 压缩问题。传统方法遵循同一范式:为每个缓存 token 计算一个重要性分数 s(例如基于累积注意力或未来预测损失),然后保留得分最高的 token。Random Attention 提出完全不同的策略:保留 prompt 中的所有 token,然后在每个 注意力头 内部对推理过程中产生的 token 进行均匀随机驱逐,不计算任何分数。

输入:一个 Transformer 模型在长推理任务中生成的 KV cache,包含 prompt 部分和逐步生成的推理 trace。

关键模块:

  1. Prompt 保护:始终保留最初输入 prompt 对应的所有 KV 项,因为实验证明 prompt 是缓存中最脆弱的部分,任何选择信号的差异大多取决于是否碰巧保住了 prompt。
  2. 逐头独立随机驱逐:对于超出预算 K 的推理 token,在每个 attention head 内独立地按均匀分布随机选择要驱逐的 token(或等价地随机保留 K 个 token),不同 head 的选择相互独立。这利用了推理 trace 在跨注意力头层面存在冗余:每个 head 都保有自己的 trace 副本,因此即使随机保留一部分,多个 head 中总有足够副本覆盖模型后续所需的内容。
  3. 无评分计算:方法完全跳过重要性估计步骤,因此省去了推理时的额外前向或统计开销。

输出:压缩后的 KV cache,直接用于后续生成,在 vLLM 部署中可实现比最强先验驱逐器(TriAttention)高 32–43% 的吞吐量。

跟同类方法的差异点:与所有基于重要性评分的驱逐器不同,Random Attention 完全放弃选择信号,靠随机性和推理 trace 的内在冗余实现同等任务质量,同时显著降低计算开销。

实验

实验设计

论文在四个模型、六个推理任务上对比 Random Attention 与现有 KV cache 驱逐方法(如 TriAttention 等基于评分的选择器)。采用周期性驱逐预算 K 和缓冲 r,基线是不同评分函数 s 的选择策略。效率评测在 vLLM 分页服务下进行。

关键发现

  • 在数学和科学推理任务上,选择信号没有带来收益:没有 selector 显著拉开差距。竞争数学上无一个 selector 胜出;代码推理上由于 prompt 更长,大多基于信号的 selector 崩坏。
  • 压缩压力加大时差距扩大,且在各方法族中均如此。
  • 为什么选择信号作用甚微?因为 prompt 是缓存中脆弱的部分,而推理轨迹通过冗余自我保护:文本层面模型重述所需内容,注意力头层面各头保留自己的轨迹副本。一旦 prompt 安全,随机抽取足以保留足够副本,无需评分。
  • 效率上,分页服务下 Random Attention 比 TriAttention 快 32–43%;同等内存下所有驱逐器都有收益,评分传递开销决定其余差异。

与基线对比的深度解读

现有方法基于 token 重要性评分保留 top-scoring,但随机驱逐在精度上匹配甚至超越,说明评分信号在推理场景中并非必要。真正的瓶颈在于 prompt 保留和利用推理轨迹的冗余性。这一发现挑战了缓存压缩中“必须评分”的固有假设,为工程实现提供了更简单的选择:跳过评分计算可显著提升服务吞吐,且不牺牲下游任务表现。

行业影响

落地场景

Random Attention 针对长推理链导致 KV cache 显存爆炸的 LLM 服务场景。典型产品包括:多步骤数学/代码求解、金融研报推理问答、企业知识库问答、AI agent 任务规划与执行。这些场景中 prompt 通常较长,且模型会生成冗长的思考过程,显存瓶颈直接影响可服务的并发量和响应延迟。

商业价值

  • 降本:用均匀随机逐出替代逐 token 打分,去除额外的注意力计算或价值估计过程,推理吞吐提升 32-43%,节省 GPU 租赁成本。
  • 体验提升:在相同显存预算下支持更高并发数,降低长任务排队等待;同时保留 prompt 可避免关键上下文丢失,维持推理准确率。
  • 风险可控:论文在四个模型、六个推理基准上验证,随机策略匹配最强的选择式逐出方法,不会带来明显精度损失。

与现有产品/工作流的接口

集成方式非常轻量:直接在 vLLM 等 paged attention 推理框架中替换 KV cache 管理模块。无需额外的打分网络或模型,只需实现以下逻辑:

  1. 永久保留 prompt 对应的 KV 缓存;
  2. 对生成 token 在每个 attention head 内均匀随机采样逐出;
  3. 沿用现有的逐出周期 K 和 buffer r 配置。

改动集中在 cache eviction 接口,不影响模型权重或训练流程,适合作为默认策略或通过 A/B 测试切换。对于已有定制打分逐出器的团队,可快速替换以验证吞吐收益。

具体 use case:

  1. 代码辅助平台:用户提交长上下文代码库问答,模型生成多步推理。采用随机逐出后,每 GPU 可同时处理的请求数提升,P99 延迟下降。
  2. 电商智能客服:处理订单、优惠、退换货政策的组合推理。长对话 prompt 与长推理 trace 叠加,随机逐出能显著缓解显存压力,使更多工单并行处理。

局限

  • **Random Attention** 主要针对推理任务(数学、科学、代码)设计,这些任务的推理痕迹具有高度冗余。对于需要精确检索的长上下文理解任务(如长文档问答、多轮对话、事实核查),随机驱逐可能丢失关键证据 token;且此类任务中 prompt 本身可能非常长,保留完整 prompt 会占据大量缓存,导致可驱逐的推理 token 很少,压缩收益明显下降。论文未在非推理任务上验证,适用性存疑。
  • 方法的核心假设是推理痕迹在文本和跨注意力头两个层面都有足够冗余,但这一假设依赖于模型在推理过程中会不断复述所需信息。对于需要精确数值计算或引用独特中间结果的推理任务,复述可能不充分,随机驱逐会导致正确率波动。论文虽报告了 run-to-run variability,但未深入分析在极低预算下的稳定性,极端压缩时生成质量可能剧烈抖动。
  • 与选择性驱逐方法相比,**Random Attention** 省去了评分开销,但当 prompt 占据大部分缓存时,它能驱逐的推理 token 数量有限,压缩比上限受 prompt 长度制约。此外,论文未讨论与其他 KV cache 压缩技术(如量化、稀疏化、跨层共享)的协同,仅与基于评分的驱逐方法比较,未覆盖更广泛的优化空间。
论文Heng Wang2026-09-03原文

相关内容