PagedAttention and KV Cache Management: The Architecture Behind High-Throughput LLM Serving

Serving Large Language Models (LLMs) in production is fundamentally a battle against memory bandwidth and memory capacity. While training is dominated by raw compute (FLOPS), autoregressive inference is severely constrained by the rate at which GPUs can stream data from High Bandwidth Memory (HBM) to compute cores.

At the center of this bottleneck lies the Key-Value (KV) Cache. During the generation (decoding) phase, every token generated requires reading the cached key and value vectors of all preceding tokens in the sequence. In naive serving systems, managing this dynamic cache causes 60% to 80% of GPU memory to be wasted due to internal and external memory fragmentation, severely capping concurrent request throughput.

The breakthrough that revolutionized high-concurrency LLM serving is PagedAttention—pioneered by the UC Berkeley team behind vLLM. By applying classic Operating System virtual memory paging principles to GPU memory management, PagedAttention slashes memory waste to near zero (<4%), enables seamless Copy-on-Write prefix caching, and unlocks 2x to 4x throughput improvements across production clusters.

In this comprehensive systems guide, we dissect the mathematical and architectural foundations of the KV Cache, analyze why traditional allocation fails, explore PagedAttention block tables and CUDA kernel mechanics, and implement a full-featured PagedAttention block allocator in Python.

The Physics of LLM Inference: Prefill vs. Decode Phase

To understand why KV cache management dictates inference throughput, we must separate an LLM request into its two distinct operational phases:

The KV Cache Size Equation

In standard multi-head self-attention, for every token in a sequence, we store its Key and Value vector representations across every layer of the transformer. The exact memory footprint (in bytes) of a single request's KV Cache is governed by the formula:

Memory_KV (bytes) = 2 × n_layers × n_kv_heads × d_head × seq_length × bytes_per_element

Where:

• $2$: Accounts for both the Key and Value tensors.
• $n_{\text{layers}}$: Number of transformer layers (e.g., 32 in LLaMA-3-8B, 80 in LLaMA-3-70B).
• $n_{\text{kv\_heads}}$: Number of Key-Value attention heads (in Grouped-Query Attention, $n_{\text{kv\_heads}} \ll n_{\text{q\_heads}}$, typically 8).
• $d_{\text{head}}$: Dimension of each attention head (typically 128).
• $\text{seq\_length}$: Total number of tokens (prompt + generated tokens).
• $\text{bytes\_per\_element}$: 2 bytes for FP16/BF16, 1 byte for FP8 / INT8.

Real-World Memory Sizing Example (LLaMA-3-70B)

Let's calculate the KV Cache requirement for LLaMA-3-70B ($n_{\text{layers}} = 80$, $n_{\text{kv\_heads}} = 8$, $d_{\text{head}} = 128$, in 16-bit precision):

Per-token KV Cache size = 2 × 80 × 8 × 128 × 2 bytes = 327,680 bytes ≈ 320 KB / token

If we serve a batch of 64 concurrent user requests, each with an average context length of 4,096 tokens:

Total KV Cache = 64 requests × 4,096 tokens × 320 KB = 83.88 GB

For context, an 80GB NVIDIA A100 / H100 GPU hosting a quantized 70B model has only ~40GB of free HBM remaining after loading the model weights. The KV Cache quickly becomes the dominant consumer of GPU memory, directly dictating how many concurrent users an inference node can serve.

The Memory Fragmentation Crisis in Traditional Serving

Before PagedAttention, mainstream serving frameworks (such as Hugging Face Accelerate, naive PyTorch, and early TensorRT engines) stored the KV Cache of each sequence in contiguous GPU memory buffers. Because the engine cannot predict how many tokens a model will generate before encountering an <EOS> (end-of-sequence) token, it allocated memory based on the worst-case scenario:

flowchart TD
    subgraph Traditional Allocator: Contiguous Pre-allocation
        R1[Request 1: Reserved 4096 Tokens] --> U1[Used: 250 Tokens]
        R1 --> W1[Wasted / Reserved Space: 3846 Tokens]
        R2[Request 2: Reserved 4096 Tokens] --> U2[Used: 120 Tokens]
        R2 --> W2[Wasted / Reserved Space: 3976 Tokens]
        R3[Request 3: Reserved 4096 Tokens] --> U3[Used: 500 Tokens]
        R3 --> W3[Wasted / Reserved Space: 3596 Tokens]
    end
    style W1 fill:#e11d48,stroke:#fda4af,color:#fff
    style W2 fill:#e11d48,stroke:#fda4af,color:#fff
    style W3 fill:#e11d48,stroke:#fda4af,color:#fff
    style U1 fill:#059669,stroke:#6ee7b7,color:#fff
    style U2 fill:#059669,stroke:#6ee7b7,color:#fff
    style U3 fill:#059669,stroke:#6ee7b7,color:#fff

