PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.08097v1cs.DC

TL;DR

Long-context and reasoning workloads make KV-cache capacity and movement a major limit on LLM decode serving. OasisKV stores the full cache off-GPU while using speculative-decoding lookahead to prefetch the sparse KV blocks needed next. It stays within 0.7 accuracy points of full attention while improving throughput across reasoning, long-context, and disaggregated settings.

  • Problem

    Long-context decode serving is constrained by HBM capacity and bandwidth, while full KV transfer also burdens prefill-decode disaggregation.

  • Method

    OasisKV keeps the full KV cache in host or remote memory and uses speculative-decoding lookahead with an asynchronous pipeline to prefetch needed KV blocks into HBM.

  • Results

    Within 0.7 accuracy points of full attention, OasisKV improves decode throughput by 1.69× on reasoning workloads, up to 2.1× on multi-GPU long-context serving, and 2.1–2.3× under PD disaggregation.

  • Takeaways & Limitations

    Sparse KV staging can expand effective in-decode memory capacity while improving throughput and reducing KV admission and decode-node host-memory requirements under disaggregation.

  • Takeaways & Limitations

    On Qwen3-235B, sparsity saves less GPU memory while per-step overhead remains, making OasisKV slower than dense serving at low batch sizes.

Abstract

from arXiv · show

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.69\times$ over dense vLLM on the reasoning workload at 0.1 points of accuracy loss, and up to $2.1\times$ on multi-GPU long-context serving. Under prefill--decode disaggregation, OasisKV reaches about $2\times$ dense throughput while admitting each request with $6.5$--$9.7\times$ less KV and holding $2.2$-$2.6$ less decode-node host memory than full KV transfer.

1 Introduction

Long-context agentic workloads make HBM capacity and bandwidth a decode bottleneck. OasisKV addresses this memory wall by predicting and prefetching only the KV blocks needed for future sparse attention.

  • Motivation: Over 10× longer contexts and decode-dominated latency make limited HBM capacity and bandwidth a central throughput bottleneck.Long-context agentic workloads require substantially more context than chatbot-era queries, while decoding remains the dominant latency component.
  • Prior approaches: Sparse attention reduces KV reads but retains the full KV cache in HBM, whereas retrieval reduces resident capacity but can place transfer latency on the decode critical path.On-demand retrieval becomes especially costly as batch size and active KV-cache size increase.
  • OasisKV: OasisKV uses speculative-decoding lookahead tokens to predict future important KV blocks and asynchronously prefetches them from off-GPU memory into HBM.The design combines lookahead-driven sparsity prediction, off-GPU prefetch overlap, and sparse staging rather than transferring the full cache.
  • Contributions: The paper combines sparse attention, KV-cache prefetching, and production-engine integration to expand effective in-decode memory capacity.Its contributions include a vLLM prototype supporting multi-GPU deployment and prefill-decode disaggregation.

2 Background

LLM decoding must retain growing KV caches, making long-context serving fundamentally capacity-constrained. Existing sparse and hierarchical approaches reduce computation or resident GPU data but leave important latency and capacity trade-offs.

  • KV-cache pressure: Each decoding query attends to prior-token KV states, so the cache must remain available throughout generation.The KV footprint grows with batch size, context length, layers, KV heads, head dimension, and bytes per element.
  • KV-cache pressure: The KV footprint grows linearly with context length, making long-context decoding fundamentally capacity-constrained.For a representative 32B-class GQA model, one KV token occupies 256 KiB under FP16/BF16 storage.
  • KV-cache pressure: 8.6 GB of KV cache per request at 32.7K average context limits an optimistic 80 GB KV-only GPU to nine requests.Real batch capacity is lower because HBM also stores weights, activations, CUDA graphs, and runtime workspace.
  • Existing approaches: Production systems such as vLLM and SGLang support KV offloading or hierarchical caching, especially for reused prefixes in multi-round and agentic workloads.These mechanisms expand the effective KV working set without eliminating the underlying movement and capacity trade-offs.
  • Existing approaches: Sparse attention selects important historical tokens to reduce attention computation, but conventional designs still retain the full KV cache in GPU memory.Hierarchical systems offload the full cache and retrieve selected blocks, placing transfer on the decode critical path.

3 Motivation

Production sparse KV serving must predict next-step blocks accurately, hide retrieval within decode latency, and avoid exhausting bandwidth, HBM, or decode-node memory. OasisKV addresses these constraints with lookahead prediction and sparse staging across memory tiers.

  • Accurate prediction is pivotal because missed blocks either lose attention context or require critical-path corrective fetches.Prefetching only hides latency when predicted blocks match the next step’s attended blocks.
  • Existing predictors trade portability for accuracy: dedicated modules predict well but require training, while previous-token proxies are unreliable across layers.OasisKV therefore targets a training-free, low-cost signal available a full step ahead.
  • 118 tokens per request per step is the practical PCIe budget for newly active KV on Qwen3-8B with a 2K active context.The estimate assumes ≈17 ms decode steps and a ≈64 GB/s PCIe link.
  • Sparse serving must jointly optimize prediction and placement because fewer fetched blocks protect bandwidth but can reduce attention context, while larger resident sets consume HBM.The supported operating point depends on future access patterns and resident-block choices.
  • Full KV staging in disaggregated serving increases admission latency and decode-node DRAM pressure, requiring sparsity in both network transfers and materialized host KV.For Qwen3-235B-A22B at 100k average context, full staging yields a theoretical ceiling of roughly 26 requests on a 1 TB decode node before other allocations.

