Paper

OasisKV 论文速读:KV Cache 不必全待在 HBM

5 min read ·

OasisKV: Scaling In-Decode KV Cache Beyond HBM with Lookahead Sparse Prefetching 是 Microsoft 团队在 8 月 8 日发布的推理论文,也进入了 Hugging Face Daily Papers。它讨论的是一个所有长上下文服务都会遇到的问题:生成 token 时,算力不是唯一瓶颈,KV cache 占用的 HBM 才是越来越贵的资源。

论文摘要给出的思路很清楚:decode-time attention 天然稀疏,不必把全量 KV 都常驻 HBM。OasisKV 把完整 KV 存储和 HBM 里的活跃 KV 集合解耦,利用 speculative decoding 产生的 lookahead token 预测未来重要 KV 块,再从主机或远端内存提前预取到 HBM。

这篇速读关注工程意义:它为什么重要,和 KV 量化、PagedAttention、prefill-decode disaggregation 有什么关系,以及推理平台团队该从中学什么。

KV Cache 为什么会成为瓶颈

Transformer 解码时,每生成一个 token,都要读取历史 token 的 key 和 value。上下文越长、batch 越大、并发越高,KV cache 越大。

在短对话时代,权重加载和矩阵乘可能更显眼。到了长上下文、长思考、agent 轨迹、代码仓库阅读之后,KV cache 会快速吞掉显存。

典型现象是:

模型权重还能放下
单请求也能跑
一提高并发就 OOM
一拉长上下文吞吐急剧下降
GPU 算力没有打满,但 HBM 容量先爆

这就是 OasisKV 要处理的场景。它不是单纯追求某个 benchmark 分数,而是在问:能不能把 HBM 从“全量 KV 仓库”改成“高速工作集缓存”?

核心思想:HBM 只放重要工作集

传统做法里,decode 阶段会假设历史 KV 都可能被访问,所以尽量放在 GPU HBM。PagedAttention 改善了 KV 管理和碎片问题,但并没有改变“活跃请求的 KV 大量占用 HBM”的根本压力。

OasisKV 的判断是:注意力访问并不均匀。很多 token 的 KV 在下一步解码里贡献很小。既然如此,可以把完整 KV 放在更大容量层级,比如主机内存或远端内存;HBM 只保留下一段时间真正需要的 KV 块。

这听起来像 CPU cache,但难点在于预取。GPU 解码很快,如果等注意力需要某个 KV 块时再去搬,延迟就会压垮吞吐。因此系统必须提前知道哪些块重要。

Lookahead Sparse Prefetching

OasisKV 借用了 speculative decoding 的 lookahead token。投机解码本来是为了加速生成:小 drafter 先猜一些未来 token,大模型验证。OasisKV 利用这些未来 token 估计接下来注意力会关注哪些历史位置。

流程可以简化为:

drafter 生成 lookahead token
系统用 lookahead 预测未来重要 KV 块
后台 pipeline 从主机或远端内存预取 KV
真正 decode 时 HBM 已经有相关块
注意力只在稀疏重要集合上计算

这个设计的漂亮之处在于,它把投机解码从“只减少模型计算”扩展为“帮助内存系统提前调度”。换句话说,未来 token 不只是候选输出,也是内存预取信号。

和 KV 量化的关系

KV 量化问的是:每个 KV 能不能更便宜。比如从 FP16 降到 FP8 或更低精度。

OasisKV 问的是:哪些 KV 需要在 HBM,哪些可以不在。

两者可以叠加。量化降低每块 KV 的大小,OasisKV 降低 HBM 里常驻块的数量。对长上下文服务来说,这两类优化都重要,但风险不同。

KV 量化的风险主要是精度损失和硬件 kernel 支持。OasisKV 的风险主要是预取错误、稀疏注意力漏掉关键信息、跨层级搬运带来的延迟抖动。

和 Prefill-Decode Disaggregation 的关系

