Inference & Cost · Chapter IC3

The KV-cache & attention memory

The KV-cache makes generation tractable but dominates GPU memory and scales with context and concurrency. PagedAttention and prompt caching are the levers that tame it.

⏱️ ~2 hours🧪 1 lab🎯 Advanced

⚙️ To run this for realThe code here is complete and correct as written. To actually execute it you'll need:

an Anthropic API key (ANTHROPIC_API_KEY) + pip install anthropic

Reading and learning works without any of this — run when you're ready. IDs/ARNs in examples are placeholders; swap in your own.

Learning objectives

Explain what the KV-cache stores and why it dominates memory during decode.
Reason about how context length drives memory and cost quadratically-ish.
Explain PagedAttention and why it was a breakthrough for serving.
Use prompt caching as a first-class cost lever.

▶ Runnable companionThe code in this lesson is also saved under code/ic3-kv-cache/ in the course, with a README. Run the scripts or copy the configs directly.

What the KV-cache is essential

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.

🗺️ How to read this diagram

This shows one decode step and why the KV-cache exists. To pick the next word, the model looks back at every earlier token. Instead of re-deriving those tokens each step, it keeps a running notebook — the KV-cache — of two things per token: a key (K) and a value (V).

Token t (new token) — the latest token the model is processing this step.
Compute K,V (this step) — for that one new token, the model works out its key and value. Think of K as "what this token is about" and V as "the information it carries."
Append to cache (grows each step) — those K and V get added to the notebook. This is why the cache grows every single step, and why long chats use more and more memory.
Attend to all (over cached K,V) — to choose the next word, the model compares the new token against all the stored keys/values at once. Because they're cached, it doesn't recompute the earlier tokens — that's the whole speed win.
Read the arrows as one step's flow; the loop repeats for every token, and the cache only gets bigger.

In short: KV-cache = the model remembering the keys and values of past tokens so it never redoes that work. It makes decode fast but is the main thing eating GPU memory — the rest of the chapter is about shrinking and reusing it.

The memory mathKV-cache size ≈ 2 × layers × heads × head_dim × sequence_length × batch × bytes. It scales with context length AND concurrency — which is why long contexts and many users are expensive, and why serving a big context window is a memory problem before a compute one.

PagedAttention essential

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

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.

Example code for learning — review, test, and adapt it before running against real or production systems. Commands can create, change, or delete resources. See the Terms & Disclaimer.

Lab IC3.1 · Prompt caching on the Anthropic API

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

▶ How this works

This shows prompt caching: when many requests share the same long opening text (here a policy manual in the system prompt), you tell the API to remember its computed KV once so later calls skip re-processing it — cheaper and faster.

BIG_SYSTEM = open("policy_manual.txt").read() loads a long block of text that every request will reuse. Re-reading this on every call is the waste we want to avoid.
Inside system=[{...}], the line "cache_control": {"type": "ephemeral"} is the magic word: it marks this block as cacheable. The API stores its KV so it isn't recomputed next time.
The messages list holds the part that actually changes per request — the user's real question. Only this short piece needs fresh work on later calls.
resp.usage reports two counters: cache_creation_input_tokens (tokens written to the cache — happens on the first call) and cache_read_input_tokens (tokens served from the cache — cheap, on later calls).

What the output means: On the first run you'll see a big cache write and near-zero read. Run the same script again and the cache read jumps up while the write drops — that's the discount kicking in.

Try this: Point BIG_SYSTEM at any large text file and call the script twice. The second call's cache_read_input_tokens is the money you save by not re-prefilling the shared prefix.

This is the cheapest optimization you haveBefore quantizing or re-architecting, look for reused prefixes. A long system prompt or RAG context served thousands of times is pure waste without caching. It ties back to Ch 6 and Bedrock, which cache the same way.

Exercise IC3.1 — Cache a shared prefix

Context: Prompt caching is the highest-leverage cost win in the track, and the savings only land on the second and later calls that reuse the cached prefix.

