Search papers, labs, and topics across Lattice.
This paper introduces an elastic key-value (KV) cache mechanism for large language model (LLM) serving that dynamically reallocates memory between decode and prefill phases, addressing the inefficiencies of a static KV cache. The authors conducted experiments to evaluate the impact of chunk sizes on prefill latency and found that the performance penalty of smaller chunk sizes was minimal, suggesting that reclaiming the reserve may not yield significant benefits. Ultimately, they provide a reusable userspace elastic-VMM allocator that enhances memory management without requiring modifications to the attention kernel or risking out-of-memory events.
Reclaiming memory from idle KV cache during decode phases yields negligible latency improvements, challenging assumptions about prefill chunk sizes in LLM serving.
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.