越来越多推理平台把 prefill 和 decode 拆到不同节点。prefill 负责吃长 prompt,decode 负责逐 token 生成。这个结构能提高资源利用率,但也带来 KV 转移问题:prefill 产生的大量 KV 要如何交给 decode?

OasisKV 摘要中特别提到在 prefill-decode disaggregation 场景下,可以用更少 KV admission 和更少 decode-node host memory 达到接近两倍 dense throughput 的效果。这里的工程意义很大:如果 decode 节点不必接收全量 KV,就能承载更多请求。

对平台团队来说,这意味着调度器需要理解 KV 的冷热程度,而不只是理解请求长度。

生产系统怎么借鉴

即使今天还不能直接 pip install oasiskv,也可以提前调整推理平台指标。

第一,监控 HBM 中 KV 占比。不要只看 GPU 利用率和 tokens per second,要按模型、上下文长度、batch、并发拆解 KV 占用。

第二,区分 prefill 和 decode 成本。很多服务把总延迟混在一起看,导致不知道瓶颈是在首 token 之前还是生成阶段。

第三,记录长上下文任务的注意力或检索模式。未来做稀疏 KV、块缓存、KV eviction,都需要知道真实请求是否有局部性。

第四,把 speculative decoding 看成系统信号,而不是单一加速技巧。drafter 输出可以帮助调度、预取、提前分配内存。

第五,准备多层内存架构。HBM、主机内存、NVMe、远端内存的层级会越来越重要。长上下文推理不是单卡显存问题,而是内存系统问题。

可能的局限

第一,稀疏注意力假设不一定适合所有任务。代码审查、数学证明、法律文本可能在远距离 token 上有关键依赖。

第二,lookahead 预测质量取决于 drafter。若 drafter 猜错未来 token,预取的 KV 也可能错。

第三,后台预取 pipeline 对系统工程要求高。带宽、调度、kernel、流同步、跨 GPU 通信都可能成为新瓶颈。

第四,准确率需要按任务测。摘要里的精度损失很小,但生产任务不能只看通用 benchmark。

对开发者的直接影响

应用开发者不一定直接实现 OasisKV,但会感受到它代表的趋势。

未来长上下文 API 的定价和限额可能更细,不只按 token,还会按上下文驻留时间、并发、KV 复用、cache 命中计费。Agent 产品也会更关心轨迹压缩,因为长轨迹不是只增加 prompt 成本,还占用 decode 阶段 KV 资源。

如果你在写 agent,应该避免无意义地把全部历史塞回上下文。保存结构化状态、摘要、证据索引和可检索轨迹,会比无限拼接聊天记录更适合未来推理系统。

结论

OasisKV 的价值在于把 KV cache 从“只能塞进 HBM 的负担”重新定义为“可以被预测、分层、预取和稀疏化的工作集”。这和过去几年数据库、操作系统、CPU cache 的思路很像,只是现在搬到了 LLM decode 阶段。

长上下文模型越普及,推理优化就越不像单纯 kernel 优化,而更像内存系统设计。OasisKV 给出的方向很明确:下一代 LLM serving 平台要懂 token,也要懂 KV 的生命周期。

Frequently asked questions

OasisKV 解决什么瓶颈?
它解决长上下文和长推理解码阶段 KV cache 占满 HBM、限制 batch size 和吞吐的问题。
它是不是普通 KV 量化?
不是。KV 量化降低单个 KV 的存储成本,OasisKV 更关注哪些 KV 块需要常驻 HBM,哪些可以放到主机或远端内存。
为什么 speculative decoding 有帮助?
lookahead token 可以提前暴露未来注意力可能访问的位置,系统就能在真正解码前把重要 KV 块预取回 HBM。
这会不会损失准确率?
论文摘要报告在一定 KV 预算下准确率接近 full attention,但任何稀疏和预取系统都需要按业务任务重新测精度。
普通团队现在能直接用吗?
如果没有相关实现,还不能当作现成参数打开。但它给 vLLM、SGLang 和长上下文服务设计提供了明确方向。
// next.txt ›

Some outbound links in this post are affiliate links — see disclosure.