基于顿悟感知的KV缓存淘汰方法(无需注意力矩阵)

arXiv cs.LG 论文

摘要

本文介绍了EpiKV,一种基于内部表征变化(顿悟分数)而非注意力权重来评估token重要性的KV缓存淘汰方法,无需具体化注意力矩阵。该方法在推理基准测试中取得了具有竞争力的性能,同时支持长达16倍的上下文长度。

arXiv:2606.26472v1 公告类型:新 摘要:随着推理模型输出动辄数万token的思维链,KV缓存日益成为部署瓶颈。现有缓存淘汰方法根据注意力权重对token进行排序,但在长推理轨迹中,注意力权重是一种有噪声的重要性代理,并且通过迫使模型具体化注意力矩阵,阻碍了生产推理中使用融合内核。在本文中,我们转而使用一种称为“顿悟分数”的指标对token进行评分:该分数直接从前向传播中读取模型内部表征的变化,无需注意力矩阵且额外状态可忽略不计。由此产生的缓存淘汰方法EpiKV无需训练、分类器或定制内核,可直接在现有的FlashAttention推理堆栈中使用——相较于基于注意力的评分,可将可行上下文长度扩展至16倍。上层中间层(负向)并利用因果滚动z分数消除位置趋势。在4096 token缓存下,EpiKV在MATH-500上达到72%,与最强的注意力基线(ThinKV 71%,H2O 67%)持平;在8192 token下,一种延迟归一化KV变体在AIME-2024上达到37%,优于最佳基线(33%),速度提升至2.8倍。
查看原文
查看缓存全文

缓存时间: 2026/06/26 05:20

# 无需注意力矩阵的顿悟感知 KV 缓存驱逐策略

来源:https://arxiv.org/html/2606.26472

Steven Kolawole  
语言技术研究所  
卡内基梅隆大学  
[email protected]

&Virginia Smith  
机器学习系  
卡内基梅隆大学  
[email protected]

###### 摘要

随着推理模型产生数万 token 长的思维链,KV 缓存日益成为部署瓶颈。现有缓存驱逐方法通过注意力权重对 token 进行排序,但在长推理轨迹中,注意力权重是一个有噪声的重要性代理,并且迫使模型物化注意力矩阵,从而阻碍了生产推理中融合核的使用。在这项工作中,我们转而使用一种称为*顿悟*分数的指标对 token 进行评分:该分数是模型内部表示的变化,可直接从前向传播中读取,无需注意力矩阵且几乎不增加额外状态。由此产生的缓存驱逐方法 EpiKV 无需训练、分类器或自定义内核,可直接在 FlashAttention 推理栈中无需修改地使用——其可行上下文长度比基于注意力评分的方案长 16 倍。在 4096 token 缓存下,EpiKV 在 MATH-500 上达到 72%,与最强的基于注意力的基线(ThinKV 71%,H2O 67%)持平;其滞后归一化 KV 变体在 AIME-2024 上以 8192 token 缓存达到 37%,优于最佳基线(33%),速度提升高达 2.8 倍。

# 无需注意力矩阵的顿悟感知 KV 缓存驱逐策略

Steven Kolawole  
语言技术研究所  
卡内基梅隆大学  
[email protected]

Virginia Smith  
机器学习系  
卡内基梅隆大学  
[email protected]

## 1 引言

