The KV-cache & attention memory

What the KV-cache is essential

Here's the puzzle that the KV-cache quietly solves. A language model generates text one token at a time, and to pick each new token it has to look back at everything written so far. The naive way is brutal: at token 1,000 you'd re-read and re-process all 999 tokens before it, then at token 1,001 you'd do it all again for 1,000 tokens, and so on. That's the same work done over and over — the classic O(n²) trap you meet in the complexity track. On a long response it would make every model painfully slow, and it's exactly the waste the cache exists to kill.

Think of it like a student writing a long essay. Each new sentence has to stay consistent with everything already written — but a sane student doesn't re-read the entire essay from the top before every sentence. They keep a running set of notes in the margin: the gist of each paragraph, ready to glance at. The KV-cache is that margin of notes. For every token the model has already seen, it stores two little tensors — a key (K, roughly "what this token is about," used to decide how relevant it is) and a value (V, "the information this token contributes"). Once a token's K and V are computed, they never change, so the model computes them once and files them away. Every later step just reads the notes instead of re-deriving them.

The common mistake is to picture the cache as an optional speed tweak — a nice-to-have you could turn off. It isn't: without it, autoregressive generation is quadratic and effectively unusable at long context, so every serious inference stack keeps a KV-cache and the real questions are all about managing it. The second surprise, and the reason this chapter exists, is where the cost lands. People assume the model's weights are the memory hog. They're fixed and often smaller than you'd think; it's the cache that balloons. It grows by two fresh tensors per layer, per token, and it grows again for every concurrent request the server is juggling. A single long conversation can consume a gigabyte or more of GPU memory in cache alone — which is why, past a certain scale, the KV-cache (not the weights) becomes the dominant consumer of GPU memory and the thing that decides how many users you can serve at once. Now the precise version:

At each decode step the model attends to every previous token. Recomputing their keys and values every step would be wasteful, so they're stored — the KV-cache. It's what makes generation tractable, but it grows with every token and every concurrent request, and it quickly becomes the dominant consumer of GPU memory.

PagedAttention essential

Imagine running a car park where every driver, on arrival, is forced to book a whole row of spaces — just in case they later show up with a bus. Almost nobody brings a bus. The rows sit half-empty, you turn cars away at the gate, and yet the lot is "full." That absurd policy is exactly how the first generation of LLM servers handled the KV-cache: because the cache grows token by token and you don't know in advance how long a reply will be, each request reserved one big contiguous slab of GPU memory sized for its maximum possible length. A request that ends after 100 tokens still hogs space reserved for 4,000. The unused tail is pure waste — the memory equivalent of empty parking rows — and it's called internal fragmentation.

Why did they reserve contiguous memory at all? Because the older attention kernels assumed a token's keys and values sat in one unbroken run, so you could stride through them cleanly. Convenient for the math, disastrous for utilization: fragmentation routinely wasted the majority of the cache, so a card that could have held dozens of conversations held only a handful. The mistake to unlearn is thinking the fix is "buy more VRAM." More memory just buys you slightly bigger half-empty rows; the leak is the reservation scheme itself.

PagedAttention is the fix, and it borrows a fifty-year-old trick from operating systems: virtual memory. An OS doesn't hand each program a single contiguous block of physical RAM; it chops memory into fixed-size pages and hands them out on demand, keeping a page table that maps "where the program thinks its data is" to "where the pages actually live." PagedAttention does the identical thing to the KV-cache — it splits the cache into small fixed-size blocks (pages) and allocates them to a request only as it actually generates tokens, tracking them through a block table. Nothing needs to be contiguous, so there's no giant up-front reservation and no wasted tail; the only slack left is at most a partial final page per request. Reclaiming all that fragmented memory means the same GPU can now hold far more concurrent sequences — and because pages can even be shared between requests with a common prefix, it dovetails with the prompt caching in the next section. This one idea, first shipped in vLLM (IC6), is the single biggest reason modern serving engines pack so many users onto a card.

Naively, each request reserves contiguous memory for its maximum possible cache — hugely wasteful. PagedAttention (the idea behind vLLM) manages the KV-cache like virtual memory: non-contiguous pages allocated on demand. That single change lets a server fit far more concurrent requests in the same GPU — the biggest reason vLLM (IC6) is fast.

Prompt caching: reuse the prefill intermediate

Before a model can write a single token of its answer, it has to read your whole prompt. That reading phase — computing the K and V for every input token so the cache is primed — is called prefill, and it isn't free: it's compute you pay for on every call, and on hosted APIs it's billed per input token. Now picture a support agent whose every request begins with the same 5,000-token policy manual, followed by a one-line question. The manual is identical call after call, yet the naive setup re-prefills all 5,000 tokens every single time. You are paying, over and over, to teach the model the exact same thing it already computed a millisecond ago. That's the waste prompt caching removes.

