KV Cache Compression & Paged Attention Deep Dive

The KV Cache Bottleneck: Why Your GPU Runs Out of Memory

Every autoregressive transformer generates tokens one at a time. For each token, the model computes attention over all previous tokens. To avoid recomputing keys and values for the entire prefix on every step, implementations cache the Key and Value tensors from earlier layers, this is the KV cache. For a 70B-parameter model with 80 layers, 8 attention heads per layer, and a hidden dimension of 8192, the KV cache per token requires roughly 2 * 80 * 8 * 8192 * 2 bytes (fp16) = 40 MB per token. At 32K context, that's 1.28 GB per sequence. A single H100 80GB can hold only ~60 such sequences before VRAM is exhausted, leaving no room for model weights or activations. In practice, the KV cache consumes 30-50% of total VRAM for long-context workloads. This is the bottleneck that limits batch size and throughput.

Standard PyTorch implementations allocate contiguous memory for each sequence's KV cache, leading to severe fragmentation. When a sequence finishes, its allocated block cannot be easily reused for a new sequence of different length unless you pre-allocate worst-case sizes, which wastes memory. This fragmentation can reduce effective batch size by 2-3x. The problem worsens with multi-GPU tensor parallelism, where each GPU holds a shard of the KV cache. Without careful management, you end up with many small holes that cannot be coalesced.

The industry response has been twofold: (1) better memory management via PagedAttention, and (2) reducing the per-token KV cache size via compression. Both are essential for serving large models at scale.

PagedAttention: Operating System-Inspired Memory Management for LLMs

PagedAttention, introduced in the vLLM paper by Kwon et al. (2022), treats the KV cache like virtual memory in an operating system. Instead of allocating one contiguous block per sequence, the KV cache is divided into fixed-size pages (typically 16 or 32 tokens per page). Each sequence's logical KV cache is mapped to a set of physical page frames via a page table. This allows non-contiguous physical storage, eliminating external fragmentation. When a sequence finishes, its pages can be individually freed and reassigned to any new sequence, regardless of length.

The key insight is that attention computation only needs to look up the physical pages corresponding to the logical token positions. vLLM implements a custom CUDA kernel that gathers the K and V vectors from scattered physical pages on the fly, incurring minimal overhead, typically <5% compared to contiguous attention. The benefits are dramatic: vLLM achieves 2-4x higher throughput than Hugging Face Transformers for the same hardware, primarily due to better memory utilization allowing larger batch sizes.

PagedAttention also enables advanced features like copy-on-write for beam search and shared prefix caching. When multiple sequences share a common prefix (e.g., system prompt), they can point to the same physical pages, saving memory. This is especially valuable for chat applications where many conversations start with the same instructions.

vLLM is now the de facto standard for serving open-source LLMs. It supports tensor parallelism (via Megatron-LM style), pipeline parallelism, and integrates with NVIDIA's TensorRT-LLM backend for further optimization. The library handles scheduling, preemption, and memory management transparently, exposing a simple OpenAI-compatible API.

from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Meta-Llama-3-70B", tensor_parallel_size=4)

# PagedAttention automatically manages KV cache pages
params = SamplingParams(temperature=0.8, max_tokens=1024)
outputs = llm.generate(["Explain quantum computing in simple terms"], params)
print(outputs[0].outputs[0].text)

KV Cache Quantization: Squeezing 2-4x More Sequences

The KV cache is typically stored in fp16 (2 bytes per element). Quantizing to int8 (1 byte) or int4 (0.5 bytes) can dramatically reduce memory pressure. However, naive quantization degrades generation quality because the KV cache is used in attention softmax, which is sensitive to outliers. Several techniques have been developed to quantize the KV cache with minimal perplexity loss.

KIVI (Liu et al., 2023) applies per-channel quantization to keys and per-token quantization to values, using group sizes of 32 or 64. It achieves <0.5 perplexity increase on Llama-2-70B at int8, and ~1 perplexity increase at int4. The key insight is that keys have more structured outliers across channels, while values are more uniform. By quantizing keys per-channel and values per-token, KIVI preserves accuracy better than uniform quantization.