4 System Design

OasisKV keeps a sparse KV working set in GPU HBM while storing full KV off-GPU, using draft-query lookahead and asynchronous per-layer prefetching to stage future blocks. Its prediction agrees with the true next-token top-K set above 98.2% in every layer and 98.74% on average.

  • OasisKV stores full KV in local CPU DRAM or remote memory while retaining only a sparse working set, draft state, and compressed summaries in HBM.This frees HBM capacity for larger batches while preserving access to the full cache off-GPU.
  • Look-ahead attention computes normal and draft queries together over resident KV, then scans compressed per-block key summaries to rank and prefetch future blocks.The prototype uses coordinate-wise minimum and maximum key summaries, requiring two vectors per block instead of the full block keys.
  • 98.74% average agreement is achieved between propagated-draft top-K predictions and the true next-token set, exceeding 98.2% in every layer.The profile uses Qwen3-8B on GSM8K with an EAGLE-3 draft model.
  • The per-layer background chain is top-K prediction, KV selection, and KV transfer, with independent workers overlapped across layers and decode steps.The foreground waits only for the relevant layer’s transfer before its next-step attention.
  • At most C blocks per KV head are transferred per layer, bounding off-GPU traffic regardless of how many predicted top-K positions change.Selection preserves resident predicted blocks and pairs admitted nonresident blocks with least-recently-selected eviction targets.
  • The prototype retains each request’s full KV cache in prefill-node host DRAM until completion, while the decode side keeps only a demand-filled subset.The same approach is described as extensible to remote memory servers or SSD storage.

5 Evaluation

OasisKV improves decode throughput across single-GPU, multi-GPU, reasoning, and disaggregated serving by staging sparse KV blocks while managing transfer overhead. Its gains depend on bounded KV working sets and fetch caps, with accuracy remaining close to dense attention.

  • Single- and Multi-GPU serving: 2.1× throughput on Qwen3-8B at 16K context reaches 1,398 versus 676 tok/s at maximum concurrency 128.At concurrency 16, OasisKV improves both throughput and TPOT: 836 versus 649 tok/s and 17.7 versus 23.5 ms.
  • PD Disaggregation: 2.1–2.3× throughput in PD-disaggregated serving reaches 1,204–1,210 tok/s at 24K and 884–888 tok/s at 32K, while dense saturates early.Dense throughput is limited by HBM capacity, whereas OasisKV configurations continue scaling with offered request rate.
  • PD Disaggregation: 2.2× and 2.6× lower per-request decode-node host occupancy result from remote partial fetching at 24K and 32K.Aggregate occupancy also falls from 161 to 76 GiB at 24K and from 126 to 46 GiB at 32K.
  • End-to-end performance: 1.69× decode throughput on AIME24 reasoning reaches 2,083 versus 1,235 tok/s for dense vLLM, with only a 0.1-point accuracy loss.The run uses a 0.05 per-step fetch cap and FlashAttention-3 as the dense baseline.
  • Eviction Strategy: 2.6× throughput variation accompanies only a 74.9-to-77.4 accuracy change across the fetch-cap sweep, making 0.05 the best setting.At 0.05, accuracy is within 0.1 point of dense and throughput is 2.5× fetch-all throughput.
  • Remote Partial Fetching: 6.5× and 9.7× reductions in admission traffic at 24K and 32K context fall to 1.33× and 1.50× end-to-end savings after decode-time misses.Fetch caps reduce drift and can yield 1.9× and 2.2× total traffic reductions at 32K.

6 Related Work

Existing KV-retrieval and KV-prefetch systems move or predict nonresident KV blocks, but lack a production-grade cross-tier manager coordinating placement and transfer across memory tiers.

  • KV retrieval: KV-retrieval systems stage query-dependent active KV subsets on GPUs while retaining access to the full context.Nonresident KV may reside in CPU memory, with some systems using low-rank keys, offloaded values, or SSD storage.
  • KV prefetch: KV-prefetch systems identify future active blocks early enough to overlap transfers with ongoing computation.Existing predictors may require model-specific training or remain tied to native sparse-attention architectures.
  • System gap: Existing designs lack a production-grade cross-tier memory manager coordinating sparse KV placement and transfer across GPU HBM, host DRAM, and remote memory.This limits their applicability to high-throughput disaggregated serving.

7 Conclusion

OasisKV keeps full KV caches off GPU memory and stages only needed blocks in HBM, using speculative-decoding lookahead and asynchronous prefetching. It preserves near-full-attention accuracy while improving throughput and reducing disaggregated-serving memory demands.

  • System design: OasisKV keeps the full KV cache off GPU and stages only the blocks needed by each decode step into HBM.This design decouples full KV storage from GPU memory during decoding.
  • Lookahead prefetching: Speculative-decoding draft tokens predict future KV blocks one step ahead, while an asynchronous pipeline hides transfer behind the forward pass.The approach is compatible with speculative decoding and supports multi-tier memory access.
  • Conclusion: OasisKV preserves accuracy close to full attention while substantially improving decode throughput across reasoning and large-batch long-context workloads.Under PD disaggregation, it also reduces per-request KV admission and decode-node host-memory requirements.
Loading 2608.08097v1…