KVMem: Virtualizing Million-Token Agent Workspaces on a Consumer GPU
Di Chai, Leye Wang, Zeshen Su, Zhiguo Xia, Zhihang Yu
cs.LG
2026-09-04
KVMem 把溢出的 Agent 历史保存成分页 KV，跨 GPU、主机内存和 NVMe 虚拟化。24GB 笔记本上撑到百万 token；DeepSWE 上成功率从压缩方案的 43.8% 提到 48.4%。
长跑 Agent 的工作区会越积越大。读文件、调工具、改代码、看日志，历史很快顶上两道硬墙。一道是 GPU 上能放下的 KV 量：权重、投机解码状态、运行时缓冲都在抢同一块显存，消费级卡上实际能用的上下文往往远小于模型标称窗口。另一道是模型自己的上下文上限。Qwen3.6-27B 标称 256K，长程 Agent 轨迹已经到几十万甚至上百万 token。
现有系统几乎都走文本中心。旧历史压成摘要，不够再把原文片段检索回来。摘要必须在未来还没发生时决定哪些细节值得留；一条工具输出里的数字、一条用户约束，当时看着不重要，后面可能就是成败关键。检索能补回细节，但这些 token 模型早就算过一遍 KV，现在又要整段 prefill。保真度和算力绑死在一起。
KVMem 把 Agent 工作区当成可虚拟化的系统资源，地址空间和 GPU 驻留量脱钩，类比操作系统的虚拟内存。溢出的历史以分页 KV 躺在 GPU、主机内存和 NVMe 上；每一步只把当前相关的块装进模型原生窗口里的执行视图。
要同时回答三个问题：何时召回、召回什么、怎么装回去。
步级调度来自注意力统计。8 条 OpenHands SWE-bench Lite rollout 上，相邻 128-token 窗口的注意力 KL，同一步内平均 0.070 bit，跨步边界 2.59 bit，差 37.3 倍。于是每个 Agent 步只更新一次工作集：当前步 prefill 之后、decode 之前；decode 期间工作集冻结，新生成的 KV 往尾巴长。
查询条件检索把粒度从 token 收到块。默认 32-token 一块，对每个 layer、每个 KV head 存一份去掉 RoPE 之后的 Mean-K，用当前 query 在注意力空间里打分。分析窗口平均 103.1 块历史，top-8 已经占 66.5% 注意力质量，top-16 占 77.0%。不必扫整段 KV 张量。sink 和最近窗口始终保留，其余预算按分数填，装入前再按时间顺序排好。
分层搬运把仓库视图和执行视图拆开。仓库记下每块在 GPU、主机内存还是 NVMe；执行视图只装被选中的块，按时间顺序排成紧凑窗口。K 带着 RoPE 位置，块换了逻辑位置就不能直接用，系统在 GPU 外保留一份位置无关的 raw K，装入时按新位置重做 RoPE，V 原样拷。连续两步重叠的页直接复用，热块优先留主机内存。miss 的块打包传输，gather、H2D、scatter 加重 RoPE 流水重叠。
底层引擎叫 QW3，自研 C++/CUDA。vLLM、SGLang 按单调增长的前缀位置恢复 KV，撑不住每步重排稀疏历史块。Mean-K 索引整份放主机内存，GPU 上只流过固定大小的 tile。backing store 装满时原型会退回文本压缩，触发粒度比窗口溢出粗得多。
受控长历史评测用 Qwen3.6-27B，同一套 QW3 后端，只换上下文策略。活动窗口按规模设成 32K、64K 或 100K。
| 设置 | Full Context | Compact-only | Compact+RAG | KVMem |
| LongMemEval-S 准确率 / 时延 | 86.60% / 0.30s | 45.60% / 18.92s | 86.20% / 26.63s | 85.60% / 0.48s |
| MemoryAgentBench（>256K）总分 | 无 | 27.54 | 34.80 | 40.99 |
| AgentLongBench ≤256K 成功率 | 59.54% | 15.84% | 47.49% | 60.87% |
| AgentLongBench 1M 成功率 / 时延 | 无 | 32.00% / 380s | 42.00% / 416s | 50.00% / 0.73s |
相对 Compact+RAG 去掉摘要生成之后的检索时延，KVMem 仍低 11.4 到 53.8 倍。512K 的 AgentLongBench 上 KVMem 成功率 53.0%，Compact+RAG 54.0%，几乎打平，但时延是 0.62s 对 264s。原生窗口内 KVMem 略超 Full Context（60.87% 对 59.54%）。论文的解释是长上下文本身会退化，只给查询相关子集可能更干净。
消费级对照很直白。同一台 24GB RTX 5090 笔记本，vLLM 大约 10K 上下文，llama.cpp 大约 80K；KVMem 执行视图同样 80K，地址空间拉到 1M，单会话约 50 token/s。服务器上执行视图钉死 64K，工作区从 256K 扩到 10M，GPU 占用几乎不动（33.9 到 34.9 GiB），NVMe 从 8.5 GiB 涨到 324 GiB，TTFT 从 0.43s 到 1.60s。Mean-K 索引从 0.25 GiB 涨到 9.5 GiB，始终在主机侧。
完整轨迹换成 Qwen3.8-27B，挂 Claude Code harness，DeepSWE v1.1 规范顺序前 16 题，每题 4 次采样。
| 配置 | Pass@1 | Pass@4 | 平均任务时长 | prefill |
| Compaction-only | 43.8% | 81.3% | 52.5 min | 211.5 s |
| KVMem | 48.4% | 93.8% | 45.8 min | 95.0 s |
16 题里 KVMem 有 15 题至少过一次，压缩方案是 13 题。解码 token 从 154.4K 降到 125.6K，省的不只是输入侧。
本地长跑 Agent 现在卡在两件事上：窗口写满就压缩，压缩就丢细节；想找回细节就重新 prefill。KVMem 给出第三条路：历史已经付过的计算费，别扔掉。对在笔记本上跑 Qwen 级编码 Agent 的人，这是能直接感知的能力差。同一块 24GB 卡，从 80K 可执行上下文变成 1M 可寻址工作区，生成速度还在交互档。
它没有把一次 attention 的窗口拉到百万。每一步仍然只看见有限的执行视图。省下的是再算一遍已经算过的历史。跟 MemGPT、Mem0 比，检索信号来自服务模型自己的注意力空间，召回物是 KV。跟 vLLM PagedAttention、LMCache 比，难点是工作集每步都在换，还要重映射 RoPE 位置。
工程落地有门槛：必须能改 KV 分配、位置编码和搬运路径，黑盒 API 套不上去。云厂商若做成服务，缓存命中和新鲜 prefill 的价差可能比单机更有经济账。
论文写得很清楚。虚拟化的是可寻址工作区，不是单次联合注意力长度。复用历史 KV 不等于在新拼出来的上下文里重算一遍：K 做了 re-RoPE，V 原样拷，当初的因果上下文已经变了。检索还可能漏块。资源允许时加大执行视图，能少依赖检索。
存储账很重。10M token 工作区要 324 GiB NVMe，Mean-K 索引 9.5 GiB 主机内存。KV 比文本大一个数量级以上，这是用空间换可复用计算。backing store 顶满仍会退回文本压缩。当前实现也做不了多用户争抢同一块 GPU 和盘。
DeepSWE 只跑了规范顺序前 16 题、每题 4 个 seed，不是全量榜。公开模型对照用 mini-swe-agent，本地实验用 Claude Code，不能当成模型实力横评。QW3 是自研引擎，和 vLLM、llama.cpp 的容量对比也受各自优化目标影响：vLLM 面向吞吐，并不以单会话最大上下文为优先。