KV Cache (Key-Value Cache)

The KV Cache stores the computed key (K) and value (V) vectors for all previous tokens in a sequence during autoregressive generation. For each new token generated, the self-attention mechanism only computes the K and V vectors for that new token. It then retrieves the cached K and V vectors for all prior tokens, avoiding the quadratic recomputation of the entire sequence's attention matrix. This transforms the computational complexity of generating a sequence of length n from O(n²) to O(n), leading to drastic latency reductions, often by 10x or more for long sequences.

The primary trade-off of the KV Cache is a linear increase in GPU memory consumption with sequence length. For a model with n layers, h attention heads, and a hidden dimension d, the cache size for a sequence of length L is approximately 2 * n * h * L * (d/h). This memory footprint can become the bottleneck for long-context models, limiting the practical batch size or maximum sequence length that can be processed on a given GPU. Efficient cache eviction policies and quantization of cached tensors are essential for managing this overhead.

Without a KV Cache, processing sequences longer than the training context length is computationally prohibitive. The cache enables techniques like:

• Sliding Window Attention: Only a fixed window of recent tokens is kept active in the cache, allowing the model to process streams longer than its full context.
• Streaming Context: New tokens can be incrementally encoded by appending their K/V vectors to the cache, supporting real-time, continuous interaction.
• Context-Aware Batching: Sequences in a batch can share cached prefixes, optimizing throughput for conversational or multi-turn applications.

For very long interactions or constrained memory, not all cached vectors can be retained. Cache eviction algorithms determine which parts of the history to discard. Common strategies include:

• Least-Recently-Used (LRU): Discards the K/V vectors for the tokens attended to the longest time ago.
• Heuristic-based Eviction: Prioritizes keeping cached vectors for system prompts, recent user turns, or summarization anchors.
• Selective Caching: Only caches K/V vectors for certain layers or attention heads to reduce memory footprint. These policies directly influence a model's effective context freshness and long-range reasoning ability.

High-performance inference servers like vLLM, TGI (Text Generation Inference), and TensorRT-LLM implement sophisticated KV Cache management. Key features include:

• PagedAttention: Analogous to virtual memory in operating systems, it allows non-contiguous storage of cached K/V blocks, drastically reducing memory fragmentation and waste from padding.
• Continuous Batching: Dynamically groups incoming requests of varying lengths, sharing GPU memory for the cache across a batch to maximize hardware utilization.
• Quantized Caches: Storing K/V vectors in FP8 or INT4 precision to halve or quarter memory use with minimal accuracy loss.

FlashAttention is a complementary optimization that speeds up the core attention computation itself by reducing memory reads/writes between GPU SRAM and HBM. When combined with a KV Cache, the system benefits from two orthogonal optimizations:

1. FlashAttention: Optimizes the computation of attention for the current forward pass.
2. KV Cache: Eliminates redundant computation across forward passes during generation. Together, they form the foundation for state-of-the-art inference latency, enabling the deployment of large models in production environments with strict throughput and cost requirements.

The KV cache is implemented within the self-attention mechanism of decoder-only transformer blocks. During autoregressive generation, for each new token, the model computes:

• Query vectors for the new token.
• Key and Value vectors for the new token, which are then appended to the cache.

The attention scores are computed using the new Query against all cached Keys from previous steps. The output is a weighted sum of the cached Values. This avoids the O(n²) recomputation of K, V for the entire sequence at each step, reducing complexity to O(n) for the attention calculation per token.

Frameworks like PyTorch manage this via stateful past_key_values tensors passed between generation steps.

The KV cache grows linearly with sequence length, consuming significant GPU memory. Production systems implement eviction policies:

• Static Allocation: Pre-allocate a fixed cache size (e.g., 2048 slots). Generation stops or triggers a full recompute if exceeded.
• Dynamic Eviction: Algorithms like Least-Recently-Used (LRU) identify and discard the K,V pairs for tokens deemed least important to current generation.
• Windowed Attention: A hybrid approach where the cache maintains only a sliding window of the most recent tokens (e.g., last 4096). This enforces a hard limit but may lose long-range context.

Memory Footprint: For a model with n_layers layers, n_heads attention heads, and d_head dimension, caching a sequence of length L requires ~ 2 * n_layers * n_heads * d_head * L * dtype_size bytes.

In high-throughput inference servers, continuous batching (also known as iteration-level scheduling) is essential. Unlike static batching, it:

• Dynamically adds new requests to a running batch as others finish.
• Requires per-request KV cache management.

When a request finishes generation, its allocated cache blocks are instantly freed for new requests. Systems like TGI (Text Generation Inference) and vLLM coordinate this by maintaining a scheduler that:

1. Tracks the live sequences and their cache blocks.
2. Packs the active queries for the next forward pass.
3. Reclaims memory without stopping the server, achieving high throughput and low latency for variable-length requests.

To reduce the memory footprint of the KV cache, quantization is applied:

• KV Cache Quantization: Storing the Key and Value tensors in a lower precision format (e.g., FP8 or INT8) instead of FP16/BF16.
• This can reduce cache memory by 50-75% with minimal accuracy loss.
• Implementations often use per-tensor or per-channel quantization scales.

Trade-off: Dequantization adds a small computational overhead during attention score computation. Frameworks like AWQ (Activation-aware Weight Quantization) and GPTQ offer integrated pathways for quantizing both weights and the KV cache. The choice of quantization is crucial for serving large models on memory-constrained hardware.

Major frameworks expose the KV cache through specific APIs:

• PyTorch / Transformers: The past_key_values tuple in Hugging Face's model.generate() and lower-level model() calls. The use_cache=True argument enables it.
• OpenAI API: Managed transparently for the user via the stream parameter and stateful chat completions.
• NVIDIA TensorRT-LLM: Provides highly optimized, static graph compilation with pre-allocated KV cache buffers for maximum performance on NVIDIA GPUs.
• JAX/Flax: Uses explicit stateful models where the cache is part of the mutable model state, passed and updated in each decoding step.

Developer Control: Advanced users can manually manipulate the cache for techniques like context shifting or speculative decoding, but standard inference is fully abstracted.

The attention mechanism is the core computation in a transformer model that determines the relevance or "attention" each token should pay to all other tokens in a sequence. It calculates a weighted sum of value vectors, where the weights are derived from the compatibility between a token's query vector and the key vectors of all other tokens. The KV Cache stores these pre-computed key and value vectors to avoid recalculating them for every new token during autoregressive generation.

A context window is the fixed-size, contiguous block of tokens that a transformer-based language model can process in a single forward pass. It represents the model's working memory for a given input and output sequence. The KV Cache is a critical optimization within this window; it stores the intermediate states for tokens already processed, allowing the model to generate new tokens without reprocessing the entire context from scratch, thus staying within the computational bounds of the window.

Autoregressive generation is the sequential process where a language model produces output one token at a time, with each new token being conditioned on all previously generated tokens (and the original input). This is the primary inference mode for text generation. The KV Cache provides its major speed-up here: instead of recomputing attention over the entire growing sequence for each step, the model simply appends the new token's key-value pairs to the cache and performs attention with the cached history.

Incremental decoding is an inference strategy synonymous with autoregressive generation that leverages the KV Cache. At each decoding step, only the activations for the newest token are computed, while the key-value pairs for all previous tokens are fetched from the cache. This reduces the computational complexity of generating a sequence of length N from O(N²) to O(N) for the attention layers, leading to drastic reductions in latency, especially for long outputs.

The transformer architecture is a neural network design based entirely on attention mechanisms, dispensing with recurrence and convolutions. It consists of an encoder and/or decoder stack of layers, each containing multi-head self-attention and feed-forward networks. The KV Cache is an optimization applied specifically to the decoder (or decoder-only) transformers during inference. It exploits the fact that the self-attention computation for a given layer and head can be decomposed into static key and value projections that are independent of the query token being generated.

Cache eviction refers to the algorithms used to manage the KV Cache when generating sequences longer than the available GPU memory can hold. Common policies include:

• Least-Recusently Used (LRU): Discards the key-value pairs for tokens that have not been attended to recently.
• FIFO (First-In, First-Out): Evicts the oldest tokens in the cache.
• Strategic Eviction: Removes tokens from less critical middle segments of long contexts. These policies are essential for enabling infinite-length generation or working with extremely long context windows without running out of memory.