arXiv'26 | OasisKV：KV Cache 超出 HBM 后，如何用 Lookahead Sparse Prefetch 把等待藏起来
arXiv’26 | OasisKV：KV Cache 超出 HBM 后，如何用 Lookahead Sparse Prefetch 把等待藏起来
原文：OasisKV: Scaling In-Decode KV Cache Beyond HBM with Lookahead Sparse Prefetching
1. 先说结论：这不是“把 KV Cache 挪到 CPU”这么简单
长上下文 LLM 推理最容易被低估的地方，是生成阶段（decode）其实更像一个内存系统 benchmark，而不是一个 FLOP benchmark。每生成一个 token，模型新增的矩阵计算并不大，却要反复读取几千、几万甚至几十万历史 token 的 Key/Value（KV）向量。
当 KV Cache 放不进 GPU 的 HBM（高带宽显存）时，最直觉的方案是把它放到 CPU 内存，算 attention 前再把需要的 block 搬回来。问题也正是从这里开始：如果“选 block + PCIe/网络传输”在关键路径上，稀疏 attention 节省的计算会被等待时间吃掉。
OasisKV 的切入点很朴素，但需要把几个环节同时做对：用一个便宜的 draft token 预测下一步会访问哪些 KV block；当前 token 做正常 attention 的同时，后台完成选择和预取；下一步到来时，GPU 只处理已经准备好的 sparse working set。论文把这套机制叫 Lookahead Sparse Prefetching。
这篇工作的价值不在于又发明了一个 top-K 算法，而在于它把“稀疏性、分层存储、异步调度”放进同一条 decode pipeline 里，直接对准了远端 KV 的核心矛盾：容量在远端，带宽在本地，延迟不能暴露给前台。
2. 先把背景讲清楚：LLM decode 为什么越来越像内存系统问题
2.1 Prefill 和 Decode 是两种完全不同的工作
LLM 推理通常分两段。
- Prefill：把用户输入的整段 prompt 一次性读进去，计算每个 token 的表示。
- Decode：模型已经开始生成后，每次只追加一个 token，再生成下一个 token。
Prefill 更像大矩阵乘法，GPU 的计算单元容易吃饱。Decode 则相反：每一步新增的计算量不大，却要反复读取历史上下文。历史 token 在每一层注意力里的 Key 和 Value 会被缓存下来，这就是 KV Cache。
可以把 KV Cache 想成模型的“工作记忆”。生成第 10000 个 token 时，模型要拿当前 Query 去和前面许多 token 的 Key 做匹配，再用匹配结果加权 Value。上下文越长，历史 KV 越大。
2.2 为什么 KV Cache 会成为瓶颈
粗略地说，一个 batch 中长度为 $L$ 的请求，其 KV Cache 内存量近似为：
\[M_{KV} \approx B \times L \times N_{layer} \times 2N_{KV}d_h \times s\]这里 $B$ 是 batch size，$N_{layer}$ 是层数，$N_{KV}$ 是 KV head 数量，$d_h$ 是 head dimension，$s$ 是每个元素的字节数，乘以 2 是因为同时保存 Key 和 Value。
这条式子没有什么神秘之处：上下文翻倍，KV 就翻倍；并发请求翻倍，KV 也翻倍。模型权重可以通过量化和共享来压缩，但 KV Cache 是跟着每个请求、每个历史 token 增长的。
于是 GPU 里出现了一个很实际的分层结构：HBM 容量小但带宽高，CPU 内存容量大但访问慢，远端内存或网络容量更大但延迟和带宽更差。问题从“GPU 能不能算”变成了“GPU 需要的数据能不能及时到”。
这张 roofline 图先把坐标轴说清楚：横轴是 Token-KV intensity，即每搬运一个 KV byte 能生成多少 token；纵轴是 decode throughput。黑、蓝、红三条斜线分别对应 HBM、CPU-GPU、远端内存带宽上限。往右走代表 KV reuse 更高或 attention 更稀疏，系统才有机会从“内存受限”爬回计算上限。图里的 OasisKV 星标沿着绿色箭头移动，含义不是凭空增加 HBM，而是通过更高的有效 intensity 和预取，把工作点从远端带宽限制区推向更高吞吐。
2.3 稀疏注意力解决了什么，又没有解决什么
全注意力会扫描整个历史 KV。稀疏注意力则只选择最重要的一小部分 KV block。一个 block 可以理解成连续的一小段历史 token，系统先估计每个 block 对当前 Query 的重要性，再只对 Top-K block 做精确 attention。
这样做同时减少了计算量和需要读取的数据量，但它有一个隐含前提：Top-K 选择必须足够准。如果选错，模型的答案可能变差；如果每一步选出来的 block 都要从 CPU 搬到 GPU，搬运也可能成为新的瓶颈。
这就是 OasisKV 和普通 sparse attention 的分界线：它关心的不只是哪一块 KV 重要，还关心下一步会重要哪一块，以及能不能在计算当前 token 的时候把下一步的 KV 搬过来。
3. 现有方法的问题：算得少，不代表等得少
论文在 Qwen3-32B、H100、32K context 上给了一个很有杀伤力的 latency breakdown：on-demand KV retrieval 把选择、PCIe transfer 和 attention 串在一起，单步 TPOT（time per output token，生成一个 token 的平均耗时）里，传输直接变成大头。也就是说，把 full KV 放到 CPU 只解决了容量，不解决 decode critical path。
先看四种 decode 流程的差别。
全注意力把所有 KV 放在 GPU 上，计算慢、容量也受限。稀疏注意力只计算 Top-K，但 KV 仍然需要在 GPU 上。KV retrieval 允许大部分 KV 放在 CPU，却在每一步先 fetch，再 attention，于是容量问题缓解了，延迟问题没有。
KV prefetch 的目标是把 fetch 和当前 attention 重叠起来：当前步骤做前景计算，后台同时准备下一步数据。只要预取足够早、预测足够准，下一步 attention 就不必再停下来等 PCIe 或网络。
4. OasisKV 的关键观察：下一个 Query 可以提前预测
4.1 Draft token 提供了一个便宜的 lookahead
自回归生成的第 $t+1$ 个 token，通常由第 $t$ 步的模型输出决定。OasisKV 利用这个顺序关系，在当前 token 的正常 forward 之外，构造一个轻量的 draft token，然后用 draft token 对下一步的 Query 做预测。这里的 draft 不是凭空训练出来的另一个模型，而是借用了 speculative decoding 里已经比较成熟的 MTP（multi-token prediction，多 token 预测）信号；论文实验采用 EAGLE-3 draft head。它的作用类似“先用一个便宜的小头猜后续 token”，主模型仍然负责最终输出。
更细一点看，draft token 先在当前层得到 draft hidden state，随后沿着 transformer layer 传播出 draft Query。这个传播过程不需要为每一层重新读取 full KV：GPU 只维护每个 block 的 compressed key 摘要，draft Query 对摘要做一次便宜的 Top-K ranking。真正完整的 Key/Value 只在 block 被选中后才从 CPU pool 搬运。于是 lookahead 的成本主要是两条短路径：共享一次 QKV/权重计算，再扫描每个 block 的两个压缩向量。
预测并不需要把下一步完整算完。论文观察到，用 lookahead token 做 KV block 预测，比用 previous token 更准：所有层上的 Top-20 KV-block 预测准确率，前者平均约 96.0%，后者约 83.9%。这里的准确率是预测集合与下一步真实 Top-K 集合的重合程度，越高意味着预取的浪费越少、cache miss 越少。
这个差别很关键。预取的本质是“提前下注”：猜对，传输可以隐藏在计算后面；猜错，GPU 要么白搬数据，要么下一步仍然要等待。
4.2 Look-ahead attention 做两件事
OasisKV 的 look-ahead attention 可以拆成前景和后台两条路径：
前景路径使用当前 Query，在当前的 sparse KV working set 上完成真正的 attention，并继续生成下一个 token。
后台路径使用 draft Query，先扫描 GPU 上更小的 compressed key pool，预测下一步的 Top-K block；然后从 CPU KV pool 选出这些 block，发起 KV transfer。
compressed key pool 是怎么来的？CPU 上保留完整 KV，GPU 上则额外保存每个 block 的压缩 Key。论文使用 block 内 Key 的 min/max pooling 来构造紧凑表示。它不负责最后的精确 attention，只负责快速判断“这个 block 可能重要吗”。真正的 Value 和完整 Key，仍然只对入选 block 搬运和计算。这个设计的关键是把“候选筛选”和“精确 attention”拆开：前者必须便宜到能和前台并行，后者才值得消耗带宽。
图中左侧是 CPU 的 full KV 如何映射成 GPU 上的 compressed key 和 sparse KV 两种缓存；右侧是正常 Query 做 attention 的同时，draft Query 扫 compressed key，预测下一步需要的 block。
4.3 预取流水线为什么能隐藏延迟
它不是简单地把“预测”挪到前面，而是把跨两个 decode step 的依赖排成流水线：
在 step $t$，模型一边执行当前层的 QKV projection、attention、FFN，一边为 step $t+1$ 做 Top-K prediction、KV selection 和 KV transfer。下一步开始时，被预测的 KV 应该已经抵达 GPU。
这里有一个工程上的取舍：预取不能无限大。预取太少，命中不了下一步所需的 block；预取太多，又会浪费 PCIe/网络带宽和 GPU 内存。OasisKV 因此提供 fetch-cap，并结合 cache eviction 和 request scheduling 管理有限的工作集。
论文的 ablation 很说明问题：在 Qwen3-8B 的 AIME24 workload 上，fetch ratio 设为 0.05 时，吞吐约 2,083 tok/s，准确率比 dense 只低 0.1 个百分点；如果把 cap 放开到“几乎全取”，吞吐反而掉到约 824 tok/s，而准确率只从 74.9 小幅涨到 77.4。多搬 6 倍数据，换不到同等比例的正确率。这也是系统论文和单纯 sparse attention 论文的区别：预算不是一个离线超参数，而是每一步都在和互联带宽竞争的在线资源。
4.4 为什么需要 head-wise 映射
GPU 的 sparse working set 和 CPU 的 full KV 并不是同一张 page table。OasisKV 为每个 KV head 维护额外的 GPU-logical-block 到 CPU-logical-block 映射；原有 page table 仍然负责请求内的逻辑 token 到物理页转换。这样做看起来像 bookkeeping，但它决定了“命中一个 GPU block”到底对应 CPU 上哪一段连续 KV，也决定了 eviction 后能否无锁地复用 GPU page。
举个具体例子：假设 block size 是 16，某个 head 的 GPU logical block 7 当前映射到 CPU logical block 31，那么 token 120（$120//16=7$）在 GPU 上查到 block 7 后，miss 时要从 CPU 的 block 31 取 offset 8，而不是重新扫描整条 KV。这个额外映射把“稀疏选择”变成真正可执行的 block transfer。
5. 系统架构：三层 KV Pool 加两个平面
OasisKV 的架构可以按“数据放在哪里”和“谁负责调度”来理解。
数据侧有三类 pool：
- GPU 上的 Sparse KV Pool：当前真正参与 attention 的稀疏工作集。
- GPU 上的 Compressed Key Pool：用于快速预测下一步重要 block。
- CPU/远端内存上的完整 KV Pool：容量足够大，保存完整历史。
系统侧有三个平面：
- Compute plane：执行前景 forward 和后台 lookahead attention。
- Control plane：由 scheduler 和 pool manager 负责请求调度、缓存准入、淘汰和映射。
- Memory plane：负责 GPU、CPU 和远端内存之间的数据传输。
这个拆分的价值在于，KV 不再被当成“必须全驻留 GPU 的一块连续张量”，而是被当成可按 block 管理、按 head 访问、可异步迁移的分层数据。
6. 远端内存场景：不必每次都传完整 block
当 KV 放在另一台机器上时，瓶颈更明显：网络往返和传输时间可能比一次 sparse attention 还长。OasisKV 进一步设计了 remote partial fetching。
在 admission 阶段，系统可以先把部分 KV 从远端搬入 CPU/GPU 层级；在 decode 阶段，如果某些 block 已经在本地，就只通过网络取缺失部分。对 draft KV 也可以做部分传输，减少不必要的网络流量。具体来说，decode worker 初次接收请求时只拿三样东西：第一步 working set、compressed-key cache、draft token 的 KV state。prefill worker 可以在 host DRAM staging 的同时释放自己的 GPU page，不必等整条 full KV 传完。这一点对 prefill-decode disaggregation 很重要，否则 decode 节点的 host memory 会被“还没开始生成的请求”提前吃满。
这背后的原则是：访问粒度、缓存粒度和传输粒度不一定要完全相同。把整条 KV Cache 作为一个不可分割的对象，系统很难做容量控制；拆成 block 后，才能做 Top-K 选择、部分预取和按需淘汰。
7. 实验结果：准确率只是入场券，重叠才决定速度
OasisKV 的实验重点不是单独证明某一个 kernel 更快，而是验证整条链路是否成立：
- lookahead token 能否准确预测下一步的 KV block；
- 预取是否真的和前景计算重叠；
- GPU HBM 不够时，额外容量是否能换来更好的吞吐和并发；
- 放到 disaggregated serving 后，网络传输是否仍然可控。
论文报告的端到端结果是：在 Qwen3-8B 的 AIME24 reasoning workload 上，OasisKV 达到 2,083 tok/s，而 dense vLLM 是 1,235 tok/s，即约 1.69x throughput，准确率只低 0.1 个百分点；multi-GPU long-context serving 的最高收益是 2.1x。prefill-decode disaggregation 场景下，系统吞吐是 dense 的约 2.1–2.3x。它不会先搬完整条 KV，而是把 admission 时的网络传输量降到 full transfer 的 1/6.5–1/9.7；对应地，decode 节点的每请求 host-memory 占用可降到约 1/2.2–1/2.6。这里的 accuracy point 指任务准确率的绝对百分点差，不是相对百分比，论文的意思是用很小的质量代价换吞吐。
但这些数字不能脱离前提看：lookahead token 的预测集合要足够准，异步 transfer 要能被当前层的 attention/FFN 覆盖，而且请求调度不能让多个 request 同时把 PCIe 或网络打爆。换句话说，预测准确率是入场券，真正决定收益的是 transfer 是否被重叠。
但这里也要看清边界：OasisKV 不是把 HBM 变成了同等速度的“大内存”。如果工作负载的访问模式变化剧烈，或者网络/PCIe 带宽不足以覆盖预取窗口，仍然会出现 cache miss 和等待。它做的是把可预测的远端访问尽可能提前，并把剩余等待压缩到最小。
8. 我对这篇工作的理解
OasisKV 最有意思的地方，是它把 sparse attention 的问题从“选哪些 token”推进到了“什么时候开始搬、搬多少、放在哪一层”。这一步很重要，因为真实推理系统里的延迟通常不是某个算子孤立决定的，而是计算、内存和调度共同决定的。
它也提醒我们，KV Cache 优化有三个不同层次：
- 表示层：量化、压缩，减少每个 KV 元素的字节数；
- 访问层：稀疏注意力，只访问少数重要 block；
- 时间层：预取和流水线，把访问延迟藏到当前计算后面。
只做其中一层，往往会把瓶颈推给下一层。OasisKV 的贡献在于把三层放到一个 serving architecture 里共同设计。
当然，lookahead 也不是免费午餐。它需要额外的 draft/预测计算，需要更复杂的 cache manager 和 scheduler，还要处理预测错误、淘汰冲突和多请求之间的带宽竞争。未来真正落地时，模型架构、量化格式、GPU 到 CPU 的互联，以及请求长度分布都会影响收益。
9. 总结
如果只记住一句话：OasisKV 用当前 decode 的“草稿未来”预测下一步会访问哪些 KV block，再把这些 block 在后台提前搬到 GPU，从而让超出 HBM 容量的 KV Cache 也能参与低等待生成。
它不是简单地把显存换成内存，而是在做一个面向 LLM decode 的分层缓存系统：GPU 保存稀疏工作集和压缩索引，CPU/远端保存完整 KV，预测器和传输线程负责把下一步需要的数据提前准备好。
原文链接：OasisKV: Scaling In-Decode KV Cache Beyond HBM with Lookahead Sparse Prefetching
顺带一提，我们把 AutoML、NAS 以及 LLM 压缩的相关积累整理成了一本《动手学 AutoML：从 NAS 到大语言模型优化实战》。这篇文章讲的是 KV Cache 的系统层搬运，书里更偏模型和搜索层：第 8 章讨论 LLM 剪枝、量化和模型融合，第 11 章有 LLM 后训练剪枝的完整实战。角度不同，但核心都是同一件事——在有限的计算和内存预算下，把每一份资源花在真正影响结果的地方。
