PagedAttention & KV-Cache Optimization: How vLLM Handles Large Scale Inference

Scenario context

Large Language Model (LLM) serving is heavily bound by GPU memory capacity due to the Key-Value (KV) cache. Traditional serving frameworks allocate contiguous memory based on maximum sequence length, leading to severe fragmentation and wasted space. This guide explores how PagedAttention, inspired by operating system virtual memory, partitions the KV cache into small logical blocks to enable dynamic, non-contiguous…

This walkthrough focuses on Concept: GPU Memory Bottlenecks in LLM Inference.

Before making a design choice, clarify:

• When running inference on large transformer models, the primary bottleneck is not computing capability, but memory bandwidth and memory capacity.
• This is particularly true during the autoregressive generation phase, where the model outputs tokens one by one.
• At each step of generation, the self-attention mechanism requires the Query vector of the current token to be multiplied by the Key vectors of all previous tokens in the sequence, followed by soft-max scaling and multip…

Concepts in play: caching, quantization-and-llm-inference, deep-learning-foundations, database-indexing

Stage 1 — Mechanics: Key-Value (KV) Cache and Memory Fragmentation

• To understand why contiguous allocations degrade serving throughput, we must look at how memory is utilized over the lifetime of an active generation query.
• When a prompt is first sent to the inference engine (the prefill phase ), the model processes all prompt tokens in parallel, generating the initial KV cache tensors.
• Once the prefill phase is complete, the model enters the decode phase , generating one token at a time.
• Each generated token adds a new vector to the Key cache and another vector to the Value cache.

Observable outcome: Each generated token adds a new vector to the Key cache and another vector to the Value cache.

Stage 2 — Deep Dive: PagedAttention Architecture and Block Mapping

• PagedAttention solves these fragmentation issues by decoupling logical sequence coordinates from physical memory locations.
• Instead of requesting a single contiguous memory block for a sequence, the engine divides the KV cache of each sequence into small, fixed-size logical blocks .
• Internals: Virtual Memory Analogy and Page Tables Similar to virtual memory in operating systems, PagedAttention represents a sequence's KV cache as a series of logical blocks.
• Each logical block contains the Key and Value tensors for a fixed number of tokens (typically 16 or 32).

Logical Space (Sequence):
[ Block 0 (Tokens 0-15) ] -> [ Block 1 (Tokens 16-31) ] -> [ Block 2 (Tokens 32-47) ]
      |                             |                             |
Block Table Mapping:
Logical Block 0 -------------> Physical Block 104 (HBM)
Logical Block 1 -------------> Physical Block 87  (HBM)
Logical Block 2 -------------> Physical Block 212 (HBM)

Observable outcome: Each logical block contains the Key and Value tensors for a fixed number of tokens (typically 16 or 32).

Stage 3 — Advanced Concepts: Copy-on-Write (CoW) and Prefix Sharing

• Beyond eliminating fragmentation, decoupling logical and physical memory blocks enables advanced memory-sharing patterns, such as Prefix Sharing and Parallel Sampling .
• When prompting an LLM with common system prompts, long instructions, or document contexts, multiple requests share the same initial prefix.
• In a contiguous allocation system, this prefix must be evaluated and stored separately for each request.
• With PagedAttention, the physical memory blocks storing the KV cache of the shared prefix can be mapped to the block tables of multiple requests simultaneously.

Sequence A Block Table: [ Block 10 (Prefix) ] -> [ Block 11 (Prefix) ] -> [ Block 45 (Private A) ]
Sequence B Block Table: [ Block 10 (Prefix) ] -> [ Block 11 (Prefix) ] -> [ Block 78 (Private B) ]

Observable outcome: With PagedAttention, the physical memory blocks storing the KV cache of the shared prefix can be mapped to the block tables of multiple requests simultaneously.

Stage 4 — Flow: Token Generation Lifecycle with PagedAttention

• The flowchart below traces how the inference engine processes incoming tokens, maps virtual indices using page lookup structures, and dynamically assigns GPU physical memory during execution:

flowchart TD
    Start[New Token Request Received] --> Prefill[Generate Initial KV Cache for Prompt]
    Prefill --> MapBlocks[Divide KV Cache into Blocks of Size B]
    MapBlocks --> Allocate[Allocate Physical Blocks from Free Pool]
    Allocate --> MapTable[Update Block Table: Logical -> Physical]
    MapTable --> Decode[Generate Next Token]
    Decode --> Check{Exceeds Current Block Space?}
    Check -->|No| Store[Store KV Tensors in Current Physical Block]
    Check -->|Yes| RequestBlock[Request New Block from Free Pool]
    RequestBlock --> AllocSuccess{Free Block Available?}
    AllocSuccess -->|Yes| Assign[Assign Block & Update Block Table]
    Assign --> Store
    AllocSuccess -->|No| Evict[Preempt/Swap Out Lowest Priority KV Cache]
    Evict --> Assign
    Store --> NextToken{End of Sequence?}
    NextToken -->|No| Decode
    NextToken -->|Yes| Free[Release Physical Blocks to Free Pool]

Observable outcome: The flowchart below traces how the inference engine processes incoming tokens, maps virtual indices using page lookup structures, and dynamically assigns GPU physical memory during execution:

Stage 5 — Applications: Production Inference Engines and Model Serving

• Enterprise LLM APIs: Frameworks like vLLM, TensorRT-LLM, and TGI use PagedAttention to serve thousands of concurrent users efficiently.
• Multi-Agent Systems: Agent frameworks generate multiple reasoning paths in parallel.
• Shared prefix mapping speeds up these parallel runs.
• Retrieval-Augmented Generation (RAG): When injecting large context files, shared document caches reduce GPU memory usage.

Observable outcome: Retrieval-Augmented Generation (RAG): When injecting large context files, shared document caches reduce GPU memory usage.

Stage 6 — Practical Implementation: Simulated Block Allocation Lookup Table in Python

• Below is a complete, runnable Python script simulating a virtual page manager that allocates token KV tensors to logical blocks, updates lookup tables, and maps virtual sequence addresses:

import math
from typing import Dict, List, Optional

class KVBlockManager:
    def __init__(self, block_size: int, total_physical_blocks: int):
        self.block_size = block_size
        self.total_blocks = total_physical_blocks
        self.free_blocks: List[int] = list(range(total_physical_blocks))

        # Maps sequence_id -> List of physical block IDs
        self.block_tables: Dict[str, List[int]] = {}
        # Maps sequence_id -> current token count
        self.sequence_token_counts: Dict[str, int] = {}

    def register_sequence(self, seq_id: str, prompt_len: int) -> bool:
        if seq_id in self.block_tables:
            raise ValueError(f"Sequence {seq_id} is already registered.")

        needed_blocks = math.ceil(prompt_len / self.block_size)
        if len(self.free_blocks) < needed_blocks:
            print(f"Memory allocation failed: Not enough free blocks for prompt. Required: {needed_blocks}, Free: {len(self.free_blocks)}")
            return False

        allocated = []
        for _ in range(needed_blocks):
            allocated.append(self.free_blocks.pop(0))

        self.block_tables[seq_id] = allocated
        self.sequence_token_counts[seq_id] = prompt_len
        print(f"Registered sequence {seq_id} ({prompt_len} tokens). Allocated blocks: {allocated}")
        return True

    def append_token(self, seq_id: str) -> bool:
        if seq_id not in self.block_tables:
            raise ValueError(f"Sequence {seq_id} is not registered.")

        current_tokens = self.sequence_token_counts[seq_id]
        current_blocks = self.block_tables[seq_id]

        # Check if we need to allocate a new physical block
        if current_tokens % self.block_size == 0:
            if not self.free_blocks:
                print(f"Memory Full! Cannot allocate block for sequence {seq_id}. Triggering preemption...")
                return False

            new_block = self.free_blocks.pop(0)
            current_blocks.append(new_block)
            print(f"Allocated new block {new_block} for sequence {seq_id} at token {current_tokens}")

        self.sequence_token_counts[seq_id] = current_tokens + 1
        return True

    def free_sequence(self, seq_id: str):
        if seq_id not in self.block_tables:
            return

        released_blocks = self.block_tables.pop(seq_id)
        self.sequence_token_counts.pop(seq_id)
        self.free_blocks.extend(released_blocks)
        # Keep physical blocks sorted to minimize memory fragmentation
        self.free_blocks.sort()
        print(f"Released blocks {released_blocks} from sequence {seq_id}. Free blocks remaining: {len(self.free_blocks)}")

    def get_physical_address(self, seq_id: str, token_index: int) -> Optional[tuple]:
        if seq_id not in self.block_tables:
            return None

        if token_index >= self.sequence_token_counts[seq_id]:
            return None

        logical_block_index = token_index // self.block_size
        block_offset = token_index % self.block_size

        physical_block_id = self.block_tables[seq_id][logical_block_index]
        return (physical_block_id, block_offset)