The insight is small once you've internalised the earlier sections: the KV of a token depends only on that token and the tokens before it. So the KV for a fixed prefix is always the same — it doesn't matter what question follows. If the model computes it once and keeps it, every later request that starts with that same prefix can reuse the cached KV and jump straight to prefilling only the part that actually changed (the user's question). Think of it like a lecturer who films the standard introduction once and replays the recording each session, spending live time only on the new Q&A. The recording is the cached prefix; the live Q&A is your variable suffix.

Two mistakes trip people up here. First, order matters: caching only helps if the shared text sits at the front, because a cache hit needs an exact-matching prefix — put the reused manual before the changing question, never after, or there's no common prefix to reuse. Second, the economics are lopsided in a way worth stating plainly: the first call pays a little extra to write the cache, and the payoff arrives only on the second and later calls that read it — so caching a prefix used once is a net loss, while caching one reused thousands of times is the cheapest optimization in this entire chapter. On hosted APIs that reuse shows up as a billing discount (cache reads are charged at a small fraction of a normal input token); when you self-host, it's a direct compute and latency saving because the prefill work simply doesn't happen. With that mental model, the precise definition:

If many requests share a prefix — a long system prompt, a big RAG context, a few-shot preamble — you can cache its KV and skip re-prefilling it every time. This is a first-class cost lever: on hosted APIs it's a billing discount; when self-hosting it's a real compute saving.

prompt_cache.pyfrom anthropic import Anthropic
client = Anthropic()

BIG_SYSTEM = open("policy_manual.txt").read()   # a long, reused preamble

resp = client.messages.create(
    model="claude-opus-4-8", max_tokens=300,
    system=[{
        "type": "text", "text": BIG_SYSTEM,
        "cache_control": {"type": "ephemeral"},   # cache this block's KV
    }],
    messages=[{"role": "user", "content": "Given the policy, can I refund order 123?"}],
)
u = resp.usage
print("cache write:", u.cache_creation_input_tokens)   # first call
print("cache read: ", u.cache_read_input_tokens)        # cheap on later calls

def kv_bytes(layers, heads, head_dim, seq_len, batch=1, bytes_per=2):
    return 2 * layers * heads * head_dim * seq_len * batch * bytes_per

b = kv_bytes(layers=32, heads=32, head_dim=128, seq_len=2048)
print(f"{b/1024**2:.1f} MB")    # 1024.0 MB  (1 GiB for a single 2k sequence!)

def kv_gb(seq_len, batch, layers=32, heads=32, head_dim=128, bpp=2):
    return 2*layers*heads*head_dim*seq_len*batch*bpp / 1024**3

for seq in (1024, 8192, 32768):
    row = [f"{kv_gb(seq, b):.1f}GB" for b in (1, 16)]
    print(f"seq={seq:>6}: batch1={row[0]:>8}  batch16={row[1]:>8}")
# seq=  1024: batch1=  0.5GB  batch16=  8.0GB
# seq=  8192: batch1=  4.0GB  batch16= 64.0GB
# seq= 32768: batch1= 16.0GB  batch16=256.0GB

reqs = [("a", 300, 4096), ("b", 1200, 4096), ("c", 80, 4096)]  # id, used, max
per_tok = 0.5  # MB per token of KV (toy)

# Naive: reserve max for every request
naive = sum(mx * per_tok for _, _, mx in reqs)
used  = sum(u  * per_tok for _, u, _  in reqs)
print(f"naive reserved: {naive:.0f} MB, actually used: {used:.0f} MB")
print(f"wasted: {naive-used:.0f} MB ({(naive-used)/naive*100:.0f}%)")

# Paged: allocate in pages of 128 tokens, only as needed
PAGE = 128
paged = sum(-(-u // PAGE) * PAGE * per_tok for _, u, _ in reqs)  # ceil-div pages
print(f"paged allocated: {paged:.0f} MB (waste now < one page/req)")

def max_concurrency(free_gb, seq_len, layers=32, heads=32, head_dim=128, bpp=2):
    kv_per_req = 2*layers*heads*head_dim*seq_len*bpp / 1024**3
    return int(free_gb // kv_per_req), round(kv_per_req, 2)

for seq in (2048, 8192):
    n, per = max_concurrency(free_gb=40, seq_len=seq)
    print(f"seq={seq}: {per} GB/req -> {n} concurrent requests")
# seq=2048: 1.0 GB/req  -> 40 concurrent
# seq=8192: 4.0 GB/req  -> 10 concurrent

# --- savings math (runnable) ---
prefix_tok = 5000
calls = 100
# cached reads are ~10% the price of a normal input token
normal = prefix_tok * calls
cached = prefix_tok * 1 + prefix_tok * (calls - 1) * 0.10
print(f"input tokens billed: {normal} -> {int(cached)}")   # 500000 -> 54500
print(f"~{(1-cached/normal)*100:.0f}% cheaper on the prefix")  # ~89% cheaper

# --- real Anthropic prompt caching (needs the SDK) ---
from anthropic import Anthropic
client = Anthropic()
resp = client.messages.create(
    model="claude-sonnet-4-5",
    max_tokens=256,
    system=[{
        "type": "text",
        "text": LONG_STABLE_DOCS,                 # the reused 5k-token prefix
        "cache_control": {"type": "ephemeral"},   # cache this block
    }],
    messages=[{"role": "user", "content": "Summarize section 3."}],
)
print(resp.usage)   # cache_creation_input_tokens / cache_read_input_tokens

layers, heads, head_dim, bpp = 32, 32, 128, 2
def kv_gb(seq, batch):
    return 2*layers*heads*head_dim*seq*batch*bpp/1024**3

free = 30  # GB after weights
print("1 req @ 8k :", round(kv_gb(8192, 1),1), "GB -> fine")   # 4.0
print("12 req @ 8k:", round(kv_gb(8192,12),1), "GB -> OOM")    # 48.0 > 30

fixes = [
    ("cap max context to 4k", "halves KV/req -> ~24 concurrent", "free, config"),
    ("switch to vLLM/PagedAttention", "reclaims fragmentation, packs more", "free-ish, eng time"),
    ("add a GPU", "linear more room", "$$$ ongoing"),
]
for f in fixes: print("-", f[0], "|", f[1], "|", f[2])

Knowledge check check yourself