工作日早上 7 点读 AI，周日早上 8 点读周报订阅邮件

METAL LAB

OasisKV: Scaling In-Decode KV Cache Beyond HBM with Lookahead Sparse Prefetching

arXiv:2608.080972026-08-07

OasisKV把大模型生成文字时用的键值缓存大部分移出GPU，靠提前预测下一步真正需要的部分来预取,让解码吞吐量提升了近一倍

大语言模型逐词生成文本时,每一步都要回顾之前所有词元留下的键值缓存(KV cache),而这个缓存会迅速占满GPU容量有限的高速显存(HBM),尤其在长上下文场景下更为严重。OasisKV利用投机解码(speculative decoding)本来就会生成的草稿词元,不需要额外训练就能提前一步准确预测下一步真正会用到的KV块,再通过后台流水线从主机内存或远程内存把这些块预取到GPU显存中。基于vLLM实现后,在保持与全量注意力相比精度损失不超过0.7分的同时,把解码吞吐量最多提升了约2倍。

METAL LAB 解读图

OasisKV的提前预测与KV预取流程

证据状态已报告实测结果

  1. 问题所在长上下文解码需要把不断增长的KV缓存全部留在GPU显存中,限制了批处理规模和吞吐量
  2. 提前预测利用投机解码产生的草稿词元,在不训练专门预测器的情况下预测下一步会用到的KV块
  3. 筛选与限量传输将预测结果与GPU上已有的块比较,只拉取缺失的块,并按步骤设置上限以保证传输量在PCIe带宽预算之内
  4. 异步后台流水线预测、筛选、传输三个阶段在各层之间重叠执行,让数据传输隐藏在GPU计算过程之后
  5. 结果基于vLLM实现后,吞吐量最高提升约2.1倍,精度损失控制在0.7分以内,预填充解码分离场景下达到约2倍吞吐量
这是 METAL LAB 制作的解读图，并非论文作者提供的原图。

他们做了什么

  1. 问题背景:随着智能体类任务把上下文长度和推理长度都推得更长,KV缓存成为占用GPU高速显存(HBM)最多的部分,限制了批处理规模和吞吐量。
  2. 方法:OasisKV复用投机解码本就产生的草稿词元作为无需训练的信号,提前一步准确预测下一个解码步骤真正需要的KV块,再通过预测、筛选、传输三阶段的后台流水线,把这些块从主机内存或远程内存预取到GPU显存中。
  3. 实现:基于vLLM搭建原型,新增按注意力头的逻辑到逻辑映射机制管理分页KV池,并用限额淘汰策略控制每步允许拉取的KV块数量,使传输能被计算时间掩盖。
  4. 结果:在2048词元的KV预算下,精度与全量注意力(保留全部KV)相比差距不超过0.7分,在一个推理类工作负载上吞吐量比原生vLLM提升1.69倍(精度损失0.1分),在多GPU长上下文场景下最高提升2.1倍。
  5. 扩展:在预填充与解码分离的服务架构下,OasisKV相比传输全部KV的方式,每个请求所需传输的KV量减少6.5至9.7倍,吞吐量达到约2倍,解码节点主机内存占用减少2.2至2.6倍。
