Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Paged Attention › The PagedAttention Algorithm

Michael BrenndoerferJanuary 8, 202610 min read

Language AI HandbookAll chapters

Part 42 of 67Inference Optimization

Chapter 315Paged Attention › The PagedAttention Algorithm

Reading section 4 of 13 · The PagedAttention Algorithm · 10 minSection 4 of 13 · 10 min

Section 4 of 13 · 10 min read

Paged Attention › The PagedAttention AlgorithmOne section · Sections in this chapter4 / 13 ↓

Sections in this chapter13

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

With KV cache data potentially scattered across non-contiguous memory locations, we need to modify the attention computation to work with this new layout. The PagedAttention algorithm handles this by fetching keys and values according to the block table during attention calculation. This modification preserves the mathematical correctness of attention while accommodating the scattered physical layout.

The modification is conceptually straightforward: instead of indexing into a contiguous key-value tensor using the token position directly, we first compute the logical block number and within-block offset, look up the physical block from the block table, and then index into the KV cache using the physical block address. The attention score computation and the softmax normalization are identical to standard attention; only the data access pattern changes.

Think of the change as the difference between reading a book straight through versus reading it by consulting an index first. The content you read is identical, and you understand it the same way. The only difference is that you now follow a lookup step before accessing each page. The lookup is fast, and the content is exactly what standard attention would have read from a contiguous layout.

The key insight is that attention only needs the keys and values for a specific sequence, in order. It does not care what their physical addresses are, only what values they hold. The block table gives us exactly what we need: a translation from the logical "token 25 in sequence 3" to the physical "slot 9 in GPU memory block 47." Once we have the value, the computation proceeds identically.

Standard Attention vs PagedAttentionLink Copied

In standard attention, keys and values are stored in contiguous tensors. For a sequence of length $n$ , the attention mechanism computes a weighted sum of values based on the similarity between the query and keys:

\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V

where:

$Q$ : the query matrix, typically for the single current token being generated (shape $1 \times d_k$ for autoregressive decoding)
$K$ : the key matrix containing keys for all tokens in the context (contiguous tensor of shape $n \times d_k$ )
$V$ : the value matrix containing values for all tokens in the context (contiguous tensor of shape $n \times d_v$ )
$d_k$ : the dimension of the key vectors, used for scaling
$d_v$ : the dimension of the value vectors
$n$ : the sequence length (all previously generated tokens plus the prompt)
$\text{softmax}$ : the function that normalizes raw similarity scores into a probability distribution summing to 1

The formula consists of three computational stages. First, the dot product $QK^T$ measures the similarity between the query and each key. Intuitively, this operation asks: "how relevant is each previous token to the token we are currently generating?" Keys that are more similar to the query receive higher scores, showing greater relevance. The result is a vector of raw similarity scores, one for each position in the sequence.

The second stage applies the $\text{softmax}$ function to these raw scores. This normalization turns the arbitrary similarity values into a proper probability distribution that sums to 1. The softmax operation accentuates differences: high scores become relatively higher, and low scores become relatively lower. The output is the attention weights, which tell us how much to attend to each previous position. A weight of 0.8 on token 5 means we draw heavily from that token's value vector; a weight of 0.001 means that token contributes almost nothing to the output.

The third stage multiplies these attention weights by the value matrix $V$ , computing a weighted average of the value vectors. Tokens with higher attention weights contribute more to the final output, while tokens with lower weights contribute less. This mechanism allows the model to selectively focus on the most relevant information from the entire context window, regardless of how long that window is.

The division by $\sqrt{d_k}$ in the formula ensures numerical stability. Without this scaling factor, the dot products between queries and keys tend to grow large as the dimension $d_k$ increases. Large dot products cause the softmax function to saturate, creating attention distributions that are nearly one-hot (all attention concentrated on a single token). The scaling factor keeps the dot products in a reasonable range, letting the softmax to produce more fine-grained attention patterns and improving gradient flow during training.

In PagedAttention, $K$ and $V$ are stored in non-contiguous blocks rather than contiguous tensors. The attention computation must gather the relevant blocks according to the block table. While the mathematical operation remains identical, the implementation must navigate the scattered physical layout:

