LLM DAILYA field guide to language models

Day 64 / 8 min read

PagedAttention — virtual memory for the KV cache

Yesterday we explored prefix caching: when two requests share an identical prompt prefix, an inference engine can reuse the corresponding key/value (KV) states instead of recomputing them.

That raises a practical question:

How does a serving system store millions of KV vectors when requests have different lengths, grow one token at a time, and sometimes share prefixes?

The answer in many modern inference engines is paged KV-cache management, popularized by PagedAttention.

The intuition is borrowed from operating systems:

\[ \boxed{\text{give each request a logical sequence of KV blocks, but store those blocks wherever memory is free}} \]

The model still performs the same attention calculation. What changes is how the runtime organizes its memory.

The problem with reserving one big KV buffer

Imagine three conversations:

Request A:  120 tokens
Request B: 1800 tokens
Request C:  430 tokens

All three may grow as the model generates new tokens. Their final lengths are not known in advance.

A straightforward allocator could reserve enough contiguous GPU memory for each request's maximum allowed length:

GPU KV memory:

[A: 2048-token reservation              ]
[B: 2048-token reservation              ]
[C: 2048-token reservation              ]

That is simple, but much of the reserved space may sit unused.

Alternatively, the system could allocate only the currently needed space and repeatedly enlarge it. But growing contiguous allocations may require copying data, and freeing requests of different sizes leaves awkward holes.

These are familiar memory-management problems: overreservation, fragmentation, and relocation.

They matter because KV memory competes with model weights and active batches for scarce GPU memory.

A quick reminder: what attention needs

At decode position \(t\), one attention head computes something like:

\[ \operatorname{Attention}(q_t,K_{1:t},V_{1:t}) = \operatorname{softmax}\left( \frac{q_tK_{1:t}^{\top}}{\sqrt{d_h}} \right)V_{1:t}. \]

The new query \(q_t\) needs access to all relevant previous keys and values.

Crucially, the equation says which logical tokens the query attends to. It does not require those tokens' KV vectors to occupy consecutive physical GPU addresses.

That distinction creates the opportunity for paging.

Divide the KV cache into fixed-size blocks

Suppose each physical block holds KV states for 16 token positions.

Instead of reserving one contiguous region per request, the runtime maintains a pool of blocks:

Physical GPU blocks:

Block 0  [16 token slots]
Block 1  [16 token slots]
Block 2  [16 token slots]
Block 3  [16 token slots]
Block 4  [16 token slots]
...

Each request gets a block table mapping its logical blocks to physical ones.

For example:

Request A logical blocks:   0    1    2
                            │    │    │
                            ▼    ▼    ▼
Physical block IDs:        17    4   29

The physical blocks do not need to be adjacent.

As Request A grows past 48 tokens, the engine assigns another free block and extends its table.

This is analogous to virtual memory: a program sees a contiguous logical address space even though its pages may be scattered across physical RAM.

How attention reads scattered blocks

A paged-attention kernel uses the block table to locate the relevant KV data:

new query token
       │
       ▼
logical positions 0 ... t
       │
       ▼
request's block table
       │
       ▼
physical blocks 17, 4, 29, ...
       │
       ▼
attention computation

The kernel gathers or streams keys and values from the mapped blocks while computing attention.

Mathematically, the result is intended to match ordinary attention over the same sequence. The physical layout is an implementation detail.

This is why PagedAttention is primarily a memory-management and kernel-design technique, not a new attention rule that changes which tokens can attend to which others.

Worked example: how much memory can paging save?

Use the illustrative model from Day 63:

\[ L=32,\quad H_{kv}=8,\quad d_h=128,\quad b=2\text{ bytes}. \]

Its KV bytes per token are approximately:

\[ m_{\text{token}}=2LH_{kv}d_hb =2(32)(8)(128)(2) =131{,}072\text{ bytes}. \]

That is 128 KiB per token, across all layers, under these assumptions.

A 16-token block therefore holds approximately:

\[ 16\times128\text{ KiB}=2\text{ MiB} \]

of KV data across the model's layers. Actual engines may allocate separate physical blocks per layer; the 2 MiB figure is the aggregate equivalent.

Now imagine 500 active requests, each currently 1,500 tokens long, but each given a 2,048-token contiguous reservation.

Unused reserved positions per request:

\[ 2048-1500=548. \]

Across 500 requests, that is:

\[ 500\times548\times128\text{ KiB} \approx33.4\text{ GiB} \]

of reserved-but-unused KV capacity.

With 16-token blocks, a 1,500-token request needs:

\[ \left\lceil\frac{1500}{16}\right\rceil=94 \]

blocks, or 1,504 slots. Only four slots are unused per request:

\[ 500\times4\times128\text{ KiB} =250\text{ MiB}. \]

In this simplified comparison, paging reduces unused reserved capacity from about 33.4 GiB to 250 MiB.

This is not a universal measured saving: real allocators differ, requests have varying lengths, and block tables and kernel machinery add overhead. The example isolates the benefit of allocating near actual sequence length rather than reserving a fixed maximum.

Internal versus external fragmentation

Paging does not eliminate all waste.

If a request has 33 tokens and blocks hold 16, it needs three blocks:

Block 1: 16 used
Block 2: 16 used
Block 3:  1 used, 15 empty

Those 15 unused slots are internal fragmentation.

But a freed block can be returned to the shared pool and reused by any other request. We no longer need to find a large contiguous hole for an entire sequence, greatly reducing the external fragmentation problem.

Smaller blocks reduce internal waste but increase the number of block-table entries and potentially the overhead of address lookup and memory access.

Larger blocks reduce bookkeeping but waste more space in partially filled final blocks.

So block size is a systems tradeoff, not an arbitrary constant.

Why this fits continuous batching

Recall continuous batching: an inference server admits new requests and retires completed ones without waiting for an entire fixed batch to finish.

Imagine the active set changes:

step 1: A, B, C
step 2: A, B, C, D
step 3: A, C, D       (B finished)
step 4: A, C, D, E

With paged KV storage, when B finishes, its blocks return to the free pool. D or E can reuse them without relocating A or C.

When A generates another 16 tokens, the runtime can simply append a new block mapping.

This flexibility makes it easier to keep GPU memory occupied by useful work rather than stranded reservations.

How paging connects to prefix caching

Yesterday's prefix-caching lesson becomes more concrete with blocks.

Suppose two requests start with the same 32-token system prompt:

Request A: [shared 16][shared 16][A-specific ...]
Request B: [shared 16][shared 16][B-specific ...]

Their block tables can point to the same physical prefix blocks:

A: [block 7][block 12][block 20]...
B: [block 7][block 12][block 31]...

Blocks 7 and 12 are shared; later blocks differ.

The runtime tracks references so shared blocks are not freed while either request still needs them.

For branching or beam-search-like workflows, copy-on-write can preserve shared blocks until a branch needs to modify a block. For ordinary append-only decoding, new suffix blocks can simply be allocated separately.

Paging supplies the physical-memory mechanism; prefix caching supplies the policy for identifying reusable computation. They are complementary, not synonymous.

Why paging does not magically solve every KV problem

Three caveats matter.

First, paging does not compress KV data. A million stored token positions still require roughly a million positions' worth of KV bytes unless you also use techniques such as KV quantization, fewer KV heads, or eviction.

Second, the block table introduces indirection. The attention kernel must translate logical positions into physical locations, and scattered blocks can make memory access less straightforward. Good kernels are essential.

Third, sharing or reusing blocks must preserve correctness and isolation. A block from another tenant or an incompatible model configuration must never be treated as a valid cache hit merely because it happens to occupy the same memory address.

The win comes from better allocation and reuse, not from changing the attention mathematics.

A concrete mental model: a library with movable shelves

Imagine a library that gives every reader a private row of 2,048 book slots, even if most readers borrow only a few hundred books. The building quickly fills with empty reserved shelves.

A better library stores books in standardized crates of 16:

Reader A's catalog: crate 17 → crate 4 → crate 29
Reader B's catalog: crate  8 → crate 6

The crates can live anywhere in the warehouse. Each reader's catalog records their logical order.

When a reader needs more books, the library assigns another crate. When they leave, their crates become available to someone else.

If two readers need the same reference volume, they can even point to a shared crate, provided the library manages ownership and access correctly.

For an LLM service:

reader's catalog = request block table
crate            = physical KV block
book order       = logical token positions
warehouse        = GPU KV memory pool
shared reference = reusable prefix block

The reader sees an ordered collection of books; the warehouse is free to store the crates wherever space exists.

The deeper lesson

Recent lessons have explored three different ways to avoid inference waste:

speculative decoding → avoid some sequential target-model steps
prefix caching       → avoid recomputing identical prompt prefixes
PagedAttention       → avoid wasting KV memory on rigid allocations

These optimizations reinforce one another. Better memory utilization allows more concurrent requests; more concurrency improves batching; and prefix sharing can reduce both prefill computation and duplicated KV storage.

Today's main principle is:

\[ \boxed{\text{logical sequence order need not equal physical memory order}} \]

That is an old operating-systems idea with enormous consequences for modern LLM serving. The Transformer specifies what attention must compute; the inference runtime decides how to store the state efficiently enough to compute it at scale.

LLM Daily Last updated