Figure 1. KV-cache memory roofline. Decode throughput is capped by Token-KV intensity × KV-cache bandwidth (three-tier slope-1 roofs) and by the dense/sparse throughput ceilings. Higher Token-KV intensity means more KV reuse (e.g., GQA (1), MLA (5)) or more KV sparsity (e.g., NSA (42), DSA (4)). It is defined as the tokens generated per decode pass per KV-cache byte transferred over HBM or off-GPU IOs. Two throughput ceilings are determined by the maximum batch size and minimum time per output token (TPOT) bounded by compute and weight-loading time.
Figure 1. KV-cache memory roofline. Decode throughput is capped by Token-KV intensity × KV-cache bandwidth (three-tier slope-1 roofs) and by the dense/sparse throughput ceilings. Higher Token-KV intensity means more KV reuse (e.g., GQA (1), MLA (5)) or more KV sparsity (e.g., NSA (42), DSA (4)). It is defined as the tokens generated per decode pass per KV-cache byte transferred over HBM or off-GPU IOs. Two throughput ceilings are determined by the maximum batch size and minimum time per output token (TPOT) bounded by compute and weight-loading time.
Figure 2. The per-output-token decode latency (TPOT) breakdown of on-demand KV retrieval over PCIe (Qwen3-32B, one H100 SXM-HBM3 GPU, context length 32K, BF16). All latencies are computed from a roofline model, assuming 10% of the attended KV is fetched from CPU DRAM each step. The percentage above each bar denotes the KV retrieval overhead.
Figure 2. The per-output-token decode latency (TPOT) breakdown of on-demand KV retrieval over PCIe (Qwen3-32B, one H100 SXM-HBM3 GPU, context length 32K, BF16). All latencies are computed from a roofline model, assuming 10% of the attended KV is fetched from CPU DRAM each step. The percentage above each bar denotes the KV retrieval overhead.
Table 1. Accuracy under the same 2,048-token KV budget. Each retrieval method is read against the full-attention anchor of its own stack: Δ is the difference from that stack’s Full column (Quest, FreeKV vs. Full (HF); Ours vs. Full (vLLM)), never across stacks.
HuggingFace Transformers stackvLLM stack
Dataset / SubsetMetricFullQuestΔFreeKVΔFullOursΔ
Long input — Llama-3.1-8B-Instruct, LongBench v2
Overall29.6229.42−0.2029.03−0.5930.2329.62−0.61
Short34.4434.440.0035.00+0.5635.5633.89−1.67
Medium28.3728.84+0.4726.51−1.8627.4427.440.00
Long24.0722.22−1.8524.070.0026.8526.850.00
Long input — Qwen3-8B, LongBench v2
Overall32.2131.61−0.6031.01−1.2033.6033.20−0.40
Short36.6737.22+0.5536.670.0039.4438.33−1.11
Medium29.3026.98−2.3227.44−1.8629.3030.70+1.40
Long30.5631.48+0.9228.70−1.8632.4129.63−2.78
Long output — Qwen3-8B, reasoning
Overallpass@k81.4178.86−2.5677.91−3.5078.8478.18−0.66
avg@k69.4866.65−2.8366.85−2.6367.6367.28−0.35
AIME24pass@886.6786.670.0080.00−6.6783.3383.330.00
avg@877.5075.42−2.0872.08−5.4279.1776.67−2.50
AIME25pass@883.3376.67−6.6680.00−3.3380.0080.000.00
avg@870.8364.17−6.6668.75−2.0865.4267.97+2.55
GPQA-Diamondpass@474.2473.23−1.0173.74−0.5073.2071.21−1.99
avg@460.1060.36+0.2659.72−0.3858.3057.20−1.10
Figure 3. Decode pipelines of dense attention, sparse attention, KV retrieval, and KV prefetching.
Figure 3. Decode pipelines of dense attention, sparse attention, KV retrieval, and KV prefetching.
Figure 4. Top-20 KV-block prediction accuracy across model layers using the previous token vs. a lookahead token.
Figure 4. Top-20 KV-block prediction accuracy across model layers using the previous token vs. a lookahead token.
Table 2. Fetch-cap ablation on Qwen3-8B (AIME24, LRU eviction). The cap bounds the blocks fetched per step by a fetch ratio. Bold marks our default operating point. Dense attention scores 76.04 / 90.00.
Fetch RatioFetchBWTPSAIME24
(GB/step)(GB/s)(tok/s)avg@32pass@32
0.010.305.02,17874.9090.00
0.020.609.82,06675.1090.00
0.051.4923.82,08375.9486.67
0.102.8731.41,42176.7786.67
0.204.3434.01,03577.4093.33
Fetch all5.0533.582476.4686.67
Figure 5. Overview of the OasisKV architecture.
Figure 5. Overview of the OasisKV architecture.
Figure 6. Per-layer agreement between the top-K set predicted by the propagated draft query and the exact set of the true next-token query. The profile uses Qwen3-8B on GSM8K with Tengyunw/qwen3_8b_eagle3 as the EAGLE-3 draft model.
Figure 6. Per-layer agreement between the top-K set predicted by the propagated draft query and the exact set of the true next-token query. The profile uses Qwen3-8B on GSM8K with Tengyunw/qwen3_8b_eagle3 as the EAGLE-3 draft model.

