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
cs.DC
2026-08-08
微软 OasisKV 用投机解码的草稿 token 提前一步预测下一解码步要用的 KV 块,从 CPU/远端内存异步预取进 HBM,Qwen3-8B 推理吞吐到 1.69 倍、精度几乎不掉。
长上下文与长推理链把 LLM 推理从「算力受限」推成了「显存受限」。KV cache 随上下文线性增长:一个 32B 的 GQA 模型在 32.7K 上下文(生产级 coding agent 的实测均值)下,单条请求就要 8.6 GB KV;即便乐观地把 80 GB HBM 全留给 KV,也只能并 9 条请求。HBM 容量稀缺又贵,直接卡死了 batch size 和吞吐。
已有的三条路都各有缺口。稀疏注意力减少了 KV 读取,但完整 KV 仍占着 HBM,容量问题没解决;KV 检索把完整 KV 挪到 CPU 内存、只把活跃部分搬回 GPU,容量解决了,但检索卡在解码关键路径上,直接拖慢每一步;KV 预取想用预测把搬运与计算重叠起来,但已有方案被钉死在 CPU-GPU 的 IO 带宽墙上,既够不到生产级吞吐,也撑不住跨节点的显存解耦。
OasisKV 的核心观察是:解码时的注意力天然稀疏,而且「下一步要用的 KV 块」可以提前一步准确预测。预测信号白送:投机解码(speculative decoding)本来就会用 EAGLE-3 这类草稿模型起草未来 token,这些草稿 token 就是天然的、免训练的 lookahead。把草稿 token 在当前 token 的稀疏 KV 上做一次前向,它预测出的 top-K 块与真实下一步查询的命中重合度,逐层平均 98.74%。
围绕这个信号,系统做了几件事:
精度(同样 2048-token KV 预算,各方法对各自全注意力锚点比):OasisKV 在长输入 LongBench v2 上 Qwen3-8B 只掉 0.4、Llama-3.1-8B 掉 0.6;长输出推理(AIME24/25、GPQA)pass@8 掉 0.66、avg@8 掉 0.35,是所有方法里最接近全注意力的,明显好于 Quest(掉 2.6 至 3.5)和 FreeKV。
吞吐才是重头戏:
| 场景 | 对比 dense vLLM | 备注 |
| Qwen3-8B 单卡 16K | 1398 vs 676 tok/s(2.1×) | 并发 128 |
| Qwen3-8B AIME24 真实推理 | 2083 vs 1235 tok/s(1.69×) | 精度仅掉 0.1 |
| Qwen3-235B TP8 | 至 1.9× | MoE,KV 占比小增益略低 |
| PD 解耦 | 2.1 至 2.3× | 远端部分搬运省 2.2 至 2.6× 解码侧主机内存 |
2048-token 的 KV 上限让 16K 下并发从 22 提到 90 至 95。一个关键消融把瓶颈说得很直白:fetch cap 从 0.01 放开到全量搬运,每步流量从 0.30 GB 涨到 5.05 GB,吞吐从 2178 塌到 824 tok/s(2.6 倍下降),而精度只在 74.9 到 77.4 之间晃。解码吞吐是被 PCIe 带宽、不是算力卡住的;默认的 0.05 cap 在精度仅差 0.1 的前提下拿到全量搬运 2.5 倍的吞吐。
对跑长上下文、推理链、agentic 负载的 serving 团队,OasisKV 把 KV 容量从 HBM 上解耦,直接换 batch 和吞吐,而且复用了已有投机解码的草稿 token,不需要单独训练预测器。实现挂在 vLLM 上,支持多卡和解耦部署,是少见的奔着生产级去的稀疏 KV 方案。
作者承认:草稿 token 目前只当 lookahead 用、强制拒绝接受,还没有跟投机解码的加速叠加(计划中的未来工作);在 Qwen3-235B 这类 MoE 上增益偏小,因为 KV 占比小、额外每步开销在低 batch 时反而让 TPOT 高于 dense;原型还要在 prefill 节点主机内存里为每条请求留完整 KV 直到结束,且暂不支持 prefix caching(只做了 TTFT 的解析建模)。补一条:整套方案依赖每个模型都有现成的 EAGLE-3 草稿头,虽然这类头越来越多作为可复用制品发布,但冷门模型仍要自己训。