Paged attention and prefix caching.

Before a model writes a single word of its answer, it reads your whole prompt and builds a set of notes about it in GPU memory, one entry per word-piece. Prefix caching means keeping those notes after the answer is done, so the next request that starts with the same text can pick up where the last one stopped instead of reading it all over again.

This matters because almost nothing you send is new. Your agents resend the same long instructions on every step, and every chat turn resends the whole conversation. Without a cache, your Spark server spends about a second rereading text it saw a moment ago, every single time. It is also why Anthropic and OpenAI charge less for input they have seen recently, and part of why Fireworks can look faster than you on repeated prompts.

You already know where the notes come from (The KV cache, by hand) and how an operating system chops memory into small numbered pieces it can hand out and share (Virtual memory). Here those two ideas meet, inside the server from Baseline: my Spark vs Fireworks. After this, the memory managers inside vLLM and SGLang will look familiar, because you will have built a small one.

Today

• Put a byte count on the notes one request keeps in memory (its KV cache), and explain why reserving room for the longest possible answer, for every request, leaves most of the GPU empty.
• Explain how servers cut that memory into small fixed-size blocks, with a lookup table per request saying where each block lives (vLLM calls this PagedAttention), and why it is the same trick as virtual memory pages.
• Explain how blocks get shared: inside one request (several answers to one prompt point at the same blocks, and a block is copied only when one of them needs to write to it) and across requests (each block gets a fingerprint of all the text up to it, looked up like a dictionary).
• Build that cache into your server: the fingerprints, the lookup, a count of who is using each block, and two rules for throwing blocks out (least recently used goes first, and anything idle too long expires). Measure the wait for the first word with a shared 2,000-token system prompt, and how often the cache hits as you change its size.
• Say exactly what this speeds up (the wait before the first word of the answer) and what it does not (each word after that), and how the same idea becomes sending requests to the machine that already holds their prompt, and the “prompt caching” line on an API bill.

You need the server, harness and SCOREBOARD.md from Baseline: my Spark vs Fireworks, plus pages, page tables and copy-on-write from Virtual memory and the OOM killer. The KV cache, by hand and Batching help, but everything needed from them is recapped here. Plan on two evenings: one for the idea and the cache, one for wiring it in and measuring.

The Problem

Your agent has a system prompt. Tool descriptions, rules, a few worked examples: about 2,000 tokens, and it never changes. Each step of the agent sends that same system prompt, plus the conversation so far, plus one new line.

Watch what your server does with it. Every request, it runs prefill over all 2,000 tokens of system prompt, computes exactly the same keys and values it computed a second ago, uses them for one answer, and throws them away. From the baseline lesson you have a rough number for this: prefill of 2,000 tokens on the Spark is on the order of a second. So every agent step pays about a second of time to first token to rediscover something it already knew.

It gets worse with chat. Turn 10 of a conversation resends turns 1 to 9. The prompt grows every turn, and the work your server repeats grows with it.

And there is a second, quieter problem, about memory. When you batch requests (the next step after the naive server), each one needs room for its KV cache. How much room? You don’t know how long the answer will be. So the obvious design reserves room for the longest possible one. Most requests use a fraction of it. The GPU fills up with reserved, empty memory, and you can fit far fewer users than the hardware could hold.

Both problems have the same fix: stop treating a request’s KV cache as one private, contiguous lump. Cut it into small fixed-size blocks, keep a table of where each block lives, and let requests point at the same block when their tokens are the same.

The Idea

How big is one request’s KV cache?

Recap. A transformer layer turns each token into a query, a key and a value. To produce the next token, the new token’s query is compared against the keys of every earlier token, and the matching values are mixed together. The KV cacheKV cacheThe stored keys and values of every token processed so far, for every layer. Kept so that each new token only computes its own K and V instead of recomputing everyone’s. keeps those keys and values around, so each decode step only computes one new token’s worth.

Now the bytes, for Llama 3.1 8B in bf16. It has 32 layers. It uses grouped-query attentiongrouped-query attentionGQA. Many query heads share a smaller set of key/value heads. Llama 3.1 8B has 32 query heads but only 8 KV heads, which cuts the KV cache 4 times. with 8 key/value heads, each 128 numbers wide. Per token:

2 (K and V) × 32 layers × 8 heads × 128 dims × 2 bytes = 131,072 bytes = 128 KiB per token

