PagedAttention: Optimizing KV Cache Utilization in Large Language Models | FromZero2AI Blog

When serving Large Language Models (LLMs) in multi-user environments, memory capacity is the main bottleneck. During generation, the model saves the Key-Value (KV) history of prior tokens to avoid recomputing them. This is known as the KV Cache.

The size of the KV cache increases with the sequence length and concurrent request counts. In standard systems, KV cache memory is allocated contiguously. This introduces two major sources of waste:

Internal Fragmentation: Allocating memory for the maximum possible sequence length up front, even if the request completes early.
External Fragmentation: Memory gaps between requests that cannot be utilized due to size differences.

As a result, up to 60% to 80% of KV cache memory can be wasted. PagedAttention solves this by managing KV cache memory in non-contiguous pages, similar to virtual memory in operating systems.

PagedAttention Memory Management PagedAttention Memory Management

The Mechanics of PagedAttention

PagedAttention divides the KV cache for each request into logical blocks. Each block contains the KV vectors for a fixed number of tokens (e.g. 16 tokens).

Logical Blocks: The KV cache is represented as a sequence of logical blocks from Block 0 to Block N.
Physical Blocks: The actual GPU memory is allocated as non-contiguous physical blocks.
Page Table: A lookup table maps each request's logical blocks to their corresponding physical blocks.

During generation, as new tokens are generated, the system allocates physical blocks dynamically. The KV vectors are written to these blocks even if they are scattered across different parts of GPU memory.

The Math: Block Attention Calculation

When calculating attention for query token q_i, the system queries the page table to locate the physical addresses of the key vectors K_j. The attention score for block b is:

A_{i,b} = \text{Softmax}\left( \frac{q_i K_b^T}{\sqrt{d_k}} \right)

where K_b represents the key vectors stored within physical block b.

Implementing PagedAttention with vLLM

vLLM uses PagedAttention under the hood to achieve high serving throughput. Here is a Python script to start a vLLM engine:

# Simple inference engine with vLLM PagedAttention
# from vllm import LLM, SamplingParams
# 
# llm = LLM(model="facebook/opt-125m")
# prompts = ["Once upon a time in a galaxy far away,"]

Serving Advantages

Zero Waste: Memory is allocated dynamically per block, reducing fragmentation to under 4%.
Shared Memory (Beam Search / Multi-User): Multiple requests can share the physical blocks of a common prompt prefix (e.g. system instructions), reducing memory usage in multi-turn chats.
Increased Throughput: By reclaiming wasted memory, vLLM can process 2x to 4x more concurrent requests than standard serving frameworks.

Conclusion

PagedAttention is a key optimization for LLM serving. Bringing virtual memory paging concepts to GPU KV Cache management eliminates fragmentation, enables memory sharing, and significantly improves serving throughput.