KV Caching

The standard Transformer self-attention mechanism recomputes attention scores for every token against all previous tokens in the sequence for each new decoding step. This results in O(n²) computational complexity in the sequence length. For a prompt of length P and generating G new tokens, a model without caching performs ~(P+G)² operations. KV caching reduces this to ~P*G + G² by storing prior computations.

During the initial prefill or prompt phase, the model processes the entire input prompt in one forward pass. For each layer and attention head, it computes:

• Key (K) Tensor: A projection of the input representing what the token "contains."
• Value (V) Tensor: A projection representing the token's "output" value. These K and V tensors for all prompt tokens are stored in the KV Cache. This is the only time the full O(P²) attention computation occurs.

During autoregressive generation, to produce token t, the model only computes the new K_t and V_t for the single new token. It then concatenates these new vectors with the cached K and V tensors from all previous tokens [0:t-1]. The attention mechanism performs a query (Q_t) against this combined cache. This avoids recomputing keys and values for the entire prior context, turning an O(t²) operation into O(t).

KV caching trades increased memory consumption for drastically reduced compute latency. The memory footprint is substantial: for a model with L layers, H attention heads, hidden dimension d_h, and precision b bytes, caching a sequence of length N requires ~2 * L * H * N * d_h * b bytes. For a 70B parameter model with 80 layers and 64 heads, caching a 4096-token context in FP16 (b=2) consumes ~2.5 GB of GPU memory. This trade-off is central to inference optimization.

In production serving systems like vLLM, KV caching must be managed across multiple concurrent requests in a continuous batching setup. Key challenges include:

• Dynamic Sequence Lengths: Each request has a unique, growing cache.
• Memory Fragmentation: Allocating and deallocating variable-sized cache blocks leads to waste.
• Cache Sharing: For requests with identical prompts (e.g., in chatbot multi-turn conversations), the prompt-phase cache can be shared, eliminating redundant computation. Solutions like PagedAttention treat the KV cache as non-contiguous, pageable memory to solve fragmentation.

KV caching has inherent limits that drive further optimization:

• Memory Bound: Generation becomes limited by the bandwidth to read the large KV cache from GPU memory, not by compute (the "memory wall").
• Long Context Degradation: As the cache grows, attention over very long sequences remains slow and memory-intensive.
• Optimizations: Techniques like Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) reduce the cache size by sharing keys and/or values across attention heads. Sliding Window Attention implements a fixed-size cache that discards old tokens, trading context length for constant memory usage.

Advanced implementations involve manual cache tuning.

• Cache Quantization: Storing KV tensors in fp16, bf16, or even int8 to halve or quarter memory use, often with minimal perplexity impact.
• Cache Eviction Policies: For infinite-length conversations, strategies like Least Recently Used (LRU) evict old cache blocks to make room for new tokens.
• Selective Caching: Not caching certain attention layers (e.g., early or late layers) based on profiling to trade memory for compute.
• These optimizations are often framework-specific and require deep performance profiling.

PagedAttention is the memory management algorithm that makes KV caching practical at scale. It treats the KV cache as non-contiguous 'pages' in memory, similar to virtual memory in operating systems.

• Eliminates Memory Fragmentation: Allows for efficient sharing and dynamic allocation of cache space.
• Enables Longer Contexts: Supports context windows far exceeding GPU memory by paging cache blocks in and out.
• Core to vLLM: This algorithm is the foundational innovation behind the vLLM inference engine's high throughput.

Continuous batching (or iterative batching) is a scheduling technique that groups incoming inference requests of varying sequence lengths into a single batch to maximize GPU utilization.

• Dynamically Batches Requests: Unlike static batching, it doesn't wait for a fixed batch size, improving throughput.
• Works with KV Cache: Efficiently manages the KV caches for multiple concurrent requests within the same batch.
• Reduces Latency: By keeping the GPU busy, it lowers the average time-to-first-token (TTFT) for users.

FlashAttention is an IO-aware algorithm that speeds up the core attention computation, which directly feeds into the KV cache.

• Optimizes Memory Access: Recomputes attention scores on-the-fly in fast SRAM to avoid reading/writing the massive attention matrix to slow HBM memory.
• Accelerates Prefill: Dramatically reduces the time spent on the initial prompt processing phase, where the initial KV cache is computed.
• Enables Longer Sequences: Makes it feasible to run attention on context lengths of 128K+ tokens.

Speculative decoding is a latency-reduction technique that uses a faster model to 'draft' tokens which are then verified in parallel by the target model.

• Leverages KV Cache Efficiency: The target model's verification step is highly parallelizable and benefits from cached computations.
• Accelerates Generation: Can provide 2-3x speedups in decoding by reducing the number of serial calls to the large target model.
• Preserves Output Distribution: Mathematically guarantees the same output as the original autoregressive model.

Model quantization reduces the numerical precision of model weights and activations (e.g., from FP16 to INT8) to decrease memory footprint and increase compute speed.

• Reduces KV Cache Size: Quantizing the KV cache (e.g., to FP8) can halve its memory consumption, allowing larger batches or longer contexts.
• Increases Throughput: Lower precision arithmetic is faster on modern hardware (GPUs, NPUs).
• Applied via PTQ/QAT: Can be done post-training (PTQ) or with fine-tuning (Quantization-Aware Training) for better accuracy.

vLLM is a high-throughput, memory-efficient open-source inference and serving engine for LLMs. It is the canonical implementation that popularized PagedAttention for KV cache management.

• PagedAttention Implementation: Its core innovation is treating the KV cache as paged memory.
• Native Continuous Batching: Dynamically batches requests to maximize GPU utilization.
• Production Serving: Includes features like optimized CUDA kernels, streaming outputs, and OpenAI-compatible APIs.