Your task: Take a workload with a long shared system prompt or RAG context, measure cost/latency without caching, then with prompt caching, and report the cache-read savings on the 2nd+ call.

Requirements:

Use a workload with a long stable shared prefix
Measure cost and latency without caching
Enable prompt caching and measure again
Report the cache-read savings on the second and later calls
Frame it as usually the biggest cost win in the track

💡 Hint: The first call pays to write the cache; the win shows up on every call after, which read the prefix at a fraction of the price.

🪜 Practice ladder beginner → industry

Six graded exercises, easy to real-world. Try each before opening its solution.

Exercise 1 · KV-cache size for one sequenceBeginner

Context: The KV-cache stores per-token keys and values so decode stays fast, and a single long sequence can eat ~1 GB — which is why the cache, not the weights, limits concurrency.

Your task: Compute the KV-cache size (MB) for one 2,048-token sequence on a 32-layer, 32-head, head-dim-128, FP16 model.

Requirements:

Size = 2 × layers × heads × head_dim × seq_len × batch × bytes
Use 2 bytes (FP16) and batch 1
Convert bytes to MB
Explain the leading 2× stores both K and V
Note one ~1 GB request is why KV-cache bounds concurrency

💡 Hint: The leading 2 is K and V; plug the given dimensions straight into the formula and divide to MB.

Show solution

Plug straight into the lesson's formula:

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!)

The 2x is for K and V. One 2k-token request already eats ~1 GB — which is why the KV-cache, not the weights, is what limits concurrency.

Exercise 2 · Cache grows with context AND concurrencyIntermediate

Context: The KV-cache grows linearly on two axes at once — context length and concurrency — so long windows and many users are each a memory problem before a compute one.

Your task: Tabulate KV-cache GB for sequence lengths {1k, 8k, 32k} at batch sizes {1, 16} and explain why each axis is expensive.

Requirements:

Reuse the KV formula, converting to GB
Sweep seq_len × batch and print a table
Show the cache is linear in both context and batch
Explain long context is a memory problem before a compute one

💡 Hint: Both context length and batch multiply into the same formula, so doubling either doubles the cache — the table just makes that linearity visible.

Show solution

Sweep both terms of the formula:

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

KV-cache is linear in both context length and batch, so a 32k context served to 16 users needs 256 GB just for the cache — long windows are a memory problem before a compute one.

Exercise 3 · PagedAttention beats internal fragmentationAdvanced

Context: Naive serving reserves each request's full max context contiguously and wastes the unused tail; PagedAttention allocates fixed pages and reclaims almost all of it.

Your task: Model the waste: given actual vs max lengths per request, compute reserved-but-unused memory, then show paging in fixed pages reclaims most of it.

Requirements:

Naive reservation = sum of max × per-token
Compute wasted MB and percentage vs actually-used memory
Model paging with fixed (e.g. 128-token) pages via ceil-division
Show waste drops to at most a partial final page per request
Connect reclaimed memory to packing more concurrent sequences

💡 Hint: Ceil-divide each request's used tokens up to a whole number of fixed pages — the only waste left is the slack in the last page.

Show solution

A deterministic sim of the fragmentation PagedAttention removes:

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)")

Naive reservation wastes ~87% here because most requests never reach max length. Paging allocates in small blocks on demand, so waste drops to at most a partial final page per request — the key to fitting far more concurrent sequences.

Exercise 4 · How many concurrent requests fit?Expert

Context: Concurrency and context length trade directly against each other through the KV-cache: the free VRAM after weights, divided by per-request KV, is the ceiling a scheduler is bound by.

Your task: Given free VRAM after weights and per-request KV growth, compute the max concurrent sequences at a target context length.

Requirements:

Compute per-request KV in GB from the formula
Max concurrency = free_gb // kv_per_req
Show a 4× longer context quarters concurrency
Frame it as admission-control math for the scheduler (IC4)

💡 Hint: It's the KV formula turned into a division — free memory over per-request cost is how many sequences fit at once.

Show solution

The admission-control math a serving engine runs:

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

Quadrupling the context window quarters how many users you can serve on the same card. Concurrency and context length trade directly against each other through the KV-cache.

Exercise 5 · Prompt caching as a cost lever (real, needs API)Professional

Context: Prompt caching reuses the prefill of a stable prefix (system prompt, docs) across calls, and it's usually the single biggest cost win in the whole track.

Your task: Show the real Anthropic cache_control shape and compute the savings when a 5k-token system prompt is reused across 100 calls.

Requirements:

Normal billing = prefix_tokens × calls
Cached = one full write + (calls−1) reads at ~10% price
Show the reuse is ~90% cheaper
Place the reused prefix in a system block with cache_control: {type: ephemeral}
Read cache_creation_input_tokens / cache_read_input_tokens off resp.usage

💡 Hint: A stable prefix becomes a one-time write plus cheap reads — mark it with cache_control and let usage report the split.

Show solution

The savings math (runnable) plus the real API shape (needs pip install anthropic + ANTHROPIC_API_KEY):

# --- 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

Reusing the prefill turns a repeated 5k-token prefix into a one-time cost plus cheap cache reads — a first-class cost lever whenever a large prefix is stable across requests.

Exercise 6 · Diagnose a 'server runs out of memory under load' incidentIndustry scenario

Context: A classic prod incident: one request is fine but the server OOMs as concurrency climbs — the KV math proves it's cache growth, not a leak.

Your task: Given the KV math, write the reasoning that identifies the OOM cause and rank three fixes (shorter max context, PagedAttention/vLLM, more VRAM) by cost-effectiveness.

Requirements:

Show one request fits but N concurrent requests exceed free VRAM
Conclude the cause is KV growth with concurrency, not a memory leak
Rank capping max context (free/config, halves KV per request) first
Then switching to vLLM/PagedAttention (eng time, reclaims fragmentation)
Then adding VRAM ($$$, only defers the same math)

💡 Hint: Do the KV arithmetic at the failing concurrency — when N requests exceed free memory but one doesn't, the cache is the culprit and config is the cheapest fix.

Show solution

Turn the incident into KV arithmetic, then rank fixes:

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])

The root cause is KV-cache growth with concurrency, not a leak. Cap the context window first (free, immediate), adopt PagedAttention next (biggest structural win), and buy hardware last — the expensive option that only defers the same math.

✓ Checkpoint — you can move on when you can…

Explain what the KV-cache stores and why it dominates memory.
Do the KV-cache memory math and explain the context-length cost.
Explain PagedAttention and why it boosts concurrency.
Use prompt caching to skip re-prefilling a shared prefix.

Knowledge check check yourself

✓ Knowledge check

What does the KV-cache store, why does it exist, and why does it dominate GPU memory during decode?

Show answer

It stores the key (K) and value (V) tensors for every previous token so the model can attend to them without recomputing them each decode step — that's the speed win. It grows with every generated token and every concurrent request (size ≈ 2 × layers × heads × head_dim × seq_len × batch × bytes), so it scales with context length and concurrency and becomes the dominant memory consumer.

✓ Knowledge check

What problem does PagedAttention solve and how, and why does prompt caching help a workload with a long shared system prompt?

Show answer

Naively each request reserves contiguous memory for its maximum possible cache, which is hugely wasteful; PagedAttention manages the KV-cache like virtual memory with non-contiguous pages allocated on demand, letting a server fit far more concurrent requests. Prompt caching stores the KV of a shared prefix so it isn't re-prefilled on every call — a billing discount on hosted APIs and real compute savings when self-hosting.

Educational content, provided as-is and without warranty. Code samples are examples — review, test, and adapt them before using in production. See the Terms of Use & Disclaimer. Use at your own risk.

Learning objectives

What the KV-cache is essential

PagedAttention essential

Prompt caching: reuse the prefill intermediate

Exercise IC3.1 — Cache a shared prefix

🪜 Practice ladder beginner → industry

✓ Checkpoint — you can move on when you can…

Knowledge check check yourself

More in Inference & cost