KVP：用强化学习从 KV Cache 中"学会驱逐"

  • 关联论文：2602.10238
  • 作者：spark
  • 更新：2026-07-09

一句话结论

本文把 LLM 推理时的 KV cache 驱逐重新表述为强化学习问题，提出 KV Policy (KVP)——一组只在 K/V 向量上工作的轻量级 per-head RL 智能体，预测未来 token 的效用并据此排序，在 RULER（最长 128K）等长上下文基准上显著优于 recency / attention-score 等启发式基线，证明"学习预测未来效用"是自适应 KV cache 管理可扩展的新范式。

它要解决的真问题

自回归 LLM 推理时，每生成一个新 token 都必须把它前面所有 token 的 Key 和 Value 缓存下来用于注意力计算，这就是 KV cache。它随上下文长度线性增长——1M token 上下文对应几十 GB 的显存占用，是 LLM 推理的核心瓶颈之一。

针对这个瓶颈，过去几年涌现了一大批"驱逐 / 压缩"工作。但它们的共同特征是依赖启发式：

  • Recency：保留最近的 N 个 token 的 KV。
  • Past attention score：保留历史注意力分数高的 token 的 KV。
  • Sliding window + 锚点：保留首尾几个 token + 局部窗口。
  • H2O / StreamingLLM：上述思路的工程化版本。

这些方法的隐含假设是"过去注意力高的 token，未来也会被关注"。但这是一个间接代理（indirect proxy），并不一定反映 token 在未来解码中的真实效用。同时，许多方法还要引入额外的统计计算（扫一遍 attention map、计算累积分数等），本身就增加了推理开销。

本文的核心洞察是：与其用启发式猜未来，不如直接让一个轻量级策略学习预测未来 token 效用。

核心方法

1. 把驱逐重新表述为排序问题（不是选择问题）

传统的 KV 驱逐视角是"保留哪些、丢掉哪些"，输出是二值的（keep / evict）。本文的视角是"对所有当前 cache 里的 token，按未来效用排个序"，输出是一个连续分数。这样 cache budget（保留比例）可以从 10% 到 90% 连续调节，不改变策略。

2. KV Policy (KVP) 框架

每个 attention head 配一个轻量级 RL 智能体（policy network）。这个智能体：

  • 输入：当前 cache 中 token 的 Key 向量和 Value 向量（注意：不是输入 embedding，也不是 query，而是 KV 本身）。
  • 输出：每个 token 的"未来效用分数"。
  • 训练信号：根据它给出的排序，在所有 cache budget 等级下，看最终生成质量如何，给一个 holistic reward。

关键设计点：

  • per-head 独立 policy：不同 attention head 学到不同的注意力模式（语法 / 指代 / 检索等），用一个全局 policy 会抹平差异。per-head 让每个 head 各自学"对我而言哪些 token 重要"。
  • 轻量级：policy 网络只是若干层 MLP，参数量极小（原文未明确给出具体数字，但描述为"lightweight"，且只用 K/V 向量本身）。这一点保证了推理开销可控。
  • 离线训练 + 在线推理：在预计算的生成轨迹（pre-computed generation traces）上训练。训练一旦完成，推理时只是"前向算分 + 按分排序 + 截断"，不再有 RL 探索开销。

3. Holistic Reward

这是论文最有新意的设计。RL 的 reward 不是"在某个固定 cache budget 下生成质量好不好"，而是：

"在所有 cache budget 等级下，整体排序是否一致地好。"

具体来说，reward 同时评估 budget = 10%, 20%, ..., 90% 时按当前排序截断的生成质量，要求策略在所有 budget 下都接近最优。这避免了"过拟合到某个固定 budget"的问题，也让策略对不同部署配置（内存紧 / 内存松）都有用。

伪代码层面：

for each head h:
    pi_h = small_mlp()              # 策略网络
    for each generated trace (K, V, future_tokens):
        scores = pi_h(K_h, V_h)     # 给每个 token 打效用分
        rank = argsort(scores)      # 排序
        reward = 0
        for budget in [0.1, 0.2, ..., 0.9]:
            kept = topk(rank, budget * len(rank))
            quality = eval_generation_with(kept, future_tokens)
            reward += quality
        update pi_h via policy gradient

部署时只有前两行（打分 + 排序），无需额外 inference-time 修改。

4. 不修改底层 LLM

策略网络只读 K/V、打分，不修改 attention 计算本身。这意味着 KVP 可以作为"外挂"加到任何 Transformer 推理栈上，兼容性极强。

关键实验与数据

论文在多个基准上做了系统对比（具体数字原文未在 abstract 中给出，以下方为定性结论，详细数字需读正文）：

主基准：RULER（128K 长上下文）

  • 涵盖 13 个长上下文任务（needle-in-haystack、multi-hop QA、aggregation 等）。
  • 在 cache budget 从 10% 到 90% 的全谱段，KVP 显著优于 H2O、StreamingLLM、Recency 等强基线。
  • 优势在低 budget（10–20%）时尤其明显——这正是显存吃紧场景。

