长上下文推理 KV 缓存撑爆显存怎么办?OasisKV 前瞻稀疏预取把吞吐提升最高 2.1×
OasisKV 利用解码阶段注意力的天然稀疏性,借助投机解码的 lookahead tokens 预测未来重要 KV 块,把完整 KV 存储从 HBM 解耦到高层内存并异步预取。在 2048-token 预算下精度仅损失 0.7 点,多 GPU 长上下文服务吞吐最高提升 2.1×。
长上下文、长思维链(long-CoT)推理的兴起,正在把大模型推理服务从"算力受限"推向"显存受限"。在 decode(逐 token 生成)阶段,KV cache 会随着上下文变长而线性膨胀,既持续吞噬 HBM 显存容量,又要在每一步生成时被完整读取一遍,成为显存流量与带宽的双重瓶颈。HBM 容量稀缺且昂贵,直接卡死了推理的 batch size 与系统吞吐。2026 年 8 月 8 日,Can Xiao、Sukmin Cho 等研究者(提交人 Sukmin Cho)在 arXiv 提交论文《OasisKV: Scaling In-Decode KV Cache Beyond HBM with Lookahead Sparse Prefetching》(arXiv:2608.08097),提出一种"内存为中心"的推理系统设计:把完整 KV cache 存储从 HBM 上解耦下来,只在 HBM 里保留注意力真正需要的那一小撮 KV,其余下沉到容量更大、成本更低的 host / remote 内存,再用"前瞻稀疏预取"把下一步要用的重要 KV 提前搬回 HBM。
OasisKV 的核心洞察有两点。其一,decode 阶段的注意力是"天然稀疏"的——任意一步生成,真正被高频访问的 KV 条目只是全量中的一小部分,绝大多数历史 token 对当前步的贡献微乎其微,把全部 KV 钉在 HBM 里本质上是一种浪费。其二,接下来几步会"重要"的 token,其实可以提前预知:OasisKV 借助投机解码(speculative decoding, SD)已经起草出来的 lookahead tokens(草稿 token),去预测未来哪些 KV 块会在后面的注意力里被频繁命中。换句话说,SD 顺手"看一眼"未来要生成的 token,就能反推它们会去读哪些历史 KV——这套前瞻信号,正是传统 KV 卸载方案一直缺少的"预取方向标"。
机制上,OasisKV 在 vLLM 之上实现了一套高效的"注意力后台流水线":先用 lookahead 预测锁定未来重要的 KV block,再把这些块从高层内存(host 内存或远端内存)异步预取到 HBM,确保在下一步 decode 真正需要它们之前就已就位;注意力计算只在 HBM 内的"活跃 KV 子集"上进行,从而把宝贵的 HBM 容量让给真正起作用的少量条目。于是完整 KV 仍保留在高容量内存层(保证可恢复、信息不丢),而 HBM 只承担"当前最相关"的薄层,显存压力被大幅卸载。整个设计是 training-free、drop-in 式的,无需改动模型权重或重训,直接套在现有推理引擎上即可生效。
在 vLLM 上的实测相当扎实(arXiv:2608.08097)。给定 2048-token 的 KV 预算,OasisKV 把精度控制在与 full attention 仅差 0.7 个点的范围内;在此基础上,它在推理负载上相对 dense vLLM 取得 1.69× 吞吐(精度损失仅 0.1 点),在多 GPU 长上下文服务场景下最高达到 2.1×。把 prefill 与 decode 分离部署时,OasisKV 达到约 2× 的 dense 吞吐,同时每个请求占用的 KV 量减少 6.5–9.7×,decode 节点 host 内存占用仅为全量 KV 传输的 2.2–2.6×。换句话说,同样的 GPU,能扛下更长的上下文、更大的 batch,单位算力产出近乎翻倍。
对思陌大模型微调而言,OasisKV 直接对应「推理部署服务」这条业务线,也反哺「基座微调 / 领域适配」的交付效率。其一,长上下文 + 推理密集(如法律/医疗长文档问答、代码仓库级理解、Agent 多轮长程任务)正是企业客户的高价值场景,却常常被显存成本卡住;OasisKV 这类"把 KV 搬出 HBM + 前瞻预取"的思路,正是思陌给客户做部署方案时可以评估、集成的降本增效手段——在不换模型、不降精度的前提下,用系统层优化把单卡 GPU 的服务能力翻一倍。其二,它再次印证了一个判断:微调解决的"能力"问题,最终要落到"部署成本"上才能转化为客户价值;思陌在交付 LoRA / 全参数微调成果时,应同步把推理引擎(vLLM 类)层面的 KV 管理、量化、批处理优化纳入交付标准,把"训得出来"和"跑得省钱"打包给客户。当然,OasisKV 当前验证以通用推理负载为主,其在强领域模型、超长思维链(如 o1/R1 类深度推理)下的稀疏性与预取命中率,仍值得在真实业务负载上做进一步评测,这也为思陌后续的「模型评估优化」留下了用武之地。
参考文献
📄 论文原文信息
| 论文标题 | OasisKV: Scaling In-Decode KV Cache Beyond HBM with Lookahead Sparse Prefetching |
| 作者 | Can Xiao, Sukmin Cho, Junbong We, Zhixiong Niu, Jianyi Cheng, Yiren Zhao, Youngjin Kwon, Yongqiang Xiong, Rui Ma, Junyi Liu |
| 期刊 | arXiv (cs.DC) |
| 发表时间 | 2026-08-08 |
| DOI | 10.48550/arXiv.2608.08097 |
论文原文信息地址:https://www.simolm.com/blog/2026-08-19-oasiskv-lookahead-sparse-prefetching/