The Three Forms of Memory Waste

1. Internal Fragmentation (Over-Reservation): To avoid expensive GPU re-allocations and memory copies during token generation, the server statically reserves a contiguous chunk for the maximum potential length (e.g., max_seq_len = 4,096). If the user prompt is 200 tokens and the answer is 100 tokens, 92.6% of the allocated memory sits completely idle.
2. External Fragmentation: Over time, as requests of arbitrary lengths start and finish, contiguous free memory becomes chopped into small, disconnected segments. Even if the GPU has 10GB of total free memory, it cannot admit a new request requiring a 4GB contiguous buffer because no single continuous slice exists.
3. Reservation Waste (Future Token Padding): Systems that grow buffers dynamically must reserve headroom for the next batch of tokens. Memory allocated for slot $N+100$ cannot be utilized by any other concurrent request while waiting for slots $1 \dots N$ to be generated.

In empirical benchmarks across real-world workloads, only 20% to 40% of the allocated KV cache memory held actual token data. The remaining 60% to 80% was pure waste, starving the GPU and throttling batch concurrency.

PagedAttention: OS Virtual Memory Architecture for GPUs

PagedAttention solves fragmentation by borrowing the oldest and most proven concept in computer architecture: Virtual Memory Paging.

Instead of demanding contiguous physical memory for each request, PagedAttention partitions the sequence's KV cache into fixed-size Logical Blocks (typically holding $B = 16$ or $32$ tokens). The serving engine maintains a global pool of Physical Blocks in GPU HBM and manages a per-request Block Table (analogous to an OS Page Table).

flowchart LR
    subgraph Logical Space (Request View)
        L0["Logical Block 0 (Tokens 0-15)"]
        L1["Logical Block 1 (Tokens 16-31)"]
        L2["Logical Block 2 (Tokens 32-47)"]
    end

    subgraph Block Table (Page Table Mapping)
        BT["Block Table
L0 -> Physical #7
L1 -> Physical #2
L2 -> Physical #19"]
    end

    subgraph Physical GPU HBM (Non-contiguous)
        P2["Physical Block #2 (Tokens 16-31)"]
        P7["Physical Block #7 (Tokens 0-15)"]
        P19["Physical Block #19 (Tokens 32-47)"]
        PFree["Free Physical Blocks Pool..."]
    end

    L0 --> BT
    L1 --> BT
    L2 --> BT
    BT --> P7
    BT --> P2
    BT --> P19
    
    style P7 fill:#0284c7,stroke:#38bdf8,color:#fff
    style P2 fill:#0284c7,stroke:#38bdf8,color:#fff
    style P19 fill:#0284c7,stroke:#38bdf8,color:#fff
    style BT fill:#7c3aed,stroke:#c084fc,color:#fff

How Block Paging Operates During Generation

1. On-Demand Allocation: When a request starts, the engine only allocates physical blocks necessary for the initial prompt (e.g., a 25-token prompt with $B=16$ allocates exactly 2 physical blocks).
2. Boundary Allocation: As the model generates tokens, it writes into the currently active block. Only when block boundary $B$ is crossed does the block manager pop a new physical block from the free list and record its index in the request's block table.
3. Zero External Fragmentation: Because all physical blocks share identical byte sizes ($B \times 2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}} \times \text{dtype\_size}$), any free block can satisfy any request. External fragmentation is mathematically eliminated.
4. Minimal Internal Fragmentation: Memory waste is strictly restricted to the unwritten slots in the final block of each sequence (at most $B - 1$ tokens, typically < 3% waste).

Under the Hood: The PagedAttention CUDA Kernel

In standard FlashAttention or PyTorch SDPA, the query vector $q_i$ computes dot products with key vectors $K$ stored in a contiguous array in GPU shared memory. With PagedAttention, keys and values are scattered across non-contiguous physical pages in GPU global memory.

To execute attention without copying pages into a contiguous staging area (which would obliterate performance), the PagedAttention CUDA kernel accepts the Block Table directly as a kernel argument:

Within the CUDA thread block, each warp computes attention over a slice of logical blocks. By looking up physical_block_id = block_table[seq_id, logical_block_id], the thread loads the $K$ and $V$ vectors directly into registers/SRAM using hardware memory coalescing, computes intermediate Softmax reductions, and outputs the final attention context vector.

Advanced Memory Sharing: Prefix Caching & Copy-on-Write

Beyond solving fragmentation for independent requests, PagedAttention unlocks a monumental performance optimization: Complex KV Cache Sharing.

1. Shared System Prompts & Prefix Caching

