~/blog

PagedAttention - Explained

Jul 19, 202615 min readBy Mohammed Vasim
PagedAttentionvLLMLLM InferenceGPU MemoryKV Cache

Introduction

If you have ever tried deploying Large Language Models (LLMs) in a high-throughput production environment, you know the harsh reality: training an LLM is a compute-bound problem, but serving an LLM is a memory-bound nightmare.

When training, we maximize hardware efficiency by stuffing dense, static matrices into giant batches. But during inference, requests arrive unpredictably, sequence lengths fluctuate wildly, and the engine must output tokens one by one in an autoregressive fashion.

The hidden bottleneck that chokes production serving systems is the Key-Value (KV) Cache. To avoid recomputing past context at every single token step, engines store historical keys and values in high-bandwidth memory (HBM). However, the traditional way engines managed this cache forced contiguous allocation on the GPU, creating staggering levels of memory waste. This memory fragmentation meant GPUs ran out of memory long before their compute engines reached full utilization.

In this architectural deep dive, we will move past the high-level hand-waving of the original vLLM paper. We will build an intuitive understanding of why traditional systems choke, explore how virtual memory concepts inspired a modern serving paradigm, walk through the actual layout of the CUDA kernels executing PagedAttention, and look at how this technology fits into today's modern inference landscape.


The Transformer Recap (Only What We Need)

To understand why the KV Cache dominates memory, let's trace exactly what happens inside a standard Multi-Head Attention (MHA) block during generation.

An input sequence of tokens is projected into three distinct matrices per layer: Queries (), Keys (), and Values (). The mathematical formulation of scaled dot-product attention is defined as:

During autoregressive generation, the model generates tokens one by one. To predict token , the query vector must attend to the key vectors of all previous tokens , and multiply the resulting weights against all previous value vectors .

Without caching, generating token would require re-running the entire attention projection for tokens to . This scales quadratically () in compute cost per generation step.

To bypass this redundant calculation, we introduce the KV Cache. Once a token's and vectors are calculated, they are saved in GPU memory. For every subsequent token, the system only projects the single new token into its , , and states, appends the new and to the historical cache, and computes attention across the entire historical record.


Understanding the KV Cache

What does this look like under the hood? The KV cache stores tensor data for every single layer, every attention head, and every token across a batch.

Cache Layout across Layers

For a single request, the KV cache is stored as a 5D tensor across the whole system:

The multiplier 2 accounts for storing both the Key tensor and the Value tensor.

Memory Consumption Formula

The exact memory footprint in bytes can be expressed mathematically as:

Where:

  • = Number of layers
  • = Number of attention heads (or KV heads if using Grouped-Query Attention)
  • = Head dimension ()
  • = Sequence length (tokens)
  • = Precision data type size in bytes (e.g., for FP16/BF16, for FP8)

Concrete Example: Llama-3-8B

Let’s calculate the KV Cache memory footprint for a single context window filled to 8,192 tokens using a standard Llama-3-8B architecture running at Native BF16 precision.

  • Layers (): 32
  • KV Heads (): 8 (Note: Llama-3 uses Grouped-Query Attention, so it features 8 KV heads instead of 32 Query heads)
  • Head Dimension (): 128
  • Sequence Length (): 8,192
  • Data type (): 2 bytes (BF16)

At a standard production batch size of 64 concurrent streams:

On an 80GB NVIDIA H100 or A100 GPU, the weight matrix of Llama-3-8B occupies roughly 16 GB of space. When you layer a 68.48 GB traditional KV cache on top of that, you completely saturate the remaining GPU VRAM, leaving almost zero overhead for workspace activation tensors.


The Real Problem: Why Traditional KV Cache Fails

Before systems like vLLM emerged, frameworks like Hugging Face Transformers or naive FasterTransformer implementations allocated the KV cache as contiguous memory chunks inside the GPU.

Because an LLM generation loop doesn't know in advance when a user request will hit a stop token, or if it will run all the way to its max capacity (e.g., 4096 tokens), the engine had to statically pre-allocate a continuous slice of memory matching the maximum possible sequence length for every active request slot.

This layout results in three catastrophic types of memory waste:

  1. Internal Fragmentation (Reserved for Future Use): If a request has generated 50 tokens but could run up to 4,096 tokens, the memory for the remaining 4,046 tokens is locked down and completely unusable by other concurrent streams.
  2. Over-allocation (Internal Waste): Memory allocated for tokens that are never actually generated because the model hits an early <|endoftext|> token.
  3. External Fragmentation: As different requests finish at different times, memory spaces become chopped up. Even if a GPU has enough total free gigabytes to service a new request, it cannot do so because those gigabytes are scattered across non-contiguous memory segments.
text
Traditional Contiguous Allocation Visualized:
Sequence A (Active): [████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░] (Reserved to max length)
Sequence B (Short):  [██░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░] (95% wasted memory space)
Sequence C (Long):   [██████████████████████████████░░░░░░░░░] 
                     └───────── Total Pre-Allocated Space ───┘

Empirical data reveals that traditional serving systems suffer from 60% to 80% memory waste due to these structural realities. This artificial memory barrier severely suppresses batch sizes, throttling overall system throughput.


Operating System Memory Paging (The Inspiration)

To solve the constraint of contiguous physical allocation, the creators of vLLM looked at a foundational concept developed in operating systems architecture during the 1960s: Virtual Memory Paging.

When an operating system runs a process, it presents the process with a illusion of a massive, contiguous block of address space called Virtual Memory. Under the hood, the operating system splits this virtual space into fixed-size chunks called Pages.

Physical memory (RAM) is similarly divided into fixed-size blocks called Page Frames. The operating system maintains a central ledger called a Page Table that maps a process's sequential logical pages into arbitrary, non-contiguous physical frames scattered across raw hardware.

text
Logical Pages (Sequential) -> [ Page 0 ] [ Page 1 ] [ Page 2 ]
                                   |          |          |
Page Table Mapping --------->      ▼          ▼          ▼
Physical Frames (Scattered)-> [ Frame 4] [ Frame 1] [ Frame 9]

When a program writes data sequentially, the underlying physical allocations can be completely fragmented. The operating system handles the resolution completely transparently, entirely eliminating the requirement for large, continuous blocks of physical RAM.


Introducing PagedAttention

PagedAttention translates this operating system paradigm directly to LLM tensor management. Instead of allocating a massive contiguous tensor for the entire maximum lifetime of an individual request, the system breaks the KV cache down into small, fixed-size chunks called Physical Blocks.

The Core Architecture

  • Logical Blocks: The sequence of KV items generated chronologically by the model. These appear as contiguous blocks of tokens (e.g., 16 tokens per block).
  • Physical Blocks: Non-contiguous regions of GPU memory reserved across a global block pool managed by vLLM's central memory manager.
  • Block Tables: A ledger mapping each individual request's sequential Logical Blocks to arbitrary Physical Blocks in GPU memory.

As shown above, when new tokens are generated, they don't look for contiguous free space behind the preceding token tensor. Instead, they fill up the current fixed-size Logical Block. Once that block is full, the block manager queries its pool of free physical space, grabs a new block frame from anywhere on the GPU, and adds an entry to the request's individual Block Table.

Because block allocations are small and fixed-size, external fragmentation is completely eliminated. The only waste that occurs is inside the very last block of a running sequence (internal fragmentation). Because block sizes are tiny (typically set to 16 tokens), this maximum possible waste is restricted to a negligible fraction of a percent of overall GPU memory.


How PagedAttention Works Internally

Let's trace the exact lifecycle of a request as it moves through a PagedAttention execution pipeline.

Step-by-Step Execution Sequence

  1. Block Allocation: Triggered at Request Initialization. When a prompt arrives, vLLM counts the incoming tokens, computes the number of blocks needed (), and reserves those physical blocks out of the global block manager pool.

  2. Autoregressive Sequence Growth: Executed on every Token Step. As the model generates a new token, the KV vectors for that specific single token are written to the exact slot index inside the current physical block. If the slot index hits the block size boundary, the manager links a new block from the free pool into the table.

  3. Block Table Lookup: Pre-Attention Kernel Phase. Before the attention kernel fires, the system passes an array of integers representing the physical block IDs mapped to the active sequence into the custom CUDA kernel configuration parameters.

  4. Non-Contiguous Attention Kernel Execution: The Custom Core Logic. The custom PagedAttention CUDA kernel launches. Instead of running a standard strided matrix multiplication across continuous memory addresses, GPU threads use the block table to dynamically look up the address offsets of the physical blocks, fetch the keys and values on the fly, and perform the attention math.

The Custom CUDA Kernel Overview

The mathematics of the attention equation do not change. PagedAttention computes the exact same values as standard Multi-Head Attention. The critical innovation lies entirely within the CUDA memory gather logic.

In a traditional attention kernel, computing involves reading a continuous array of memory:

In the PagedAttention kernel, the memory lookup requires an extra layer of indirection:

GPU warps execute highly parallelized gather operations across these scattered block addresses. By utilizing shared memory within the GPU Streaming Multiprocessors (SMs), the kernel hides the latency of these scattered non-contiguous reads, matching or exceeding the raw compute efficiency of standard contiguous memory implementations.


Prefix Sharing and Copy-on-Write (CoW)

One of the most powerful downstream advantages of PagedAttention is its native support for Prefix Sharing.

In scenarios like multi-turn chat agents, parallel text generation, or few-shot in-context learning, multiple distinct requests share an identical starting sequence (e.g., a massive system prompt or a long source document).

In a traditional serving setup, if 100 users are chatting with an agent using a 2,000-token system prompt, that prompt's KV cache must be copied 100 times in memory. With PagedAttention, the logical blocks of all 100 requests point to the exact same physical blocks in GPU memory.

text
Prefix Sharing Memory Layout:
User Request 1 Logical Blocks: [Block 0 (System)] -> [Block 1 (System)] -> [Block 2 (User 1 Dev)]
                                      │                     │
                                      ▼                     ▼
Physical Blocks in VRAM:       [Physical X]          [Physical Y]
                                      ▲                     ▲
                                      │                     │
User Request 2 Logical Blocks: [Block 0 (System)] -> [Block 1 (System)] -> [Block 3 (User 2 Dev)]

Reference Counting & Copy-on-Write (CoW)

To prevent one user's generated response from corrupting the shared prompt history of another, the block manager uses a reference counting ledger:

  1. Each physical block tracks how many logical entries point to it (RefCount).
  2. For the shared system prompt blocks, the RefCount is equal to the number of active user streams sharing that prompt.
  3. When a specific user request enters the generation phase and needs to write its own unique tokens into a block that is shared, the engine intercepts the write, clones the target physical block into a new, independent block frame, decrements the original block's RefCount, and modifies the individual user's block table entry.

This classic engineering pattern ensures absolute isolation between concurrent users while extracting maximal memory efficiency out of identical system workflows.


Continuous Batching

Traditional deep learning systems rely on Static Batching: requests are grouped together, run through the network, and the entire batch blocks execution until the single longest sequence finishes its generation run. If one user requests a 5-token response and another requests a 500-token response, the first user's slot sits completely idle for 495 cycles, wasting valuable hardware resources.

PagedAttention acts as the foundation for modern Continuous Batching (or cellular batching). Because memory is allocated dynamically at the block level, the serving engine can safely schedule requests at the level of individual iteration steps.

  • Dynamic Arrival: New requests join the running batch at the next token iteration boundary without waiting for ongoing generations to complete.
  • Dynamic Departure: As soon as an individual stream encounters a stop token, its entry is dropped from the batch, and its allocated physical blocks instantly return to the free global manager pool.

This level of granular control keeps the GPU continuously saturated, optimizing overall serving infrastructure costs.


Performance Gains

By combining PagedAttention with continuous batching, vLLM changes the economy of LLM serving. The data reported across production deployments consistently highlights major leaps in scale:

MetricTraditional Serving (Hugging Face / Naive)vLLM Serving (PagedAttention)
KV Cache Memory Waste60% – 80%< 4%
Concurrent Batch Size CapacityLower (e.g., Batch Size 8-16)High (e.g., Batch Size 32-64+)
System Throughput Under LoadBaseline (1x)2x – 4x Throughput Improvement
GPU Utilization ProfileHighly volatile, compute under-saturatedConsistently high saturation

By reclaiming the massive footprint previously locked away by internal and external fragmentation, engines can increase concurrency limits until the hardware hits its physical compute processing cap.


PagedAttention vs. FlashAttention

A common point of confusion for engineers entering this space is distinguishing between FlashAttention and PagedAttention. They sound similar, but they address entirely separate bottlenecks inside the system architecture.

FeatureFlashAttentionPagedAttention
Primary Bottleneck TargetedCompute & SRAM/HBM Read-Write LatencyVRAM Capacity & Fragmentation
Operational PhaseTraining & Inference (Prefill heavy focus)Inference Generation (Decoupled Decoding)
Core MechanismTiling, online softmax, recomputation to avoid writing intermediate attention matrices to HBM.Memory virtualization using fixed-size logical-to-physical block tables for storing the KV Cache.
Core BenefitDramatic speedups in raw attention computation math.Enormous increases in concurrent batch capacities.

They Are Complementary

These two innovations do not compete; they operate together. Modern serving frameworks utilize FlashDecoding (an extension of FlashAttention tailored for inference) to optimize the raw mathematical execution speeds of the attention step, while using PagedAttention to control how those tensors are mapped and laid out across global GPU memory.


PagedAttention vs. Other Modern Serving Techniques

The LLM serving landscape has evolved into an advanced ecosystem of optimization frameworks. Here is how PagedAttention tracks against alternative modern strategies:

  • TensorRT-LLM (NVIDIA): NVIDIA’s highly optimized, proprietary inference engine. It uses its own native implementation of paged cache tracking alongside specialized hardware acceleration hooks. It delivers extreme performance but ties workloads directly to NVIDIA architectures.
  • SGLang (RadixAttention): An advanced frontend and execution engine that manages prefix caching using a hierarchical graph structure called a Radix Tree. While PagedAttention manages raw memory blocks, RadixAttention automates the eviction and reuse of complex multi-turn prompt structures over time.
  • vAttention: A modern paradigm that shifts the paging logic out of user-space Python runtime code and pushes it directly down into the system's CUDA Virtual Memory Management (VMM) APIs. By letting the GPU's hardware memory management unit (MMU) handle address virtualization, it achieves physical paging behavior without requiring custom non-contiguous attention kernels.
  • Hugging Face TGI (Text Generation Inference): A production-ready architecture that natively incorporates PagedAttention and FlashAttention layers to provide an out-of-the-box system for enterprise APIs.

Simplified Python Simulation

To make these concepts concrete, let’s build a functional, step-by-step toy Python simulation of a PagedAttention Memory Manager. This code demonstrates the internal mechanics of block tables, allocation rules, sequence growth, and reference-counted copy-on-write actions.

python
import math

class PagedKVBlockManager:
    def __init__(self, total_blocks: int, block_size: int = 16):
        self.block_size = block_size
        self.total_blocks = total_blocks
        # Physical block pool tracking: free blocks list
        self.free_blocks = list(range(total_blocks))
        # Reference counting mapping block_id -> count
        self.block_ref_counts = {i: 0 for i in range(total_blocks)}
        # Central routing register request_id -> list of block_ids
        self.block_tables = {}
        
    def allocate_request(self, request_id: str, prompt_tokens: int):
        """Initial allocation for an incoming prompt."""
        needed_blocks = math.ceil(prompt_tokens / self.block_size)
        if len(self.free_blocks) < needed_blocks:
            raise MemoryError("Out of Physical VRAM Block Pool Capacity!")
            
        allocated = []
        for _ in range(needed_blocks):
            block_id = self.free_blocks.pop(0)
            self.block_ref_counts[block_id] = 1
            allocated.append(block_id)
            
        self.block_tables[request_id] = allocated
        print(f"[Allocated] Req '{request_id}': Tokens={prompt_tokens} -> Blocks={allocated}")

    def append_token(self, request_id: str, current_sequence_length: int):
        """Simulates growth of a sequence by one token step."""
        block_table = self.block_tables[request_id]
        
        # Check if we need to cross the block boundary to open a new block frame
        if current_sequence_length % self.block_size == 0:
            if not self.free_blocks:
                raise MemoryError("VRAM saturated during token generation!")
            
            new_block = self.free_blocks.pop(0)
            self.block_ref_counts[new_block] = 1
            block_table.append(new_block)
            print(f"[Growth] Req '{request_id}' reached boundary. Added new Physical Block {new_block}")
        else:
            # Check for Copy-On-Write situation on the active writing target block
            target_block_idx = len(block_table) - 1
            target_block_id = block_table[target_block_idx]
            
            if self.block_ref_counts[target_block_id] > 1:
                # Trigger Copy-on-Write
                if not self.free_blocks:
                    raise MemoryError("VRAM saturated during Copy-on-Write execution!")
                
                new_block = self.free_blocks.pop(0)
                self.block_ref_counts[target_block_id] -= 1
                self.block_ref_counts[new_block] = 1
                block_table[target_block_idx] = new_block
                print(f"[CoW Triggered] Req '{request_id}' cloned shared Block {target_block_id} into independent Block {new_block}")

    def share_prefix(self, source_request_id: str, target_request_id: str):
        """Simulates a second request sharing the prefix cache of the first request."""
        source_blocks = self.block_tables[source_request_id]
        # We assume the entire current length is shared prefix
        self.block_tables[target_request_id] = list(source_blocks)
        for b_id in source_blocks:
            self.block_ref_counts[b_id] += 1
        print(f"[Prefix Shared] Req '{target_request_id}' now maps to identical blocks as '{source_request_id}'")

    def free_request(self, request_id: str):
        """Cleans up memory structures when a request completes."""
        if request_id not in self.block_tables:
            return
        block_table = self.block_tables.pop(request_id)
        for b_id in block_table:
            self.block_ref_counts[b_id] -= 1
            if self.block_ref_counts[b_id] == 0:
                self.free_blocks.append(b_id)
        print(f"[Released] Req '{request_id}' resources returned to free pool.")

# Execution Walkthrough Simulation
if __name__ == "__main__":
    # Initialize a tiny manager with 10 total physical memory blocks
    manager = PagedKVBlockManager(total_blocks=10, block_size=16)
    
    # 1. Allocate Request 1 (Prompt length = 24 tokens -> requires 2 blocks)
    manager.allocate_request("User_1", prompt_tokens=24)
    
    # 2. Setup Request 2 sharing the exact same initial prefix blocks
    manager.share_prefix("User_1", "User_2")
    print(f"Global Reference counts: {manager.block_ref_counts}")
    
    # 3. User 1 generates a token while at 32 tokens total (boundary edge condition)
    manager.append_token("User_1", current_sequence_length=32)
    
    # 4. User 2 generates a token, causing a write inside a shared block zone (CoW)
    manager.append_token("User_2", current_sequence_length=24)
    
    # 5. Terminate and free jobs
    manager.free_request("User_1")
    manager.free_request("User_2")
    print(f"Final Free Block Pool State: {manager.free_blocks}")

Key Takeaways

  • Algorithmic Preservation: PagedAttention does not alter the core mathematical formulas of attention. The computed probabilities, outputs, and model quality remain completely identical. It is strictly an engineering optimization targeting memory layout management.
  • Operating System Inspiration: The architecture directly adapts standard operating system virtual memory concepts, swapping logical pages and physical frames for token blocks and GPU memory segments.
  • Elimination of Fragmentation: By fixing block lengths to tiny static increments (e.g., 16 tokens), it drives system-wide external memory fragmentation down to nearly zero.
  • Throughput Multiplication: Reclaiming previously wasted space allows frameworks to scale concurrent batch processing limits up significantly, offering major efficiency gains for high-demand AI platforms.

Appendix

A. KV Cache Memory Formula Complete Derivation

The general memory footprint calculation for storing standard Transformer KV tensors maps directly to tensor dimension sizes. Because each token requires tracking vectors for both Key () and Value (), we begin with a multiplier of 2:

If a model uses standard Multi-Head Attention, the number of KV Heads matches the number of Query attention heads. If the architecture utilizes Grouped-Query Attention (GQA), the value represents the smaller, downsampled head count value (typically or of the total hidden layer size mapping).

B. Glossary of Core Concepts

  • KV Cache: Saved Key and Value activation vectors designed to minimize computational repetitions throughout autoregressive decoding phases.
  • Logical Block: A sequence of contiguous token states grouped together based on generation timestamps.
  • Physical Block: A static chunk of raw physical GPU VRAM managed within a global pool framework.
  • Block Table: A lookup vector ledger mapping structural logical indices to arbitrary physical VRAM coordinates.
  • Fragmentation: Memory space that is locked away but underutilized or unaddressable due to boundary misalignments.
  • Prefix Cache: Shared memory architecture designed to map identical system instructions across separate individual user pipelines.
  • Copy-on-Write (CoW): An allocation paradigm that delays physical tensor cloning operations until an active write command targets a shared resource block.
  • Continuous Batching: A dynamic system optimization loop that permits real-time addition and eviction of active generation tasks at token boundary cycles.

Stay in the loop

New posts on ML, AI engineering, and building things — straight to your inbox. No spam.

Comments (0)

No comments yet. Be the first to comment!

Leave a comment