无需妥协的遗忘：固定预算下流式KV-Cache驱逐的Nexus采样

# 在固定预算下进行流式 KV 缓存逐出的 Nexus 采样 来源：https://arxiv.org/html/2606.23961 ## 遗忘但不妥协：固定预算下流式 KV 缓存逐出的 Nexus 采样 Duc Duong同等贡献。联系方式：[email protected] 和 [email protected]。  
格林内尔学院计算机科学系  
Jianwen Xie  
Lambda, Inc  
Anshumali Shrivastava  
莱斯大学计算机科学系  
Zhaozhuo Xu  
Workato  
###### 摘要 长上下文和智能体式 LLM 工作负载将 KV 缓存推至任何固定内存预算之外，迫使推理堆栈在持续推理流的每一步中永久性地*逐出*令牌。现有方法都遵循相同的模板：每一步直接注意力分数后跟*确定性 top-KK* 选择，这会将一个低于阈值的步骤转换为不可逆转的判定，并永久擦除任何直接注意力无法从噪声中区分出来的微妙重要令牌。为了解决这一挑战，我们提出了 Nexus 采样，一种无需训练的方法，它将*Nexus 评分*（一种迭代遍历直接注意力以暴露桥接令牌的方法）与*加权蓄水池采样*（以包含概率保留令牌，代替确定性 top-KK）配对。理论上，我们证明 Nexus 采样在微妙重要令牌的长期存留方面优于确定性 top-KK。实验上，在 80% 的 KV 缓存逐出率下，Nexus 采样在 LongBench 上与全注意力匹配（差距约 1 分），而在检索密集型任务上优于 top-KK 基线，且每个序列的缓存内存最多小 10 倍。 ## 1 引言 GPU 内存为 LLM 推理堆栈在长上下文生成期间能够持有的内容设定了硬性上限，然而键值（KV）缓存随上下文长度线性增长，并迅速超出此预算[19 (https://arxiv.org/html/2606.23961#bib.bib20), 5 (https://arxiv.org/html/2606.23961#bib.bib17)]。在长寿命、智能体式部署中，如多轮助手、持续推理循环[18 (https://arxiv.org/html/2606.23961#bib.bib21)]以及仓库级编码智能体[1 (https://arxiv.org/html/2606.23961#bib.bib22), 12 (https://arxiv.org/html/2606.23961#bib.bib23)]，有效上下文趋于无限流，差距进一步扩大。稀疏注意力方法[20 (https://arxiv.org/html/2606.23961#bib.bib6), 14 (https://arxiv.org/html/2606.23961#bib.bib24), 24 (https://arxiv.org/html/2606.23961#bib.bib9)]通过*读取*完整缓存来提高效率，部分解决了这个问题。但一旦缓存不再适合内存，高效读取就不够了：推理堆栈必须永久性地*逐出*令牌。越来越多的研究[25 (https://arxiv.org/html/2606.23961#bib.bib5), 22 (https://arxiv.org/html/2606.23961#bib.bib7), 3 (https://arxiv.org/html/2606.23961#bib.bib10), 8 (https://arxiv.org/html/2606.23961#bib.bib11), 16 (https://arxiv.org/html/2606.23961#bib.bib12), 9 (https://arxiv.org/html/2606.23961#bib.bib13)]执行了这种*KV 缓存逐出*，保留在当前（或最近）查询下得分最高的令牌，并丢弃其余部分。现有方法都与稀疏注意力共享相同的每步设计：在每一步逐出时，它们通过*确定性 top-KK* 选择保留得分最高的 KK 个令牌。然而，我们认为，**逐出与稀疏注意力有根本不同**：稀疏注意力是对仍然完整的缓存执行一次性读取，而逐出是一个*流式、固定预算*问题，其中令牌连续到达，逐出是不可逆转的，并且策略必须反复决定保留哪些令牌以应对尚未看到的未来查询。在这种观点下，每步确定性 top-KK 是一个次优原语：它孤立地处理每一步，本质上将暂时的边缘性与永久的不重要混淆，因此一个分数仅间歇性高于阈值的令牌在第一个边缘性步骤上就被丢弃，即使其跨步骤的长期平均重要性可能很高。确定性 top-KK 错过了重要性波动的令牌。  
KV 缓存中的每步重要性分数是嘈杂且随时间变化的。现代 LLM 中的注意力是重尾分布的[25 (https://arxiv.org/html/2606.23961#bib.bib5), 14 (https://arxiv.org/html/2606.23961#bib.bib24)]：一小部分头令牌占据了大部分注意力质量，而缓存的其余部分处于近乎均匀的低分状态，令牌之间的相对排名由噪声主导（图 1 (https://arxiv.org/html/2606.23961#S1.F1)）。一个对查询 t 重要的令牌可能在 t+1 时看起来边缘化，而在 t+10 时再次重要。然而，top-KK 将这些每步分数视为最终判决：一个在任何一步中落在低于阈值位置的令牌，与深埋尾部的令牌一样被最终丢弃，无论其分数离阈值有多近。长上下文会产生许多这样的判决，并且每步误差会累积。结果是*单调边缘侵蚀*：每个分数仅间歇性高于阈值的令牌都会悄然且永久地丢失。  
参见说明  
图 1：注意力分数是重尾分布的：约 80% 的累积质量集中在约前 15% 的令牌位置（黑色曲线），具有一个高分令牌的小头部（左上角集群）和一个长而嘈杂的边缘分数令牌尾部（底部带），其中确定性 top-KK 逐出被迫在近乎相同的分数之间选择。  
蓄水池采样是应对这种设置的流式算法答案。加权蓄水池采样[21 (https://arxiv.org/html/2606.23961#bib.bib14), 7 (https://arxiv.org/html/2606.23961#bib.bib15)]以与其当前分数成正比的概率保留每个令牌，使得令牌跨许多步骤的长期存留追踪其*跨步骤的平均重要性*，而不是在其分数恰好低于阈值的第一个步骤上崩溃。我们认为这是 KV 缓存逐出的正确选择原语候选，但它继承了其采样的分数校准。然而，现有方法使用的直接注意力量级是局部短视的：它只计数当前查询关注的内容，而忽略了我们称之为*桥接令牌*的内容——这些令牌没有单个近期查询强烈关注，但它们在上下文中将一个强连接集群的互相关注令牌维系在一起，以至于移除一个会切断集群的内部连接。桥接令牌是边缘排名令牌的典型实例：确定性 top-KK 在每一步都会逐出它们，即使它们维系在一起的集群可能高度重要。因此，我们将蓄水池的输入权重从这种直接注意力分数提升为*Nexus 评分*，它在任何选择发生之前就暴露桥接令牌。  
我们提出了 Nexus 采样，一种无需训练的 KV 缓存逐出方法，围绕两个新机制构建，共同解决了上述两个失败模式，并统一应用于预填充和解码阶段。上游环节，*Nexus 评分*通过一个短的迭代行走递归传递直接注意力分数，在任何选择发生之前暴露桥接令牌。下游环节，*加权蓄水池采样*以与该行走增强权重成比例的包含概率抽取被保留的 KK 个令牌，取代了所有先前逐出方法共享的确定性 top-KK 选择步骤。每个修复都恢复了现有模板无法看到的一类重要性：行走恢复了直接注意力 top-KK 遗漏的桥接令牌，而蓄水池恢复了任何确定性 top-KK 侵蚀的边缘排名令牌。理论上，Nexus 的长期令牌存留衰减为每步包含概率的乘积（跨步骤的均值），而确定性 top-KK 的则在第一个低于阈值的步骤上崩溃为零（详情见附录 A (https://arxiv.org/html/2606.23961#A1)）。贡献如下：  
- • 我们识别了确定性 top-KK KV 缓存逐出的一个根本限制：它将暂时的边缘性转化为永久损失，在流式、固定预算机制中产生*单调边缘侵蚀*，这本质上是 KV 缓存逐出所固有的。
- • 我们提出了 Nexus 采样，一种无需训练的 KV 缓存逐出方法，统一应用于预填充和解码，结合了加权蓄水池采样与*Nexus 评分*，后者暴露*桥接令牌*，即锚定窗口中互相关注令牌的强连接集群的令牌。
- • 我们建立了理论保证，表明蓄水池采样的长期令牌存留是*跨步骤的乘积*，而确定性 top-KK 的是*跨步骤的最小值*，并在 20% 密度下实验证明，Nexus 采样在 LongBench 上与全注意力相差约 1 分，同时在检索密集型长上下文任务上优于 top-KK 基线，每个序列的缓存内存比密集 FlashAttention-2 小最多 10 倍。

我们介绍 Nexus 采样，一种与解码并行运行的 KV 缓存逐出方法：在每个逐出步骤 t（当解码流中达到预算时），Nexus 通过两个按顺序应用的组件产生保留集 S_t：*Nexus 评分*后跟*加权蓄水池选择*。Nexus 评分通过一个短的迭代行走递归传递每块注意力分数，该递归暴露*桥接令牌*（锚定窗口中互相关注令牌的强连接集群的令牌），并产生反映间接重要性的每块采样权重 w^(t)。然后蓄水池步骤在加权无放回法则下以与 w^(t) 成比例的包含概率抽取被保留的 KK 个块，取代了所有先前逐出方法共享的确定性 top-KK 选择。名称反映了每个组件在流中每一步的两个角色：Nexus 评分识别维系集群的*枢纽*块，而蓄水池从由此产生的权重中*采样*，使得每个正权重块保持正包含概率，而不是在第一次处于边缘排名时被确定性地判定。  
### 2.1 预备知识与符号  
流式视角。我们将 KV 缓存逐出视为推理步骤上的一个流式问题。在每个步骤 t，LLM 接收一个新查询 q_t ∈ R^{1×D}（预填充时是提示令牌，解码时是解码令牌），关注当前缓存 K_t, V_t，并追加一个新的（键，值）对，之后缓存可能超出内存预算。每当达到预算时，就会发生*逐出步骤*：保留一个大小为 B 的子集 S_t ⊆ {1, ..., T_t} 的位置，其余被丢弃，推理继续使用 K_{S_t}, V_{S_t}。因此，逐出在两个阶段中都是相关的：长提示可能在预填充结束时触发一次性逐出，而长生成会在整个解码过程中产生长串的逐出步骤。在这两种情况下，逐出都是不可逆转的：在步骤 t 逐出的位置在任何后续步骤 t' > t 中都无法恢复。  
块粒度。我们以*键块*而不是单个令牌的粒度工作，将每个决策的成本分摊到 b 个令牌上（b=32）；块由 j ∈ {1, ..., N_k} 索引，其中 N_k = ⌈T_k / b⌉，T_k 是缓存的键数量。  
每块基础分数。单查询注意力是嘈杂的，且由即时令牌主导，因此遵循标准实践[16 (https://arxiv.org/html/2606.23961#bib.bib12), 9 (https://arxiv.org/html/2606.23961#bib.bib13)]，我们根据长度为 W 的近期查询观察窗口 Q ∈ R^{W×D} 对块进行评分，使用头部平均表示。每块*基础分数* a ∈ R^{N_k} 是以下矩阵的行平均值：  
Ŝ = rownorm(BlockSum_b(softmax(QK^⊤ / √D)))  
a = 1/W ∑_{w=1}^W Ŝ_{w,·}, ∑_j a_j = 1,   (1)  
符号。在本节的其余部分，我们使用 W 表示观察窗口长度，Ŝ ∈ R^{W×N_k} 表示每行块分布，a ∈ R^{N_k} 表示基础分数（公式 1），c ∈ R^{N_k} 表示行走状态，w ∈ R^{N_k} 表示组合采样权重，B 表示总保留块预算，n 表示蓄水池平均计数。完整的符号表见附录 A.1。  
参见说明  
图 2：一个逐出步骤中 Nexus 采样的概览。*Nexus 评分*（顶部）结合了直接重要性项（窗口化每查询注意力）和通过在相同查询行上迭代行走获得的桥接项。*流式蓄水池采样*（底部）然后以与结果权重成比例的包含概率抽取保留的块，取代了所有先前逐出方法使用的确定性 top-KK。  
### 2.2 Nexus 评分  
Nexus 评分给每个块分配一个两项*Nexus 分数*  
s_j = a_j + λ c̃_j,   (2)  
它同时捕捉了块 j 的*直接*重要性 a_j（近期查询关注它的程度）和*间接*重要性 c̃_j，后者衡量块 j 作为*桥接令牌*的作用强度，即它在窗口中锚定一个互相关注令牌集群的强度，即使没有单个近期查询强烈关注它。这两项涵盖了单查询注意力分数的互补失败模式：a_j 对当前查询有理由查看的块进行评分；c̃_j 对当前查询有理由*记住*的块进行评分，因为它们在近期上下文中的作用。混合权重 λ ≥ 0 平衡了两项。  
两项都从 Ŝ 的相同每查询行计算得出。将行记为 a^(q) ∈ R^{N_k}，因此 a^(q) = Ŝ_{q,·} 是窗口查询 q 的块分布。*直接*重要性是均匀的行平均值：  
a = 1/W ∑_{q=1}^W a^(q),   (3)  
它将每个窗口查询视为平等投票。相比之下，*桥接*分数通过对相同行进行对齐加权的迭代行走来聚合。从 C = 0 开始，行走运行 H 步；在每一步，我们插入一个新的窗口查询行 a^(q) 并更新：  
C ← C + γ_q a^(q), γ_q = 1 + ⟨a^(q), C⟩,   (4)  
然后 c̃ = C / ‖C‖_1。提供给蓄水池的每块权重是 Nexus 分数加上一个小的近期平局打破项：  
w_j = s_j + ε_tie r_j, r_j = j / (N_k - 1),   (5)  
其中 r_j 是从 0（最旧块）到 1（最新块）的线性斜坡，ε_tie 足够小以至于不会翻转任何真实信号，但大到足以确定性地打破平局，倾向于更新近的块。  
### 2.3 加权蓄水池块选择  
给定每块权重 w ∈ R^{N_k}，我们