Key-Value (KV) Cache

During autoregressive generation, a transformer model processes an input sequence token-by-token. For each new token, it recomputes the self-attention scores against all previous tokens. The KV Cache stores the computed key (K) and value (V) tensors for all previous positions in the sequence. When generating the next token, the model retrieves these cached tensors instead of recomputing them, avoiding redundant matrix multiplications for the historical context. This transforms the computational complexity of generating a sequence of length N from O(N²) to O(N), which is the primary source of its speedup.

The KV Cache's main value is drastically reducing per-token latency during text generation. By eliminating the need to reprocess the entire growing context for each new token, inference becomes significantly faster. For example, generating a 100-token response with a model that has a 32-layer transformer might require ~3,200 fewer layer-wise attention computations for the historical context if caching is used. This directly translates to higher throughput (tokens/second) and lower response latency for end-users, which is a critical metric for production LLM APIs.

The KV Cache introduces a significant, dynamic memory overhead. The cache size is proportional to:

• Batch Size (B): Number of parallel requests.
• Sequence Length (L): Total context length (prompt + generation).
• Number of Layers (N_l): Transformer blocks in the model.
• Hidden Dimension (D) & Attention Heads (H).

For a large model (e.g., Llama 3 70B with 80 layers, 64 heads, 128 head-dim), caching a 4096-token sequence for a batch of 8 can require ~60 GB of VRAM just for the cache. This necessitates advanced memory management techniques like PagedAttention (used in vLLM) or continuous batching to handle memory fragmentation and out-of-memory errors efficiently.

In production servers, the KV Cache is managed in tandem with continuous batching (also called iterative batching). This optimization allows new requests to be added to a running batch as previous requests finish generation. Each request in the batch has its own, logically separate KV Cache. The serving engine must:

• Allocate and deallocate cache memory dynamically per request.
• Maintain correct attention masking so requests within a batch don't attend to each other's caches.
• Handle variable sequence lengths efficiently across the batch. This combination is what enables high GPU utilization and throughput in servers like vLLM and TGI, as the GPU is kept constantly busy processing tokens from multiple requests at different generation stages.

When serving models fine-tuned with Parameter-Efficient Fine-Tuning (PEFT) methods like LoRA or Adapters, the KV Cache interacts uniquely with the inference process. For a merged model (where LoRA weights are fused into the base model), the KV Cache operates normally. However, in multi-adapter serving scenarios, where a single base model dynamically switches between different adapter sets:

• The base model's KV Cache computation remains largely unchanged, as the frozen base weights are used to compute the initial K and V projections.
• The adapter-specific modifications typically affect the query or later projections, but the cached K and V tensors from the base model can still be reused efficiently.
• This allows for efficient adapter switching without invalidating the entire cache, supporting multi-tenant or multi-task serving from a single GPU instance.

Managing the KV Cache involves key operational decisions:

• Cache Eviction: For very long conversations or document processing, the cache may exceed available memory. Systems implement eviction policies (e.g., FIFO, discarding oldest tokens) or context window sliding, which degrades performance as recomputation is needed.
• Precision: Caching in FP16 vs BF16 vs INT8 involves a trade-off between memory savings and potential precision loss affecting output quality.
• Warm-up: A model warm-up phase pre-populates the cache with a dummy forward pass to ensure memory allocators and kernels are primed, preventing high latency for the first real request (cold start).
• Observability: Monitoring cache hit rates, memory usage, and per-request cache allocation is crucial for performance debugging and capacity planning in production.

During autoregressive decoding (generating one token at a time), a transformer's self-attention mechanism recomputes key (K) and value (V) tensors for all previous tokens in the sequence for each new token. The KV Cache stores these computed K and V tensors in a memory buffer. For each subsequent generation step, the model retrieves cached tensors for previous tokens and computes them only for the new token, transforming a quadratic computational complexity into a linear one. This is the primary driver behind the significant speed-up in text generation for models like GPT, Llama, and Claude.

The KV Cache's memory consumption is substantial and grows linearly with:

• Batch Size: Number of sequences processed in parallel.
• Sequence Length: Number of tokens generated so far.
• Model Size: Number of layers and attention heads.

The formula is roughly: Memory ≈ 2 * batch_size * seq_len * n_layers * n_heads * d_head * bytes_per_param. For a large model serving many long conversations, this can consume tens of gigabytes, often exceeding the memory required for the model weights themselves. This creates the central serving challenge: managing this dynamic, growing memory allocation efficiently across many concurrent requests.

PagedAttention is a breakthrough optimization algorithm introduced by the vLLM inference engine. It solves the problem of memory fragmentation caused by variable-length sequences. It draws an analogy from operating system memory management:

• The KV Cache is divided into fixed-size blocks.
• The cache for each sequence is represented as a list of these blocks, which don't need to be contiguous in physical memory.
• This allows for non-contiguous storage, eliminating fragmentation waste.
• It enables highly efficient memory sharing for identical prompts in sampled decoding (e.g., in beam search). This technique is why vLLM achieves near-optimal memory utilization (over 90%) and high throughput.

Continuous Batching (or iterative batching) is an inference scheduling paradigm that optimizes how the KV Cache is utilized across requests. Unlike static batching, it continuously adds new requests to a running batch as other requests finish generation.

• The scheduler must manage the lifecycle of the KV Cache for each request (allocation, growth, and eviction).
• Finished sequences free their cache blocks, making them available for new sequences.
• This leads to significantly higher GPU utilization and throughput, as the GPU is kept constantly busy, and cache memory is recycled efficiently. Systems like TGI (Text Generation Inference) and vLLM implement this.

To reduce the memory pressure of the KV Cache, several quantization and compression techniques are applied:

• KV Cache Quantization: Storing the K and V tensors in a lower precision format (e.g., FP8, INT8) instead of FP16/BF16. This can halve the memory footprint but may require calibration to minimize accuracy loss.
• Selective Caching: Not caching keys/values from certain layers or attention heads based on profiling that shows they are less critical for performance.
• Compressed Caching: Applying lossy compression algorithms to older tokens in very long contexts, as their influence on the next token often diminishes. These techniques are essential for serving long-context models (e.g., 128K+ tokens) cost-effectively.

Serving models fine-tuned with LoRA or Adapters introduces complexity for the KV Cache. The cache is computed from the active model weights. When using multi-adapter serving:

• The KV Cache is adapter-specific. The cached keys and values are a function of the base weights plus the active adapter's weights.
• Simply switching adapters on a shared base model makes the existing cache invalid, as the internal representations change.
• Solutions involve either:
  • Cache recomputation on adapter switch (costly).
  • Adapter-aware caching, where the cache is keyed by a (sequence_id, adapter_id) pair, logically partitioning the memory.
  • Using merged weights for inference, which creates a standalone model with a single, consistent KV Cache.

Also known as iterative batching, this is an advanced inference optimization for autoregressive models like transformers. Unlike static batching, it allows new requests to be added to a running batch as previous requests finish generating tokens. This technique maximizes GPU utilization and throughput by keeping the computational hardware constantly occupied, and it works in tandem with the KV Cache to manage the state of multiple concurrent sequences efficiently.

The core neural network component that the KV Cache optimizes. In the transformer architecture, self-attention computes a weighted sum of values (V) based on the compatibility between a query (Q) and a set of keys (K). For a generated token at position i, the attention scores are computed against all previous tokens 1...i. The KV Cache stores the pre-computed K and V tensors for these previous tokens, eliminating the need to recompute them for every new token, which is the source of its significant speedup.