In production LLM applications (customer support bots, coding assistants, RAG pipelines), thousands of distinct user requests share identical system instructions (often 1,000 to 4,000 tokens long). With PagedAttention, the engine stores the system prompt once in physical blocks and maps all incoming requests' block tables to the same physical blocks:

flowchart TD
    Sys["System Prompt (Physical Blocks #1, #2) RefCount = 2"]
    
    ReqA["Request A Block Table
Block 0 -> #1
Block 1 -> #2
Block 2 -> #8 (User Prompt A)"]
    ReqB["Request B Block Table
Block 0 -> #1
Block 1 -> #2
Block 2 -> #14 (User Prompt B)"]

    ReqA --> Sys
    ReqB --> Sys
    
    style Sys fill:#059669,stroke:#6ee7b7,color:#fff
    style ReqA fill:#0284c7,stroke:#38bdf8,color:#fff
    style ReqB fill:#0284c7,stroke:#38bdf8,color:#fff

2. Parallel Sampling & Beam Search via Copy-on-Write (CoW)

When generating $N$ candidate responses for a single prompt (e.g., n=4 in code generation or tree-of-thought exploration), traditional engines duplicate the prompt's KV cache 4 times. PagedAttention forks the child requests with zero memory copying by simply pointing their block tables to the parent's physical blocks and incrementing the block Reference Count ($ref\_count$).

When one branch generates a unique token that modifies a shared block, the engine executes a Copy-on-Write (CoW): it allocates a fresh physical block, copies only the $B$ tokens from the shared block, decrements the shared block's $ref\_count$, and updates the child's block table.

Continuous Batching vs. Static Batching

Traditional deep learning batching is Static: $N$ requests are grouped together, and the batch runs until the single longest request finishes. Shorter requests that finish early output padding tokens, wasting GPU compute cycles.

PagedAttention pairs seamlessly with Continuous Batching (Iteration-Level Scheduling) (first introduced in Orca):

• Iteration-Level Dispatch: At each forward decode step ($t$), the scheduler inspects the running batch. Completed sequences are immediately evicted, their physical blocks returned to the free pool.
• Dynamic Request Insertion: Waiting requests from the queue have their prefill step injected into the running iteration without waiting for existing requests to complete.
• Preemption & GPU Swapping: If a sudden surge of long-output requests causes GPU HBM exhaustion, the block manager can selectively preempt low-priority sequences by swapping their physical blocks to CPU DRAM or discarding them for re-computation later.

Hands-On Implementation: Building a PagedAttention Block Allocator in Python

To master how block tables, reference counting, and Copy-on-Write work in production runtimes like vLLM, let's build a complete, runnable PagedAttention Block Manager in Python:

import math
from typing import List, Dict, Optional

class PhysicalBlock:
    """Represents a physical memory block in GPU HBM."""
    def __init__(self, block_id: int, block_size: int = 16):
        self.block_id = block_id
        self.block_size = block_size
        self.ref_count = 0
        self.tokens: List[int] = []

    @property
    def is_full(self) -> bool:
        return len(self.tokens) >= self.block_size

    def append_token(self, token_id: int):
        if self.is_full:
            raise RuntimeError(f"PhysicalBlock {self.block_id} is full!")
        self.tokens.append(token_id)

    def reset(self):
        self.ref_count = 0
        self.tokens.clear()

class BlockManager:
    """Manages the pool of physical GPU blocks and page-table mappings."""
    def __init__(self, num_blocks: int, block_size: int = 16):
        self.block_size = block_size
        self.num_blocks = num_blocks
        self.free_blocks: List[PhysicalBlock] = [
            PhysicalBlock(i, block_size) for i in range(num_blocks)
        ]
        self.all_blocks: Dict[int, PhysicalBlock] = {
            b.block_id: b for b in self.free_blocks
        }
        # Maps request_id -> List of PhysicalBlock IDs (Block Table)
        self.block_tables: Dict[str, List[int]] = {}

    def allocate_block(self) -> PhysicalBlock:
        if not self.free_blocks:
            raise MemoryError("GPU Out of Memory: No free physical blocks available!")
        block = self.free_blocks.pop(0)
        block.ref_count = 1
        return block

    def free_block(self, block_id: int):
        block = self.all_blocks[block_id]
        block.ref_count -= 1
        if block.ref_count == 0:
            block.reset()
            self.free_blocks.append(block)

    def allocate_sequence(self, request_id: str, prompt_tokens: List[int]):
        """Allocates initial logical and physical blocks for a new prompt."""
        num_blocks = math.ceil(len(prompt_tokens) / self.block_size)
        table: List[int] = []

        for i in range(num_blocks):
            block = self.allocate_block()
            start_idx = i * self.block_size
            end_idx = min((i + 1) * self.block_size, len(prompt_tokens))
            for tok in prompt_tokens[start_idx:end_idx]:
                block.append_token(tok)
            table.append(block.block_id)

        self.block_tables[request_id] = table

    def append_token_to_seq(self, request_id: str, token_id: int):
        """Appends a new generated token during the decode phase."""
        table = self.block_tables[request_id]
        last_block_id = table[-1]
        last_block = self.all_blocks[last_block_id]

        if last_block.ref_count > 1:
            # Copy-on-Write: If block is shared with other requests, fork it
            new_block = self.allocate_block()
            new_block.tokens = list(last_block.tokens)
            last_block.ref_count -= 1
            table[-1] = new_block.block_id
            last_block = new_block

        if last_block.is_full:
            # Allocate a new physical block when crossing page boundary
            new_block = self.allocate_block()
            new_block.append_token(token_id)
            table.append(new_block.block_id)
        else:
            last_block.append_token(token_id)

    def fork_sequence(self, src_req_id: str, dst_req_id: str):
        """Forks a sequence for parallel sampling using Copy-on-Write."""
        src_table = self.block_tables[src_req_id]
        dst_table = []
        for block_id in src_table:
            block = self.all_blocks[block_id]
            block.ref_count += 1
            dst_table.append(block_id)
        self.block_tables[dst_req_id] = dst_table

    def free_sequence(self, request_id: str):
        """Frees all physical blocks associated with a completed request."""
        if request_id not in self.block_tables:
            return
        for block_id in self.block_tables[request_id]:
            self.free_block(block_id)
        del self.block_tables[request_id]

# --- Simulation Demo ---
if __name__ == "__main__":
    manager = BlockManager(num_blocks=10, block_size=4)
    print(f"Total Physical GPU Blocks: {manager.num_blocks} (Capacity: {manager.num_blocks * 4} tokens)")

    # 1. Allocate Request A (Prompt of 6 tokens -> requires 2 blocks)
    prompt_a = [101, 2045, 1037, 2795, 102, 300]
    manager.allocate_sequence("req_A", prompt_a)
    print(f"\n[Req A Allocated] Block Table: {manager.block_tables['req_A']}")
    print(f"Free Blocks Remaining: {len(manager.free_blocks)}")

    # 2. Fork Request A into Request B (Parallel Sampling / CoW)
    manager.fork_sequence("req_A", "req_B")
    print(f"[Req B Forked] Block Table: {manager.block_tables['req_B']}")
    print(f"Free Blocks Remaining (Zero-Copy): {len(manager.free_blocks)}")

    # 3. Generate divergent tokens for Req A and Req B
    # Appending to Req A triggers CoW on Block #1 (the partially full block)
    manager.append_token_to_seq("req_A", 9999)
    print(f"\n[Req A Generated Token 9999] Block Table: {manager.block_tables['req_A']}")
    print(f"Req B Block Table after A's CoW: {manager.block_tables['req_B']}")
    print(f"Free Blocks Remaining: {len(manager.free_blocks)}")

    # 4. Clean up
    manager.free_sequence("req_A")
    manager.free_sequence("req_B")
    print(f"\n[Cleaned Up] Free Blocks Restored: {len(manager.free_blocks)} / {manager.num_blocks}")

Execution Output of the Allocator Simulation

Total Physical GPU Blocks: 10 (Capacity: 40 tokens)

[Req A Allocated] Block Table: [0, 1]
Free Blocks Remaining: 8

[Req B Forked] Block Table: [0, 1]
Free Blocks Remaining (Zero-Copy): 8

[Req A Generated Token 9999] Block Table: [0, 2]
Req B Block Table after A's CoW: [0, 1]
Free Blocks Remaining: 7

[Cleaned Up] Free Blocks Restored: 10 / 10

Notice how forking req_B from req_A consumed 0 additional physical blocks. When req_A added token 9999, Copy-on-Write transparently duplicated only the active partially full block (Block 1 $\rightarrow$ Block 2) while continuing to share Block 0 without overhead!

Production Comparison: Serving Architecture Ecosystem

Key Architectural Takeaways

• LLM Inference is Memory Bound: Generation throughput is constrained by GPU HBM memory bandwidth and capacity, not raw compute.
• Fragmentation Was the Invisible Killer: Static contiguous allocation wasted upwards of 80% of GPU memory on unutilized buffer padding.
• Paging Brings Order to GPU Memory: PagedAttention breaks KV caches into fixed-size logical blocks mapped dynamically to physical pages, cutting memory waste below 4%.
• Zero-Cost Sharing: By treating memory blocks as reference-counted pages, prefix caching and branching operations (parallel sampling, speculative decoding, multi-turn chat) share GPU memory with zero duplication until write divergence.