PagedAttention by hand

Table of Contents

• Motivation
• What is PagedAttention
• Setup
• Forward pass (PagedAttention)
• Why divide memory into physical blocks?
• Code

Motivation

PagedAttention (the core algorithm behind the vLLM inference engine) is widely used to serve Large Language Models (LLMs) with high throughput. While libraries like vLLM provide out-of-the-box support for it, understanding how PagedAttention is implemented by hand may help with:

• Building intuition: understand how it maps logical token sequences to non-contiguous physical memory blocks, mimicking operating system virtual memory.
• Understanding the memory-capacity bottleneck: understand why the number of sequences you can batch is set by how much KV cache fits in GPU memory, not by how fast the GPU computes, and why fragmentation limits it.
• Optimizing deployments: the trade-off between internal fragmentation and bandwidth in choosing the optimal block_size.

This article can serve as a complementary study guide for understanding the vLLM paper: Efficient Memory Management for Large Language Model Serving with PagedAttention.

This article is written with the assistance of AI.

What is PagedAttention

When an LLM generates a response, it produces tokens one by one. To avoid recomputing the attention states for all past tokens at every step, the model caches the Key and Value vectors for all previously generated tokens. This is called the KV cache. I have a worked example of a KV cache with a single-layer model of embedding dimension d = 2 here as reference.

However, the KV cache has unique characteristics that make it a nightmare for contiguous memory allocators:

1. It grows dynamically over time.
2. Its final length is unknown ahead of time.

For systems which pre-allocate contiguous chunks of memory for the maximum possible sequence length, this will lead to three distinct wastes:

Internal vs external fragmentation

Internal fragmentation is waste inside an allocation owned by a sequence, e.g. the sequence asked for a 2048-slot chunk and only used 60. The other 1988 slots are sitting idle and no other sequences can use them.

External fragmentation is waste between allocations. The memory is free, but it is in pieces of the wrong size or in the wrong places, e.g. you can have 10 GB free in total and still fail a 2 GB sequence because no single contiguous run is 2 GB long.

Using the parking lot as an analogy: internal fragmentation is being assigned a bus bay and parking a motorbike in it. External fragmentation is a lot where cars parked at random leave plenty of empty tarmac but no single gap wide enough for a bus.

How PagedAttention solves reservation and fragmentations

PagedAttention solves this by borrowing the concept of virtual memory with paging from operating systems:

1. Partitions Memory: It divides the KV cache into fixed-size "blocks" (pages), where each block contains the keys and values for a fixed number of tokens.
2. Non-contiguous Storage: These blocks do not need to be stored next to each other in physical GPU memory.
3. Block Tables: It uses a "Block Table" (like a page table) to map the continuous logical sequence of tokens to their scattered physical blocks.

Uniform blocks solve external fragmentation because with only one block size, there are no holes of the wrong size. The free list is a stack of interchangeable IDs, and any free block satisfies any sequence.

Why not set B = 1 to solve internal fragmentation?

The paper tests this (§7.2, Fig. 18b of paper) and finds it among the worst-performing settings. The reason is that a block is not only an allocation unit, but also the unit of a coalesced memory read. vLLM assigns one GPU warp per block (§5.1 of paper), so at B=16 those 32 threads fetch 16 tokens of adjacent key vectors in a single wide transaction. At B=1 you get one narrow scattered read per token, and since decoding is memory-bandwidth-bound, that costs far more than the ≤15 slots saved. Block size is a trade between fragmentation and bandwidth.

Estimating size of a block byte

block bytes = B tokens × d numbers × 2 bytes
            = 16 × 128 × 2
            = 4,096 bytes

At B=1 that is 128 × 2 = 256 bytes. 128 is the length of a single token's key vector, and it stays 128 regardless of what B is. Changing B changes how many of those 256-byte vectors sit side by side.

Setup

To keep the scope of the article manageable, this article discusses PagedAttention in the context of autoregressive generation phase of LLM inference (serving).

We will walk through the attention computation for a single attention head generating a new token. Let's assume our LLM has a tiny hidden dimension of d = 2.

Our sequence currently has 5 tokens, and we are generating the 6th token.
Instead of a contiguous tensor, our KV cache is managed by PagedAttention with a Block Size (B) of 2 tokens per block.