So a 2,000-token system prompt is about 256 MiB of KV cache. A full 128k-token context is 16 GiB, the same size as the weights. How much memory does a model need? does this carefully; for today, remember 128 KiB per token, and that one prompt’s KV is hundreds of megabytes.

Reserving the worst case

There are three ways memory gets wasted when each sequence gets one contiguous slab:

• Reserved for later. Room for tokens that haven’t been generated yet, and may never be, because the answer stops early.
• internal fragmentationinternal fragmentationMemory handed out but not used, inside an allocation. A sequence that reserved 8,192 slots and used 1,000 wastes 7,192 inside its own slab.. The leftover inside a slab when the answer does stop early.
• external fragmentationexternal fragmentationFree memory split into gaps that are each too small to use. Enough free space in total, but no single hole big enough for the next request.. Gaps between slabs. When requests of different sizes come and go, free memory ends up in holes, and a new request that needs one big contiguous slab can’t fit in any of them.

You have seen this exact problem before. In Virtual memory, physical RAM is cut into 4 KiB frames and every process gets a page tablepage tableA per-process map from virtual page numbers to physical frames. The CPU’s MMU walks it on every memory access (with the TLB as a cache) so a process can treat scattered frames as one contiguous range. that maps its pages to frames anywhere. No process needs a contiguous chunk of physical RAM. A frame is handed out only when a page is first touched. Pages can be shared between processes, and copied only when someone writes.

PagedAttention is that, for KV caches.

PagedAttention: pages for attention

Cut the KV memory into fixed-size KV blocksKV blocksFixed-size chunks of KV cache memory, each holding the keys and values of a fixed number of tokens (16 by default in vLLM) for every layer. The unit of allocation, sharing and eviction., each holding the keys and values for 16 tokens (vLLM’s default). For Llama 3.1 8B a block is 16 × 128 KiB = 2 MiB. At startup, the server allocates one big pool of blocks with whatever memory is left after the weights, and keeps a list of the free ones.

Each sequence then has a block tableblock tableA small array per sequence: entry i is the physical block that holds the sequence’s tokens 16i to 16i+15. The KV-cache version of a page table.. Logical block 0 (tokens 0 to 15) lives in physical block 7; logical block 1 (tokens 16 to 31) in physical block 1; and so on. When a sequence fills its last block, the server grabs any free block and appends one entry to the table. No reservation in advance, no need for neighbours.

The part that needed new code is the attention computation itself. The usual attention kernel assumes a sequence’s keys sit in one contiguous tensor. The PagedAttention kernel takes the block table as an extra input and, for each block of keys, looks up where it lives before reading it. That costs a little indirection per block, the same deal the CPU makes with page tables. In return, memory waste drops to less than one block per sequence: the paper reports under 4%.

The translation is direct:

Sharing inside one request

Once a sequence is a list of pointers, two sequences can point at the same block.

The first place this pays off is parallel samplingparallel samplingAsking for several independent answers to one prompt (the n parameter in OpenAI-style APIs). All samples start from the same prompt, so their prompt KV is identical.. Ask for four samples of a 2,000-token prompt and the naive server stores four identical copies of 256 MiB of prompt KV. With blocks, all four block tables point at the same prompt blocks, and each block carries a reference countreference countA counter on a shared object: how many users currently point at it. Incremented when someone starts using it, decremented when they stop. At zero, nobody needs it. of 4.

What happens when a sample writes? The prompt’s last block is usually only partly full: 2,000 tokens is 125 full blocks exactly, but a 2,008-token prompt leaves 8 tokens in block 125 and 8 empty slots. Each sample’s first generated token wants to go into those empty slots, and each sample generates a different token. They can’t all write into one shared block.

So the server does what the kernel does after fork: copy-on-writecopy-on-writeSharing a block read-only and copying it only at the moment one of the sharers needs to write to it. Only the blocks that actually differ get duplicated.. Before a sample writes into a block whose reference count is above 1, the server copies that block to a fresh one, points the sample’s table at the copy, and drops the old block’s count by one. Every full prompt block stays shared forever.

beam searchbeam searchA decoding method that keeps the k most likely partial answers at every step, extends each, and keeps the best k again. Beams constantly split and die, so they share long common histories. is the extreme case. Beams branch off each other constantly, so their histories overlap almost entirely. With reference counts, a beam that dies just decrements its blocks’ counts; a block whose count reaches zero goes back to the free list. The paper reports large memory savings for beam search from this sharing alone.

Sharing across requests: prefix caching

Here is the leap. Within one request, the samples share a prompt because they are the same prompt. But two different requests from two different users can share a prompt too: the same system prompt, the same few-shot examples, the same document, the same conversation history.

Can request B reuse request A’s KV blocks? Only for a prefixprefixThe first n tokens of a sequence. Two prompts share a prefix of length n if their first n tokens are identical., and here is why. The key and value for token i at layer 2 depend on the output of layer 1 at token i, which, through attention, depends on every token before it. And with RoPERoPERotary position embeddings. Each key is rotated by an angle that depends on its position, so the same word at position 5 and at position 500 has different keys., each key also depends on its position. So token i’s KV is a function of tokens 0 to i, at positions 0 to i. Nothing after it matters, because of the causal maskcausal maskThe rule in decoder attention that a token may only attend to itself and tokens before it, never after. It is why generation works left to right.. Nothing else does either.

That gives a crisp rule: if two prompts agree on their first n tokens, their KV for those n tokens is identical. Not similar. Identical, up to floating-point rounding. Request B can skip prefill for those n tokens and only run the rest.

Naming a block: the hash chain

To reuse a block, the server needs to find it. It needs a name for “the KV of these 16 tokens, given everything before them”, which it can compute from the new request’s tokens alone, without running the model.

vLLM’s answer: hash each full block together with the hash of the block before it.

key_0 = H(root,  tokens 0-15)
key_1 = H(key_0, tokens 16-31)
key_2 = H(key_1, tokens 32-47)

Because each key includes the previous key, key_2 depends on tokens 0 to 47, exactly like the KV it names. The same 16 tokens after a different past get a different key. It is the same chaining idea as a git commit hash, which covers the parent commit’s hash and so the whole history.

The cache is then just a hash tablehash tableA dictionary: a structure that finds a value from its key in roughly constant time, by turning the key into an array position. Python’s dict is one. from key to physical block. A cache lookupcache lookupAsking a cache “do you have this?” and getting back either the stored thing (a hit) or nothing (a miss), in which case you compute it yourself. for a new request:

1. Tokenize the prompt. Cut it into full 16-token blocks; ignore the partial tail.
2. Compute the keys left to right.
3. Walk them: if key_0 is in the table, that’s a hit; try key_1; keep going until the first miss. Everything after the first miss is a miss even if it happens to be in the table, because a block is only usable if every block before it is too. (In practice, with chained keys, a later block can’t be present without its parents unless they were evicted.)
4. Point the request’s block table at the hit blocks. Run prefill only on the tokens after them.

One catch: the model has to run on at least one prompt token, because the last prompt position’s output is what produces the first new token. If the entire prompt is cached, the server still recomputes the last token. vLLM does exactly this.

A few details matter in production, and each is a good interview answer:

• What else goes into the key. Anything that changes the KV besides the tokens: which LoRA adapter is active, the hashes of any images in the prompt, and, for multi-tenant servers, a per-tenant salt (more below). vLLM’s design docs list these as “extra keys” mixed into each block’s hash.
• Collisions. If two different prefixes ever hash to the same key, the second request silently gets the first request’s KV and produces garbage that looks plausible. With a 256-bit cryptographic hash like SHA-256 this effectively never happens by accident. With a fast 64-bit hash an attacker who can choose prompts might find one on purpose. vLLM added a SHA-256 option for this reason; check its docs for the current default.
• Timing leaks. A cache hit is faster. If users share a cache, someone could send a guess at another user’s prompt and time the response to learn whether it was cached. Salting the hash per tenant (each user only ever hits their own entries, plus explicitly public ones like a shared system prompt) closes this. vLLM accepts a cache_salt per request for exactly this; check the current API.

The tree version: RadixAttention

A hash table knows “this block exists”. It doesn’t know the shape of what is cached: which prefixes branch from which. SGLang’s RadixAttentionRadixAttentionSGLang’s prefix cache. Stores all cached prefixes as a radix tree whose edges are runs of tokens, so the longest cached prefix of a new prompt is found by walking down the tree. keeps that shape explicitly as a radix treeradix treeA compressed prefix tree (trie). Each edge is labelled with a run of tokens rather than one token, and a node exists only where two cached sequences branch apart..

Lookup is walking down from the root, matching tokens, as far as they agree. Insertion is adding or splitting an edge. The tree gives two things the flat table doesn’t make easy:

• Evict leaves first. A leaf is a suffix only one conversation used. A node near the root is shared by everything below it. Evicting the least recently used leaf keeps the valuable shared stems. SGLang does exactly that: LRU over leaves whose reference count is zero.
• Schedule by prefix. When many requests are waiting, the scheduler can see which ones share long cached prefixes and run them together, before the shared nodes get evicted. SGLang calls this cache-aware scheduling.

vLLM’s flat hash table gets close to the first point with a trick you will build yourself: when a request finishes, it puts the request’s blocks onto the free list tail first, so the end of a chain is always older than its head and gets evicted first. Same effect, no tree.

Eviction: who leaves when the pool is full

The pool is finite. A cache that never forgets would eventually hold every prefix ever seen and have no room for running requests. So when a request needs a new block and the free list is empty, something cached has to go. Three rules decide what.

Rule 1: never evict a block in use. A block a running request points at holds part of that request’s attention context. Free it and hand it to someone else, and the running request reads another request’s keys on its next step. That is not a slow request; it is a wrong one. So a block can be evicted only when its reference count is zero: nobody running is using it. If every block is pinned, the new request waits. (In vLLM, a running request can also be preempted: its blocks freed and recomputed later, a deliberate choice rather than eviction.)

Rule 2: among unpinned blocks, evict the least recently used. LRULRULeast recently used: when full, evict the item whose last use is oldest. Bets that what was used recently will be used again soon. is the default because traffic is bursty: the system prompt that served a request a second ago will probably serve another one in the next second. It is cheap to implement too: an ordered list of unpinned blocks, move a block to the back when it is released, pop from the front to evict.

Rule 3 (optional): drop anything idle longer than a TTL. A TTLTTLTime to live: how long an entry may stay in a cache without being used. After that it is treated as gone, even if nothing else needed its space. answers a different question from LRU. LRU answers “who leaves when we are full?” TTL answers “how long may an idle thing stay?” A pure LRU cache that is never under pressure keeps a prefix from last Tuesday forever. A TTL frees that memory on a schedule, which matters when the memory could be doing something else (more concurrent requests), or when the cache is a promise to a customer who is paying for it. Hits refresh the timer, so a prompt in steady use never expires.

Now play with it. Sixteen blocks of 16 tokens. Prompt A and prompt B are system prompts of four blocks each. The chat turn reuses A and grows one block of history per turn. The one-off prompt is five blocks nobody will ever send again.

Things to notice:

• A, then A again. Four hits, 64 prompt tokens of prefill skipped. The second request costs nothing but the tail.
• Chat turns. Each turn hits the system prompt and every earlier turn, and computes only one new block. That is why chat and agent loops are where prefix caching shines: the reusable prefix grows every turn.
• One-off prompts. They fill the pool with blocks that will never be hit again and push out useful ones. A real server sees plenty of these. Cache pollution is the main reason hit rates fall short of what the workload “should” give.
• Send one-offs without pressing Finish. Eventually every block is pinned and a new request can’t start at all. Reference counts are doing their job: nothing in use was evicted.
• Evictions take tails first. When A was idle and space was needed, A3 went before A0. The next A request still hits A0 to A2.
• TTL. Advance the clock and idle blocks expire even with empty slots all around. A hit refreshes the timer.

What it buys: TTFT, not decode

The precise statement: prefix caching skips prefill compute for the cached tokens. Prefill is the compute-bound phase, its cost grows with prompt length, and it is most of TTFT for long prompts. Decode is the bandwidth-bound phase from Why ~17 tokens per second?: each step reads the weights plus the KV, and none of that changes.

There are two indirect wins on a busy batched server. Prefill that doesn’t run frees compute for other requests, so throughput goes up when prompts are long and shared. And shared blocks are stored once, so more concurrent requests fit in the pool. On a one-at-a-time server like yours today, you will only see the TTFT win.

Many replicas: cache-aware routing

Run eight replicas of your server behind a load balancer. Each replica has its own KV pool, and its own cache. Round robin sends turn 2 of a conversation to a different replica than turn 1 seven times out of eight. The prefix is recomputed, and now seven replicas waste pool space on copies of it.