In[16]:

Code

import torch
import torch.nn.functional as F


def paged_attention_reference(
    query: torch.Tensor,  # Shape: (batch, num_heads, 1, head_dim) for single token
    key_cache: torch.Tensor,  # Shape: (num_blocks, block_size, num_heads, head_dim)
    value_cache: torch.Tensor,  # Shape: (num_blocks, block_size, num_heads, head_dim)
    block_tables: torch.Tensor,  # Shape: (batch, max_blocks) - physical block indices
    context_lens: torch.Tensor,  # Shape: (batch,) - actual sequence lengths
    block_size: int,
    scale: float,
) -> torch.Tensor:
    """
    Reference implementation of PagedAttention for understanding.
    Not optimized - real implementations use custom CUDA kernels.
    """
    batch_size, num_heads, _, head_dim = query.shape
    output = torch.zeros_like(query)

    for b in range(batch_size):
        seq_len = context_lens[b].item()
        num_blocks_used = (seq_len + block_size - 1) // block_size

        # Gather keys and values from physical blocks
        keys_list = []
        values_list = []

        for block_idx in range(num_blocks_used):
            physical_block = block_tables[b, block_idx].item()

            # Calculate how many tokens to use from this block
            start_pos = block_idx * block_size
            end_pos = min(start_pos + block_size, seq_len)
            tokens_in_block = end_pos - start_pos

            # Fetch from physical memory location
            keys_list.append(key_cache[physical_block, :tokens_in_block])
            values_list.append(value_cache[physical_block, :tokens_in_block])

        # Concatenate gathered keys and values
        # Shape: (seq_len, num_heads, head_dim)
        keys = torch.cat(keys_list, dim=0)
        values = torch.cat(values_list, dim=0)

        # Reshape for attention computation
        # keys: (num_heads, seq_len, head_dim)
        keys = keys.permute(1, 0, 2)
        values = values.permute(1, 0, 2)

        # Query shape: (num_heads, 1, head_dim)
        q = query[b]  # (num_heads, 1, head_dim)

        # Compute attention scores
        # (num_heads, 1, head_dim) @ (num_heads, head_dim, seq_len) -> (num_heads, 1, seq_len)
        attn_scores = torch.matmul(q, keys.transpose(-2, -1)) * scale
        attn_probs = F.softmax(attn_scores, dim=-1)

        # Apply attention to values
        # (num_heads, 1, seq_len) @ (num_heads, seq_len, head_dim) -> (num_heads, 1, head_dim)
        attn_output = torch.matmul(attn_probs, values)
        output[b] = attn_output

    return output

This reference implementation shows the core idea: before computing attention, we gather the keys and values from their physical locations according to the block table. The attention computation itself remains unchanged. The gathering process iterates through logical block indices, looks up the corresponding physical block in the block table, and fetches the data from that physical location. Once all the keys and values are assembled into contiguous tensors, the standard attention formula applies exactly as before.

Demonstration with Non-Contiguous BlocksLink Copied

Let's verify that PagedAttention produces correct results even when blocks are scattered:

In[17]:

Code

# Set up a small example
torch.manual_seed(42)

num_blocks = 8
block_size = 4
num_heads = 2
head_dim = 8
batch_size = 2

# Physical KV cache - blocks can be used by any sequence
key_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
value_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)

# Block tables - note non-contiguous physical blocks
# Sequence 0 uses physical blocks [0, 3, 5] (scattered!)
# Sequence 1 uses physical blocks [1, 7]
block_tables = torch.tensor(
    [
        [
            0,
            3,
            5,
            0,
            0,
            0,
            0,
            0,
        ],  # Sequence 0: 3 blocks (12 tokens max, using 10)
        [
            1,
            7,
            0,
            0,
            0,
            0,
            0,
            0,
        ],  # Sequence 1: 2 blocks (8 tokens max, using 7)
    ]
)

# Actual sequence lengths
context_lens = torch.tensor([10, 7])

# Query for current token being generated
query = torch.randn(batch_size, num_heads, 1, head_dim)
scale = 1.0 / (head_dim**0.5)

