PagedAttention

Overview

PagedAttention is a breakthrough memory management technique for Large Language Model (LLM) inference that dramatically reduces memory fragmentation and enables efficient batch processing of requests with variable sequence lengths.

• Introduced: June 2023
• Paper: "Efficient Memory Management for Large Language Model Serving with PagedAttention"
• Authors: Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, et al. (UC Berkeley)
• Impact: 10-40x throughput improvement, 50% memory savings
• Adopted by: vLLM, HuggingFace TGI, and other inference engines

-

The Problem: Memory Fragmentation in LLM Inference

Traditional KV Cache Management

When generating tokens, transformer models store Key-Value (KV) cache for each token to avoid recomputation:

Sequence: "The quick brown fox jumps"
 ↓
Token 1 "The": K1, V1 → stored
Token 2 "quick": K2, V2 → stored (K1, V1 reused)
Token 3 "brown": K3, V3 → stored (K1, V1, K2, V2 reused)
Token 4 "fox": K4, V4 → stored (K1-K3, V1-V3 reused)
Token 5 "jumps": K5, V5 → stored (K1-K4, V1-V4 reused)

Cache size grows: O(sequence_length)

The Bottleneck: Memory Waste

Scenario: Batch Processing with Variable Lengths

Batch Processing Timeline:

Time 0:
Request 1: Allocated 4096 tokens (uses 256 actually) → 3840 wasted
Request 2: Allocated 4096 tokens (uses 512 actually) → 3584 wasted
Request 3: Allocated 4096 tokens (uses 128 actually) → 3968 wasted
Request 4: Allocated 4096 tokens (uses 1024 actually) → 3072 wasted
─────────────────────────────────────────────────────────────────
Total GPU Memory: 16,384 tokens allocated
Actual Use: 1,920 tokens
Fragmentation Loss: ~88% wasted!

Time 1:
Request 1 completes (frees 4096)
Request 5 arrives but needs 2048 tokens - MUST WAIT (fragmentation!)

Why This Happens

1. Pre-allocation: Each request pre-allocates max possible memory
2. Variable lengths: Different requests need different amounts
3. Contiguous memory: Traditional approach requires contiguous allocation
4. No sharing: Memory freed by completed requests is fragmented

Impact on Throughput