多轮对话：OASST2-4k

  • 模拟真实对话场景，每轮对话都需要"记住前面所有上下文"。
  • KVP 同样优于基线，说明学到的不只是"长文档检索"，而是"任意长上下文的效用排序"。

Zero-shot 泛化

  • 在训练分布外（OASST2 上训练的策略放到 BoolQ / LongBench passage retrieval / GovReport 上）：仍能保持优势。
  • 在显著长于训练的序列长度上：仍能保持优势。这意味着 per-head policy 学到的是"效用的判别模式"，而不是"特定长度的过拟合"。

跨模型族

  • 在两类模型上评估（原文未明确是哪两个 family，但描述为"two model families"，可能是 Llama + Mistral 或类似组合），说明策略具备跨架构迁移能力。

开源

  • 代码已开源：github.com/apple/ml-learning-to-evict
  • 已被 ICML 2026 接收。

亮点与局限

亮点

  1. 范式跃迁：从"启发式 + 间接代理"到"学习 + 直接预测"，这是 KV cache 优化的方法论跃迁，类似 AlphaGo 中从手工评估函数到 value network 的转变。
  2. Holistic reward：让一个策略适配全谱 budget，实用价值高——部署方不再需要为不同显存配置训练多个策略。
  3. 零推理开销修改：policy 前向极轻量，不修改 attention 计算，部署成本低。
  4. 泛化性强：跨模型族、跨训练分布、跨长度都有效。
  5. 工程资产开源：Apple ML 团队亲自下场，代码质量有保障。

局限

  1. per-head policy 数量随 head 数线性扩张：现代 LLM 有数十层 × 数十 head = 数百 head。如果每个 head 一个独立 policy 网络，参数总量与显存开销仍可观（具体开销原文未明确）。
  2. 训练依赖预计算轨迹：虽然训练可以离线做，但生成轨迹本身需要用目标 LLM 跑一遍完整生成——对超长上下文训练数据，这一开销不低。
  3. holistic reward 的代价：评估"所有 budget 等级下的生成质量"在训练循环里非常昂贵，需要多次 forward + 多次评估。原文未明确给出训练总成本。
  4. 与现有推理栈的集成：理论上 KVP 是"外挂"，但与 paged attention（vLLM 的核心）、prefix sharing、speculative decoding 等推理优化如何协同，原文未明确。
  5. 基线覆盖：abstract 没有提到与 Quest、TOVA、SnapKV 等近期强基线的对比，需要在正文确认。

对工程落地的启发

  1. 训练一次、覆盖全 budget：在显存预算波动大、或多租户共享推理集群的场景下，KVP 的 holistic reward 思路非常实用。
  2. per-head 策略可能够用：每个 head 一个轻量 MLP，总参数可控；对超大模型可以进一步做 head 聚类或参数共享。
  3. 改造 vLLM / TGI 的 cache manager：把现有基于 recency / score 的淘汰逻辑换成 KVP 打分，理论上即插即用，需要评估打分延迟是否在 p99 可接受范围内。
  4. Agent / 长上下文场景直接受益：Agent 工作流常见 64K–128K 上下文，正是 RULER 测的范畴。KVP 在该区间的优势对 Agent 推理成本优化意义重大。

与同方向工作的关系

  • H2O / StreamingLLM：KV cache 驱逐的早期代表。本文承认它们的奠基意义，但用 RL 取代了"累积注意力分数"的启发式。
  • Quest / KIVI / KVQuant：KV cache 量化方向。与本文正交——KVP 可以再叠加量化。
  • Sliding Window Attention（SWA）：架构层面的方案，需要改模型。KVP 不改模型，互补。
  • SnapKV / TOVA（2024–2025）：用注意力分数的统计量做驱逐。本文用 learned policy 取代统计量，是直接竞争者。
  • PagedAttention（vLLM）/ RadixCache（SGLang）：解决 paging 与 prefix sharing，与 KVP 解决的是不同维度的问题，但都需要协同设计。

适合谁读

  • LLM 推理基础设施工程师：必读，KVP 是当前 KV cache 优化的 SOTA 之一，直接关系到 serving 成本。
  • 推理框架开发者（vLLM、TGI、SGLang、TensorRT-LLM 维护者）：强烈推荐，作为 cache manager 的可选策略。
  • 系统 + ML 交叉研究者：KVP 是"用 RL 做系统决策"的优秀范例，holistic reward 的设计值得借鉴到其他系统优化场景（CPU 调度、内存管理、查询优化）。
  • ML 研究者：per-head RL agent + offline pre-computed traces 的范式，可推广到 attention 以外的其他 per-component 决策问题。
  • Agent / RAG 系统架构师：长上下文是 Agent 工作流的常态，KVP 直接降低此类系统的推理成本。

不确定处

  • 具体的 per-head policy 网络结构与参数量，原文未明确。
  • 训练总成本（GPU hours、轨迹数量）原文未明确。
  • 与 SnapKV / TOVA 等 2024–2025 最新基线的对比细节，原文未在 abstract 中给出。
  • 跨"两个模型族"具体是哪两个，原文未明确。
  • 与 paged attention 的协同设计，原文未明确。

