Paged KV cache

Deep dive

How it works

In: token ids of a request and a pool of fixed-size KV blocks · Out: a physical slot for each token's K/V and the block table the attention kernel reads · State: free list, per-block refcounts, one block table per sequence, prefix-hash map · Cost: O(1) per block alloc/free; one extra table lookup per 16 tokens in attention

physical pool (per layer): [blk0][blk1][blk2][blk3][blk4] ...   each = 16 tokens × kv_heads × head_dim (K and V)

seq A  block table: [3, 0, 4]      tokens 0-15 → blk3, 16-31 → blk0, 32-40 → blk4 (partially filled)
seq B  block table: [3, 1]         forked from A after token 15: shares blk3 (refcount 2)

slot for token t of seq s = table[s][t / 16] * 16 + t % 16

1. Prefill: allocate ⌈prompt_len/16⌉ blocks and write K/V through the slot mapping (the reshape_and_cache kernel).
2. Decode step: if the last block is full, allocate one new block. If it is shared (refcount > 1), copy-on-write it first.
3. Attention: the kernel receives block_tables[batch, max_blocks] and context_lens, loads K/V block by block (one gather per 16 tokens), and runs an online softmax as in FlashAttention.
4. Finish or preempt: decrement refcounts and return blocks to the free list. Under memory pressure, preempt a sequence by freeing it and recomputing later, or swap its blocks to CPU.
5. Prefix caching: key each full block by a hash of (parent block hash, its 16 token ids). A new request whose prefix hashes match reuses the blocks without recomputing them. SGLang keeps the same information in a radix tree.

Code

Block manager with refcounts, copy-on-write and slot mapping: the CPU side of a paged KV cache.

pub const BS: usize = 16;                               // tokens per block

pub struct BlockPool { free: Vec<u32>, refcnt: Vec<u32> }
#[derive(Clone, Default)]
pub struct Seq { pub table: Vec<u32>, pub len: usize }

/// What the GPU must do for this append: write K/V at `slot`, maybe copy a block first.
pub struct Append { pub slot: usize, pub cow_copy: Option<(u32, u32)> } // (src, dst)

impl BlockPool {
    pub fn new(n: u32) -> Self { Self { free: (0..n).rev().collect(), refcnt: vec![0; n as usize] } }
    pub fn num_free(&self) -> usize { self.free.len() }
    fn alloc(&mut self) -> Option<u32> {
        let b = self.free.pop()?;
        self.refcnt[b as usize] = 1;
        Some(b)
    }
    fn decref(&mut self, b: u32) {
        let r = &mut self.refcnt[b as usize];
        *r -= 1;
        if *r == 0 { self.free.push(b); }
    }
    /// Reserve a slot for one more token. None = out of blocks: preempt someone.
    pub fn append(&mut self, s: &mut Seq) -> Option<Append> {
        let off = s.len % BS;
        let mut cow_copy = None;
        if off == 0 {
            let b = self.alloc()?;
            s.table.push(b);
        } else {
            let last = *s.table.last().unwrap();
            if self.refcnt[last as usize] > 1 {          // shared partial block: copy-on-write
                let nb = self.alloc()?;
                self.decref(last);
                *s.table.last_mut().unwrap() = nb;
                cow_copy = Some((last, nb));
            }
        }
        let slot = *s.table.last().unwrap() as usize * BS + off;
        s.len += 1;
        Some(Append { slot, cow_copy })
    }
    /// Parallel sampling / beam search: share all blocks, copy nothing.
    pub fn fork(&mut self, s: &Seq) -> Seq {
        for &b in &s.table { self.refcnt[b as usize] += 1; }
        s.clone()
    }
    pub fn release(&mut self, s: &mut Seq) {
        for b in s.table.drain(..) { self.decref(b); }
        s.len = 0;
    }
}

fn main() {
    let mut pool = BlockPool::new(8);
    let mut a = Seq::default();
    for _ in 0..20 { pool.append(&mut a).unwrap(); }        // 20-token prompt → 2 blocks
    assert_eq!((a.table.len(), pool.num_free()), (2, 6));
    let mut b = pool.fork(&a);                              // shares both blocks
    let ap = pool.append(&mut b).unwrap();                  // b writes into the shared partial block
    assert!(ap.cow_copy.is_some());                         // → it gets its own copy
    assert_eq!(b.table[0], a.table[0]);                     // full prefix block stays shared
    assert_ne!(b.table[1], a.table[1]);
    pool.release(&mut a);
    pool.release(&mut b);
    assert_eq!(pool.num_free(), 8);
    println!("paged KV ok");
}

Complexity and constants

• KV size per token: Llama-2-13B (40 layers, 40 heads × 128, fp16) = 2·40·5120·2 = 819 KB. Llama-3-8B (32 layers, 8 KV heads × 128) = 128 KB. Llama-3-70B (80 layers, 8 × 128) = 320 KB.
• Capacity example: an 80 GB H100 holding a 16 GB 8B model leaves about 60 GB for KV, about 460k tokens, i.e. about 110 concurrent 4k-token sequences. Contiguous max-length reservation at 32k would fit only 14.
• Waste: at most BS−1 slots per sequence (about 0.5 block on average) plus whatever the kernel launch reserves. In vLLM's measurements, existing systems used only 20-38% of KV memory for actual tokens, while vLLM reached about 96%.
• Kernel overhead of the indirection is reported at about 20-26% for the attention kernel alone, and small end to end.

Efficiency

• Cost: KV is 2 · layers · kv_heads · head_dim · bytes per token: 320 KB for Llama-3-70B in fp16 (computed above). Block alloc/free is O(1) and waste is under one block per sequence. The attention kernel pays 20-26% more latency for the block-table indirection, small end to end (Kwon et al.). Bound: HBM capacity (it caps the batch) and HBM bandwidth in decode.
• Levers:
  • paged blocks instead of max-length reservation — KV memory holding real tokens rises from 20-38% to about 96%, 2-4x throughput over FasterTransformer and Orca (Kwon et al. 2023) — variable-length serving
  • prefix caching by block hash — shared prompt blocks are neither recomputed nor stored twice (vLLM, SGLang) — shared system prompts, few-shot examples, multi-turn chats
  • copy-on-write forks — sibling samples and beams share prompt blocks (Kwon et al.) — parallel sampling, beam search
  • fewer KV bytes per token (GQA, MLA, fp8 KV) — capacity scales inversely (Llama-3-8B: 128 KB/token vs 819 KB for 13B MHA, above) — when accuracy allows
  • CPU/SSD/remote KV tiers — keep long-lived prefixes (LMCache, Mooncake) — long contexts with high reuse
• Naive → standard → state of the art:
  • contiguous max-length reservation · 20-38% of KV memory holds tokens (Kwon et al.)
  • 16-token blocks, block tables, preemption · about 96% usage
  • paged blocks plus hash or radix prefix cache and KV offload tiers · reuse across requests; gain depends on prefix overlap (order of the shared fraction of each prompt)

Implementation notes

• Block size trade-off: smaller blocks waste less and share prefixes more finely; larger blocks mean fewer table lookups and better memory coalescing. 16 is vLLM's default; FlashInfer and TensorRT-LLM support other sizes.
• Lay out K as [num_blocks, kv_heads, head_dim/x, BS, x] (with x = 16 B worth of elements) so that a warp's loads are 16-byte vectors. Look up the block table once per block, not per token.
• Allocate the whole pool up front: profile peak activation memory with a dummy max batch, then give the rest (e.g. 90% of the GPU via gpu_memory_utilization) to KV blocks.
• Admission control must reserve blocks for the prompt plus a margin. Otherwise every step ends in preemption churn.
• Prefix-cache eviction: LRU over blocks with refcount 0, evicting leaves before their parents.

Pitfalls

• Forgetting copy-on-write on the partial last block after a fork corrupts sibling sequences silently.
• Speculative decoding: draft tokens append KV slots that must be rolled back (truncate len, free trailing blocks) on rejection.
• Prefix hashes must include everything that changes K/V: model, LoRA adapter, image/multimodal inputs, and position-encoding offsets for sliding windows.
• Sliding-window or chunked attention layers need per-layer block freeing, or memory grows anyway.
• Swapping to CPU over PCIe (about 25-50 GB/s) is slower than recomputation for short sequences. vLLM recomputes by default.

Variants and alternatives

In production

vLLM (origin of PagedAttention), SGLang, TensorRT-LLM (paged KV), Hugging Face TGI, FlashInfer kernels, LMDeploy, llama.cpp server (unified KV cache with sequence ids).

Implementations

Sources

• Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023. paper · code
• Zheng et al., SGLang: Efficient Execution of Structured Language Model Programs, NeurIPS 2024. paper · code
• Prabhu et al., vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention, ASPLOS 2025. paper
• Ye et al., FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving, MLSys 2025. paper · code