Introduction to PagedAttention
Large language models (LLMs) have transformed how developers build applications, but serving them efficiently at scale remains a significant engineering challenge. One of the biggest bottlenecks is the Key-Value (KV) cache, the memory used to store intermediate attention states during autoregressive generation. PagedAttention, introduced by the vLLM team in 2023, is an attention algorithm inspired by operating system virtual memory paging that dramatically reduces KV cache waste and improves throughput.
In this tutorial, you will learn what PagedAttention is, why it matters, how it works under the hood, and how to integrate it into your own LLM serving stack using practical code examples.
What Is the KV Cache?
During autoregressive decoding, an LLM generates one token at a time. At each step, the model computes attention over all previously generated tokens. Recomputing the Key and Value projections for every past token at every step would be prohibitively expensive. Instead, these projections are cached in what is called the KV cache.
For a transformer with L layers, hidden size d, and sequence length n, the KV cache stores 2 * L * n * d floating point values per request. For a model like LLaMA-2 70B, a single request with a 4K token context can consume several gigabytes of memory. When you serve many concurrent requests, the KV cache quickly dominates GPU memory usage.
The Memory Fragmentation Problem
Traditional serving systems allocate a contiguous block of memory for each request's KV cache, sized for the maximum possible sequence length. This leads to two major problems:
- Internal fragmentation: If a request is allocated space for 2048 tokens but only generates 200, the remaining memory is wasted. Measurements show 60-80% of KV cache memory can be wasted this way.
- External fragmentation: Allocating and freeing variable-sized contiguous blocks over time fragments the memory pool, making it impossible to fit new requests even when enough total free memory exists.
These inefficiencies cap the number of concurrent requests a GPU can handle, directly limiting throughput and increasing per-token cost.
How PagedAttention Works
PagedAttention borrows the classic idea of virtual memory paging from operating systems. Instead of requiring a contiguous block of memory for each sequence's KV cache, it partitions the KV cache into fixed-size blocks, each holding the keys and values for a small number of tokens (typically 16).
Block Tables and Non-Contiguous Storage
Each sequence maintains a block table that maps logical token positions to physical block locations in GPU memory. The blocks for a single sequence do not need to be contiguous. This means:
- Memory is allocated on demand, only as a sequence actually generates tokens.
- Blocks from completed or stopped sequences can be returned to a free pool and reused immediately.
- Fragmentation is virtually eliminated because all allocations are the same fixed size.
During the attention computation, the kernel walks the block table, loads the relevant blocks, and computes attention over the concatenated keys and values. This adds a small amount of indirection overhead, but the memory savings and throughput gains far outweigh it.
Shared Prefix Caching
Because blocks are independent units, PagedAttention enables a powerful optimization: shared prefix caching. If multiple requests share the same system prompt or few-shot examples, their KV cache blocks for the shared prefix can reference the same physical blocks via copy-on-write semantics. This can dramatically reduce memory usage and recomputation for prompt-heavy workloads.
Why PagedAttention Matters
The original vLLM paper reported that PagedAttention improves LLM serving throughput by 2-4x compared to naive implementations, and by similar margins over systems like HuggingFace Transformers with continuous batching. The gains come from three sources:
- Higher batch sizes: Less wasted memory means more concurrent requests fit on the GPU.
- Lower latency: Better memory locality and on-demand allocation reduce stalls.
- Prefix reuse: Shared system prompts across requests are computed once and reused.
For production LLM APIs where cost per token directly affects margins, these improvements are transformative. PagedAttention has become a foundational technique adopted by vLLM, TensorRT-LLM, SGLang, and other modern inference engines.
Using PagedAttention with vLLM
The most accessible way to use PagedAttention is through vLLM, an open-source LLM inference engine. The following examples walk through installation, basic serving, and programmatic usage.
Installation
pip install vllm
vLLM bundles CUDA kernels that implement PagedAttention. You need a CUDA-capable GPU (compute capability 7.0 or higher) and a recent PyTorch build.
Offline Batched Inference
The simplest way to use PagedAttention is through vLLM's LLM class, which handles KV cache paging automatically:
from vllm import LLM, SamplingParams
# Initialize the engine. vLLM allocates a paged KV cache
# sized to fill available GPU memory by default.
llm = LLM(model="meta-llama/Llama-2-7b-chat-hf")
# Define sampling parameters for all requests.
sampling_params = SamplingParams(
temperature=0.7,
top_p=0.9,
max_tokens=512,
)
# A batch of prompts. vLLM will schedule these together,
# using PagedAttention to share GPU memory efficiently.
prompts = [
"Explain quantum entanglement in one paragraph.",
"Write a Python function to reverse a linked list.",
"Summarize the plot of Romeo and Juliet.",
"What are the benefits of containerization?",
]
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
prompt = output.prompt
generated = output.outputs[0].text
print(f"Prompt: {prompt}")
print(f"Generated: {generated}\n")
Behind the scenes, vLLM allocates a pool of fixed-size KV cache blocks. Each prompt is assigned blocks on demand as tokens are generated. When one sequence finishes, its blocks are immediately returned to the pool and reused by the next queued request.
Starting an OpenAI-Compatible Server
For production deployments, you typically run vLLM as a server. PagedAttention is enabled by default:
python -m vllm.entrypoints.openai.api_server \
--model meta-llama/Llama-2-7b-chat-hf \
--tensor-parallel-size 1 \
--gpu-memory-utilization 0.9 \
--max-model-len 4096 \
--port 8000
You can then send requests using the standard OpenAI client library:
from openai import OpenAI
client = OpenAI(
base_url="http://localhost:8000/v1",
api_key="not-needed",
)
response = client.chat.completions.create(
model="meta-llama/Llama-2-7b-chat-hf",
messages=[
{"role": "system", "content": "You are a helpful coding assistant."},
{"role": "user", "content": "Write a function to check if a string is a palindrome."},
],
temperature=0.7,
max_tokens=512,
)
print(response.choices[0].message.content)
Enabling Automatic Prefix Caching
To take advantage of PagedAttention's shared prefix caching, enable the enable_prefix_caching flag. This is especially valuable when many requests share a long system prompt or common context:
python -m vllm.entrypoints.openai.api_server \
--model meta-llama/Llama-2-7b-chat-hf \
--enable-prefix-caching \
--port 8000
With prefix caching enabled, the first request computes the KV cache for the shared system prompt. Subsequent requests with the same prefix reuse those blocks via copy-on-write, reducing time-to-first-token and memory consumption.
Programmatic Control Over KV Cache
For advanced use cases, you can inspect and tune how vLLM manages the paged KV cache. The EngineArgs object exposes relevant parameters:
from vllm import LLM, SamplingParams
from vllm.engine.arg_utils import EngineArgs
engine_args = EngineArgs(
model="meta-llama/Llama-2-7b-chat-hf",
# Block size in tokens. 16 is the default and usually optimal.
block_size=16,
# Fraction of GPU memory to reserve for KV cache.
gpu_memory_utilization=0.9,
# Swap space (in GiB) for offloading KV cache to CPU.
swap_space=4,
# Enable prefix caching for shared prompts.
enable_prefix_caching=True,
# Maximum number of sequences to schedule at once.
max_num_seqs=256,
)
llm = LLM(**engine_args.to_dict())
prompts = ["Hello, world!"] * 10
sampling_params = SamplingParams(max_tokens=64, temperature=0.0)
outputs = llm.generate(prompts, sampling_params)
for o in outputs:
print(o.outputs[0].text)
The block_size parameter controls the granularity of paging. Smaller blocks reduce waste for short sequences but increase block table overhead. The default of 16 tokens is a well-tuned balance for most workloads.
Implementing a Simplified Paged KV Cache
To build intuition, here is a simplified Python implementation of a paged KV cache manager. This is not production code, but it illustrates the core data structures:
import torch
from dataclasses import dataclass, field
from typing import List, Dict
@dataclass
class Block:
"""A fixed-size block holding KV cache for `block_size` tokens."""
block_id: int
block_size: int
# Shape: (block_size, num_heads, head_dim)
keys: torch.Tensor = field(init=False)
values: torch.Tensor = field(init=False)
ref_count: int = 0 # For copy-on-write prefix sharing
def __post_init__(self):
self.keys = torch.zeros(self.block_size, 8, 128)
self.values = torch.zeros(self.block_size, 8, 128)
class PagedKVCache:
def __init__(self, num_blocks: int, block_size: int = 16):
self.block_size = block_size
# Pre-allocate a pool of blocks.
self.blocks: Dict[int, Block] = {
i: Block(block_id=i, block_size=block_size)
for i in range(num_blocks)
}
self.free_block_ids: List[int] = list(range(num_blocks))
# Each sequence maps to a list of physical block IDs.
self.block_tables: Dict[int, List[int]] = {}
# Track how many tokens are written in the current block.
self.seq_lengths: Dict[int, int] = {}
def allocate_sequence(self, seq_id: int):
"""Register a new sequence with no blocks yet."""
self.block_tables[seq_id] = []
self.seq_lengths[seq_id] = 0
def append_token(self, seq_id: int, key: torch.Tensor, value: torch.Tensor):
"""Append one token's K and V to the cache, allocating blocks as needed."""
length = self.seq_lengths[seq_id]
position_in_block = length % self.block_size
# Need a new block if we are at the start of a fresh block slot.
if position_in_block == 0:
if not self.free_block_ids:
raise RuntimeError("KV cache exhausted: no free blocks available.")
new_block_id = self.free_block_ids.pop(0)
self.blocks[new_block_id].ref_count += 1
self.block_tables[seq_id].append(new_block_id)
current_block_id = self.block_tables[seq_id][-1]
block = self.blocks[current_block_id]
block.keys[position_in_block] = key
block.values[position_in_block] = value
self.seq_lengths[seq_id] += 1
def free_sequence(self, seq_id: int):
"""Return all blocks used by a finished sequence to the free pool."""
for block_id in self.block_tables[seq_id]:
block = self.blocks[block_id]
block.ref_count -= 1
if block.ref_count == 0:
self.free_block_ids.append(block_id)
del self.block_tables[seq_id]
del self.seq_lengths[seq_id]
def get_keys_values(self, seq_id: int):
"""Gather all keys and values for a sequence across its blocks."""
block_ids = self.block_tables[seq_id]
length = self.seq_lengths[seq_id]
keys = torch.stack([self.blocks[bid].keys for bid in block_ids])
values = torch.stack([self.blocks[bid].values for bid in block_ids])
# Reshape to (length, num_heads, head_dim)
keys = keys.reshape(-1, *keys.shape[2:])[:length]
values = values.reshape(-1, *values.shape[2:])[:length]
return keys, values
# Demonstration
cache = PagedKVCache(num_blocks=100, block_size=16)
cache.allocate_sequence(seq_id=1)
# Simulate appending 40 tokens (requires 3 blocks: 16 + 16 + 8).
for i in range(40):
k = torch.randn(8, 128)
v = torch.randn(8, 128)
cache.append_token(seq_id=1, key=k, value=v)
keys, values = cache.get_keys_values(seq_id=1)
print(f"Sequence 1 cached {keys.shape[0]} tokens across "
f"{len(cache.block_tables[1])} blocks.")
# Free the sequence; blocks return to the pool for reuse.
cache.free_sequence(seq_id=1)
print(f"Free blocks after freeing: {len(cache.free_block_ids)}")
This example demonstrates the three core ideas: a fixed-size block pool, per-sequence block tables, and on-demand allocation. In a real implementation, the attention kernel itself is modified to traverse the block table on the GPU, avoiding the need to gather blocks into contiguous memory before computing attention.
Best Practices
Choose the Right Block Size
The default block size of 16 tokens works well for most models and workloads. Smaller block sizes (such as 8) reduce waste for short sequences but increase block table lookup overhead. Larger block sizes (such as 32) improve kernel efficiency but waste more memory when sequences end mid-block. Benchmark with your specific traffic patterns before deviating from the default.
Tune GPU Memory Utilization
The gpu_memory_utilization parameter controls how much GPU memory vLLM reserves for the KV cache pool after loading model weights. Setting it too low wastes capacity; setting it too high risks out-of-memory errors during spikes. A value between 0.85 and 0.92 is a good starting point for dedicated inference GPUs.
Enable Prefix Caching for Prompt-Heavy Workloads
If your application uses long shared system prompts, few-shot examples, or retrieval-augmented context that repeats across requests, always enable prefix caching. The first request pays the full prefill cost, and subsequent requests reuse cached blocks. This can cut time-to-first-token by 50% or more for shared prefixes of several hundred tokens.
Use Continuous Batching
PagedAttention is most effective when combined with continuous batching (also called iteration-level scheduling). Instead of waiting for all sequences in a batch to finish, the scheduler inserts new requests at every decoding step as old ones complete. PagedAttention's on-demand block allocation makes this practical because freed blocks are immediately reusable. vLLM enables continuous batching by default.
Monitor KV Cache Utilization
Track metrics like KV cache block usage, prefix cache hit rate, and number of running versus waiting sequences. vLLM exposes Prometheus metrics that make it easy to observe these in Grafana. If you see frequent "KV cache exhausted" errors, either increase gpu_memory_utilization, reduce max_model_len, or add more GPU capacity.
Consider Tensor Parallelism for Large Models
For models too large to fit on a single GPU, use tensor parallelism. PagedAttention works across tensor-parallel ranks, with each rank managing its own shard of the KV cache block pool. Use the --tensor-parallel-size flag to split the model across multiple GPUs while retaining all the memory efficiency benefits of paging.
Conclusion
PagedAttention is one of the most impactful innovations in LLM inference engineering. By applying operating system paging principles to the KV cache, it eliminates the memory fragmentation that plagues naive implementations, enables much higher batch sizes, and unlocks optimizations like shared prefix caching. Whether you are deploying an open-source model with vLLM or building a custom inference stack, understanding and leveraging PagedAttention is essential for achieving production-grade throughput and cost efficiency. Start with vLLM's sensible defaults, enable prefix caching if your workload benefits from it, and monitor KV cache utilization to squeeze the most value out of every GPU.