KV Cache

During the decode phase, the primary computational cost shifts from matrix multiplication to loading the KV cache from GPU memory. The system enters a memory-bound regime, where performance is limited by memory bandwidth, not arithmetic throughput. This makes cache size and access patterns critical for latency.

• Impact: Reading cached keys/values dominates runtime.
• Optimization Goal: Minimize cache footprint and maximize data locality.

The KV cache stores the computational state of the transformer's attention mechanism for all previously generated tokens. Each new decoding step appends a single new key-value pair per layer and head while reading the entire cached history. This stateful design is what enables efficient autoregressive generation but also creates linear memory growth with sequence length.

• Core Function: Enables O(1) computation for attention over history.
• Consequence: Cache size scales as O(batch_size * num_layers * num_heads * seq_len * head_dim).

Transformer inference has two distinct phases governed by KV cache usage:

• Prefill Phase: The initial, compute-intensive parallel processing of the input prompt. All prompt tokens are processed simultaneously, and their key/value tensors are computed and stored, populating the initial cache.
• Decode Phase: The iterative, memory-intensive token generation. For each new token, the model reads the entire existing cache and appends one new KV pair.

This dichotomy necessitates different optimization strategies for each phase.

The KV cache consumes GPU memory that grows linearly with both batch size and sequence length. For a model with parameters L (layers), H (heads), and D (head dimension), the cache size for one sequence is approximately 2 * L * H * D * seq_len * dtype_size. This linear scaling is the fundamental constraint on context window length and batch size.

• Example: A 70B model with 80 layers, 64 heads, and 128-dim heads at FP16 requires ~2.5MB of cache per token.

The structure and size of the KV cache are directly determined by the model's attention architecture:

• Multi-Head Attention (MHA): Each query head has a dedicated key and value head, resulting in the largest cache.
• Multi-Query Attention (MQA): All query heads share a single key and value head, reducing cache size by a factor of H.
• Grouped-Query Attention (GQA): A tunable hybrid where groups of query heads share a KV head, offering a memory/quality trade-off.

Choices like Rotary Positional Embedding (RoPE) also affect caching efficiency by enabling relative position calculations without storing absolute positions.

Efficiently managing the KV cache in production involves solving several systems challenges:

• Continuous Batching: Dynamically adding/removing sequences from a shared batch requires non-contiguous, variable-length cache allocation.
• Cache Eviction: Policies (e.g., LRU) for discarding cached tokens when the context exceeds physical memory.
• Offloading & Sharding: Moving cache to CPU/NVMe (offloading) or splitting it across GPUs (sharding) to overcome memory limits.
• Quantization & Compression: Storing cache in lower precision (e.g., INT8) to reduce footprint.

Algorithms like PagedAttention (from vLLM) treat the cache as virtual memory with pages to solve fragmentation and sharing problems.

The prefill phase is the initial, compute-intensive stage of transformer inference where the model processes the entire input prompt in parallel. During this phase, the model computes the initial KV cache for all prompt tokens, which is then stored for reuse. This phase is characterized by high FLOP utilization but only occurs once per sequence, setting the stage for the efficient decode phase.

The decode phase is the token-by-token autoregressive generation stage. It relies heavily on the KV cache to avoid recomputing key and value tensors for previous tokens. In this phase, the model:

• Reads the cached K and V tensors for all prior tokens.
• Computes attention scores with the new token's query.
• Generates the next token's logits. This phase is typically memory-bandwidth bound, as performance is limited by the speed of reading the large KV cache from GPU memory.

These are attention variants designed explicitly to reduce KV cache memory footprint.

• Multi-Query Attention (MQA): All query heads share a single key head and a single value head. This drastically reduces the size of the K and V projections that must be cached.
• Grouped-Query Attention (GQA): A generalization of MQA where query heads are grouped, and each group shares a single K and V head. It provides a tunable trade-off between the cache efficiency of MQA and the model quality of standard Multi-Head Attention. Models like Llama 2 and Llama 3 use GQA to optimize inference cost.

PagedAttention is a memory management algorithm, popularized by the vLLM inference engine, that organizes the KV cache into non-contiguous, fixed-size blocks or 'pages'. This approach:

• Eliminates internal fragmentation caused by variable-length sequences in a continuous batch.
• Enables efficient memory sharing for common prefixes between requests (e.g., in a shared system prompt).
• Allows for dynamic allocation and cache eviction at the block level. It is analogous to virtual memory paging in operating systems and is foundational to high-throughput serving systems.

Continuous batching is an inference scheduling technique that dynamically groups incoming requests into a shared batch. It is critically dependent on efficient KV cache management because:

• New requests can join the batch as soon as GPU resources are available.
• Completed sequences can exit the batch, freeing their allocated cache.
• It requires the KV cache for each request to be managed independently within the same physical memory space. This maximizes GPU utilization and throughput compared to static batching, where the entire batch must finish before a new one begins.

Techniques to reduce the memory footprint of the KV cache itself.

• Quantization: Storing the cached key and value tensors in a lower numerical precision (e.g., FP8 or INT8) instead of FP16/BF16. This reduces memory bandwidth consumption and capacity needs, often with a minimal impact on output quality.
• Compression: A broader category including pruning (removing less important cache entries) and other algorithms to represent the cache in a more compact form. These methods directly address the memory-bound nature of the decode phase, enabling longer context windows or larger batch sizes on fixed hardware.