1M 上下文塞进显存有多难？有人让 AI 学会自己"扔掉没用的记忆"

  • 关联论文：2602.10238

为什么这事值得关心

大模型每吐一个字，都要把之前所有字的"笔记"留着——这就是 KV Cache。它随上下文线性膨胀，1M token 的上下文动辄吃掉几十 GB 显存，是长上下文推理最硬的一道成本墙。

过去几年的解法基本是启发式：留最近的（Recency）、留注意力分数高的（Attention Score）、首尾加滑窗（H2O / StreamingLLM）。它们背后都藏着一个假设——过去被关注的，未来还会被关注。可这个假设是间接代理，不是未来真实效用。办法是猜，那就有猜错的可能。

苹果 ML 团队的 KVP 换了个思路：别猜了，直接学。

核心机制：把"扔什么"变成"排个队"

三件事让它跟以往工作不一样。

第一，输出从"留/扔"变成连续分数。 传统驱逐是二值的，本文把所有 token 按"未来效用"排序。这样缓存预算能从 10% 调到 90% 而不用换策略——一条策略适配所有显存配置，这是它最实用的地方。

第二，每个 attention head 一个独立的轻量 MLP。 输入不是 query、也不是输入 embedding，而是 K/V 向量本身；输出就是每个 token 的效用分。为什么要 per-head？因为不同 head 干的事完全不同，有的管语法、有的管指代、有的管检索，用一个全局策略会把差异抹平。

第三，Holistic Reward——这是全文最巧的设计。 奖励不是"在某个固定预算下生成质量好不好"，而是"在 10%、20%……90% 所有预算档位下，排序是不是一致地好"。训练一次，全谱段可用。

训练走的是离线路线：先用预生成的轨迹离线训策略，部署时只剩"算分 → 排序 → 截断"三步，不改 attention 计算本身，等于外挂。

证据边界：数字长什么样

在 RULER（最长 128K，13 类长上下文任务）上，KVP 全预算段优于 H2O、StreamingLLM、Recency，低预算（10–20%）区间优势最明显——正是显存最紧的场景。多轮对话基准 OASST2-4k 上同样领先，说明学到的不是"长文档检索"而是通用的效用排序。Zero-shot 迁移（OASST2 训、BoolQ/LongBench/GovReport 测）、跨模型族、跨更长序列，优势都保持。

代码开源于 github.com/apple/ml-learning-to-evict，论文被 ICML 2026 接收。

但要泼几盆冷水：具体数值原文摘要没给，得读正文；与 Quest / TOVA / SnapKV 等近两年强基线的对比也需在正文确认；"两个模型族"具体是哪两个、per-head policy 参数量、训练总成本，原文均未明确。

工程局限：落地前先看这四条

head 数量会撑大开销。 32 层 × 32 head = 1024 个 MLP，即便每个只有 8K 参数，全量也是几十 MB 级别；更大的模型会更可观。而且每个 decode step 都要读一遍全量 K/V，128K 上下文下这延迟不便宜。经验做法：先用 recency 粗筛到 16K，再用 KVP 细筛。

训练数据是真成本。 每条 trace 是一次完整 forward（注意力 O(n²)），10K 条轨迹比单纯做推理贵一个数量级。可以先用小模型生成轨迹再做行为克隆，或者复用推理框架里已积累的 KV dump。

vLLM 里有个粒度坑。 vLLM 是 block（通常 16 token）粒度的 paged KV，KVP 排的是 token 粒度。落地时得先把分数聚合到 block 级别（取块内平均分），会引入近似误差，16 / 32 / 64 不同 block size 值得实测。TensorRT-LLM 更受限——静态分配不支持动态驱逐，KVP 只能在 prefill 阶段做静态预筛。

打分延迟必须过 p99。 A100 上 7B 模型 8K 上下文做一次 KVP 打分，建议控制在 5ms 以内，否则 decode 直接被它拖住。

一句话总结：KVP 把 KV Cache 管理从"手写启发式"推到了"学出来的策略"，holistic reward 那套"一条策略通吃所有预算"的思路值得借鉴到别的系统优化场景；但 per-head 的规模效应、训练开销、推理栈集成这三块，仍是落地前必须自己算清楚的账。

三个标题变体

  1. 别再靠"猜"来决定扔掉哪些上下文：有人教模型自己学
  2. 1M 上下文塞不进显存？苹果把 KV Cache 驱逐变成了一道 RL 题
  3. 一条策略通吃所有显存预算：KV Cache 管理的学习范式来了

小红书卡片文案

正文：

上下文一长，显存就不够用——这是每个做大模型应用的人都撞过的墙。

以前大家扔缓存的思路很朴素：留最近的、留注意力分数高的。听着合理，但它们都建立在一个假设上——「过去被看过的，未来还会被看」。这假设对不对？没人验证过。

苹果 ML 的做法是：别猜了，让模型自己学。

每个注意力头配一个小小的打分网络，直接读 K/V 向量，输出每个 token 的「未来会有多有用」。然后按分数排队，该留的留。最妙的是它的训练目标：不是「在某个内存占用下表现多好」，而是「从 10% 到 90% 所有占用档位下都排得好」。一次训练，全部配置通吃。

结果是 128K 长文本基准上全面超过 H2O、StreamingLLM 这些老办法，而且越是在内存紧张的时候优势越大——刚好像是大家最需要它的场景。

不过要说实话：具体多少分要翻正文；跟最新那批基线的对比还得再核；真正落地时，1024 个打分网络的开销、128K 上下文下每步的读取延迟、跟 vLLM 的 block 粒度怎么对齐，都是要自己算的账。

「用强化学习做系统决策」的教科书案例。#大模型 #AI推理 #长上下文 #KVCache #显存优化 #人工智能

标题： 大模型内存不够用？苹果让 AI 学会自己扔缓存 封面文案： 「过去被看的，未来还会被看」——这个假设该换换了 标签： #大模型 #AI推理 #长上下文 #KVCache #显存优化 #机器学习