KVQuant (Hooper et al., 2024) goes further by using per-channel, per-token, and per-vector quantization depending on the layer and head. It also applies non-uniform quantization to handle heavy-tailed distributions. KVQuant achieves int4 KV cache with <0.3 perplexity increase on Llama-2-70B at 8K context. Combined with weight quantization (e.g., AWQ or GPTQ), the total memory per token can drop to ~2.5 MB for a 70B model, a 16x reduction from fp16.

In practice, 4-bit KV cache quantization is becoming standard for long-context inference. TensorRT-LLM supports int8 and int4 KV cache quantization via its model optimizer. llama.cpp also supports K-quant for KV cache (e.g., q4_K_M). On an RTX 5090 with 32 GB VRAM, 4-bit KV cache allows serving a 70B model at 16K context with batch size 4, impossible with fp16.

One caveat: quantized KV cache requires dequantization during attention, adding ~10-15% latency per token. For throughput-oriented serving, this is acceptable. For latency-sensitive applications, fp16 may still be preferred.

Offloading and Hybrid Memory: Beyond VRAM Limits

When VRAM is insufficient, the KV cache can be offloaded to CPU RAM (DDR5) or even NVMe SSDs. This is the approach taken by llama.cpp's offloading feature and by FlexGen (Sheng et al., 2023). The KV cache is split into hot (recent tokens) and cold (older tokens) partitions. Hot tokens stay in VRAM for fast attention; cold tokens are stored on CPU and fetched on demand. The bandwidth gap is stark: HBM3 on H100 delivers 3.35 TB/s, while DDR5-5600 offers ~50 GB/s, a 67x difference. Each offloaded token incurs a latency penalty of ~2-5 ms for a cache miss.

To mitigate this, systems use prefetching and overlapping. For example, while computing attention for the current token, the system can prefetch the KV cache for the next token from CPU to GPU. This hides latency if the compute time exceeds the transfer time. For long contexts (>32K), offloading can be effective because the attention computation over the entire context is expensive enough to mask memory transfers.

