Inference Systems · Serving Architectures

Paged KV Cache Allocation and the Block Table

Why reserving contiguous cache for a worst-case output length caps batch size, and how fixed-size blocks with an indirection table remove most of that waste.

Decode-time attention needs the keys and values of every earlier position in the sequence. Those tensors are produced once and read on every subsequent step, so they are kept resident for the life of the request. The cache is therefore not a cache in the replacement-policy sense: nothing is evicted while a request is alive, and the total footprint is the sum over resident sequences of their current length times the per-token cost. A serving loop that admits requests against this footprint is doing memory management, and the allocation policy determines how much of the accelerator’s capacity is actually usable.

The naive policy allocates one contiguous region per request, sized for the maximum output length at admission. Its waste is structural, not incidental.

Three ways contiguous reservation loses memory

Internal fragmentation comes first. A request admitted with a 2,048-token reservation that produces 90 tokens holds 2,048 tokens’ worth of memory for 90 tokens of data. The unused tail cannot be given to anyone else, because the region was reserved as one extent. In a batch of sequences with mixed output lengths, the average occupancy of a reservation is roughly the mean output length over the maximum, and that ratio is well below one for any heavy-tailed distribution.

External fragmentation comes second. Even if reservations were exact, contiguous extents of varying size leave holes between them. A hole large enough for a 300-token request is useless to a 3,000-token request, so the allocator must either compact, which means copying live cache, or refuse admission while free memory exists in aggregate. Long-running servers accumulate a scattered free list.

Preemption cost comes third, and it is the one that shapes behaviour under pressure. When memory runs out, a contiguous allocator cannot shrink a resident request; it can only evict it entirely and later restore it. Restoration means either copying the cache out to host memory and back, or discarding it and recomputing the prompt. Both options cost far more than the marginal token that triggered the shortage.

Blocks, tables, and the indirection

The fix borrows the virtual-memory arrangement wholesale: split the cache into fixed-size blocks, allocate them on demand, and reach them through a per-sequence table. A block holds a fixed number of token slots, a value in the range of 16 for many deployments, chosen so that a block is a manageable unit for the attention kernel and large enough that table overhead stays small.

Each sequence owns a block table: an array mapping logical block index to physical block number. The table is short, proportional to the sequence length divided by the block size, and it lives on the host until the attention kernel needs it. Physical blocks come from a global free list and need not be adjacent. The last block of a sequence is partially filled, and its occupancy is tracked separately.

def append_token(self, seq, k, v):
    slot = seq.length % self.block_size
    if slot == 0:
        # the sequence has outgrown its table: take one more physical block
        block = self.free_list.pop()
        seq.table.append(block)
    physical = seq.table[-1]
    offset = physical * self.block_size + slot
    self.k_cache[offset] = k
    self.v_cache[offset] = v
    seq.length += 1

Waste now comes only from two sources. The tail block of each sequence is at most block_size - 1 slots empty, so with 16-slot blocks the expected waste is under 8 slots per sequence regardless of how long the sequence eventually becomes. Free memory is a set of blocks, and any block fits any request, so there is no external fragmentation at all. Reservation disappears: a request is admitted with one block and grows as it generates.

The arithmetic is stark. Take a model whose per-token cache cost is 512 KiB, as computed in Iteration-Level Batching. A request reserving 2,048 tokens claims 1 GiB. If actual output averages 200 tokens, the reservation is ten times the need, and a 40 GiB cache holds forty such requests. With 16-token blocks the same 40 GiB holds roughly 400 sequences at 200 tokens each, less one tail block apiece. The batch is limited by work available rather than by the worst case.

Sharing blocks between sequences

The block table is a level of indirection, and indirection admits aliasing. Two sequences whose prompts begin with the same tokens produce identical keys and values for those positions, because attention at position i depends only on positions up to i and the model is deterministic in inference mode. Those blocks can be referenced by both tables instead of copied, with a reference count on each physical block.

Reads need no coordination. A shared block is immutable while it is shared, so any number of sequences may attend to it concurrently. The moment one sequence writes past the shared prefix, it must allocate a fresh block and copy the affected slots into it before appending, which is copy-on-write with a block granularity rather than a page granularity.

Beam search and parallel sampling are the natural beneficiaries. A request that explores eight continuations from one prompt produces eight sequences whose first L tokens are identical, so the shared prefix costs one copy of cache rather than eight. If the sampler later prunes a branch, its blocks are released and the count on the shared prefix drops by one. The saving is bounded by how quickly branches diverge, which for sampling with a moderate temperature is within a few tokens, so the gain concentrates in the prompt rather than in the continuation.

The reference count is what makes reclamation correct. Retiring a sequence decrements the count of every block in its table and returns only those that reach zero. A block with two referrers survives one of them, and the free list never receives a block that another table still names. Getting this wrong produces a use-after-free that manifests as corrupted attention scores rather than as a crash, since the physical block is still valid memory but no longer holds the expected values.

What the design costs

The indirection is not free in the kernel. Contiguous cache allows attention over a sequence to be expressed as one strided access pattern; paged cache requires gathering blocks, either by rewriting the attention kernel to walk the table or by copying non-contiguous blocks into a scratch buffer before the kernel runs. The first option costs kernel complexity and a table load per block; the second costs bandwidth proportional to the cache it moves. Both are acceptable because attention over long contexts is bandwidth-bound anyway, but a naive implementation that materializes a contiguous copy per layer per step can lose much of the gain.

Block size is a genuine trade with no dominant answer. Larger blocks reduce table length and the number of gather operations, but increase tail waste per sequence and the amount copied on a copy-on-write fault. Smaller blocks do the reverse and add per-block bookkeeping. A block of 16 tokens keeps tail waste near 8 slots per sequence; a block of 128 makes tail waste up to 127 slots, which matters when many sequences are short.

Finally, the scheme does not remove the need for admission control. It converts a reservation problem into a rate problem: a sequence that keeps generating eventually consumes all blocks, and the server must decide whether to stop admitting, to swap a sequence’s blocks to host memory, or to recompute. Those decisions belong to the scheduler, and the deadline-aware policies in Deadline-Aware Batching and Tenant Fairness choose between them under explicit latency targets.

One consequence deserves emphasis: because the table is per-sequence and small, it can be copied, serialized, and moved between processes cheaply, while the blocks it names cannot. That asymmetry is what makes cache migration between machines plausible at all, and it underlies the separation of prefill from decode described in Prefill and Decode Disaggregation.

References

  1. Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., ACM SOSP 2023
  2. Operating System Concepts — Silberschatz, Galvin and Gagne, Wiley
  3. Orca: A Distributed Serving System for Transformer-Based Generative Models — Yu et al., USENIX OSDI 2022