def run_paged_attention_simulation():
    print("--- Starting PagedAttention Simulator ---")
    # Initialize manager: 16 tokens per block, 10 physical blocks total
    manager = KVBlockManager(block_size=16, total_physical_blocks=10)
    print(f"Initial Free Blocks: {manager.free_blocks}\n")

    # 1. Register two sequences
    seq_a = "req_001"
    seq_b = "req_002"

    # Prompt A has 24 tokens (needs 2 blocks)
    manager.register_sequence(seq_a, prompt_len=24)
    # Prompt B has 35 tokens (needs 3 blocks)
    manager.register_sequence(seq_b, prompt_len=35)
    print(f"Free Blocks remaining: {len(manager.free_blocks)}\n")

    # 2. Map logical token index to physical memory block coordinates
    target_token = 20
    address = manager.get_physical_address(seq_a, target_token)
    if address:
        block_id, offset = address
        print(f"Token {target_token} of {seq_a} is mapped to Block {block_id} at offset {offset}.")

    # 3. Simulate generation steps (appending tokens)
    print("\n--- Simulating Decode Iterations ---")
    # Generate 10 tokens for Sequence A (from 24 to 34 tokens)
    for _ in range(10):
        if not manager.append_token(seq_a):
            break

    # Trace mapping after generation
    new_target_token = 33
    address = manager.get_physical_address(seq_a, new_target_token)
    if address:
        block_id, offset = address
        print(f"After generation: Token {new_target_token} of {seq_a} mapped to Block {block_id} at offset {offset}.")

    # 4. Clean up sequences
    print("\n--- Cleaning Up Resources ---")
    manager.free_sequence(seq_a)
    manager.free_sequence(seq_b)
    print(f"Final Free Blocks: {manager.free_blocks}")

if __name__ == "__main__":
    run_paged_attention_simulation()

Observable outcome: Below is a complete, runnable Python script simulating a virtual page manager that allocates token KV tensors to logical blocks, updates lookup tables, and maps virtual sequence addresses:

End-to-end outcome

After working through the stages, you should be able to explain:

• Tune the Block Size Carefully: Small blocks (e.g., 8 tokens) reduce fragmentation but increase block lookup overhead.
• Large blocks (e.g., 64 tokens) reduce lookup latency but increase internal fragmentation.
• Memory Bottleneck: LLM generation throughput is bound by GPU memory capacity, driven by the linear growth of the Key-Value (KV) cache.
• PagedAttention Solution: Divides the KV cache into fixed-size blocks (16 or 32 tokens), mapping logical blocks to non-contiguous physical memory via page tables.

Alternatives considered

Decision Guide: Custom Memory Allocations vs. Page Tables

The table below contrasts PagedAttention against other memory management techniques for LLM serving:

Failure and degraded modes

Trade-offs and Failure Modes

• Detection: Block Lookup Overhead: Accessing non-contiguous memory blocks in HBM introduces minor latency penalties compared to reading flat contiguous memory.
• User or operator impact: Block Lookup Overhead: Accessing non-contiguous memory blocks in HBM introduces minor latency penalties compared to reading flat contiguous memory.