Paged Key-Value Residency and Block Tables

A cache allocated as one contiguous extent per sequence has three problems that look like different problems and are the same one. Memory is reserved for the maximum length before the length is known. A sequence that stops early holds its reservation until the end of the request. And when reservations of differing sizes are handed out and returned, the pool develops holes that cannot be filled by a request needing a large extent even though total free bytes suffice. Paging solves all three with one change: allocate in fixed-size blocks, and record the mapping from sequence positions to blocks in a per-sequence table.

The virtual-to-physical indirection

A sequence’s cache is addressed as if it were contiguous, position $t$ mapping to a byte offset $t \times c$ where $c$ is the per-token cost. That virtual layout is what the attention kernel wants. The physical layout is a set of blocks drawn from a global pool in arbitrary order. The block table is the translation between them: an array of block identifiers, indexed by the block number within the sequence.

With a block holding $P$ tokens, position $t$ lives in logical block $\lfloor t/P \rfloor$ at slot $t \bmod P$. The attention kernel reads the logical block index, looks up the physical block identifier in the table, and computes the physical address. The table is small: a 4,096-token sequence at $P = 16$ needs 256 entries, and at $P = 256$ needs 16. It lives in device memory alongside the pool and is read once per attention call.

The indirection costs one extra memory load per block touched, and it changes the shape of the attention read. A contiguous cache is a linear scan; a paged cache is a gather over blocks, where each block is contiguous internally. Since attention for one query token reads the entire context, the gather is over $\lceil T/P \rceil$ blocks, and the block size is chosen so that each block is large enough to be read efficiently and small enough that the internal fragmentation stays small.

Allocation and the end of reservation

With paging, admission reserves blocks rather than extents, and blocks are allocated as the sequence grows. A request that declares a 4,096-token maximum is admitted with a promise of $\lceil 4096/P \rceil$ blocks, but only the blocks covering tokens actually produced are taken from the pool. A sequence that stops at 200 tokens holds 13 blocks instead of 256.

The saving is the difference between worst-case and actual length, which for interactive traffic is large. It converts a reservation problem into a growth problem: the sequence needs a new block every $P$ tokens, and the allocation must succeed at that moment or the sequence must be preempted. This is the same trade as demand paging in an operating system, and it carries the same consequence, that the system needs a policy for what to do when a page fault cannot be served.

Fixed-size blocks also bound external fragmentation. If every allocation is exactly one block and blocks are all the same size, the pool can never develop an unusable hole: any free block serves any need. The only waste is internal, the partially filled last block of each sequence, which averages half a block per sequence. For a pool of $M$ bytes and block size $S$, the waste is at most $B \cdot S/2$ bytes for $B$ resident sequences, which is a predictable number rather than a workload-dependent one.

Reference counts and shared blocks

Once allocation is by block and the table is a level of indirection, two sequences can point at the same physical block. The mechanism that makes this safe is a reference count per block: a block is freed only when the count reaches zero.

Sharing arises naturally in three places. Sequences that begin with the same prompt can share the prompt’s blocks until one of them diverges; parallel samples from one prompt share the entire prompt cache. Beam search branches share the prefix up to the branch point. In each case the sharing is copy-on-write at block granularity: when a sequence needs to write a token into a block that is shared, it takes a fresh block, copies the block’s contents, and updates its table entry. The copy costs one block of memory traffic, which for $P = 16$ tokens at 128 KiB per token is 2 MiB.

Copy-on-write at block granularity means the sharing decision has a cost that depends on block size. Large blocks make copies expensive and tables small; small blocks make copies cheap and tables large. The break-even is where the table read overhead per block becomes comparable to the attention work over that block, which for typical head dimensions is a few tens of tokens per block. Blocks in the low hundreds of tokens sit in the comfortable range.

The prefix-sharing case is worth its own treatment, since the residency saving there is the largest available in most workloads; prefix caching and shared residency covers the hashing and eviction side.

What the kernel does differently

Attention over a paged cache is usually written as a fused kernel that iterates logical blocks, translates each to a physical block, and accumulates. The translation is a table lookup on the device; the loop structure is:

