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

Deep Dive: 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.

Deep Dive: 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).

Deep Dive: 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.

Deep Dive: 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:

Deep Dive: 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.

Deep Dive: 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: