Source-linked AI summary
Elastic KV Cache for LLM Serving:A Working Reclamation Mechanism, and Why Chunked Prefill Already Closes the Gap
Sathishkumar Sivashanmugam
TL;DR
The paper asks whether idle prefill-activation memory can be reclaimed for the KV cache and builds a userspace elastic mechanism that lends and returns this reserve. The mechanism works, but the opportunity does not on tested workloads: small prefill chunks add little latency, so lowering max_num_batched_tokens recovers more KV at nearly equal latency.
Problem
An LLM serving engine permanently reserves memory for worst-case prefill activation, leaving that capacity idle during decode-dominant phases and motivating whether it can be reclaimed.
Method
The paper builds a userspace elastic KV mechanism that grows into the prefill reserve during decode and shrinks before prefill without changing the attention kernel or driver.
Results
The mechanism works and dynamic toggling is necessary, but a small prefill chunk raises median time-to-first-token by only about 1%, so lowering max_num_batched_tokens recovers more KV at nearly equal latency.
Takeaways & Limitations
On the tested workloads, vLLM’s chunked prefill already dissolves the reserve tradeoff, making lower max_num_batched_tokens a better lever than an elastic cache.
Takeaways & Limitations
The paper did not find a configuration where the elastic cache was the right tool; a useful regime would require small chunks to throttle prefill throughput while KV remains scarce.
Abstract
from arXiv · showhide
An LLM serving engine sizes its key-value (KV) cache once, at startup, permanently setting aside a reserve for the worst-case prefill activation. During decode-dominant phases that reserve sits idle, yet it cannot be handed to the KV pool because it is exactly the memory a large prefill needs. We ask whether this reserve is reclaimable, and build a mechanism to test it. Our elastic KV cache lends the reserve to the KV pool during decode and returns it before prefill, driven by the scheduler's one-step-ahead view of the next batch. It is pure userspace on the CUDA virtual-memory path: two physical handles mapped into one contiguous virtual range per layer, so the attention kernel is unchanged and no driver patch is required. It decommits in a few milliseconds and recommits in tens of milliseconds, works with CUDA graphs and prefix caching, and never triggers an out-of-memory event. A static commit of the same memory is unsafe, crashing on prefill bursts, which makes the dynamic toggle necessary. Having built the mechanism, we test the premise it rests on and report an honest negative result. It only pays off if a small prefill chunk size badly hurts prefill latency. In a controlled experiment injecting long prompts into a live decode load, that penalty is small (median time-to-first-token differs by about 1% between chunk sizes of 8192 and 32768 tokens), because prefill is compute bound and decode consumes only about one token per sequence per step. Simply lowering max_num_batched_tokens recovers more KV than the controller does, at nearly equal latency. The reserve also dilutes under tensor parallelism, from 16% of KV at TP1 to 2.7% at TP4. We state precisely when reclaiming the reserve could still help, and release the mechanism as a reusable userspace elastic-VMM allocator.
1 Introduction
LLM engines reserve memory for worst-case prefill activations, reducing KV capacity even though the reserve is idle during decode. The paper builds an elastic mechanism to reclaim it, then finds that chunked prefill largely removes the latency tradeoff motivating reclamation.
- The KV cache bounds concurrent requests and maximum context length, making its capacity a scarce serving resource.
- Worst-case prefill profiling permanently reserves activation memory, and larger prefill chunks further reduce KV capacity.
- The paper lends the idle reserve to KV during decode and returns it before large prefill using a userspace elastic mechanism.
- The mechanism preserves the attention kernel and driver while dynamically toggling the reserve in milliseconds.
- Static reclamation is unsafe on prefill bursts, whereas dynamic toggling is necessary to capture the reserve safely.
- The reserve dilutes from 16% to as little as 2.7% of KV under tensor parallelism.
2 Background and Motivation
vLLM allocates a fixed paged KV block pool from memory left after weights, CUDA graphs, and activation reserves. The prefill reserve grows with chunk size, while the utilization knob targets a separate headroom pool.
- vLLM profiles the largest allowed batch at startup and converts remaining memory into fixed-size KV blocks managed by an integer free list.
- KV capacity shrinks as prefill chunks grow because more memory is reserved for the activation peak.
- Increasing max_num_batched_tokens from 2048 to 32768 reduces KV capacity from about 366K to about 308K tokens.
- At a 32768-token chunk, the 3.1 GiB activation reserve is about 16% of KV capacity and remains free during decode.
- Raising gpu_memory_utilization captures stranded device headroom, not the activation reserve provisioned inside the utilization budget.
3 An Elastic KV Mechanism
The elastic KV mechanism keeps each layer’s virtual address range contiguous while changing its physical backing. A scheduler-controlled CUDA VMM toggle commits the reserve during decode and decommits it before prefill with millisecond-scale overhead.
- The design grows and shrinks KV without changing the attention kernel by preserving one contiguous virtual range per layer.
- Two physical handles back one reserved virtual range: the base remains mapped while an elastic slice is mapped or unmapped at the range’s top.
- The scheduler commits the elastic slice during decode-only phases and decommits it before a predicted prefill, using hysteresis to avoid thrashing.
- About 30 ms to commit and about 3 ms to decommit the 3 GiB reserve across 28 layers is faster than the naive whole-handle commit at about 180 ms.
4 Static Reclamation is Unsafe
Static reclamation recovers the reserve only if it can be returned before a large prefill; otherwise the borrowed memory causes an out-of-memory failure.
- The mechanism is necessary because the full reserve is reachable only through dynamic reclamation before prefill.A static commit cannot preserve enough activation memory for the largest scheduled prefill.
- A static 3 GiB expansion recovered the full KV capacity gap but OOMed on the first large prefill.The capacity increased from about 314K to about 370K tokens before the failure.
- The OOM-safe static ceiling was only about 1.2 GiB, a 6% gain, because larger commitments starved the largest prefill.A 62K-token burst required about 2.3 GiB of activation memory.
- The dynamic toggle safely decommitted the elastic slice before the same 62K-token burst, then recommitted it afterward.Decommitment completed in 4.4 ms, allowing the prefill to run cleanly.
- A caching-allocator bug caused the first recommit after a large prefill to OOM until empty_cache() was called before cuMemCreate.
5 The Decisive Test: Does a Small Prefill Chunk Actually Hurt?
The decisive test finds that small prefill chunks barely increase latency under live decode load, so lowering the chunk size recovers more KV than elastic reclamation at nearly equal latency.
- Mode A, using the small chunk, had the most KV while keeping TTFT within about 1% of the other modes.Table 1 identifies Mode A as the winner on both TTFT and KV capacity.
- The experiment compared 8192-token chunks, 32768-token chunks, and an elastic controller at 32768 tokens with the reserve lent during decode.Prefix caching, CUDA graphs, and vLLM’s asynchronous engine were enabled.
- About 1% median TTFT separated 8192- and 32768-token chunks, far below the expected roughly fourfold difference.The experiment injected six approximately 25K-token prompts into a live load of 40 background decode sequences.
- The small-chunk penalty is small because prefill is compute bound and chunking changes granularity rather than total FLOPs.
- Decode contributes about 280 tokens against an 8192-token chunk budget at the 280-sequence maximum, or about 3%.This prevents decode from meaningfully starving the prefill budget in the tested regime.
- Across harsher regimes, the authors found no case where the elastic controller beat simply lowering the chunk size.
6 Scope: When Could This Help?
The elastic cache could matter when the reserve is large and KV is scarce, but the tested workloads did not show a regime where it beats lowering the chunk size. Tensor parallelism and chunked prefill both narrow its opportunity.
- 95–99.8% decode-only wall time occurs in most tested regimes, so phase availability does not explain the technique’s weak results.Under sustained saturation, decode-only time falls to 24% with chunk 2048 and 49% with chunk 32768.
- 16% of KV at TP1 falls to 2.7% at TP4, because tensor parallelism shards the per-rank activation cost while KV remains large per rank.High tensor parallelism dominates for large models, making their reserve relatively small.
- Mode A, using the small chunk, has the most KV at about the same median TTFT, leaving little for the elastic cache to reclaim.This places the small-chunk baseline ahead on both axes in the KV-versus-latency frontier.
- The reserve is largest when tensor parallelism is low, the model is small to mid-sized, the context is long, and KV capacity is scarce.In this regime, the reserve is roughly 10–16% of KV.
- The elastic controller could help only if a small chunk throttles prefill throughput while KV remains scarce, a regime not found on this hardware.Lowering the chunk size is also cheapest for small models because they prefill quickly.
7 Comparison with Alternatives
Compared with allocator headroom, demand paging, and fixed-region KV managers, the elastic mechanism targets activation-reserve reclamation. However, the tested reserve is not worth reclaiming because chunked prefill already makes small chunks cheap.
- The utilization knob at util=0.97 OOMs on a 48-prompt prefill burst and survives with only 0.04–0.15 GiB of margin when expandable segments are enabled.It captures headroom rather than the activation reserve and is therefore not robust for this workload.
- vAttention preserves contiguity through demand paging within a fixed KV budget, but requires a replaced NVIDIA UVM driver and pinned software versions.It does not make the KV budget elastic against the activation reserve.
- Jenga sizes KV blocks within a fixed region and is orthogonal to reserve reclamation, so the two mechanisms could compose.For a single-group model such as Qwen, Jenga behaves as stock vLLM.
- The tested alternatives do not reclaim the activation reserve, but chunked prefill already makes a small chunk cheap enough that reclamation is not worthwhile.The larger chunk has more reserve and more decode-only availability, so availability is not the limiting comparison.
8 Related Work
The work extends virtual-memory-based KV management with a fine-grained, kernel-transparent elastic budget. Its negative result identifies chunked prefill as the scheduling technique that already resolves the motivating tradeoff.
- The mechanism extends the paged block-pool line with a two-handle-one-VA toggle that grows or shrinks a contiguous GPU tensor without kernel changes.Related systems manage contiguity, allocation placement, or fixed-budget KV sizing rather than changing the budget elastically.
- Chunked prefill interleaves prefill chunks with decode and, according to the paper’s negative result, already resolves the prefill-versus-KV-capacity tension.This is the scheduling finding that limits the practical value of an elastic KV budget in the tested setting.
9 Conclusion
The paper builds and releases a userspace elastic KV mechanism, then finds that its motivating opportunity is largely closed by chunked prefill. The conclusion is based on experiments on one A100-40GB system running Qwen2.5-7B-Instruct with vLLM 0.23.0.
- The userspace mechanism toggles the reserve in milliseconds without modifying the attention kernel or CUDA driver.It uses a torch pluggable allocator, a block-pool gate, and a scheduler-boundary controller.
- A static reserve commit is unsafe, while dynamic toggling captures the reserve without causing OOM on the tested prefill bursts.The dynamic design is therefore necessary for safe reclamation.
- 16% of KV at TP1 declines to 2.7% at TP4, removing the main regime where the reserve is large.The paper reports this tensor-parallelism effect as a reason large models do not improve the opportunity.
- All reported numbers come from a single A100-40GB running Qwen2.5-7B-Instruct on vLLM 0.23.0.The decisive experiments used small standalone harnesses.