You met the fix in Load balancers and state: cache-aware routingcache-aware routingRouting each request to the replica most likely to already hold its prompt prefix, while still avoiding overloaded replicas.. The simple version hashes something stable (the conversation ID, or the first few blocks of the prompt) and uses consistent hashing so the same prefix keeps landing on the same replica. The better version, used by SGLang’s router, keeps an approximate radix tree of what each replica has been sent, routes to the longest match, and falls back to the least-loaded replica when the match is overloaded. Pure affinity is dangerous: one hugely popular system prompt would send everyone to one GPU.

This is a sticky session, and it is fine, because it is a cache. If the replica dies, the next request goes elsewhere and pays the prefill. Nothing is lost but time.

The productized version: API prompt caching

Every major LLM API now sells this mechanism back to you as a line item. Under the hood it is the same idea: the provider keeps the KV for your prefix somewhere (GPU memory, and for some providers host memory or disk) and skips prefill when your next request starts with the same tokens. What changes is the contract, because now someone is paying.

• Anthropic makes it explicit. You mark the end of the reusable part of the prompt with a cache_control breakpoint (up to 4 per request), or set it at the top level to let the API place it. The default TTL is 5 minutes, refreshed every time the entry is read; a 1-hour TTL is available. Writing the cache costs more than normal input (1.25× for the 5-minute TTL, 2× for 1 hour), reading it costs about a tenth (less on some newer models). There is a model-dependent minimum length below which nothing is cached. The response’s usage shows cache_creation_input_tokens and cache_read_input_tokens. Check the prompt caching docs for current numbers.
• OpenAI makes it automatic. Prompts above a minimum length (1,024 tokens at launch) are cached without any flag, matched in increments, and the cached part is billed at a discount. Entries typically last minutes of inactivity, longer off-peak, with a longer-retention option on some models. usage.prompt_tokens_details.cached_tokens tells you how many hit. There is a prompt_cache_key parameter that influences routing, which is cache-aware routing exposed to the customer. Check their prompt caching guide for current details.
• Fireworks and other open-model hosts document prompt caching too. That matters for your scoreboard: some of Fireworks’ TTFT on repeated prompts may already have been cache hits. Look for a cached-tokens field in their usage.

Why do APIs use TTLs when your server will mostly use LRU? Because an API is a promise. “Your cache lives 5 minutes after the last use” is something a customer can plan around and a provider can budget memory for. “Your cache lives until something more popular pushes it out” is neither. The write premium pays for holding the memory; the read discount is the prefill compute you didn’t burn.

The habits that make API caching work are the same ones that make your server’s cache work: put stable content first (tools, system prompt, examples), volatile content last (the date, the user’s question), and never put a timestamp or request ID at the top.

Build It

You are adding a prefix cache to the server from the baseline lesson. If you already rewrote it for batching, the cache plugs in the same way; reference counts start mattering for real once requests overlap.

An honest note before you start. Real engines share blocks in place: the attention kernel reads the cached blocks directly through the block table, with no copying. Hugging Face’s generate expects each request’s KV as ordinary contiguous tensors. So your version stores each cached block as its own small tensors, and on a hit copies the hit blocks into a fresh cache for the request. That costs a copy of the cached prefix per request (about 256 MiB for a 2,000-token prompt, a few milliseconds of memory bandwidth plus Python overhead), and during a request the prefix lives in memory twice. You get the hashing, the lookup, the reference counting and the eviction exactly right, and the TTFT win; you do not get PagedAttention’s memory sharing. That is the trade.

inference/
  prefix_cache.py          # new: the cache (yours)
  test_prefix_cache.py     # new: the spec
  server_naive.py          # gains a cache, behind PREFIX_CACHE=1
  bench.py                 # records cached_tokens
  prompts/shared_2k.jsonl  # new
  prompts/multi_prefix_zipf.jsonl, multi_prefix_rr.jsonl   # new

# prefix_cache.py: block-hashed prefix cache. Knows nothing about the model.
import hashlib, threading, time
from collections import OrderedDict
from dataclasses import dataclass


@dataclass
class Block:
    key: bytes          # chained hash: names this block AND everything before it
    kv: object          # per layer (K, V) for these B tokens; opaque to the cache
    ref: int = 0        # running requests using it right now
    last_used: float = 0.0


class PrefixCache:
    def __init__(self, block_size=16, max_blocks=1024, ttl_s=None, salt=b"", clock=time.monotonic):
        self.B, self.max_blocks, self.ttl = block_size, max_blocks, ttl_s
        self.root = hashlib.sha256(b"prefix-cache|" + salt).digest()
        self.blocks: dict[bytes, Block] = {}
        self.free_lru: OrderedDict[bytes, None] = OrderedDict()  # ref == 0 only, oldest first
        self.lock = threading.Lock()
        self.clock = clock
        self.stats = dict(lookups=0, hit_blocks=0, miss_blocks=0, evicted=0, expired=0)

    def hash_blocks(self, ids: list[int]) -> list[bytes]:
        """One key per FULL block: key_i = sha256(key_(i-1) + bytes of the block's token ids)."""

    def lookup(self, keys: list[bytes], max_hit: int) -> list[Block]:
        """Longest cached prefix, at most max_hit blocks. Sweep expired blocks first.
        Pin each returned block: ref += 1, and take it off free_lru if it was there."""

    def insert(self, key: bytes, kv) -> bool:
        """Store a block the current request just computed, pinned (ref = 1).
        Evict until there is room. Return False (don't cache it) if everything is pinned."""

    def release(self, blocks: list[Block]) -> None:
        """The request is done with these. ref -= 1; stamp last_used; any block reaching
        ref 0 goes to the BACK of free_lru. Think hard about the order you do this in."""

    def _evict_one(self) -> bool:
        """Drop the front of free_lru. False if free_lru is empty."""

    def sweep(self, now=None) -> None:
        """TTL: drop unpinned blocks idle longer than ttl_s. They are at the front of free_lru."""

# test_prefix_cache.py   run: python -m pytest -q
from prefix_cache import PrefixCache


class Clock:
    def __init__(self): self.t = 0.0
    def __call__(self): return self.t


def fill(c, ids):
    """Simulate one request: look up, insert what was missing, finish."""
    keys = c.hash_blocks(ids)
    hit = c.lookup(keys, max_hit=len(keys))
    new = [c.blocks[k] for k in keys[len(hit):] if c.insert(k, kv=None)]
    c.release(hit + new)
    return len(hit), len(keys)


def test_hash_chain():
    c = PrefixCache(block_size=4)
    a = c.hash_blocks([1, 2, 3, 4, 5, 6, 7, 8, 9])
    b = c.hash_blocks([1, 2, 3, 4, 5, 6, 7, 0])
    assert len(a) == 2                       # the partial tail is not a block
    assert a[0] == b[0] and a[1] != b[1]
    x = c.hash_blocks([9, 9, 9, 9, 5, 6, 7, 8])
    assert x[1] != a[1]                      # same tokens, different past, different key


def test_hit_after_insert():
    c = PrefixCache(block_size=4)
    assert fill(c, list(range(12))) == (0, 3)
    assert fill(c, list(range(12))) == (3, 3)
    assert fill(c, list(range(8)) + [99] * 4) == (2, 3)


def test_pinned_blocks_are_never_evicted():
    c = PrefixCache(block_size=4, max_blocks=2)
    keys = c.hash_blocks(list(range(8)))
    for k in keys:
        assert c.insert(k, kv=None)          # pinned (ref 1), never released
    other = c.hash_blocks([7] * 4)
    assert c.insert(other[0], kv=None) is False
    assert all(k in c.blocks for k in keys)


def test_lru_evicts_tails_first():
    c = PrefixCache(block_size=4, max_blocks=3)
    fill(c, list(range(12)))                 # one chain of 3 blocks, then released
    fill(c, [50] * 4)                        # needs 1 block, so one must go
    keys = c.hash_blocks(list(range(12)))
    assert keys[0] in c.blocks and keys[1] in c.blocks and keys[2] not in c.blocks


def test_ttl():
    clk = Clock()
    c = PrefixCache(block_size=4, ttl_s=300, clock=clk)
    fill(c, list(range(8)))
    clk.t = 200; assert fill(c, list(range(8)))[0] == 2   # hit, and the timer restarts
    clk.t = 450; assert fill(c, list(range(8)))[0] == 2   # idle 250 s: still alive
    clk.t = 800; assert fill(c, list(range(8)))[0] == 0   # idle 350 s: expired

import torch
from transformers import DynamicCache