// One thread computes attention for one query vector over a paged cache.
// Layout: each physical block holds BLOCK_TOKENS keys then BLOCK_TOKENS values,
// each of HEAD_DIM floats, so a block is contiguous in memory.
__device__ void attend(const float *q, const float *k_cache, const float *v_cache,
                       const int *table, int n_blocks, int tail, float scale,
                       float *out) {
    float max_s = -INFINITY, sum = 0.0f;
    for (int d = 0; d < HEAD_DIM; ++d) out[d] = 0.0f;

    for (int lb = 0; lb < n_blocks; ++lb) {
        const int n = (lb + 1 == n_blocks) ? tail : BLOCK_TOKENS;   // ragged last block
        const float *kb = k_cache + (long)table[lb] * 2 * BLOCK_TOKENS * HEAD_DIM;
        for (int j = 0; j < n; ++j) {
            const float *k = kb + (long)j * HEAD_DIM;
            const float *v = k + (long)BLOCK_TOKENS * HEAD_DIM;     // values follow keys
            float s = 0.0f;
            for (int d = 0; d < HEAD_DIM; ++d) s += q[d] * k[d];
            s *= scale;
            const float new_max = fmaxf(max_s, s);                  // online softmax
            const float corr = expf(max_s - new_max);
            const float w = expf(s - new_max);
            sum = sum * corr + w;
            for (int d = 0; d < HEAD_DIM; ++d) out[d] = out[d] * corr + w * v[d];
            max_s = new_max;
        }
    }
    for (int d = 0; d < HEAD_DIM; ++d) out[d] /= sum;
}

The n term is where the ragged end of the sequence appears: every block is full except the last, and the kernel needs the sequence’s true length to know how many entries that block holds even though the table is block-granular. The table[lb] lookup is the only indirection, and it is what turns a linear scan of logical blocks into a gather over physical ones.

The layout inside a block is a separate choice from the block table. The snippet above places all of a block’s keys before its values, so a block is one contiguous run and the value for position $j$ sits a fixed stride after the key. The alternative keeps two arrays, one for keys and one for values, with the same block identifier indexing into both; the block stride is then uniform across the two reads and the two can be issued without a computed offset. Both work. What does not work is a per-token interleaved layout across the whole cache, because then the gather stride inside a block changes at every token and the read degenerates into a strided access pattern.

Sizing the block

The block size is the single tunable in the design, and it trades four quantities against each other.

Block size	Table size	Internal waste	Copy-on-write cost	Gather efficiency
Small (tens of tokens)	Large	Small	Low	Poorer
Large (thousands of tokens)	Small	Up to a block per sequence	High	Better

A small block wastes little at the tail and copies cheaply, but a long sequence’s table grows large enough to matter for both memory and lookup bandwidth, and each gather touches fewer contiguous bytes. A large block makes each read efficient and the table negligible, but a sequence that stops just past a block boundary wastes nearly a full block, and any copy-on-write duplicates that entire block.

The dominant consideration is usually the copy, because prefix sharing turns every branch into a copy. A second consideration is the growth granularity: a sequence allocates a new block every $P$ tokens, so $P$ also sets how far ahead the allocator must look. Very small blocks mean an allocation decision every few tokens, which is host-side work on the critical path of the step.

Failure modes

Table capacity. The table must be sized for the maximum context, not the current length, so a sequence with a huge declared maximum needs a table entry per possible block even if it never reaches them. For a 100,000-token maximum and 16-token blocks that is 6,250 entries per sequence.

Block exhaustion mid-step. When the pool empties, the next allocation fails and the sequence cannot advance. The scheduler must then preempt something, and the choice between evicting the failing sequence and evicting a different one is a policy question. Preemption by swap and by recomputation is covered in preemption by swap versus recomputation.

Fragmentation of the last block. A block shared between two sequences with different lengths has valid data only up to the shorter length. The table cannot express that, so the sharing rule must be that a shared block is exactly full, or that the length is tracked per sequence separately from the block count. Implementations that conflate the two produce attention over garbage positions.

Concurrent modification. A block being copied on write must not be read by another sequence’s kernel in the same step. Since all sequences in a step are launched together, the copy has to complete before the step’s attention runs, or the copy has to target a block that is not referenced by any sequence in the current batch.

Paging does not make the cache smaller. It makes the cache’s size match what is actually live, which is a different and usually larger gain. It also makes the cache pool a shared resource with a single allocator, which is what allows the scheduler to reason about admission, sharing, and eviction with one set of numbers rather than per-sequence extents. The cost is a level of indirection in the hottest loop in the system, and the design lives or dies on making that indirection cheap.