Since we have 5 tokens, we require $\lceil 5 / 2 \rceil = 3$ logical blocks.

Step 1: Initial Setup - Defining the Tensors and Memory

Our sequence currently has 5 tokens. We are running the decode step that will predict the 6th. The input to this step is token 5, so the query vector q is the query of token 5. Note that token 5's own key and value are already in the cache (vLLM writes them before calling the attention kernel):

We have a pre-allocated pool of 9 physical blocks in GPU memory, shared by every sequence on this worker. The system assigned our sequence to Physical Blocks 2, 0, and 1 (in that exact order) — the first three that happened to be free. Each of the other 8 blocks are either free, or already belong to another sequence.

This mapping is stored in our Block Table:
Block Table = [2, 0, 1]

Here is what is currently stored in those physical KV cache blocks (Keys and Values for past tokens):

Logical Block 0 (Tokens 1 & 2) -> Mapped to Physical Block 2:
$K_{phys_2} = \begin{pmatrix} 1 & 0 \newline 0 & 1 \end{pmatrix}$ , $V_{phys_2} = \begin{pmatrix} 1 & 1 \newline 2 & 2 \end{pmatrix}$

Logical Block 1 (Tokens 3 & 4) -> Mapped to Physical Block 0:
$K_{phys_0} = \begin{pmatrix} 1 & 1 \newline 2 & 0 \end{pmatrix}$ , $V_{phys_0} = \begin{pmatrix} 3 & 3 \newline 4 & 4 \end{pmatrix}$

Logical Block 2 (Token 5 & Empty) -> Mapped to Physical Block 1:
Notice that the second slot in this block is currently empty because we have an odd number of tokens (5).
$K_{phys_1} = \begin{pmatrix} 0 & 2 \newline 0 & 0 \end{pmatrix}$ , $V_{phys_1} = \begin{pmatrix} 5 & 5 \newline 0 & 0 \end{pmatrix}$

Note that each token contributes two vectors — one key and one value. A block with B = 2 therefore holds two rows in its K matrix and two rows in its V matrix, with row i of each belonging to the same token.

Forward pass (Paged Attention)

Standard attention calculates $Score = softmax(q K^T) V$ .
Because our K and V matrices are scattered across physical memory, PagedAttention processes them block by block. (Note: For simplicity in mental math, we will omit the standard $\frac{1}{\sqrt{d}}$ scaling factor).

Step 2: Compute Unnormalized Attention Scores ( $q K^T$ )

We iterate through the Block Table [2, 0, 1].

1. First Logical Block (Physical Block 2)

• Score_0[0] = (1*1) + (2*0) = 1
• Score_0[1] = (1*0) + (2*1) = 2

Result: [1, 2]

2. Second Logical Block (Physical Block 0)

• Score_1[0] = (1*1) + (2*1) = 3
• Score_1[1] = (1*2) + (2*0) = 2

Result: [3, 2]

3. Third Logical Block (Physical Block 1)

• Score_2[0] = (1*0) + (2*2) = 4
• Score_2[1] = (1*0) + (2*0) = 0

Result: [4, 0]

We only have 5 valid tokens. The 6th slot in Physical Block 1 is uninitialised memory, so we must stop the model attending to it. We replace the invalid score with negative infinity (-inf), which sends its softmax weight to exactly zero.

We do not leave the score at 0, because doing so means the slot absorbs about 1% of the probability mass, diluting every other token's weight and shifting the output by roughly the same amount. In a real server that slot holds whatever a previously evicted sequence left behind, so the contribution is arbitrary rather than merely small.

In production, instead of masking, we can bound the so that the CUDA kernel never reads those addresses at all and invalid slots cost zero memory traffic. Since decoding is memory-bandwidth-bound, not reading can save bandwidth on unnecessary reads.

Step 3: Apply Softmax

We concatenate our block scores into a single logical sequence of scores:
Scores = [1, 2, 3, 2, 4, -inf]

To calculate the softmax, PyTorch subtracts the maximum value (which is 4) for numerical stability, then takes the exponent, and divides by the sum of exponents.

Scores - 4 = [-3, -2, -1, -2, 0, -inf]
Exp = [0.0498, 0.1353, 0.3679, 0.1353, 1.0000, 0]
Sum of Exp = 1.6883

