PagedAttention: Detailed Technical Guide

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: 4096 - 512 = 3584 tokens per request

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: Q @ K^T / sqrt(d_k)
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