Pytorch implementation of paper: Efficient Memory Management for Large Language Model Serving with PagedAttention (2023)

The 2023 paper “Efficient Memory Management for Large Language Model Serving with PagedAttention” by Kwon et al. (UC Berkeley) introduced vLLM, a high-throughput serving engine.

The core problem it solves is KV Cache fragmentation. In standard LLM serving, memory for the Key-Value (KV) cache is allocated contiguously for the maximum possible sequence length. This leads to “Internal Fragmentation” (reserved but unused space) and “External Fragmentation” (scattered small free blocks). PagedAttention treats GPU memory like Virtual Memory in an Operating System, breaking the KV cache into non-contiguous “pages.”

---

1. The PagedAttention Logic

Instead of a single large tensor for a request, PagedAttention stores the KV cache in Physical Blocks. A Block Table maps the logical position of a token in a sequence to a physical slot in a GPU memory block.

$$KV_{logical} \rightarrow \text{BlockTable} \rightarrow KV_{physical}$$

2. PyTorch Implementation: The Kernel Logic

While the actual high-performance kernel is written in CUDA, we can implement the Logical-to-Physical mapping in PyTorch to understand how the memory is managed.

Python

import torch
import torch.nn as nn

class PagedAttentionExecutor(nn.Module):
    def __init__(self, num_heads, head_dim, block_size=16):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.block_size = block_size # Number of tokens per block
        self.scale = head_dim ** -0.5

    def forward(self, query, k_cache, v_cache, block_table, context_lens):
        """
        query: [num_queries, num_heads, head_dim]
        k_cache: [num_blocks, num_heads, head_dim, block_size]
        v_cache: [num_blocks, num_heads, block_size, head_dim]
        block_table: [num_queries, max_num_blocks_per_query]
        context_lens: [num_queries]
        """
        num_queries = query.shape[0]
        output = torch.zeros_like(query)

        for i in range(num_queries):
            q = query[i] # [num_heads, head_dim]
            cur_context_len = context_lens[i]
            cur_block_table = block_table[i]
            
            # 1. Fetch relevant blocks from the scattered physical cache
            # This avoids having one giant contiguous tensor per request
            num_blocks = (cur_context_len + self.block_size - 1) // self.block_size
            
            # Logic: Collect all K and V tensors from the mapped blocks
            keys = []
            values = []
            for j in range(num_blocks):
                block_idx = cur_block_table[j]
                keys.append(k_cache[block_idx])   # [num_heads, head_dim, block_size]
                values.append(v_cache[block_idx]) # [num_heads, block_size, head_dim]
            
            # Concatenate blocks into a temporary view for attention calculation
            K = torch.cat(keys, dim=-1)[..., :cur_context_len] # [num_heads, head_dim, context_len]
            V = torch.cat(values, dim=-2)[:, :cur_context_len, :] # [num_heads, context_len, head_dim]

            # 2. Standard Attention Calculation
            # energy = (Q @ K) * scale
            attn_weights = torch.matmul(q.unsqueeze(1), K).squeeze(1) * self.scale
            attn_weights = torch.softmax(attn_weights, dim=-1) # [num_heads, context_len]
            
            # context = weights @ V
            out = torch.matmul(attn_weights.unsqueeze(1), V).squeeze(1)
            output[i] = out
            
        return output

---

3. Key Component: The Block Manager

The “secret sauce” is the CPU-side Block Manager, which tracks which physical blocks on the GPU are free, occupied, or shared (for Copy-on-Write).

Python

class BlockManager:
    def __init__(self, num_total_blocks):
        self.free_blocks = list(range(num_total_blocks))
        self.active_requests = {} # request_id -> list of block_ids

    def allocate(self, request_id, num_tokens):
        num_blocks_needed = (num_tokens + 15) // 16
        allocated = []
        for _ in range(num_blocks_needed):
            allocated.append(self.free_blocks.pop(0))
        self.active_requests[request_id] = allocated
        return allocated

    def free(self, request_id):
        blocks = self.active_requests.pop(request_id)
        self.free_blocks.extend(blocks)

---

4. Technical Specifications & Impact

Metric	Contiguous Cache	PagedAttention (vLLM)
Memory Waste	60% – 80% (Reserved)	< 4% (Only last block)
Throughput	Bottlenecked by VRAM	2x – 4x Increase
Decoding	Sequential	Highly Parallel
Complex Sampling	Difficult (Beam Search)	Efficient (Copy-on-Write)

Why it is a 2023 Milestone

PagedAttention enabled Near-Zero Memory Waste. By allowing the KV cache to be fragmented physically but contiguous logically, it allowed servers to pack significantly more requests into the same GPU. This is the technology that powers current high-speed inference APIs and allows for efficient “Beam Search” and “Parallel Sampling” by simply sharing physical blocks between different candidate sequences.