Dividing by the sum gives us our attention probabilities:
Probs ≈ [0.0295, 0.0802, 0.2179, 0.0802, 0.5923, 0]

Step 4: Multiply by Values ( $Probs \times V$ )

We slice the probabilities back into their respective blocks and multiply them by the physical Value blocks.

1. First Logical Block (Physical Block 2)

Probs_0 = [0.0295, 0.0802]

2. Second Logical Block (Physical Block 0)

Probs_1 = [0.2179, 0.0802]

3. Third Logical Block (Physical Block 1)

Probs_2 = [0.5923, 0]

4. Final Aggregation

We sum the outputs from all blocks to get the final attention output for the new token:

This tensor is then passed into the rest of the Transformer block (Linear layers, etc.) to predict the 6th token. At the next iteration, token 6 becomes the model's input; its key and value are computed and written into the empty slot in Physical Block 1 before attention runs, and the cycle repeats.

Why divide memory into physical blocks?

Since the numerical output is exactly the same as standard Attention, why go through the headache of block tables and scattered memory?

1. Zero External Fragmentation & Minimal Internal Fragmentation

With contiguous memory allocation, if an LLM supports a max sequence of 2048 tokens, the system pre-allocates contiguous memory for 2048 tokens the moment a sequence arrives. If the user prompt is only 50 tokens, 1998 tokens are wasted to reserved fragmentation. If the model responds with 10 tokens, now 1988 slots of memory are wasted to internal fragmentation.

With PagedAttention, memory is allocated on demand. In our example, we only allocated exactly 3 blocks (6 slots) for a 5-token sequence. The only "wasted" space is the single empty slot in the last block (internal fragmentation).

§7.2 of the paper default block size is 16, which is large enough to keep the GPU's read parallelism busy and small enough to avoid meaningful internal fragmentation. That default guarantees that wasted space is bounded by at most B − 1 slots per sequence, or 15 slots in this case, no matter whether the sequence is 20 tokens or 20,000.

2. Memory Sharing (Copy-on-Write)

Imagine we ask for two samples (n=2), from the prompt prompt "The capital of France is":

• Candidate A attempts to generate: "Paris, which is..."
• Candidate B attempts to generate: "a city that..."

With contiguous memory, the KV cache for the prompt would have to be duplicated for each candidate.

With PagedAttention, both candidates share the same logical-to-physical block mapping for the prompt. The block manager tracks a reference count for each physical block. When a candidate needs to write a token into a block whose reference count is greater than 1, it performs an OS-style copy-on-write: allocate a fresh physical block, copy that one block's contents, decrement the original's reference count, and write into the copy.

Note: A reference count is a tally on the number of sequences currently using a physical block.

Example of Copy-on-Write

Prompt: 5 tokens. Logical blocks: [t1,t2] [t3,t4] [t5,_]. Both candidates map to physical blocks 2, 0, 1 with refcount 2 on each.

Now both generate token 6, and they disagree.

Afterwards: blocks 2 and 0 (holding t1–t4) are still shared, reference count is 2. Only one block was duplicated.

Compare the alternatives. Never copy: physically impossible, the numbers differ. Always copy: that's the old system, i.e. duplicate all 3 blocks per candidate, which is the waste vLLM exists to remove. copy-on-write gives you 2 blocks shared, 1 copied.

Why contigous KV cache has difficulty implementing prefix matching

If a sequence's KV cache is a contiguous tensor, the attention kernel takes a base pointer and a length, then strides through memory linearly. Under that constraint, candidate A's cache mus look like:

[prompt tokens 1-5][A's tokens 6,7,8...]

all in one unbroken run. Candidate B needs the same thing with its own tokens appended. If both point at one shared prompt region, then B's token 6 has to occupy the address immediately after the prompt, which is where A's token 6 already lives. Two different values, one address, which makes prefix matching. There are alternatives to PagedAttention, such as RadixAttention, that can also enable prefix matching.

Code

For those who prefer following the code, the manual walk-through of PagedAttention above is replicated below using PyTorch. In production, vLLM uses custom CUDA kernels to execute these block reads and matrix multiplications faster.

import torch

# --------------------------------------------------------------------------
# 1. SETUP: Define tensors and physical memory pool
# --------------------------------------------------------------------------

BLOCK_SIZE = 2   # tokens per block. vLLM's default is 16; try changing it.
HEAD_DIM = 2
NUM_PHYSICAL_BLOCKS = 9

# New Query token: shape (1, 2)
q = torch.tensor([[1.0, 2.0]])

# We simulate our GPU physical memory pool. Our sequence will only use 3 of
# these 9 blocks -- the rest are free, or belong to other sequences.
# Shape: (total_physical_blocks, block_size, head_dim) = (9, 2, 2)
#
# Note: vLLM's real cache uses a more intricate layout, reshaping the key
# cache to interleave the head dimension so each block read coalesces across
# a GPU warp (§5.1, "fused reshape and block write"). We flatten it here for
# clarity -- the algorithm is identical, only the memory layout differs.
K_cache = torch.zeros(NUM_PHYSICAL_BLOCKS, BLOCK_SIZE, HEAD_DIM)
V_cache = torch.zeros(NUM_PHYSICAL_BLOCKS, BLOCK_SIZE, HEAD_DIM)

# Fill Physical Block 2 (Logical Block 0: tokens 1, 2)
K_cache[2] = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
V_cache[2] = torch.tensor([[1.0, 1.0], [2.0, 2.0]])

# Fill Physical Block 0 (Logical Block 1: tokens 3, 4)
K_cache[0] = torch.tensor([[1.0, 1.0], [2.0, 0.0]])
V_cache[0] = torch.tensor([[3.0, 3.0], [4.0, 4.0]])

# Fill Physical Block 1 (Logical Block 2: token 5, empty slot)
K_cache[1] = torch.tensor([[0.0, 2.0], [0.0, 0.0]])
V_cache[1] = torch.tensor([[5.0, 5.0], [0.0, 0.0]])

# Blocks 3-8 remain free for other sequences. Notice our sequence's blocks are
# neither contiguous nor in order -- that is the entire point.
block_table = [2, 0, 1]
context_length = 5   # Number of valid tokens we currently have

print("--- PagedAttention Forward Pass ---")

# --------------------------------------------------------------------------
# 2. Compute Unnormalized Scores
# --------------------------------------------------------------------------
scores = []
for i, physical_block_idx in enumerate(block_table):
    # Retrieve the physical keys for this block
    K_block = K_cache[physical_block_idx]      # shape (BLOCK_SIZE, HEAD_DIM)

    # Calculate q * K^T
    score = torch.matmul(q, K_block.T)         # shape (1, BLOCK_SIZE)

    # Mask out unfilled slots in the last block
    if i == len(block_table) - 1:
        valid_in_last_block = context_length % BLOCK_SIZE
        if valid_in_last_block != 0:
            score[0, valid_in_last_block:] = float('-inf')

    scores.append(score)

# Concatenate block scores into a single sequence
scores = torch.cat(scores, dim=-1)
print(f"Unnormalized Scores: {scores}")

# --------------------------------------------------------------------------
# 3. Softmax
# --------------------------------------------------------------------------
probs = torch.softmax(scores, dim=-1)
print(f"Attention Probabilities: {probs}")

# --------------------------------------------------------------------------
# 4. Multiply by Values
# --------------------------------------------------------------------------
final_output = torch.zeros(1, HEAD_DIM)
start_idx = 0

for i, physical_block_idx in enumerate(block_table):
    V_block = V_cache[physical_block_idx]

    # Slice the probabilities corresponding to this block
    p_block = probs[0, start_idx:start_idx + BLOCK_SIZE].unsqueeze(0)

    # Multiply and accumulate
    final_output += torch.matmul(p_block, V_block)

    start_idx += BLOCK_SIZE

print(f"\nFinal Attention Output: {final_output}")

Output

--- PagedAttention Forward Pass ---
Unnormalized Scores: tensor([[1., 2., 3., 2., 4., -inf]])
Attention Probabilities: tensor([[0.0295, 0.0802, 0.2179, 0.0802, 0.5923, 0.0000]])

Final Attention Output: tensor([[4.1256, 4.1256]])

Note: our hand-calculation gives 4.1259 because we rounded probabilities to four decimals at each step. The code, working in full precision, returns 4.1256.