Latency Optimization › KV Cache Management

The key-value (KV) cache is the most consequential data structure in LLM serving. Every transformer layer computes key and value projections for each token during the forward pass. During autoregressive decoding, these values do not change for previously computed tokens. Storing them avoids recomputing the full attention operation from scratch at every step.

Without a KV cache, decoding token $t$ would require processing all $t$ tokens again. With the cache, decoding token $t$ processes only the new token and reads the cached keys and values for positions $1$ through $t-1$. The time savings are dramatic: without caching, decode time would grow as $O(N_{\text{out}}^2)$; with caching, it is $O(N_{\text{out}})$.

To understand why the KV cache is so large and why managing it carefully matters, it helps to trace what is being stored. In each transformer layer, for each token position, the model computes a key vector $k \in \mathbb{R}^{d_{\text{head}}}$ and a value vector $v \in \mathbb{R}^{d_{\text{head}}}$ for each attention head. These vectors are the inputs to the attention operation. Once computed, they never change for that position, so storing them avoids recomputing the attention keys and values from the residual stream at every subsequent decode step. The tradeoff is memory: every token in the context requires storing $2 \times L \times H$ vectors of dimension $d_{\text{head}}$, and this memory grows linearly with context length.

The KV cache size for a single sequence is:

$$\text{size}_{\text{KV}} = 2 \times L \times H \times d_{\text{head}} \times N_{\text{ctx}} \times \text{bytes\_per\_element}$$

where:

• $L$: number of transformer layers
• $H$: number of attention heads per layer
• $d_{\text{head}}$: dimension of each attention head
• $N_{\text{ctx}}$: sequence length (prompt plus all generated tokens so far)
• The factor of 2 accounts for both keys and values

For a 70-billion parameter model with 80 layers, 64 heads, and 128 dimensions per head, storing the KV cache for a single 4,096-token context at float16 requires:

$$2 \times 80 \times 64 \times 128 \times 4096 \times 2 \approx 10.7 \text{ GB}$$

This is a substantial fraction of a GPU's HBM. Managing KV cache memory is one of the primary constraints on serving throughput and the primary reason batching cannot be increased without limit.

KV cache memory scales linearly with sequence length, number of layers, and number of heads. A 70B model serving a single request with an 8,192-token context uses over 20 GB of KV cache, leaving little room on an 80 GB GPU for model weights and activations, let alone other requests. This arithmetic explains why modern decoder models increasingly use grouped-query attention (GQA) or multi-query attention (MQA), both of which reduce the number of KV heads. By sharing keys and values across multiple query heads, GQA can reduce KV cache size by 4x to 8x without meaningful quality loss. We covered the architectural details of GQA in the chapter on modern decoder models; here the practical consequence is simply that models designed with serving in mind tend to have significantly smaller KV caches than their older counterparts.

Traditional KV cache implementations pre-allocate a contiguous memory block for each request based on the maximum sequence length. This wastes memory because most requests do not fill their allocated space, and it makes it impossible to split a cache across non-contiguous memory regions.

PagedAttention (introduced with the vLLM system) borrows virtual memory concepts from operating systems. The KV cache is divided into fixed-size pages (typically 16 tokens per page), and a page table maps logical page indices to physical memory locations. When a request needs more cache space, it is allocated the next free physical page, which may not be contiguous with its existing pages. Attention kernels are modified to perform the scatter-gather needed to read non-contiguous pages correctly.

The gains are significant. With contiguous allocation, a request for a 1,000-token sequence that uses only 300 tokens wastes 700 tokens of memory. With paged allocation, only populated pages consume memory. System-wide memory utilization improves, allowing more requests to be batched simultaneously. Paging also enables prefix caching (discussed below) because multiple sequences can share the same physical pages for common prefixes without copying.

An additional benefit of paged allocation is that it enables copy-on-write semantics for beam search and parallel sampling. When a single prompt generates multiple completions (as in beam search), the prompt's KV pages can be shared across all beams until the beams diverge. Only the pages containing positions where the beams differ need to be copied. This reduces memory usage for multi-completion workloads substantially compared to the naive approach of allocating a full separate cache per beam.

The tradeoff of paged attention is implementation complexity. Attention kernels must support non-contiguous memory access, which requires custom CUDA kernels rather than standard matrix multiply routines. Libraries like vLLM and TensorRT-LLM have invested significant engineering effort in these kernels. For practitioners using these libraries, paged attention is transparent; for those building custom serving infrastructure, implementing it from scratch is a substantial undertaking.

Many production LLM deployments send the same system prompt or few-shot examples with every request. Without caching, each request must prefill the entire prompt, even when the first 2,000 tokens are identical to every previous request.

Prefix caching stores the KV states for common prompt prefixes in a shared cache. When a new request arrives with a prefix matching a cached entry, the serving system skips prefill for those tokens and loads their KV states directly from cache. This turns an expensive compute operation into a memory read, reducing TTFT proportionally to the fraction of the prompt that hits the cache.

The effectiveness of prefix caching depends on cache hit rate, which in turn depends on how repetitive the prompt patterns are. Retrieval-augmented generation systems, where document chunks are retrieved and inserted into prompts, can achieve high hit rates for repeatedly retrieved chunks. Multi-turn conversation systems benefit when system prompts are cached across turns.

Prefix caching works well when the shared prefix appears at the very beginning of every prompt, because token positions in a causal attention model must be consistent with the original sequence. A cached prefix for positions 1 through $K$ can only be reused if the new request's first $K$ tokens are byte-for-byte identical to those in the cache entry. Even a single token difference invalidates the match. This constraint means that serving systems typically hash the prefix tokens to identify matching entries, and any mutation of the system prompt (even adding a timestamp) defeats the cache entirely. Prompt engineering for cached deployments must be disciplined about keeping the shared prefix stable.

Beyond paging and prefix sharing, another approach to extending KV cache capacity is quantizing the cached tensors to lower precision. Standard KV caches store keys and values as float16 (2 bytes per element). Reducing to int8 (1 byte) halves the memory footprint; int4 reduces it to a quarter.

The accuracy impact of KV cache quantization is generally mild for int8 and acceptable for int4 in many settings, though it grows with context length. The keys and values that most affect current attention scores are those for recent tokens; distant token values can tolerate more quantization error before outputs are noticeably affected. This observation motivates mixed-precision KV caching, where recent tokens are stored at higher precision and older tokens at lower precision, though production implementations of this are still rare.