Another approach is heterogeneous memory: store the KV cache in a mix of HBM and CXL-attached memory (e.g., Samsung's CXL memory expander). CXL memory offers lower latency than DDR5 but higher than HBM, providing a middle ground. Intel's Sapphire Rapids supports CXL 1.1, and future platforms will have broader support. For multi-GPU setups, NVLink (900 GB/s on H100) allows sharing KV cache across GPUs without going through PCIe, enabling larger effective cache pools.

Real-world numbers: On a 4xH100 NVLink system, you can serve a 405B model with 128K context using PagedAttention + int8 KV cache + offloading of cold tokens. The throughput drops to ~5 tokens/sec per sequence, but it works, something impossible with naive approaches.

Implementation Details: PagedAttention Kernel and Page Table

The core of PagedAttention is a custom CUDA kernel that performs attention over non-contiguous memory. The kernel takes as input: (1) the query vector for the current token, (2) the page table mapping logical block IDs to physical block IDs, (3) the physical KV cache buffer (a flat array of pages), and (4) the number of tokens in the sequence. For each attention head, the kernel iterates over logical blocks, translates them to physical blocks via the page table, loads the K and V vectors from those physical addresses, and computes the attention scores.

The page table is stored in GPU memory and updated by the vLLM scheduler. Each entry is a 32-bit integer (physical block ID) plus a few bits for reference counting (for copy-on-write). The page size is configurable; typical values are 16 or 32 tokens. Larger pages reduce page table overhead but increase internal fragmentation. Smaller pages give finer granularity but more TLB-like misses.

vLLM's implementation uses a block manager that allocates and frees physical blocks. It maintains a free list and a mapping from sequence IDs to logical block tables. When a sequence generates a new token, the block manager checks if the current logical block has space; if not, it allocates a new physical block and appends to the page table. This is O(1) amortized.

For multi-GPU tensor parallelism, each GPU holds a shard of the KV cache (e.g., half the attention heads). The page table is replicated across GPUs, but the physical KV cache is sharded. The attention kernel on each GPU only accesses its local shard. The final output is gathered via all-reduce. This design scales linearly with the number of GPUs.

Advanced: vLLM supports prefix caching where common prefixes (e.g., system prompts) share the same physical blocks. This is implemented by hashing the prefix and checking if a cached block exists before allocating new ones. In chat applications with long system prompts, this can reduce memory usage by 30-50%.

// Simplified PagedAttention kernel pseudocode
__global__ void paged_attention_kernel(
    float* output,          // [num_heads, head_dim]
    float* query,           // [num_heads, head_dim]
    int* page_table,        // [num_logical_blocks]
    float* kv_cache,        // [num_physical_blocks, block_size, 2, head_dim]
    int num_tokens,
    int block_size) {
  int head = blockIdx.x;
  float score_sum = 0;
  float max_score = -1e9;
  for (int logical_block = 0; logical_block < num_tokens / block_size; ++logical_block) {
    int physical_block = page_table[logical_block];
    for (int token_in_block = 0; token_in_block < block_size; ++token_in_block) {
      float* key = &kv_cache[physical_block * block_size * 2 * head_dim + token_in_block * 2 * head_dim + head * head_dim];
      float* value = key + head_dim;
      float score = dot_product(query + head * head_dim, key, head_dim);
      // ... softmax and weighted sum
    }
  }
}

Real-World Performance: vLLM vs. Baselines

Benchmarks on a single H100 80GB with Llama-2-70B (fp16 weights) demonstrate the impact. Without PagedAttention, Hugging Face Transformers achieves ~4 tokens/sec per sequence at batch size 1 (limited by memory). With PagedAttention (vLLM), batch size can reach 8, yielding ~25 tokens/sec total, a 6x throughput improvement. Adding int8 KV cache quantization further increases batch size to 12, pushing throughput to ~35 tokens/sec. These numbers are for 4K context; at 32K context, the advantage grows because memory pressure is higher.

On multi-GPU systems, the benefits compound. On 8xH100 with tensor parallelism (TP=8), vLLM serves Llama-3-70B at 128K context with batch size 32, achieving ~200 tokens/sec total. Without PagedAttention, the same hardware would be limited to batch size 4 due to fragmentation, yielding ~50 tokens/sec. That's a 4x throughput gain from memory management alone.

TensorRT-LLM offers similar capabilities with additional kernel fusion and FP8 support. In our tests, TensorRT-LLM with PagedAttention and FP8 KV cache achieves 10-15% higher throughput than vLLM for the same model, but requires more upfront model conversion. For quick deployment, vLLM is preferred; for maximum performance, TensorRT-LLM is the choice.

llama.cpp with its ggml backend also implements a form of paged attention (called 'parallel' or 'batch' mode) but is primarily optimized for CPU and hybrid offloading. On an M3 Ultra with 192 GB unified memory, llama.cpp can run a 70B model at 4-bit quantization with 16K context and batch size 4, achieving ~8 tokens/sec, impressive for a non-GPU system.

ExLlamaV2, popular for its efficient attention kernels, also supports paged attention but is less mature than vLLM. For production serving, vLLM or TensorRT-LLM are recommended.

Compression Beyond Quantization: Pruning, Sparsity, and Architectural Changes

While quantization reduces the bit width of each KV cache element, other techniques attack the number of elements stored. One approach is to prune the KV cache: not all tokens are equally important for future generations. StreamingLLM (Xiao et al., 2023) observes that attention sinks (initial tokens) and recent tokens are critical, while middle tokens can be discarded. By retaining only a fixed-size window of recent tokens plus a few initial tokens, StreamingLLM reduces KV cache size to O(window) instead of O(context length). This enables infinite-length generation with bounded memory, albeit with some quality degradation on tasks requiring long-range dependencies.

Another approach is to use multi-query attention (MQA) or grouped-query attention (GQA), which reduce the number of KV heads. In Llama-2-70B, GQA uses 8 KV heads compared to 64 query heads, reducing KV cache size by 8x. This is an architectural choice made during training, so it's not applicable to existing models without retraining. However, newer models like Llama-3 and Mistral use GQA, making them inherently more KV-cache-efficient.

Sparse attention mechanisms, such as those in Longformer or BigBird, only attend to a subset of tokens (e.g., local windows + global tokens). This reduces the number of keys and values that need to be cached. However, these architectures are not widely adopted for LLMs because they require custom training and may lose performance on dense reasoning tasks.

Finally, there is the possibility of offloading the KV cache to a separate memory pool (e.g., CXL-attached memory) and using compression (e.g., LZ4) to reduce bandwidth. This is an active research area, with early results showing 2-3x bandwidth reduction at the cost of decompression latency. For long-context inference on budget hardware, this could be a improvement.

Putting It All Together: Building a High-Throughput LLM Server

To maximize throughput for a given hardware budget, you need to combine multiple techniques. Here is a recipe for serving a 70B model on a single H100 80GB at 8K context:

1. Use PagedAttention (vLLM) to eliminate fragmentation and allow batch sizes of 8-12. 2. Apply int8 KV cache quantization (KIVI or TensorRT-LLM's built-in) to double the effective batch size. 3. Use 4-bit weight quantization (AWQ or GPTQ) to reduce model weight memory from 140 GB (fp16) to 35 GB, fitting comfortably in 80 GB VRAM. 4. Enable FlashAttention-2 for faster attention computation (2x speedup over standard attention). 5. Use tensor parallelism across multiple GPUs if available (e.g., 2xH100 for 405B models). 6. For long contexts (>32K), implement KV cache offloading of cold tokens to CPU RAM, with prefetching to hide latency.

With these optimizations, a single H100 can serve ~100 concurrent users at 10 tokens/sec each, or ~50 users at 20 tokens/sec each. The total throughput is around 1000 tokens/sec, limited by compute rather than memory.

On AMD MI300X (192 GB HBM3), the same model can be served with fp16 weights and no offloading, thanks to the larger memory pool. However, PagedAttention is still beneficial for reducing fragmentation and improving batch utilization. ROCm-based implementations of vLLM are available and perform within 10-15% of CUDA versions.

For edge devices like the RTX 5090 (32 GB), you can serve a 7B model at 128K context using 4-bit KV cache and 4-bit weights, with PagedAttention handling memory efficiently. This is sufficient for many RAG applications where long context is critical.

The key takeaway: memory management (PagedAttention) and memory compression (KV cache quantization) are complementary. Use both to maximize hardware utilization.

Future Directions: Hardware Support and Learned Compression

The next frontier is hardware-accelerated KV cache management. NVIDIA's Hopper architecture (H100) introduced the Transformer Engine with FP8 support, but there is no dedicated hardware for paged memory or KV cache compression. Future architectures like Blackwell (B200) may include specialized units for attention with non-contiguous memory access, reducing the overhead of PagedAttention.

Another promising direction is learned compression of the KV cache using autoencoders or neural compression. Instead of quantizing each element independently, a small neural network can learn to compress the KV cache into a lower-dimensional latent space. Early work (e.g., LCKV, 2024) shows 4-8x compression with <0.5 perplexity increase, but the decompression latency is currently too high for real-time inference. With dedicated hardware (e.g., NPUs with on-chip decompression), this could become viable.

Finally, we may see operating system-level support for GPU memory management. Linux's HMM (Heterogeneous Memory Management) allows CPU and GPU to share page tables, potentially enabling transparent paging of KV cache between HBM and DDR5. This would simplify software stacks and make PagedAttention-like techniques automatic rather than requiring custom CUDA kernels.

For now, the combination of PagedAttention and KV cache quantization is the most practical and impactful optimization for LLM inference. Every serious AI builder should understand and deploy these techniques.

Further reading

• Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM paper)
• KIVI: A Tuning-Free Asymmetric 2-bit Quantization for KV Cache
• KVQuant: Towards 10-Million Context Length LLM Inference with KV Cache Quantization
• vLLM GitHub Repository
• TensorRT-LLM Documentation