def layer_kv(cache):
    """[(K, V)] per layer, each [batch, kv_heads, seq, head_dim]."""
    if hasattr(cache, "layers"):
        return [(l.keys, l.values) for l in cache.layers]
    return list(zip(cache.key_cache, cache.value_cache))  # older transformers


def slice_block(kv, i, B):
    """Copy out block i. clone() matters: a slice is a view that keeps the whole tensor alive."""
    return [(k[:, :, i * B:(i + 1) * B].clone(), v[:, :, i * B:(i + 1) * B].clone()) for k, v in kv]


def build_cache(blocks):
    """A fresh DynamicCache holding the hit blocks, concatenated along the sequence axis."""
    c = DynamicCache()
    for li in range(len(blocks[0].kv)):
        c.update(torch.cat([b.kv[li][0] for b in blocks], dim=2),
                 torch.cat([b.kv[li][1] for b in blocks], dim=2), li)
    return c

# ids = inputs["input_ids"][0].tolist()
# keys = cache.hash_blocks(ids)
# hit = cache.lookup(keys, max_hit=(len(ids) - 1) // cache.B)    # leave >= 1 token to compute
# past = build_cache(hit) if hit else DynamicCache()
# generate(**inputs, past_key_values=past, ...)   # in the thread, as before
#    generate sees the full input_ids, notices past already covers the first len(hit)*B tokens,
#    and only prefills the rest. It also fills `past` in place with every new token.
# after join:
#    kv = layer_kv(past); new = []
#    for i in range(len(hit), len(keys)): if cache.insert(keys[i], slice_block(kv, i, B)): new.append(...)
#    cache.release(hit + new)          # in a finally: block, so an exception can't leak pins
# usage gains "prompt_tokens_details": {"cached_tokens": len(hit) * B}

Measure

TTFT on a shared system prompt

Run shared_2k.jsonl at c = 1, --max-tokens 256 --ignore-eos, once with PREFIX_CACHE=0 and once with PREFIX_CACHE=1. Then once against Fireworks.

What to expect, and why:

• Cache off: every request pays full prefill. The first and the rest look the same, around the second-or-so you estimated for 2,000 tokens in the baseline lesson. Divide 2 × 8 billion × 2,050 by your measured TTFT to get your effective compute rate again.
• Cache on, first request: the same as cache off, plus a little for slicing and storing 128 blocks.
• Cache on, the rest: roughly a single pass over the weights plus overhead. My expectation is somewhere around 100 to 200 ms; measure it. The ratio should be large, similar in spirit to the 339 ms to 58 ms I measured on the Mac with a tiny model.
• tok/s: the same in both runs, within noise. If cache-on decode is noticeably slower, check that you didn’t leave the KV in a strange layout or dtype.
• e2e: shorter by exactly the TTFT saving and no more.

Then account for the hit overhead. Time build_cache on its own: 128 blocks × 32 layers is about 8,000 small tensors concatenated, roughly 256 MiB of copying. If it costs tens of milliseconds, try 64-token blocks: four times fewer pieces, at the cost of coarser matching. That trade (block size against match granularity and bookkeeping) is a real tuning knob in vLLM too.

Hit rate against cache size

Run the sweep on both multi-prefix files. Each 2,000-token prompt needs 128 blocks, and each request’s unique question adds about 3 more blocks, so all 8 prompts need about 1,024 blocks. For Llama 3.1 8B at 2 MiB per block, that is 2 GiB.

Before you run it on the Spark, here is what a pure-Python simulation of this exact workload gives, using a reference implementation of the cache from Task 2 (2,000 requests, 2,048-token prompts plus 50 unique tokens, 16-token blocks). Your measurements should land close to these, since the cache logic is the same; the model only adds time.

The ceiling is 2,048 / 2,098 ≈ 0.976, because the unique question at the end of every prompt can never hit. Notice that even at 64 blocks the Zipf file gets some hits: tail-first eviction keeps the heads of the hottest prompt, and a partial prefix is still a prefix.

Plot TTFT p50 against hit rate. It should be close to a straight line from full-prefill TTFT at 0 to the warm TTFT at 1, because TTFT is roughly a fixed cost plus a per-uncached-token cost.

The scoreboard row

Your scoreboard’s main columns (short prompts, c = 1) barely move: those prompts share almost nothing. That is the honest result, and worth writing down. Put the real effect in the notes:

| date | spark | prefix cache, 16-tok blocks, LRU (+TTL) | (short prompts: ~unchanged) | ~unchanged | ~unchanged | shared 2k sys prompt TTFT p50: ___ ms off → ___ ms on; Zipf-8 hit rate ___ at 1,024 blocks (2 GiB) |

Break It

Put a timestamp at the top. Add Current time: {now} as the first line of the system prompt and rerun shared_2k.jsonl. The hit rate drops to zero: the very first block differs, so every key after it differs. Move the line to the end of the prompt and the hits come back. Every API’s prompt caching guide warns about exactly this, because it is the most common way people silently lose their cache.

Release blocks head first. Change release to put blocks on the free list in order, head before tail. test_lru_evicts_tails_first fails. Rerun the sweep in simulation: at 256 blocks the Zipf hit rate falls from about 0.36 to 0.18, and the round-robin file at 1,024 blocks falls from 0.81 to 0.00. Evicting a head block makes every block after it unreachable (lookup stops at the first miss), but those orphans still occupy the pool until LRU reaches them. One line, half the hit rate.

Store views instead of copies. Remove .clone() from slice_block. Everything still works, and memory climbs request after request. A slice of a tensor is a view: it points into the original storage, so each cached 16-token block keeps the request’s entire KV tensor alive (prompt plus generated tokens, every layer). Watch the process’s memory with the cache sized at 1,024 blocks; it grows far past 2 GiB. The same bug exists in plain Python with NumPy slices. It is why “cache size in blocks” is not the same as “cache size in bytes” unless you are careful.

Use a tiny hash. Replace SHA-256 with something that keeps only 16 bits of the key. With 65,536 possible keys and thousands of blocks, collisions start within minutes of the multi-prefix run, and some requests get another prompt’s KV. Look at the outputs: they are fluent and wrong, a model confidently answering a question nobody asked, with no error anywhere. This is why a cache key must be collision-resistant, and why cache bugs are the worst kind of bug: they return an answer.

Explain It Back

Answer each out loud before opening it.

Go Deeper

• Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention” (SOSP 2023). The vLLM paper. Sections 3 and 4 are the memory-waste measurements and the block table design; the sharing and copy-on-write parts are short and clear.
• vLLM documentation, “Automatic Prefix Caching” (design docs). How the block hash is built (parent hash, tokens, extra keys), how the free queue is ordered for eviction, and the current options for hashing and cache salts. Read this right after building yours and compare.
• Zheng et al., “SGLang: Efficient Execution of Structured Language Model Programs” (NeurIPS 2024), and the LMSYS blog post “Fast and Expressive LLM Inference with RadixAttention and SGLang” (January 2024). The radix tree cache, leaf-first LRU, and cache-aware scheduling, with benchmarks on agent and few-shot workloads.
• Gim et al., “Prompt Cache: Modular Attention Reuse for Low-Latency Inference” (MLSys 2024). Goes past prefixes, reusing cached attention for prompt modules that appear in different positions, and shows what you have to give up to do it.
• Anthropic docs, “Prompt caching.” Breakpoints, the 5-minute and 1-hour TTLs, pricing, minimum lengths, and a list of what invalidates a cache. The cleanest description of a TTL cache as a product.
• OpenAI docs, “Prompt caching” guide. The automatic version: thresholds, retention, cached_tokens, and prompt_cache_key for routing.
• Hugging Face transformers docs, “KV cache strategies” (the cache classes and the prefix-reuse example). The API your server uses; check here first when DynamicCache changes.
• Arpaci-Dusseau and Arpaci-Dusseau, Operating Systems: Three Easy Pieces, chapters on paging and “Beyond Physical Memory: Policies.” Page tables, LRU and its approximations, and why LRU fails on loops bigger than the cache: the same curve as your round-robin sweep.

Connections

This lesson builds on Virtual memory (pages, page tables, copy-on-write: the same ideas, one level up), Load balancers and state (cache-aware routing is a sticky session that is safe because it’s a cache), and the server and harness from Baseline: my Spark vs Fireworks.

It connects forward to Caching and Redis, where the same questions (what is the key, how long does it live, what happens when it’s stale) come up for ordinary application data; to How much memory does a model need?, where KV bytes per token decide how big the pool can be; and to Inside vLLM and SGLang, where you will find the block manager and radix tree you just built, in production form.