研究结果

  • 在2048词元的KV预算下,精度与全量注意力相比差距保持在0.7分以内。
  • 在一个推理类工作负载上,吞吐量比原生vLLM提升1.69倍,精度损失为0.1分。
  • 在多GPU长上下文服务场景下,吞吐量最高提升2.1倍。
  • 在预填充解码分离服务中,相比全量KV传输,每个请求所需KV减少6.5至9.7倍,吞吐量达到约2倍,解码节点主机内存占用减少2.2至2.6倍。
  • 在Qwen3-8B上,将每步拉取上限设为0.05时,精度与稠密方式相比差距仅0.1分(75.94对76.04),吞吐量达到全量拉取方式的2.5倍(2083 tok/s)。
Figure 7. Look-ahead attention. Left: the CPU full KV cache maps to two GPU caches — compressed keys via per-block min/max pooling and the sparse KV working set via block-wise sparsification. Right: the attention kernel processes the normal and draft queries together over the sparse KV; the draft query then scans the compressed keys to predict the next step’s top-K blocks, which are prefetched from the CPU cache.
Figure 7. Look-ahead attention. Left: the CPU full KV cache maps to two GPU caches — compressed keys via per-block min/max pooling and the sparse KV working set via block-wise sparsification. Right: the attention kernel processes the normal and draft queries together over the sparse KV; the draft query then scans the compressed keys to predict the next step’s top-K blocks, which are prefetched from the CPU cache.
Figure 8. The asynchronous prefetch pipeline across two decoding steps. Red arrows trace one layer’s chain: the draft query at step t drives top-K prediction, KV selection, and KV transfer before that layer’s attention at step t+1.
Figure 8. The asynchronous prefetch pipeline across two decoding steps. Red arrows trace one layer’s chain: the draft query at step t drives top-K prediction, KV selection, and KV transfer before that layer’s attention at step t+1.

可应用场景

  • 为需要处理长上下文或长推理链的聊天机器人、编程智能体、网页操作智能体等服务降低GPU显存压力,提升并发请求处理能力
  • 提升多GPU部署下长上下文大模型服务的吞吐量
  • 在采用预填充与解码分离架构的大规模服务集群中,降低网络传输量和解码节点的主机内存占用
Figure 9. Head-wise mapping between the bounded GPU working set and the CPU full-KV cache. The original page tables retain their logical-to-physical translations. An additional table maps each GPU logical block to a CPU logical block for every KV head.
Figure 9. Head-wise mapping between the bounded GPU working set and the CPU full-KV cache. The original page tables retain their logical-to-physical translations. An additional table maps each GPU logical block to a CPU logical block for every KV head.
Figure 10. Remote partial fetching: a partial transfer at admission (top) and a network fetch at decode (bottom). Numbered steps are described in Sections 4.4.1 and 4.4.2.
Figure 10. Remote partial fetching: a partial transfer at admission (top) and a network fetch at decode (bottom). Numbered steps are described in Sections 4.4.1 and 4.4.2.

局限与待验证事项

  • 评估仅限于特定模型(Qwen3-8B、Qwen3-32B、Qwen3-235B)和特定基准测试(AIME24/25、GPQA-Diamond、LongBench v2、GSM8K),能否推广到其他模型或任务尚待验证。
  • 原型系统尚不支持前缀缓存,其在启用前缀缓存后对首字延迟(TTFT)的改善效果是基于分析建模估算的,并非实测结果。
  • 已报告的数值是在特定硬件条件(H100 GPU、PCIe带宽等)下测得的,在其他互联方式或内存层级配置下结果可能不同。
  • 预测仍会出现失误,需要额外的补偿性拉取,这会部分抵消解码阶段的网络流量节省效果。