DeepSeek-R1(Guo 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib19)) 等推理模型通过生成长思维链来解决难题;单个竞赛数学问题可能需要数万 token 的内部推理才能得出答案。键值 (KV) 缓存随此长度线性增长,并迅速成为部署的内存瓶颈:在 10^4–10^5 个解码 token 时,它占用绝大部分设备内存,并限制了服务器可容纳的批次大小(Kwon 等, 2023 (https://arxiv.org/html/2606.26472#bib.bib1))。KV 缓存驱逐通过仅保留 K 个 token 的预算来解决此问题,但这引出了每种方法必须回答的问题:哪些 token 重要?

现有针对推理轨迹的解码时驱逐方法通过考虑注意力权重来回答此问题(Zhang 等, 2023 (https://arxiv.org/html/2606.26472#bib.bib9); Ramachandran 等, 2026 (https://arxiv.org/html/2606.26472#bib.bib11); Hu 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib12); Su 等, 2026 (https://arxiv.org/html/2606.26472#bib.bib24))。然而,注意力权重有关键缺陷。首先,它是重要性的有噪声代理:注意力汇无论内容如何都会吸收权重(Xiao 等, 2024 (https://arxiv.org/html/2606.26472#bib.bib10)),填充 token 在生成时会吸引注意力权重,但之后永远不会被再次引用。其次,它在架构上代价高昂:读取注意力权重需要物化注意力矩阵,而 FlashAttention 等最先进方法正是为了避免这一点而构建的(Dao, 2024 (https://arxiv.org/html/2606.26472#bib.bib15))。设置 `output_attentions=True` 会强制使用 eager 内核,并在几乎每个推理轨迹长度以下耗尽 80 GB A100 内存,而在相同模型上使用 FlashAttention 前向传播可扩展一个数量级(图 1 (https://arxiv.org/html/2606.26472#S1.F1))。

![参见说明文字](图1:80 GB A100 上单次前向传播的峰值 GPU 内存与上下文长度的关系。读取注意力权重(eager)呈二次方增长,并在 8192 token 时耗尽 GPU;我们的方法读取的前向传播可扩展到 65,536 token——可行上下文长度延长 16 倍。这是无需注意力矩阵带来的架构收益(详见 §4.4 (https://arxiv.org/html/2606.26472#S4.SS4))。

我们引入*顿悟感知* KV 缓存驱逐(EpiKV),它根据模型内部表示(特定层的隐藏状态和 KV 向量)的变化而非注意力权重对 token 进行评分,并且从标准前向传播中读取,无需注意力矩阵。该名称指的是推理轨迹中的转折点(例如,一个已完成的步骤、一个已确定的洞察),在这些点上残差流变化最大,我们发现这些正是值得保留的 token。

#### 贡献。

1. 我们识别出一个 32 层推理模型中的两个带区层解剖结构:第 7–13 层(A 带)的隐藏状态变化与 token 重要性正相关,第 18–25 层(B 带)则负相关(以反事实遮挡标签衡量)。组合信号优于我们测试的所有基于注意力的信号。
2. 我们发现原始信号在轨迹内存在单调的位置趋势,因为它追踪位置和内容一样多,并且我们展示了因果滚动 z 分数可以消除这种趋势,恢复驱逐质量。
3. 在可部署的预算下,我们无需注意力矩阵的方法在 MATH-500 和 AIME-2024 上都匹配或超过了最强的基于注意力的基线(§4.3 (https://arxiv.org/html/2606.26472#S4.SS3))。
4. 我们量化了工程收益:在相同预算下,我们的方法比基于注意力的驱逐快 2.8 倍,并且避免了使基于注意力的评分在长上下文中不可行的注意力矩阵内存墙。

这些共同使得驱逐可在标准 FlashAttention 服务栈中部署(§5 (https://arxiv.org/html/2606.26472#S5.SS0.SSS0.Px2))。我们发布反事实重要性标签作为验证资源。

![transformer layers one forward pass at decode step t Band A Band B importance s(t) = z_A - z_B rolling z-score of ||Δh_l|| sink decode tokens KV cache: keep top-K, evict low-score (×) rank Attention-based scoring builds n×n matrix (O(n^2) memory); we never do. n×n attn](图2:如何在一个解码步骤计算重要性分数。对模型层进行一次前向传播,得到隐藏状态和 KV 缓存。我们读取 A 带(第 7–13 层,与重要性正相关)和 B 带(第 18–25 层,负相关)的隐藏状态变化,组合成一个因果滚动 z 分数 s(t) = z_A - z_B,并在 KV 缓存中保留前 K 个 token。与基于注意力的驱逐不同,该分数不需要 n×n 注意力矩阵,因此不增加内存,并且与融合注意力内核(如 FlashAttention)及其构建的推理栈兼容。

## 2 相关工作

#### 基于注意力的 KV 驱逐。

大多数驱逐方法通过注意力权重对 token 进行排序。H2O 保留累积注意力的“重击者”(Zhang 等, 2023 (https://arxiv.org/html/2606.26472#bib.bib9));StreamingLLM 保留注意力汇加上一个近期窗口(Xiao 等, 2024 (https://arxiv.org/html/2606.26472#bib.bib10));SnapKV 从提示末尾的观察窗口中选择上下文 token(Li 等, 2024a (https://arxiv.org/html/2606.26472#bib.bib16));PyramidKV 为较低层分配更大的预算(Cai 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib21));ChunkKV 驱逐连续块以保留局部语义(Liu 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib20))。这些方法针对长*输入*,并且都需要注意力分布;这需要物化 n×n 注意力矩阵,因此排除了生产推理所依赖的融合内核(例如 FlashAttention(Dao 等, 2022 (https://arxiv.org/html/2606.26472#bib.bib14); Dao, 2024 (https://arxiv.org/html/2606.26472#bib.bib15)))。我们直接测量这一成本(第 4 节 (https://arxiv.org/html/2606.26472#S4))。

#### 推理感知驱逐。

第二条线针对推理模型的长*生成*轨迹,其中注意力是非单调的,里程碑 token 在最后一次被关注后很久仍然重要。ThinKV 通过注意力稀疏性分类思想片段,并通过自定义内核应用按类型量化和驱逐(Ramachandran 等, 2026 (https://arxiv.org/html/2606.26472#bib.bib11))——需要注意力权重、稀疏性阈值和层子集的离线校准以及 token 块刷新窗口;RaaS 使用完全预填充保留的注意力刷新 LRU 时间戳(Hu 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib12));LongFlow 在同一模型类上以 ||softmax(scores) V||_1 评分(Su 等, 2026 (https://arxiv.org/html/2606.26472#bib.bib24));AhaKV(Gu 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib23)) 和 CAOTE(Goel 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib22)) 改进了基于注意力的分数;LagKV 针对滞后窗口对 KV 统计量进行归一化,避免注意力(Liang 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib13))。除 LagKV 外,所有方法都从注意力中导出信号。我们则使用残差流和缓存的 KV 向量中的表示变化。

#### 检索和量化。

正交方向在无需选择丢弃哪些 token 的情况下减少 KV 成本:检索保留每个 token 并按步骤获取子集(Tang 等, 2024 (https://arxiv.org/html/2606.26472#bib.bib26); Liu 等, 2026 (https://arxiv.org/html/2606.26472#bib.bib27)),SideQuest 提示模型删除过时的工具响应(Kariyappa and Suh, 2026 (https://arxiv.org/html/2606.26472#bib.bib30)),量化降低保留条目的精度(Hooper 等, 2024 (https://arxiv.org/html/2606.26472#bib.bib25); Sharma 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib31))。所有这些都是可堆叠的,与我们的信号互补。

#### 隐藏状态作为重要性信号。

网络中层承载模型的关键计算:ROME 和 MEMIT 将事实回忆定位到中层前馈模块(Meng 等, 2022 (https://arxiv.org/html/2606.26472#bib.bib17), 2023 (https://arxiv.org/html/2606.26472#bib.bib18)),这些模块充当键值记忆(Geva 等, 2021 (https://arxiv.org/html/2606.26472#bib.bib29))——正是我们发现与 token 重要性最强正相关的同一层(7–13)。投机性解码提供了收敛证据:EAGLE 从隐藏状态而非 token 嵌入进行草稿,因为它们携带更丰富的预测结构(Li 等, 2024b (https://arxiv.org/html/2606.26472#bib.bib28))。

#### 定位。

目前没有针对推理轨迹的解码时驱逐方法将非注意力重要性信号与无需注意力矩阵的评分相结合。ThinKV、RaaS 和 LongFlow 是推理感知的,但源自注意力;LagKV 是无需注意力的,但通用且仅限 KV。我们的方法两者兼具,并增加了层级别说明——重要性所在的正中层区域。

## 3 方法

### 3.1 问题设置

在自回归解码过程中,键值 (KV) 缓存每个生成的 token 每层增加一个条目。对于一个在具有 L 层、H 个维度为 d_h 的键值头的模型上的 n 个 token 的推理轨迹,缓存持有 2 L H d_h n 个标量,对于 10^4–10^5 个 token 的轨迹,这主导了设备内存。解码时驱逐将缓存限制为 K 个 token 的预算:每一步,一个策略对缓存的位置进行评分,并保留得分最高的 K 个,永久丢弃其余部分。

我们称一个策略为*FlashAttention 兼容*(FA2 兼容),如果它仅使用 (i) 已存于高带宽内存中的缓存键和值,以及 (ii) 标准 `output_hidden_states` 接口暴露的每层隐藏状态来计算其分数。这样的策略从不请求 `output_attentions`,也从不物化 n×n 注意力矩阵,因此它在 FlashAttention 前向传播内部运行,而不会强制使用 eager 回退。如果策略需要注意力矩阵(等效地,`output_attentions=True`),则称为*需要注意力*,这会禁用 FlashAttention 的分块并为评分重新引入 O(n^2) 峰值内存。图 2 (https://arxiv.org/html/2606.26472#S1.F2) 对比了两种模式。

### 3.2 隐藏状态方差信号

注意力权重是此前所有针对推理轨迹的解码时驱逐方法使用的代理,它既是有噪声的重要性信号,又在架构上难以提取(§1 (https://arxiv.org/html/2606.26472#S1));我们始终将这两个反对意见分开。

我们的信号是 token 级别残差流的变化。对于层 l 和解码位置 t,令 h_l(t) 为隐藏状态,定义 L2 差异

g_l(t) = ||h_l(t) - h_l(t-1)||_2. (1)

大的 g_l(t) 表示生成 token t 使模型在层 l 的内部状态发生了偏移,这是重要 token(中间结果、完成的步骤、从探索性推理到收敛性推理的过渡)的标志,而非流畅的填充 token。我们将这些转折点称为*顿悟* token。

#### 双带区解剖。

针对反事实重要性标签(第 3.4 节 (https://arxiv.org/html/2606.26472#S3.SS4))的逐层相关性研究识别出两个在竞赛数学问题上具有一致且相反行为的带区。*A 带*(第 7–13 层)具有一致正的 Spearman ρ:高的 g_l 标志着重要 token。*B 带*(第 18–25 层)具有一致负的 ρ:高的 g_l 标志着可丢弃的 token。我们在 §5 (https://arxiv.org/html/2606.26472#S5) 中解释这两个带区;这种划分与中层事实检索一致(Meng 等, 2022 (https://arxiv.org/html/2606.26472#bib.bib17), 2023 (https://arxiv.org/html/2606.26472#bib.bib18); Geva 等, 2021 (https://arxiv.org/html/2606.26472#bib.bib29))。

我们将这两个带区组合成一个单一分数

s(t) = ĝ_{10}(t) - ĝ_{21}(t), (2)

其中 ĝ_l(t) 是在过去 w=64 个 token 的滑动窗口上 g_l 的滚动均值。该窗口是因果的(只使用 ≤t 的位置),因此 token t 的分数从不依赖于未来 token。具有高 s 的 token 被保留。

#### 时间趋势校正。

原始分数 (2) 携带一个我们在分析中发现并作为方法论发现报告的混杂因素。在单个轨迹内,ĝ_{10} 倾向于降低,而 ĝ_{21} 倾向于随位置增加,因此 s(t) 追踪位置和内容一样多:在短轨迹中,它可能将早期(可丢弃的)token 排在晚期(承载负载的)token 之上。激励这些带区的总体 ρ 部分由跨问题结构驱动,并高估了轨迹内的排序质量。我们使用因果滚动 z 分数进行校正,

z_l(t) = (g_l(t) - μ_l(t)) / (σ_l(t) + ε), (3)

其中 μ_l(t) 和 σ_l(t) 是 g_l 在滞后窗口上的均值和标准差,并使用 z_{10}(t) - z_{21}(t) 进行评分。这将绝对幅度(受位置污染)转换为局部偏差(与位置无关),类似于滞后相对归一化(Liang 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib13)) 和分析去趋势(Gu 等, 2025 (https://arxiv.org/html/2606.26472#bib.bib23)),但应用于隐藏状态差异。去趋势变体 EpiKV 是我们的主要方法。

![参见说明文字](图3:MATH-500(左,n=100)和 AIME-2024(右,n=30)上准确率 vs. 缓存预算。)

相似文章

面向长推理的信息感知KV缓存压缩

arXiv cs.CL

本文提出InfoKV,一种熵感知的KV缓存压缩框架,结合了token级别的预测不确定性和注意力分数,以提高长上下文推理效率。实验表明,它在Llama-3.1、Llama-3.2和DeepSeek-R1上优于现有的基于注意力的方法。

来自未来的回望:基于反因果惊奇度的键值缓存管理

arXiv cs.LG

本文提出了一种KV缓存逐出策略,通过反因果惊奇度对令牌进行评分,移除那些可以从未来上下文中很好预测的过去令牌。该方法无需训练、分布内,并通过快速的单层近似实现了具有竞争力的性能。