Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Paged Attention › Worked Example: Tracing a Complete Lookup

Michael BrenndoerferJanuary 8, 20265 min read

Language AI HandbookAll chapters

Part 42 of 67Inference Optimization

Chapter 315Paged Attention › Worked Example: Tracing a Complete Lookup

Reading section 5 of 13 · Worked Example: Tracing a Complete Lookup · 5 minSection 5 of 13 · 5 min

Section 5 of 13 · 5 min read

Paged Attention › Worked Example: Tracing a Complete LookupOne section · Sections in this chapter5 / 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

To make the block table mechanism concrete, let us trace a specific attention computation step by step. This example uses small numbers to keep the arithmetic manageable, but the same procedure applies at full production scale.

Suppose we have a sequence that has generated 38 tokens so far. The block size is 16 tokens. The sequence's block table contains three entries: [7, 2, 15]. We want to compute the attention score between the query for the current (39th) token and the key for token 20.

Step 1: Identify the logical block. Token 20 falls in logical block $\lfloor 20 / 16 \rfloor = 1$ (the second block, zero-indexed). Its position within that block is $20 \bmod 16 = 4$ (the fifth slot in the block).

Step 2: Look up the physical block. We read entry 1 from the block table: $\text{block\_table}[1] = 2$ . So logical block 1 is stored in physical block 2.

Step 3: Compute the physical address. The key for token 20 lives at physical block 2, slot 4. If each block stores keys for one attention head in a contiguous sub-array of shape $[16, d_k]$ , the exact memory address is: physical block 2's base address + 4 * (size of one key vector).

Step 4: Compute the attention score. Load the key vector from that physical address. Compute the dot product with the current query vector and divide by $\sqrt{d_k}$ .

Step 5: Repeat for all 38 tokens. The same lookup applies for tokens 0 through 37. Tokens 0-15 live in physical block 7, tokens 16-31 live in physical block 2, and tokens 32-37 live in the first six slots of physical block 15. The scores for all 38 tokens are computed by following block-table lookups for each one.

Step 6: Apply softmax and weighted sum. Once all 38 scores are computed, apply softmax to get attention weights, then compute the weighted sum of value vectors (retrieved using the same physical block lookups).

The entire computation is mathematically identical to standard attention on a contiguous 38-token KV tensor. The only difference is that instead of reading keys[20] directly, we follow the two-level addressing: block table lookup, then within-block offset. On GPU hardware, this extra indirection costs only a tiny fraction of the total compute time because the block tables fit in shared memory or registers and the block lookups can be pipelined with the arithmetic.

Let us verify this numerically with a small concrete example:

In[21]:

Code

import torch

# Worked example parameters
torch.manual_seed(0)
block_size_ex = 16
d_k = 8
num_tokens = 38

# Block table: logical blocks [0, 1, 2] -> physical blocks [7, 2, 15]
# We only have 16 physical blocks in this example, so cap at 15
block_table_ex = [7, 2, 15]

# Create a small physical KV cache with 16 blocks
num_phys_blocks = 16
key_cache_ex = torch.randn(num_phys_blocks, block_size_ex, d_k)
value_cache_ex = torch.randn(num_phys_blocks, block_size_ex, d_k)

# Query for the new token
query_ex = torch.randn(d_k)
scale_ex = 1.0 / (d_k**0.5)

# Step-by-step paged attention for token 20
token_idx = 20
logical_block = token_idx // block_size_ex
within_block_offset = token_idx % block_size_ex
physical_block = block_table_ex[logical_block]
key_for_token20 = key_cache_ex[physical_block, within_block_offset]

# Gather all keys for the full sequence
gathered_keys = []
for t in range(num_tokens):
    lb = t // block_size_ex
    wb = t % block_size_ex
    pb = block_table_ex[lb]
    gathered_keys.append(key_cache_ex[pb, wb])
all_keys = torch.stack(gathered_keys)  # (38, d_k)

# Compute attention scores
scores = torch.matmul(all_keys, query_ex) * scale_ex  # (38,)
attn_weights = torch.softmax(scores, dim=0)

# Gather all values and compute output
gathered_values = []
for t in range(num_tokens):
    lb = t // block_size_ex
    wb = t % block_size_ex
    pb = block_table_ex[lb]
    gathered_values.append(value_cache_ex[pb, wb])
all_values = torch.stack(gathered_values)  # (38, d_k)

output_ex = torch.matmul(attn_weights.unsqueeze(0), all_values).squeeze(0)

Out[22]:

Console

Worked Example: PagedAttention Token Lookup
==================================================
Sequence length: 38 tokens
Block size: 16
Block table: [7, 2, 15]

Lookup for token 20:
  Logical block:      1  (= 20 // 16)
  Within-block offset:4  (= 20 % 16)
  Physical block:     2  (block_table[1])

Key vector norm at token 20: 1.6488
Max attention weight: 0.1318 at position 9
Output vector norm: 0.7807

Block coverage:
  Logical block 0 -> Physical block 7: tokens 0 to 15
  Logical block 1 -> Physical block 2: tokens 16 to 31
  Logical block 2 -> Physical block 15: tokens 32 to 37

The output confirms that the block-table indirection correctly assembles the full key and value tensors. Each token's KV data was retrieved by first mapping its position to a logical block number, then looking up the physical block, and finally reading from the correct within-block slot. The attention scores and the weighted output are computed identically to standard attention, and the result is mathematically exact regardless of how scattered the physical blocks are.