Figure 11. Synthetic decode sweep over max concurrency at 16K and 32K context, on Qwen3-8B (single H100) and Qwen3-235B (TP8, eight H100s). We compare OasisKV against dense attention on unmodified vLLM (labeled vLLM) and three hierarchical-KV baselines: ShadowKV, InfiniGen, and FreeKV. Rows: decode throughput (TPS), running batch, and per-token latency (TPOT); each request generates 2,048 tokens. Arrows mark the OasisKV/vLLM ratio at max concurrency 256.
Figure 11. Synthetic decode sweep over max concurrency at 16K and 32K context, on Qwen3-8B (single H100) and Qwen3-235B (TP8, eight H100s). We compare OasisKV against dense attention on unmodified vLLM (labeled vLLM) and three hierarchical-KV baselines: ShadowKV, InfiniGen, and FreeKV. Rows: decode throughput (TPS), running batch, and per-token latency (TPOT); each request generates 2,048 tokens. Arrows mark the OasisKV/vLLM ratio at max concurrency 256.
Figure 12. Disaggregated serving over the offered request rate, Qwen3-8B, at 24K (left) and 32K (right) context with 2,048 output tokens. Top: decode throughput. Bottom: average and peak host memory usage in the decode node.
Figure 12. Disaggregated serving over the offered request rate, Qwen3-8B, at 24K (left) and 32K (right) context with 2,048 output tokens. Top: decode throughput. Bottom: average and peak host memory usage in the decode node.

为什么重要

GPU高速显存稀缺且昂贵,长上下文负载会迅速耗尽这部分资源,严重限制服务能同时处理的请求数量;这项工作展示了一种在几乎不损失精度的前提下,用同样硬件服务更多请求的路径。这对多GPU部署和预填充解码分离的大规模服务架构尤其重要,因为它直接关系到显存与内存成本能否被有效控制。

本文术语

  • KV缓存 · 语言模型保存之前处理过的词元信息的缓存,每生成新词元时都需要参考它
  • HBM · GPU上速度很快但容量有限的高带宽显存
  • 投机解码 · 用一个小的草稿模型提前猜测未来几个词元,以加快生成速度的技术
  • 预填充解码分离 · 把处理提示词的阶段(预填充)和逐词生成的阶段(解码)分别放在不同服务器上运行的服务方式
  • 限额淘汰策略 · 限制每个解码步骤能新拉取多少KV块,以保证传输时间不超过可用的计算时间窗口

论文原文摘要（英文）

Large language model (LLM) inference serving is increasingly constrained by memory rather than compute. As long-context and long-form reasoning workloads become more prevalent, the key-value (KV) cache dominates both memory footprint and memory traffic during LLM token generation, i.e., decode. In particular, HBM capacity has become a scarce and costly resource that heavily limits inference batch size and system throughput. This paper presents OasisKV, a memory-centric LLM inference system design that alleviates HBM capacity pressure by decoupling full KV-cache storage from HBM during LLM decoding. Because decode-time attention is naturally sparse, OasisKV keeps only the KV entries of the most relevant tokens in HBMs for attention computation. We observe that future important tokens can be predicted accurately in advance using lookahead tokens drafted by speculative decoding (SD). OasisKV employs an efficient attention background pipeline to identify important KV blocks. They are then prefetched from higher-capacity memory tiers (e.g., host or remote memory) and staged in HBMs before being used in the next decode step. We implement OasisKV based on vLLM. The lookahead prediction is accurate enough to keep accuracy within 0.7 points of full attention under a 2,048-token KV budget. This lets OasisKV turn sparsity into throughput gain: 1.69times over dense vLLM on the reasoning workload at 0.1 points of accuracy loss, and up to 2.1times on multi-GPU long-context serving. Under prefill--decode disaggregation, OasisKV reaches about 2times dense throughput while admitting each request with 6.5--9.7times less KV and holding 2.2-2.6 less decode-node host memory than full KV transfer.

作者 · Can Xiao

在 arXiv 阅读

最新论文

全部论文 →

METAL LAB 最新报道

图片来源： Can Xiao et al., arXiv:2608.08097, arxiv-nonexclusive