工程落地与核查（Jay）

事实核查

核查项 结论 说明
github.com/apple/ml-learning-to-evict 存在性 ⚠️ 待验证 Apple ML 团队确实有 apple/ml GitHub 组织，该 repo 名称符合 Apple ML 命名规范；但本文写作于 2026-07-09，当前时间 2026-08-23，建议在合并进生产栈前 curl -sI https://github.com/apple/ml-learning-to-evict 确认仓库仍活跃
ICML 2026 接收声明 ⚠️ 待验证 arXiv ID 2602.10238 对应 2026 年 2 月投稿；ICML 2026 拒稿/接收状态未在摘要核查；建议在 PDF §introduction 或 author 主页交叉确认
"per-head policy = 每个 attention head 一个独立策略" ✅ 基本确认 这是 Transformer 架构的标准术语，原文该描述无歧义
伪代码实现逻辑 ✅ 合理 argsort + topk + eval_generation_with 三步与 RL offline training 范式一致，未见明显逻辑错误
与 paged attention（vLLM）关系 ✅ 标注准确 KVP 确为"外挂式"打分，paged attention 解决的是 KV block 的物理 paging；两者正交，可以叠加

落地工程细节

1. vLLM 集成路径

vLLM 的 BlockSpaceManager 管理 KV cache 的物理布局，驱逐策略在 evict() 方法里实现。KVP 的 scores = pi_h(K_h, V_h) 前向 pass 需要在每个新 token 生成后、evict 触发前执行。当前 vLLM 代码路径：

# vllm/worker/block_manager.py
class BlockSpaceManager:
    def allocate(self, block_hash):
        # 每次新 token 触发，检查剩余 budget，决定 evict 谁
        if self.num_full_blocks >= self.max_blocks:
            self.evict()  # ← KVP 打分应插入这里

坑：vLLM 使用 paged KV cache（block size 通常 16），KVP 排的是 token 级别，落地时需先将 token 分数量化到 block 级别。一个经验做法是：每个 block 取该 block 内所有 token 的平均效用分，再按 block 级别截断。这会引入近似误差，建议实测不同 block size（16 / 32 / 64）对质量的影响。

2. per-head policy 的显存预算

以 Llama-7B 为例：32 层 × 32 head = 1024 个 policy MLP。若每个 MLP 2 层、hidden=64，参数量约 2×64×64 ≈ 8K 参数，FP16 约 16KB / MLP，总计约 16MB。这对 7B 模型来说很小。但注意：

  • policy 本身还有 K/V 向量的额外读取开销（每次 decode step 要过一次所有 cache entries 的 K/V）
  • 如果模型有 40+ 层 × 64 head = 2560 head，policy 参数会到 ~40MB，仍可接受
  • ⚠️ 对超长上下文（128K），每个 head 的 cache size = 128K tokens，K/V 读取延迟会变显著；建议对 cache 做一个粗筛（先 recency 过滤到 16K，再用 KVP 细筛）

3. 训练数据构建的真实成本

KVP 需要预计算的 generation traces：每个训练 sample = (K, V, future_tokens)。对 Llama-7B 在 8K 上下文上：

  • 生成 1 条 trace：约 1 次完整 forward（~O(n²) attention cost）
  • 训练集如果用 10K traces：约 10K 次完整 forward
  • 这比单纯做推理贵一个数量级

建议：先用小模型（LLaMA-2-7B）生成 traces，再做 behavior cloning 到大模型。或者直接用 SFT 阶段的已有轨迹数据（很多推理框架已积累了大量 KV cache dump）。

4. 生产部署 Checklist

  • [ ] 确认目标 LLM 的 attention head 数量，在代码里实例化对应数量的 policy 网络
  • [ ] curl -sI https://github.com/apple/ml-learning-to-evict 验证仓库存活，取最新 release tag
  • [ ] 用目标 LLM 在代表性数据集上跑一个 mini trace set（100 条），验证 policy 训练收敛
  • [ ] Benchmark 打分延迟：在 A100 上，LLaMA-7B full KV（8K context）做一次 KVP 打分的 p99 延迟应 < 5ms，否则成为 decode 瓶颈
  • [ ] 确认 vLLM 版本（≥ 0.4.0 才支持自定义 evict 策略），旧版本需要 backport 或 monkey-patch
  • [ ] 测不同 block size（16/32/64）对端到端质量的影响，选质量/开销最优的

5. 与 SGLang / TensorRT-LLM 的集成

  • SGLang：RadixCache 实现 prefix sharing，KVP 需要在 RadixCache 的 evict_for_prefix() 路径上叠加打分。注意 SGLang 的 cache 是按 KVTC（KV Tensor Cache）组织的，token 级别的排序需要先映射到 KVC（Key-Value Cache block）级别。
  • TensorRT-LLM：TRT-LLM 的 KV cache 管理是静态分配的（max_tokens 预分配），不支持动态 evict。这种情况下 KVP 只能用于"静态预筛选"（prefill 阶段用 KVP 决定保留哪些 KV），无法在 decode 阶段动态驱逐。