OasisKV：用先行稀疏预取把解码期 KV 缓存扩展到 HBM 之外 OasisKV: Scaling In-Decode KV Cache Beyond HBM with Lookahead Sparse Prefetching
复用推测解码草稿令牌提前一步预测并异步预取关键KV块，让HBM只驻留稀疏工作集
前置知识
KV 缓存
Transformer 自回归解码时，为避免每步重算历史 token 的注意力投影，系统把每层每个 KV 头的 Key/Value 缓存下来。对 $N_{layer}$ 层、$N_{kv}$ 个 KV 头、头维 $d_h$、每元素 $s$ 字节的模型，批量 $B$、上下文 $L$ 的占用为 $M_{KV}=B\cdot L\cdot N_{layer}\cdot 2N_{kv}d_h\cdot s$。GQA/MLA 通过共享或压缩头减小这个系数。
论文的整个问题空间都围绕这块随上下文线性增长、解码期间必须可访问的缓存展开，公式 (1)(2) 直接决定最大批大小，是理解一切动机的起点。
HBM 与内存墙
HBM 是 GPU 上的高带宽显存，带宽极高但容量小、单价贵（H100 仅 80 GB）。解码阶段每生成一个 token 都要完整读一遍 KV 缓存，缓存装不下 HBM 时批大小被压死、算力闲置，这就是解码的内存墙。vLLM/SGLang 的 KV offloading 只是为前缀复用把缓存挪到 CPU，并不解除解码期的常驻要求。
本文的目标就是既把 KV 挪出 HBM（解容量）又不拖慢解码（解带宽），读懂 roofline 分析需要这个背景。
稀疏注意力（Quest 式块摘要）
注意力分数高度偏斜，每个 query 只真正关注少数历史 token。Quest 等方法把历史切成逻辑块，为每块 key 保存逐维 min/max 两个摘要向量，用 query 与摘要粗算块重要性后只对 top-$K$ 块做精确注意力，把每步读取从 $L$ 个 token 降到约 2,048 个。
OasisKV 直接复用 Quest 式压缩 key 摘要做 top-$K$ 预测，稀疏化本身也是它省 HBM 容量的前提。
KV 检索与 KV 预取
KV 检索：注意力前现查现取非驻留块，PCIe 传输位于解码关键路径，直接拉高 TPOT；KV 预取：提前预测下一步要用的块，用后台传输与前向计算重叠。预取的成败取决于预测命中率、以及每步传输字节数能否塞进 $C_{token}=B_{link}\cdot T_{decode}$ 这个每步字节预算。
本文属于预取路线，其三大设计挑战（预测准、藏延迟、控流量）全部由这条路线的固有矛盾引出。
推测解码与 MTP（EAGLE-3）
用小草稿模型一次猜出未来 token，目标模型一次前向并行验证，被接受的 token 直接输出，从而一次前向产出多个 token。EAGLE-3 是目前主流的多 token 预测（MTP）草稿头，且多以公开可复用制品的形式发布。草稿 token 若验证失败通常被直接丢弃。
OasisKV 的核心 trick 是把这些本会被丢弃的草稿 token 当作下一步注意力模式的免费预测器，这是全文的 lookahead 信号来源。
PD 分离（Prefill-Decode Disaggregation）
把算力密集的 prefill 与访存密集的 decode 拆到不同节点各自优化，代价是 KV 缓存必须在请求准入时跨网络从 prefill 节点搬到 decode 节点。现有稀疏检索系统（如 DualPath、HiSparse 类）先搬全量 KV 到 decode 节点 DRAM，造成 TTFT 恶化与内存容量瓶颈。
§4.4 的远程部分取数（RPF）专门为这个场景设计，理解 §3.3 的两个瓶颈需要 PD 分离背景。
PagedAttention 与页表
vLLM 把每条请求的 KV 缓存按固定 token 数切成逻辑块，经页表映射到物理页，按需分配。本文在 GPU 与 CPU 两套原页表之上再叠加一层按 KV 头组织的逻辑-逻辑映射，使每个 KV 头能维护不同的稀疏块集合，同时不破坏原页表与 FlashAttention 执行路径。
Fig 9 的头级映射是让稀疏解码融入生产引擎的关键工程机制，不懂页表就看不懂这篇的系统贡献。
Token-KV 强度 roofline 模型
定义为每次解码每搬运 1 字节 KV 能服务多少 token 生成，吞吐上限等于强度乘以带宽。GQA/MLA 靠复用提高强度，稀疏注意力靠少搬提高强度；带宽屋顶分 HBM、CPU-GPU IO、远端网络三层，斜率均为 1。
Fig 1 用这个模型一图定位所有已有方法：KV 检索与预取都被钉在 PCIe 屋顶，OasisKV 的目标就是打破 off-GPU 内存墙。
研究动机
智能体工作负载把上下文推向聊天时代的 10 倍以上，生产编程 agent 实测平均上下文已达 32.7K token（DualPath 追踪数据）。KV 缓存随上下文线性增长：32B 级 GQA 模型在 BF16 下每 token 占 256 KiB，32K 上下文单请求约 8.6 GB，即使把 80 GB HBM 全部留给 KV，最大批也只有 9 个请求。dense vLLM 在 Qwen3-8B、16K 上下文下并发 32 即饱和（约 676 tok/s）。分层 KV 两条路各有硬伤：KV 检索把 PCIe 传输放在解码关键路径上，roofline 估算显示每步只取回所注意力 KV 的 10%，在 batch 64、8K 活跃 token 时 TPOT 开销就高达 78%；KV 预取虽能重叠传输，但免训练预测普遍用上一 token 的 query 当代理，top-20 块命中率平均仅 83.9% 且逐层剧烈波动。PD 分离下更糟：现有系统先把全量 KV 搬进 decode 节点 DRAM，Qwen3-235B TP8 节点配 1 TB 内存、平均 100K 上下文时理论批上限只有约 26，网络传输还会主导高前缀命中率下的 TTFT。
本文的目标是构建以内存为中心的推理系统 OasisKV：解码期间全量 KV 留在主机或远端内存，GPU HBM 只驻留每步真正注意到的稀疏工作集（默认每头 2,064 token）。要求预测免训练、足够准（相对全注意力精度损失 0.7 点以内）、便宜到能领先解码一整步生成；每步跨 PCIe/网络的 KV 流量必须控制在 $C_{token}=B_{link}\cdot T_{decode}$ 预算内；并且要嵌入 vLLM 这样的生产引擎，同时支持多卡张量并行、PD 分离与多层存储。最终把稀疏性转化为吞吐：真实推理负载 1.69×、多卡长上下文最高 2.1×、PD 分离 2.1–2.3×。
与已有工作不同的是，独特切入是预测信号的来源：推测解码/MTP 反正要花一次前向草拟未来 token，OasisKV 把这些通常被丢弃的草稿 token 复用为 KV 访问模式的 lookahead 预言——草稿 query 沿当前稀疏工作集传播后，预测的 top-$K$ 块集合与真实下一步 query 的一致率逐层超过 98.2%、平均 98.74%，远超上一 token 代理的 83.9%，且无需训练任何预测器。其次，它把重叠传输、稀疏预测、稀疏-精度权衡当作联合系统问题：用 delta 选择加封顶驱逐把每步 PCIe 字节数硬性钉在预算内；用远程部分取数把 PD 分离的准入传输量从随上下文增长改为由 top-$K$ 预算封顶。相比之下，已有工作要么训练专用预测器并绑定模型（SparDA、FlashMemory），要么在 HuggingFace 小批量原型上评估、缺少面向生产引擎的跨层内存管理。
核心方法
直觉：草稿模型已经免费给出了下一步的 token，那就让它顺手当一次注意力模式预言家。系统分三个平面（Fig 5）：计算平面实现 look-ahead attention——前台执行正常稀疏前向，同时把草稿 token 分叉进一条无同步的轻量后台注意力路径；控制平面的池管理器通过页表管理 GPU 上的稀疏 KV 池、压缩 key 池与草稿 KV 池，调度器为每步派发前台与预取任务；数据平面在本地 DRAM（非分离部署）或远端内存（PD 分离）持有全量 KV。每层后台流水线串起三个阶段：top-$K$ 预测 → KV 选择 → KV 传输；层与层之间任务独立（各自维护压缩 key、驻留映射与 KV 存储），可并发执行。流水线填满后的稳态完成间隔为 $T_{pipe}=\max\{T_{pred},T_{select},T_{transfer}\}\le\Delta_{prefetch}=T_{step}/N_{layer}$，即由最慢阶段决定、仍快于每层平均可用时间，传输因此完全躲开前台关键路径。
核心创新是 lookahead 信号的产生与消费方式。与模型内嵌预测器（需端到端训练、绑定特定模型）和上一 token 代理（免训练但平均仅 83.9%、逐层不稳）本质不同，OasisKV 让草稿 token 走正常解码前向：QKV 投影把正常与草稿 query 堆叠成 $\{q_t^l; q_{t+1}^l\}$ 一次注意力调用，权重与驻留 KV 只读一遍；草稿 query 随后扫描每块的 min/max 压缩摘要（Quest 式，仅占全量 KV 的 1/16），对每个 KV 头独立给全上下文排名——不恢复任何 key 就能做全局 top-$K$，逐层一致率均超 98.2%。第二个关键区别是把流量当一等公民：KV 选择先与驻留集求交（命中块零 PCIe 流量），再把驻留集外、按最近一次被选时间排序的旧块作为驱逐目标，与非驻留新块逐一配对；每个 KV 头每步最多接纳 $C$ 块（由 fetch ratio 控制），从而给每步字节量一个确定性上界，而不是指望碰运气式的高命中率。
方法步骤详情
每步每层的完整流程：前台 QKV 投影同时产出正常 token 与草稿 token 的 query，注意力内核在驻留工作集（每头 129 块共 2,064 token，含 1 个活动本地块）上一次算出两个输出；草稿 query 提交 top-$K$ 预测 worker，扫描压缩 key 摘要排出下步块序。KV 选择 worker 把预测集与驻留集求交，命中者不动；预测集中的非驻留块按 $C$ 上限（fetch ratio）逐头配对一个按最近最少选择策略选出的驱逐目标，形成传输计划。传输 worker 用自定义 UVA gather 内核经 PCIe 把 CPU pinned 内存里的块直接写入 GPU 驻留页。三阶段作为持久 C++ worker 跑在独立 CUDA 流上，CUDA event 保证层内依赖；第 $t$ 步第 $l$ 层的传输只需在该层 $t+1$ 步注意力前完成（层内局部同步），于是预测($l{+}1$)、选择($l$)、传输($l{-}1$)三层并发。请求上下文超过预算后由稠密转稀疏：全量 KV 拷至 CPU，GPU 留选中块加本地窗口，新增的头级逻辑-逻辑映射（Fig 9）支持各头不同稀疏集、稠密/稀疏请求混批，TP 下各卡只更新本地 KV 头、无跨卡同步。PD 分离时（§4.4）：prefill 节点额外执行一步解码拿到首个 query，选出各头 top-$K$ 的块级并集，连同压缩 key 与草稿 KV 状态做部分传输，全量 KV 留在 prefill 侧 staging 池并立即释放 prefill GPU；解码中 top-$K$ 漂移导致的 miss 由同一后台流水线按层聚合、从远端 staging 池读取，每块至多过网一次，未被选中的块永不传输。
技术新颖性
技术新颖性有四点。(1) 信号复用：首个把 EAGLE-3 草稿 token 同时用作验证候选与 KV 预取预言的生产级系统，零训练、零额外前向（与正常 token 共享权重与驻留 KV 读取），98.74% 的 top-$K$ 集合一致率把免训练预测从不可用提到近乎精确。(2) 流量整形：delta 选择加封顶驱逐把预测命中率与放置策略解耦成每步确定性的字节上界，Table 2 显示这一个旋钮可在 2.6× 吞吐范围内调节而精度只波动约 2.5 点，把 PCIe 饱和（30–34 GB/s）直接变成可度量的设计约束。(3) 结构兼容：头级逻辑映射叠加在 PagedAttention 页表之上而非替换它，注意力仍走 FlashAttention 同一路径，稠密到稀疏的运行时切换、混批、TP 本地更新都不破坏引擎不变量——这正是 ShadowKV/InfiniGen/FreeKV 跑不了大批量配置的根因。(4) RPF 改变准入流量的 scaling 律：准入传输量由 top-$K$ 预算而非上下文长度决定（0.46 GiB 对 4.50 GiB），漂移部分再经 fetch ratio 二次限流。
实验结果
系统性能（Fig 11）：Qwen3-8B 单卡 16K、最大并发 128 时 OasisKV 达 1,398 tok/s 对 dense vLLM 676 tok/s（2.1×），32K 下最高 2.4×；并发 16 时吞吐与延迟双赢（836 对 649 tok/s，TPOT 17.7 对 23.5 ms）；2,048-token 预算支撑 90–95 并发而 dense 仅 22。ShadowKV/InfiniGen/FreeKV 的大批量点多因 OOM 或不支持而无法运行，且均无跨 GPU 实验。Qwen3-235B TP8 因 KV 占显存比例小，低批时 TPOT 反高于 dense，从 B=32（16K）/B=16（32K）起反超，最高 1.9×（1,102 tok/s）。PD 分离（Fig 12）：dense 因 HBM 容量在 0.27/0.19 req/s 饱和（550–554/384–386 tok/s，32K 每轮 20–33 次抢占），OasisKV 达 1,204–1,210/888–884 tok/s（2.1–2.3×）；RPF 把每请求 decode 节点 DRAM 从 3.38/4.52 GiB 降到 1.54/1.73 GiB（2.2×/2.6×），峰值聚合占用从 209/167 GiB 降到 76/46 GiB。真实负载（Fig 13）：Qwen3-8B 2,083 对 1,235 tok/s（1.69×），AIME24 avg@32 为 75.94 对 76.04（仅 −0.1 点）；235B 为 1,546 对 1,283 tok/s（1.20×），83.85 对 84.69。$K$ 扫描：8B 从 $K{=}192$ 的 1.39× 到 $K{=}64$ 的 1.89×（精度 76.77→72.40），$K\ge128$ 精度与 dense 持平，$K{=}192$ 甚至微超（76.77 对 76.04；84.90 对 84.69）。精度（Table 1）：与自家锚点比整体最接近——LongBench v2 Qwen3-8B −0.40（Quest −0.60、FreeKV −1.20）、Llama-3.1-8B −0.61；长输出推理总体 pass@$k$ −0.66、avg@$k$ −0.35，Quest/FreeKV 分别 −2.56/−3.50 与 −2.83/−2.63；AIME24、AIME25 pass@8 与全注意力打平（83.33、80.0），GPQA pass@4 −1.99。消融（Table 2）：fetch ratio 从 0.01 到全取，每步流量 0.30→5.05 GB，吞吐 2,178→824 tok/s（2.6× 暴跌）而精度仅 74.9→77.4，有效 PCIe 带宽在 0.10 后饱和于 30–34 GB/s；默认 0.05 以 2,083 tok/s 保住与 dense 差 0.1 点的精度。RPF 流量（Fig 14）：准入 0.52/0.46 GiB 对全传 3.37/4.50 GiB（6.5×/9.7×），含漂移后总节省 1.33×/1.50×，fetch ratio 0.1/0.05 把 32K 漂移从 2.53 压到 1.95/1.61 GiB（总 1.9×/2.2×）；全传链路 60% 空闲但交接时 5.3 GB/s 突发，RPF 仅 1–2% 空闲、峰值从 3.1 降到 2.1 GB/s。TTFT（Fig 15 解析模型）：90% 前缀命中时 100 Gbps 链路提速 2.0×/2.2×，400 Gbps 下仍 1.14×/1.23×。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 真实推理负载端到端服务（Qwen3-8B，AIME24 请求） | 解码吞吐（tok/s） | 2,083 | dense vLLM + FlashAttention3：1,235 | 1.69×，AIME24 avg@32 精度仅 −0.1 点（75.94 对 76.04） |
| 单卡长上下文合成负载（Qwen3-8B，16K，最大并发 128） | 解码吞吐（tok/s） | 1,398 | dense vLLM：676 | 2.1×；32K 下最高 2.4×；并发容量 90–95 对 22 |
| 多卡长上下文服务（Qwen3-235B-A22B，TP8） | 解码吞吐（tok/s） | 最高 1,102（16K）；32K 下相对加速 2.1× | dense vLLM | 大批量下最高 1.9–2.1×，低批时 TPOT 略逊于 dense |
| PD 分离服务（Qwen3-8B，24K/32K，2,048 输出） | 解码吞吐（tok/s） | 1,210 / 884 | dense vLLM：550–554 / 384–386 | 2.1–2.3×，且无 dense 的 20–33 次抢占 |
| PD 分离准入 KV 传输（32K 上下文） | 每请求传输量（GiB） | 0.46（RPF 准入） | 全量传输 4.50 | 9.7×（24K 为 6.5×）；含解码漂移后总节省 2.2×（fetch ratio 0.05） |
| 长上下文精度（LongBench v2，Qwen3-8B） | 总体准确率 | 33.20（相对全注意力 −0.40） | 全注意力锚点 33.60；Quest −0.60；FreeKV −1.20 | 三者中与全注意力最接近 |
| 长输出推理精度（AIME24/25+GPQA-Diamond，Qwen3-8B） | 总体 pass@k | 78.18（−0.66） | 全注意力锚点 78.84；Quest 78.86（−2.56）；FreeKV 77.91（−3.50） | 损失约为竞品四分之一；AIME24/25 pass@8 与全注意力完全打平 |
| 每步取回上限消融（Qwen3-8B，AIME24） | 吞吐 @ 精度 | fetch ratio 0.05：2,083 tok/s @ avg@32 75.94 | 全量取回：824 tok/s @ 76.46；dense 76.04 | 2.5× 吞吐，精度差 0.1 点以内 |
局限与改进
作者承认的局限：Qwen3-235B TP8 低批时因 KV 占显存比例小、稀疏省得少而每步额外开销固定，TPOT 反而高于 dense；当前原型把草稿 token 强制拒绝，投机解码的加速红利与草稿预取尚未联合兑现；原型不支持前缀缓存，Fig 15 的 TTFT 收益只是解析模型；prefill 侧 staging 池在请求完成前仍保留全量 KV。我自己的观察：lookahead 依赖相邻步块重要性的时间局部性（论文明说 top-$K$ 漂移靠 fetch ratio 兜底），对注意力模式突变剧烈的任务（仓库级代码跳转、多文档交叉引用）漂移 miss 会上升，且 miss 的补救是异步远程取数，可能放大长尾延迟；压缩 key 摘要占全量 KV 的 1/16 并常驻 HBM，对 MLA 这类 KV 本已极小的架构增益有限；评估覆盖 Qwen3-8B/235B 与 Llama-3.1-8B 三个模型、H100 + PCIe Gen5 + 400G RoCE 一种硬件栈，低带宽链路或非 H100 卡上的表现未知；$K$ 全局固定 128，GPQA pass@4 掉 1.99 点说明部分任务对预算敏感，缺少按层/按头的自适应预算；每个目标模型必须有配套 EAGLE-3 草稿头，模型覆盖是硬性前提。
独立分析的弱点
独立分析的弱点与改进方向：(1) 草稿 token 被强制拒绝，SD 本可带来的单请求加速被完全放弃——应让被接受的草稿 token 及其 KV 参与后续步骤，把接受长度与预取深度联合调度（作者已列为未来工作）。(2) fetch ratio 与 $K$ 全局静态：Table 2 显示 0.05 是全局折中，但不同层、不同头、不同负载阶段的边际收益差异很大（精度在 0.01–0.20 间只动 2.5 点，说明许多步取回是浪费）——可做基于 PCIe 占用率与预测置信度闭环的动态控制器，或按层自适应 $K$。(3) 时间局部性假设在突变型负载下失效，远程 miss 取数的延迟没有优先级通道——可为高漂移请求回退到更大本地窗口或为纠正取数设立抢占式高优传输。(4) 前缀缓存缺失对多轮 agent 是现实痛点，而这类负载恰恰是论文的动机场景——应把 staging 池与 vLLM APC/SGLang HiCache 统一，命中前缀只传增量。(5) MoE 大模型收益打折的根源是稀疏化没有按显存压力自适应触发——系统已有稠密转稀切换机制，可升级为按请求粒度、按实时 KV 压力决策的策略。
未来方向
作者提出的方向：联合投机解码与草稿预取，让被接受的 token 摊销草稿开销；把 prefill 侧 staging 池扩展到远端内存服务器或 SSD 层（同一方法论直接适用）；开启模型-系统协同设计空间，横跨 KV 稀疏度、长上下文能力、内存硬件扩展与系统性能。基于其成果可自然延伸：与原生稀疏注意力模型（NSA/DSA）结合，lookahead 信号可直接喂给其内建索引器；针对 MLA 架构联合设计压缩 KV 与块摘要，绕开 1/16 摘要开销；把头级稀疏工作集做成跨请求共享的分布式 KV 内存服务，进一步抽象出 KV-as-a-Memory 层；训练专门优化注意力分布预测（而非仅 token 预测）的 MTP 头，进一步提高漂移任务下的命中率；以及多租户场景下按 SLO 动态调节 fetch ratio 与驱逐策略。
复现评估
工程实现基于 vLLM v0.12.0 的 V1 引擎，扩展了稀疏注意力后端、GPU model runner、KV 缓存管理器与稠密转稀调度支持；压缩 key 更新、头级 top-$K$ 预测、稀疏页映射与 KV 传输用 C++/CUDA 编写，持久后台 worker 跑在独立 CUDA 流上，CUDA event 维护四流依赖；PD 分离传输用 NIXL 1.3.0 over UCX 1.21.0。所用模型与草稿头全部公开（Qwen3-8B、Qwen3-235B-A22B、Llama-3.1-8B-Instruct 加三个公开 EAGLE-3 草稿头），基准公开（AIME24/25、GPQA-Diamond、LongBench v2），采样参数与种子在文中交代完整。论文未明确给出开源代码链接。复现硬件门槛高：完整测试床是 8×H100 80GB + 2TB 主机内存 + 双节点 ConnectX-7 400G RoCE；但单卡配置（一张 H100）即可复现 Fig 11 上排与 Table 2 的核心结论，门槛大幅降低。总体属于可复现但费劲：页表、多 CUDA 流、UVA gather、调度器改造的系统工程量大，有 vLLM 内核开发经验的小团队预计需要数周到数月。
论文图表
用 roofline 模型刻画解码吞吐上限：吞吐 = Token-KV 强度 × KV 带宽，画出 HBM、CPU-GPU IO、远端内存三条斜率为 1 的带宽屋顶，加上 dense/sparse 两条吞吐天花板；标注 dense vLLM、稀疏注意力、KV 检索、KV 预取分别被卡在哪个屋顶，OasisKV 用绿箭头表示要打破 off-GPU 内存墙。
一图总览全文的问题定位：现有 KV 检索与预取都被钉在 PCIe 屋顶附近，远离高吞吐生产部署，这是论文的出发点和卖点。
基于 roofline 模型估算按需 KV 检索的 TPOT：每步只从 CPU DRAM 取回所注意力 KV 的 10%，开销随批量与活跃 token 数增长——4K 活跃 token 时约 52–64%，8K、batch 64 时达 78%；一旦传输时间超过基础解码时间，重叠即失效。
用数据证明按需检索在生产批量下不可行，是选择预取路线而非检索路线的直接论据。
四种解码流水线的结构对比：dense（QKV→Attention→FFN，延迟✗容量✗）；sparse attention（增加 top-$K$ 选择阶段，延迟✓但容量✗，全量 KV 仍在 HBM）；KV retrieval（注意力前插入取数，容量✓但延迟✗）；KV prefetch（提前取数并与前台计算重叠，延迟✓容量✓）。
一张图讲清四条技术路线的本质差异，说明预取是唯一能同时解决延迟与容量问题的路线，为后文设计定调。
逐层对比两种免训练预测信号的 top-20 块命中率：用当前步上一 token 的 query 平均仅 83.9%，且逐层波动明显；用 lookahead 草稿 token 平均 96.0%，逐层稳定在 90% 以上。
本文核心论点的关键实验证据：草稿 token 作为预测信号远优于既有免训练代理，直接支撑 3.1 节对准确性的要求。