Traditional serving with fragmentation:
- Serve Request 1: 8 seconds
- Serve Request 2: 8 seconds (can't start until Request 1 done)
- Serve Request 3: 8 seconds (can't start until Request 2 done)
Total: 24 seconds for 3 requests = 0.125 req/sec 

With batch processing (but high fragmentation):
- Can serve 2-3 requests concurrently
- But memory waste limits batch size to 2
- Still only 0.2-0.3 req/sec 😞

---

The Solution: PagedAttention

Core Idea: Virtual Memory for KV Cache

PagedAttention treats KV cache like operating system virtual memory:

Traditional OS Virtual Memory:

Physical Memory is divided into fixed-size PAGES
Applications see LOGICAL addresses
OS maps LOGICAL pages → PHYSICAL pages
Non-contiguous physical memory appears contiguous to app

PagedAttention applies this to GPU memory:

KV Cache is divided into fixed-size PAGES (512 tokens default)
LLM sees LOGICAL KV cache
vLLM maps LOGICAL pages → PHYSICAL GPU memory pages
Non-contiguous GPU memory appears contiguous to LLM

How PagedAttention Works: Step-by-Step

Setup: Memory Organization

Physical GPU Memory: Divided into fixed-size blocks (512 tokens each)

- ┌─────────────────────────────────────────────────────────┐
 - GPU Memory (16GB = 16384 pages of 512 tokens each) │
 - ┬──────────┬──────────┬──────────┬──────────┬──┤
 - Page 1 │ Page 2 │ Page 3 │ Page 4 │ Page 5 │..│
 - [Empty] │ [Empty] │ [Empty] │ [Empty] │ [Empty] │..│
 - ┴──────────┴──────────┴──────────┴──────────┴──┘

Request 1: "Explain quantum computing" (256 tokens)

Logical View (Request 1 sees):
- ┌────────────────────────┐
 - Logical Pages │
 - ┬───────────┤
 - Page 1 │ Page 2 │
 - (tokens │ (tokens │
 - 0-511) │ 512-1023) │
 - ┴───────────┘
 ↓ Attention Computation

Physical Mapping:
Logical Page 1 → Physical Page 1
Logical Page 2 → [EMPTY - not yet allocated]

Request generates: 256 tokens (uses only part of Logical Page 1)

Request 2 Arrives: "What is AI?" (256 tokens, in parallel)

Request 1 (continuing): needs one more page
Request 2 (new request): needs pages

Logical View (Request 1):
 - ┬───────────┬───────────┐
 - Page 1 │ Page 2 │ Page 3 │
 - (full) │ (full) │ (partial) │
 - ┴───────────┴───────────┘

Logical View (Request 2):
 - ┬───────────┐
 - Page 1 │ Page 2 │
 - (partial) │ (partial) │
 - ┴───────────┘

Physical Mapping:
Request 1 Logical Page 1 → Physical Page 1 
Request 1 Logical Page 2 → Physical Page 3 (not contiguous!)
Request 1 Logical Page 3 → Physical Page 5 (skip Page 4)
Request 2 Logical Page 1 → Physical Page 2 
Request 2 Logical Page 2 → Physical Page 4 

Physical GPU Memory (after allocation):
- ┌──────────┬──────────┬──────────┬──────────┬──────────┐
 - Page 1 │ Page 2 │ Page 3 │ Page 4 │ Page 5 │
 - Req1 │ Req2 │ Req1 │ Req2 │ Req1 │
 - Tokens │ Tokens │ Tokens │ Tokens │ Tokens │
 - 0-511 │ 0-511 │ 512-1023 │ 512-1023 │ 1024+ │
 - ┴──────────┴──────────┴──────────┴──────────┘

Key Insight: Non-contiguous pages can be used efficiently!

Request 1 Completes

Request 1 is done generating tokens

Physical Memory State:
- ┌──────────┬──────────┬──────────┬──────────┬──────────┐
 - Page 1 │ Page 2 │ Page 3 │ Page 4 │ Page 5 │
 - [Free] │ Req2 │ [Free] │ Req2 │ [Free] │
 - ┴──────────┴──────────┴──────────┴──────────┘

Request 3 Arrives: "How does photosynthesis work?" (512 tokens)

Can allocate:
- Logical Page 1 → Physical Page 1 
- Logical Page 2 → Physical Page 3 (non-contiguous, but works!)

Requests 2 and 3 run CONCURRENTLY!

Final State:
- ┌──────────┬──────────┬──────────┬──────────┬──────────┐
 - Page 1 │ Page 2 │ Page 3 │ Page 4 │ Page 5 │
 - Req3 │ Req2 │ Req3 │ Req2 │ [Free] │
 - ┴──────────┴──────────┴──────────┴──────────┘

-

Technical Deep Dive

Memory Layout with PagedAttention

Before PagedAttention (Contiguous Allocation)

# Memory allocation for each request
class KVCache:
 def __init__(self, max_seq_len, hidden_dim):
 # Pre-allocate MAX sequence length for ALL requests
 self.key_cache = torch.zeros(
 (num_heads, max_seq_len, head_dim)
) # 4096 * hidden_dim = LARGE
 self.value_cache = torch.zeros(
 (num_heads, max_seq_len, head_dim)
) # 4096 * hidden_dim = LARGE

# Even if actual sequence is 512, memory for 4096 is allocated
# Waste

With PagedAttention (Logical-Physical Mapping)

# Page-based allocation
class PagedKVCache:
 def __init__(self, page_size=512):
 self.page_size = page_size

 # Pool of physical pages
 self.physical_pages = [
 Page(id=i, is_free=True)
 for i in range(total_pages)
]

 # Logical → Physical mapping for each request
 self.page_table = {} # request_id → [physical_page_ids]

 def allocate_page(self, request_id):
 """Allocate one physical page for request"""
 for page in self.physical_pages:
 if page.is_free:
 page.is_free = False
 self.page_table[request_id].append(page.id)
 return page
 raise OutOfMemoryError()

 def free_pages(self, request_id):
 """Free all pages for request"""
 for page_id in self.page_table[request_id]:
 self.physical_pages[page_id].is_free = True

Attention Computation with Paged KV

Traditional Attention (Contiguous)

# Standard attention
def attention(Q, K, V):
 """
 Q: query (seq_len, d_k)
 K: key (seq_len, d_k) - CONTIGUOUS in memory
 V: value (seq_len, d_k) - CONTIGUOUS in memory
 """
 scores = Q @ K.T / math.sqrt(d_k) # (seq_len, seq_len)
 weights = softmax(scores)
 output = weights @ V # (seq_len, d_v)
 return output

Paged Attention (Non-contiguous)

def paged_attention(Q, K_pages, V_pages, page_table):
 """
 Q: query (seq_len, d_k)
 K_pages: list of K page blocks (not contiguous)
 V_pages: list of V page blocks (not contiguous)
 page_table: logical → physical page mapping
 """
 # Step 1: Reconstruct full K, V from pages
 # (vLLM does this efficiently in CUDA kernel)
 K_full = reconstruct_from_pages(K_pages, page_table)
 V_full = reconstruct_from_pages(V_pages, page_table)

 # Step 2: Standard attention (same as before)
 scores = Q @ K_full.T / math.sqrt(d_k)
 weights = softmax(scores)
 output = weights @ V_full

 return output

Optimization: vLLM uses custom CUDA kernels to handle the page indirection efficiently, avoiding materialization of full K, V matrices.

Memory Efficiency Calculation

Example: 7B Model with 4 requests

Without PagedAttention (waste):

Hidden dim: 4096
Num heads: 32
Head dim: 128
KV per token: 2 * 32 * 128 = 8192 bytes = 8KB

Model: Llama 7B
KV cache per token per layer: ~8KB
Num layers: 32
Total KV per token: ~256KB

Max sequence: 4096 tokens
Pre-allocated per request: 256KB * 4096 = 1GB per request

4 concurrent requests: 4GB wasted on KV cache allocation!

With PagedAttention (efficient):

Page size: 512 tokens
KV cache per page: 256KB * 512 = 131MB per page
Num pages available: 16GB / 131MB ≈ 122 pages

4 concurrent requests:
- Request 1: 256 tokens = 1 page
- Request 2: 512 tokens = 1 page
- Request 3: 384 tokens = 1 page
- Request 4: 128 tokens = 1 page
Total: 4 pages = 524MB

Memory saved: 4GB → 524MB = 7.6x improvement!

-

Implementation Details

Page Table Management

class PageTableManager:
 def __init__(self, num_pages, num_layers):
 # Page table: [num_requests, num_layers, 2]
 # Last dimension: [key_page_ids, value_page_ids]
 self.page_table = {}
 self.page_pool = list(range(num_pages))
 self.free_pages = set(self.page_pool)

 def get_or_create_sequence(self, request_id):
 if request_id not in self.page_table:
 self.page_table[request_id] = {
 'key_pages': [],
 'value_pages': [],
 'seq_length': 0
 }
 return self.page_table[request_id]

 def append_token(self, request_id):
 """Allocate new page if needed when token is generated"""
 seq = self.get_or_create_sequence(request_id)
 seq['seq_length'] += 1

 # Check if we need new page
 if seq['seq_length'] % self.page_size == 1:
 # Need new page
 if not self.free_pages:
 raise MemoryError("No free pages!")

 page_id = self.free_pages.pop()
 seq['key_pages'].append(page_id)
 seq['value_pages'].append(page_id)

 def free_sequence(self, request_id):
 """Free all pages when request completes"""
 if request_id in self.page_table:
 seq = self.page_table[request_id]
 self.free_pages.update(seq['key_pages'])
 self.free_pages.update(seq['value_pages'])
 del self.page_table[request_id]

CUDA Kernel Optimization

The real magic is in the custom CUDA kernels that vLLM uses:

// Pseudo-code for efficient PagedAttention CUDA kernel
__global__ void paged_attention_kernel(
 float* queries, // (seq_len, num_heads, head_dim)
 float* kv_pages, // Physical page pool
 int* page_table, // Logical → physical mapping
 float* output // Results
) {
 // Each thread processes one query position
 int query_idx = threadIdx.x + blockIdx.x * blockSize;

 // Load query
 float query[HEAD_DIM];
 //... load from queries...

 // Compute attention WITHOUT materializing full K, V
 float score = 0.0f;
 for (int past_token = 0; past_token < seq_len; past_token++) {
 // Figure out which page this token is in
 int page_idx = past_token / PAGE_SIZE;
 int offset_in_page = past_token % PAGE_SIZE;

 // Get physical page
 int physical_page = page_table[page_idx];

 // Access K, V directly from physical page
 float* k_page = &kv_pages[physical_page * PAGE_SIZE];
 float key[HEAD_DIM];
 //... load key from page at offset...

 // Compute attention score for this token
 float attn_score = dot_product(query, key);
 score += attn_score;
 }

 output[query_idx] = score;
}

Key optimization: Attention computation happens directly on pages without reconstructing full K, V matrices!

-

Performance Impact

Throughput Improvement

Scenario: Serving requests with variable lengths

Batch of 4 requests:
- Request 1: 256 tokens → completes at t=2s
- Request 2: 512 tokens → completes at t=4s
- Request 3: 384 tokens → completes at t=3s
- Request 4: 128 tokens → completes at t=1s

Without PagedAttention (sequential):
 - Request 1: [────────] 2s
 - Request 2: [────────────────] 4s
 - Request 3: [──────────────] 3s
 - Request 4: [──] 1s
Total: 10 seconds → 0.4 req/sec

With PagedAttention (concurrent, continuous batching):
 - Request 1: [────────] free memory at 2s
 - Request 2: [────────────────]
 - Request 3: [──────────────]
 - Request 4: [──] free memory at 1s
Timeline:
t=0-1: All 4 requests running
t=1-2: Requests 1,2,3 running (Request 4 done)
t=2-3: Requests 2,3 running (Request 1 done)
t=3-4: Request 2 running (Request 3 done)
Total: 4 seconds → 1.0 req/sec (2.5x faster!)

With larger batch + more requests:
10-40x throughput improvement typical

Memory Efficiency

Model: Llama 2 7B
KV cache size: ~14GB at max length

Without PagedAttention:
2 concurrent requests × 14GB = 28GB required (most GPUs can't handle)

With PagedAttention (same hardware):
Can serve 4-6 concurrent requests = 4-6GB KV cache
Same GPU can now handle 28-42 requests sequentially!

Latency Impact

Time to generate first token (TTFT):

Request enters system
 - KV cache allocation:
 - Without paging: 100-200ms (contiguous allocation)
 - With paging: 1-2ms (just append to page table)
 - Model forward pass: 50ms (same)
 - Total TTFT:
 Without: ~150ms
 With: ~51ms (3x faster!)

---

Comparison with Alternatives

Memory Management Approaches

Approach	Memory Efficiency	Throughput	Complexity	Notes
No paging	Low (88% waste)	0.1-0.2 req/s	Simple	Baseline
Memory pool	Medium (40% waste)	0.3-0.5 req/s	Medium	Manual management
PagedAttention	High (5-10% waste)	1-4 req/s	Complex	State-of-art
Compression	Very High	0.5-1 req/s	Very Complex	Limited sequence length

Real Benchmarks (from paper)

Model: Llama 2 70B
Hardware: 2xA100 40GB GPUs
Max sequence length: 2048
Batch size: 1-32 requests

Throughput comparison:
- ┌────────────────────────┬──────────────┬─────────────┐
 - Method │ Tokens/sec │ Speedup │
 - ┼──────────────┼─────────────┤
 - Standard PyTorch │ 50-100 │ 1x │
 - TensorRT optimized │ 200-300 │ 2-3x │
 - PagedAttention (vLLM) │ 1000-2000 │ 10-20x │
 - ┴──────────────┴─────────────┘

Memory usage:
- ┌────────────────────────┬──────────────┬─────────────┐
 - Method │ GPU Memory │ Efficiency │
 - ┼──────────────┼─────────────┤
 - Standard PyTorch │ 78GB (peak) │ 28% │
 - PagedAttention (vLLM) │ 38GB (peak) │ 58% │
 - Savings │ 40GB saved │ 51% │
 - ┴──────────────┴─────────────┘

---

Advanced Topics

Copy-on-Write for Prompt Sharing

Many requests share common prompts. PagedAttention enables efficient sharing:

Prompt: "You are a helpful AI assistant. Answer:"
This 12-token prompt appears in 100 requests

Without PagedAttention:
- Prompt stored 100 times = 100 × 12KB = 1.2MB waste

With PagedAttention + Copy-on-Write:
- Prompt stored once in pages 1-2
- All 100 requests reference same pages
- Only copy pages when request appends new tokens
- Memory: 24KB (shared) + incremental growth

Speculative Decoding with Paging

PagedAttention enables efficient speculative decoding:

Speculative decoding: Draft model generates N tokens,
main model verifies them

With PagedAttention:
1. Draft generates 5 tokens → 1 page allocated
2. Main model verifies → reuses same page layout
3. If verification passes → keep pages
4. If verification fails → free pages and regenerate

Non-contiguous pages make this efficient!

-

Practical Implications

For Users

Better throughput: 10-40x improvement Lower costs: More requests served per unit time Lower latency: Faster time to first token Better hardware utilization: Serve more concurrent requests

For Deployment

Before PagedAttention:

8x A100 GPUs needed to serve 100 requests/sec
→ $50,000/month infrastructure cost

After PagedAttention:

2x A100 GPUs enough to serve 100 requests/sec
→ $12,500/month infrastructure cost (75% savings!)

For Model Size Limits

Before PagedAttention:

Single A100 (40GB): max 7B-13B models

After PagedAttention:

Single A100 (40GB): can serve 70B models with batching
(lower peak utilization, continuous batching keeps GPU fed)

-

Implementation in vLLM

How vLLM uses PagedAttention

# From vLLM source code (simplified)
class SequenceGroup:
 """Group of sequences sharing same prompt"""
 def __init__(self):
 self.sequences = [] # Multiple sequences/requests
 self.prompt_len = 0

 def get_kv_cache_pages(self):
 """Get allocated pages for this sequence group"""
 return self.page_table # Logical → physical mapping

class Scheduler:
 """Scheduler with PagedAttention"""
 def __init__(self):
 self.page_table_manager = PageTableManager()
 self.running_sequences = []

 def schedule(self):
 """Continuous batching with PagedAttention"""
 # Add new sequences if space available
 for new_seq in waiting_queue:
 if self.page_table_manager.can_allocate(new_seq.max_len):
 self.running_sequences.append(new_seq)

 # Generate one token for all running sequences
 # Pages are allocated on-demand (no pre-allocation!)
 for seq in self.running_sequences:
 output = self.model.forward(seq)
 self.page_table_manager.append_token(seq.id)

 # Remove completed sequences (immediately free pages!)
 self.running_sequences = [
 seq for seq in self.running_sequences
 if not seq.is_completed()
]

---

Limitations & Tradeoffs

Challenges

Challenge	Impact	Mitigation
Page size selection	Too large: fragmentation, too small: overhead	Dynamic page size (future work)
Page table overhead	Additional memory for mappings	Minimal (< 1% of KV cache)
CUDA kernel complexity	Harder to implement	Standard in modern frameworks
Compatibility	Not all models supported	Universal transformer support

Current Limitations

# Some advanced attention patterns not yet optimized
# Flash attention with PagedAttention (in progress)
# Sparse attention patterns (limited support)
# Multi-head latency hiding (future optimization)

---

Future Directions

Potential Improvements

1. Swap to CPU: Overflow KV cache to CPU memory
2. Compression: Lossy compression of old tokens
3. Recomputation: Trade memory for compute
4. Hierarchical paging: Multi-level page hierarchy
5. GPU-GPU paging: Efficient multi-GPU coordination

Research Opportunities

Open questions:
- Optimal page size for different model architectures
- Interaction with sparse attention patterns
- Multi-GPU coordination efficiency
- Energy efficiency implications

---

Key Takeaways

PagedAttention solves KV cache fragmentation by treating it like virtual memory 50% memory reduction compared to contiguous allocation 10-40x throughput improvement through continuous batching Non-contiguous pages are handled efficiently by optimized CUDA kernels Massive cost savings in production LLM serving Foundation for modern LLM inference (vLLM, HF TGI, etc.)

-

Further Reading

• Original Paper: Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention" (OSDI 2023)
• vLLM GitHub: github.com/vllm-project/vllm
• Hugging Face TGI: github.com/huggingface/text-generation-inference
• Related: KV cache quantization, speculative decoding, prompt paging

-

Related Notes

• 01 Vllm
• Llm Serving Architecture
• Memory Management In Ml
• 00 Readme