# Run paged attention
output = paged_attention_reference(
    query, key_cache, value_cache, block_tables, context_lens, block_size, scale
)

# Extract block mappings for visualization
sequence_mappings = []
for i, seq_len in enumerate(context_lens):
    # Calculate number of blocks used by this sequence
    n_blocks = (seq_len.item() + block_size - 1) // block_size
    # Slice the block table to get only the used physical blocks
    blocks = block_tables[i, :n_blocks].tolist()
    sequence_mappings.append((i, seq_len.item(), blocks))

Out[18]:

Console

PagedAttention Demonstration
=============================================
Block size: 4 tokens
Number of physical blocks: 8

Block table mappings:
  Sequence 0 (len=10): uses physical blocks [0, 3, 5]
  Sequence 1 (len=7): uses physical blocks [1, 7]

Note: blocks are non-contiguous in physical memory!

Output shape: torch.Size([2, 2, 1, 8])

Optimized CUDA KernelsLink Copied

The reference implementation above gathers blocks into contiguous memory before computing attention. This defeats some of the purpose of paging, as it requires temporary memory allocation for the gathered keys and values. Production implementations like vLLM use custom CUDA kernels that compute attention directly on the paged data structure, avoiding the explicit gather step entirely.

The core challenge in writing these kernels is that standard attention kernels, including FlashAttention, assume contiguous tensors. A custom paged attention kernel must interleave block-table lookups with the attention computation, load blocks into fast shared memory, and then compute scores and apply the weighted sum, all without ever assembling a full contiguous KV tensor.

The key optimization insight is that attention can be computed incrementally as blocks are loaded, rather than requiring all data to be assembled first. This approach uses the online softmax algorithm (also known as the flash attention technique covered in the efficient attention chapter) to maintain numerical stability while processing blocks one at a time. The online softmax technique keeps a running maximum and a running sum, updating them as each new block is processed. This allows the kernel to compute the final attention output without ever needing to store all the keys and values in a contiguous buffer.

In[19]:

Code

# Pseudocode for optimized PagedAttention kernel
# This shows the conceptual approach - actual CUDA code is more complex


def paged_attention_kernel_pseudocode():
    """
    Conceptual flow of the optimized PagedAttention CUDA kernel.

    Key optimizations:
    1. No explicit gather step - compute attention directly from blocks
    2. Each thread block processes a portion of the attention
    3. Leverage shared memory for frequently accessed block table entries
    """
    kernel_description = """
    CUDA Kernel: paged_attention_v1

    Grid: (num_heads, batch_size, 1)
    Block: (NUM_THREADS,)

    For each (head, batch) pair:
        1. Load query vector for this head into registers
        2. Initialize running softmax numerator and denominator

        For each logical block in sequence:
            a. Look up physical block index from block_table
            b. Load key block from global memory (coalesced access)
            c. Compute partial attention scores for this block
            d. Update running softmax using online softmax trick

        For each logical block in sequence:
            a. Load value block from global memory
            b. Weight by final attention probabilities
            c. Accumulate to output

        Write output for this (batch, head, token)
    """
    return kernel_description


kernel_pseudocode = paged_attention_kernel_pseudocode()

Out[20]:

Console


    CUDA Kernel: paged_attention_v1

    Grid: (num_heads, batch_size, 1)
    Block: (NUM_THREADS,)

    For each (head, batch) pair:
        1. Load query vector for this head into registers
        2. Initialize running softmax numerator and denominator

        For each logical block in sequence:
            a. Look up physical block index from block_table
            b. Load key block from global memory (coalesced access)
            c. Compute partial attention scores for this block
            d. Update running softmax using online softmax trick

        For each logical block in sequence:
            a. Load value block from global memory
            b. Weight by final attention probabilities
            c. Accumulate to output

        Write output for this (batch, head, token)

The necessary optimization is computing attention incrementally as blocks are loaded, rather than gathering all blocks first. By processing blocks sequentially and maintaining running statistics about the maximum score seen so far and the running normalization denominator, the kernel achieves the same mathematical result as gathering all keys first, but with dramatically reduced memory requirements. This combination of paging and incremental attention computation is what makes PagedAttention both memory-efficient and computationally practical on real hardware.