Solving the LLM Memory Wall: A Deep Dive into PagedAttention
Solving the LLM Memory Wall: A Deep Dive into PagedAttention
In the race to deploy Large Language Models (LLMs) at scale, the bottleneck is rarely just raw compute—it's memory. Specifically, the Key-Value (KV) cache, which stores the context of a conversation to avoid redundant computations, is a notorious memory hog.
Traditional memory management for LLMs is inefficient, leading to massive waste and limiting the number of concurrent users a server can handle. Enter PagedAttention, the breakthrough mechanism powering high-throughput engines like vLLM.
In this post, we will break down how PagedAttention borrows a classic concept from Operating Systems to revolutionize LLM serving.
The Problem: The "KV Cache" Memory Crisis
To understand PagedAttention, we first need to understand the inefficiency of standard KV caching.
In autoregressive generation, the model generates one token at a time. To calculate the attention for the next token, the model needs the Key ($K$) and Value ($V$) vectors of all previous tokens. To avoid re-calculating these every time, we store them in a KV Cache.
The Traditional Approach (Contiguous Allocation)
Historically, systems allocated a contiguous chunk of GPU memory for the KV cache based on the maximum possible sequence length. This leads to two types of waste:
- Internal Fragmentation: If you allocate space for 2048 tokens but the model only generates 100, the remaining 1948 slots are wasted.
- External Fragmentation: Because requests have varying lengths, memory becomes a "swiss cheese" of small gaps that are too small to be useful for new requests.
The Intuition: Virtual Memory for LLMs
PagedAttention solves this by applying the concept of Virtual Memory and Paging from OS design.
Instead of requiring a contiguous block of memory, PagedAttention decouples the logical sequence of tokens from their physical storage. It divides the KV cache into fixed-size blocks. These blocks can be scattered anywhere in the GPU memory, and a Block Table keeps track of where they are.
The Architecture at a Glance
How it Works: The Technical Breakdown
1. The Mathematical Foundation
At its core, PagedAttention still performs the standard Scaled Dot-Product Attention. Given a query $q_i$, we compute the attention output $o_i$ as:
$$a_{ij} = \frac{q_i k_j^T}{\sqrt{d}}, \quad o_i = \sum_{j=1}^{i} \text{softmax}(a_{ij}) v_j$$
The difference isn't in the math, but in how $k_j$ and $v_j$ are retrieved from memory.
2. The Algorithmic Steps
- KV Cache Partitioning: The cache is split into blocks (e.g., 4 or 16 tokens per block).
- Non-Contiguous Allocation: Blocks are allocated from a global pool on-demand.
- Logical-to-Physical Mapping: A
BlockManagermaintains a table mapping the request's logical sequence to physical block IDs. - Paged Computation: During the forward pass, the system "gathers" the physical blocks, flattens them into a temporary sequence, and computes attention.
- Dynamic Management: When a request ends, blocks are returned to the free pool. If multiple sequences share a prefix (like in Beam Search), they can point to the same physical block, drastically reducing memory usage.
Implementation: Simulating PagedAttention in PyTorch
Below is a production-style simulation of the PagedAttention mechanism.
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Dict, List
class BlockManager:
"""Simulates the OS-style page table mapping logical blocks to physical blocks."""
def __init__(self, num_blocks: int):
self.num_blocks = num_blocks
self.free_blocks = list(range(num_blocks))
self.block_table: Dict[int, List[int]] = {}
def allocate(self, request_id: int, num_blocks_needed: int) -> List[int]:
if len(self.free_blocks) < num_blocks_needed:
raise RuntimeError("Out of physical memory blocks!")
allocated = [self.free_blocks.pop(0) for _ in range(num_blocks_needed)]
self.block_table[request_id] = allocated
return allocated
def free(self, request_id: int):
blocks = self.block_table.pop(request_id, [])
self.free_blocks.extend(blocks)
class PagedKVCache(nn.Module):
"""Physical storage pool for KV caches."""
def __init__(self, num_blocks: int, block_size: int, num_heads: int, head_dim: int):
super().__init__()
self.register_buffer("k_cache", torch.zeros(num_blocks, block_size, num_heads, head_dim))
self.register_buffer("v_cache", torch.zeros(num_blocks, block_size, num_heads, head_dim))
def update(self, block_idx: int, token_idx_in_block: int, k: torch.Tensor, v: torch.Tensor):
self.k_cache[block_idx, token_idx_in_block] = k
self.v_cache[block_idx, token_idx_in_block] = v
class PagedAttention(nn.Module):
"""Custom attention mechanism that fetches KV caches from non-contiguous blocks."""
def __init__(self, num_heads: int, head_dim: int):
super().__init__()
self.num_heads = num_heads
self.head_dim = head_dim
def forward(self, q: torch.Tensor, block_indices: List[int], kv_cache: PagedKVCache, seq_len: int):
# 1. Gather physical blocks
k_blocks = kv_cache.k_cache[block_indices]
v_blocks = kv_cache.v_cache[block_indices]
# 2. Flatten and Trim to actual sequence length
k_seq = k_blocks.view(-1, self.num_heads, self.head_dim)[:seq_len]
v_seq = v_blocks.view(-1, self.num_heads, self.head_dim)[:seq_len]
# 3. Scaled Dot-Product Attention
q = q.unsqueeze(1) # [heads, 1, dim]
k_seq_t = k_seq.permute(1, 2, 0) # [heads, dim, seq_len]
attn_weights = torch.matmul(q, k_seq_t) / math.sqrt(self.head_dim)
attn_weights = F.softmax(attn_weights, dim=-1)
v_seq_t = v_seq.permute(1, 0, 2) # [heads, seq_len, dim]
out = torch.matmul(attn_weights, v_seq_t)
return out.squeeze(1)Performance Analysis
If we run the above code with a BLOCK_SIZE of 4 and a sequence length of 7, the system allocates 2 blocks (8 slots).
- Actual tokens: 7
- Wasted slots: 1
- Memory Efficiency: $\approx 87.5%$
Compare this to a traditional system that might pre-allocate 2048 slots for the same request, resulting in an efficiency of $\approx 0.3%$.
Summary: Why This Matters
PagedAttention is a game-changer for LLM deployment because it transforms memory management from a static allocation problem into a dynamic scheduling problem.
| Feature | Traditional KV Cache | PagedAttention |
|---|---|---|
| Allocation | Contiguous, Pre-allocated | Non-contiguous, On-demand |
| Fragmentation | High (Internal & External) | Near-Zero |
| Memory Sharing | Difficult/Impossible | Easy (via Block Table) |
| Throughput | Limited by Memory Waste | Maximized GPU Utilization |
By treating GPU memory like an operating system treats RAM, PagedAttention allows us to serve more users, handle longer contexts, and reduce the cost of running the world's most powerful AI models.