Systems labAdvanced120 min estimateOriginal publication

The KV Cache Capacity Plan

Model weights are the static bill. The KV cache is the per-request bill that grows one token at a time—and often decides whether another request fits.

Authorship: InterviewsVector
Published / updated: 2026-08-11 / 2026-08-11
Review status: Artifact tests passing · primary sources recorded

Original InterviewsVector material. Code examples are covered by repository tests and primary references are recorded. No named human reviewer is claimed.

The decision in one pass

For a decoder transformer, approximate KV-cache bytes per sequence as tokens × layers × 2 × KV heads × head dimension × bytes per element. Multiply by active sequences, then add allocator overhead and headroom. Multi-query and grouped-query attention reduce KV heads; quantized cache reduces bytes per element; paging reduces fragmentation. None of those changes removes the need to budget prefill and decode separately.

Why this matters

An endpoint can pass single-request latency tests and fall over at modest concurrency because long contexts consume the memory that batching needs. Capacity planning must use request-shape distributions, not the model's maximum context as one headline number.

You will be able to

Derive the KV-cache formula from attention tensor shapes.
Distinguish attention compute, attention working memory, and persistent inference cache.
Estimate concurrency under MHA, GQA, and MQA.
Explain fragmentation, paging, prefix reuse, and eviction trade-offs.
Design a load test using prompt and output-length distributions.

Prerequisite contract

— Scaled dot-product attention
— Tensor shapes and byte units
Prefill, decode, and KV cache →

Your Vector Loop for this lab

01
Model
Treat the cache as per-sequence mutable state, separate from model weights.
02
Derive
Multiply the exact K and V tensor dimensions by precision and tokens.
03
Build
Implement a unit-safe estimator.
04
Stress
Vary context, GQA ratio, concurrency, fragmentation, and output length.
05
Operate
Translate bytes into admission control, batching, paging, and SLOs.
06
Defend
Defend a serving plan with assumptions and sensitivity analysis.

Cache the past projections, not the attention matrix

During autoregressive decoding, each new token attends to earlier tokens. Recomputing every earlier key and value projection on every step would repeat work. The server stores K and V for prior positions at every layer, appends one position per generated token, and reuses the stored tensors on the next step.

Derive bytes from the tensor shape

bytes = T × L × 2 × Hkv × Dhead × B

T is cached tokens, L layers, 2 accounts for keys and values, Hkv is the number of key/value heads, Dhead is head dimension, and B is bytes per element.

For standard multi-head attention, Hkv equals the number of query heads. Grouped-query attention shares one K/V head across a group of query heads. Multi-query attention uses one K/V head for all query heads. Query heads still affect attention compute, but KV heads determine cache width.

Variant	Query heads	KV heads	Cache effect
MHA	H	H	baseline
GQA	H	between 1 and H	reduced by grouping ratio
MQA	H	1	minimum KV-head width

Turn architecture knobs into a memory budget

KV-cache capacity planner

Adjust layers, KV heads, head dimension, precision, context, and concurrency. The planner separates raw cache, operational headroom, and capacity per accelerator.

Layers32KV heads8Head dimension128Bytes per element2Context tokens8192Concurrent sequences24Accelerator capacity (GiB)80Reserved headroom (%)25Cache overhead (%)10

Raw cache working set

24.0 GiB

With cache overhead

26.4 GiB

Usable accelerator budget

60.0 GiB

Full-sequence capacity

Fits the declared budget. Reserved headroom covers weights, activations, and temporary kernels; the overhead factor covers fragmentation and paging metadata.

Build a unit-safe estimator

kv_cache_budget.py

1from dataclasses import dataclass
2 
3@dataclass(frozen=True)
4class ModelShape:
5    layers: int
6    kv_heads: int
7    head_dim: int
8    bytes_per_element: float
9 
10def cache_bytes(shape, tokens, sequences=1, overhead=1.0):
11    if min(shape.layers, shape.kv_heads, shape.head_dim, tokens, sequences) <= 0:
12        raise ValueError("shape, tokens, and sequences must be positive")
13    raw = sequences * tokens * shape.layers * 2 * shape.kv_heads * shape.head_dim * shape.bytes_per_element
14    return raw * overhead
15 
16def gib(value):
17    return value / (1024 ** 3)
18 
19shape = ModelShape(layers=32, kv_heads=8, head_dim=128, bytes_per_element=2)
20print(round(gib(cache_bytes(shape, tokens=8192, sequences=16, overhead=1.15)), 2))

Expected output

18.4

Verify: Run python -m unittest discover courses/ai-engineering/reference-impl/kv_cache.

Why average context length lies

01Long-tail promptsA small share of long sequences can pin large cache blocks and reduce the batch that short requests need.
02Unknown output lengthAdmission control must reserve growth or risk mid-generation eviction and failure.
03Fragmented allocationContiguous per-sequence reservations waste memory when sequences finish at different lengths.
04Prefix reuseShared prefixes can reduce repeated cache state, but invalidation, tenant isolation, and hit rate determine value.
05Cancellation lagIf disconnected requests do not release cache quickly, measured concurrency understates effective occupancy.

Capacity-plan prefill and decode as different workloads

Phase	Dominant work	Useful measures	Control
Prefill	parallel compute over prompt tokens	TTFT, prompt tokens/s	chunking and prefill scheduling
Decode	memory-bound one-token steps	inter-token latency, output tokens/s	continuous batching
Cache	persistent per-sequence memory	allocated blocks, waste, evictions	admission and paging

Build a two-dimensional load distribution over prompt length and requested output length, preserve arrival bursts, include cancellation, and report latency percentiles by request shape. A single requests-per-second result conceals whether the service is compute-bound, memory-bound, queue-bound, or cache-bound.

Operate at three altitudes

Production lens

— Track cache blocks in use, fragmentation, prefix-cache hit rate, evictions, and release latency.
— Slice TTFT and inter-token latency by prompt length, output length, and batch occupancy.
— Reserve memory for weights, runtime workspaces, graph capture, and failure headroom before assigning cache capacity.
— Use request-shape-aware admission control and cancellation propagation.

Staff lens

— Ask whether the bottleneck is model compute, memory bandwidth, cache capacity, scheduler policy, or queueing before buying accelerators.
— Model capacity with distributions and sensitivity bounds, then verify with load tests.
— Keep tenant isolation and prefix-cache privacy in the same design review as throughput.

Interview defense

How would you estimate KV-cache memory for an LLM service?

I would calculate tokens × layers × two tensors × KV heads × head dimension × bytes per element per sequence, then multiply by active sequences. I would add allocator overhead and operational headroom, account for GQA or MQA, and test the estimate against prompt/output length distributions because average length hides long-tail occupancy.

Expect the interviewer to press on

— How does GQA change the estimate?
— Does FlashAttention remove KV-cache memory?
— How would you admit a request with unknown output length?

Misconceptions to remove

“The model's maximum context tells us service capacity.”

Capacity depends on the distribution of active prompt and output lengths, concurrency, cache architecture, precision, and headroom.

“FlashAttention eliminates quadratic attention problems and the KV cache.”

It reduces IO for exact attention; decode still retains prior keys and values for reuse.

“Batch size is a static deployment setting.”

Continuous batching changes membership every decode step as requests arrive, finish, and cancel.

Check your model

1. If KV heads drop from 32 to 8 with every other variable fixed, how does raw cache size change?

It drops by a factor of four because cache width is linear in the number of KV heads.

2. Why reserve expected output growth at admission time?

The cache grows during decode; admitting only on current tokens can exhaust memory after a request has already begun.

Prove the mechanism

Capacity-plan two model shapes for an 80 GiB accelerator. State every reserved-memory and request-shape assumption, then find the sensitivity to p95 output length.

Add a production constraint

Add block-size fragmentation and prefix-sharing to the estimator. Simulate arrivals and completions instead of multiplying by fixed concurrency.

Artifact: Tested KV-cache estimator

courses/ai-engineering/reference-impl/kv_cache/kv_cache_budget.py

Download reference implementation

Primary references and next links

References

1. Attention Is All You Need
Vaswani et al.. Primary transformer architecture paper.
2. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Dao et al.. Primary paper on IO-aware exact attention.
3. Efficient Memory Management for LLM Serving with PagedAttention
Kwon et al.. Primary vLLM and PagedAttention systems paper.

Continue through the graph

Prefill, decode, and KV-cache interview answer →
Rehearse the compact systems explanation.
Dynamic batching trade-off →
Connect memory to scheduler policy.
LLM p99 latency diagnosis →
Apply the capacity model to an incident.

Glossary: prefill · decode · KV cache · GQA · MQA · continuous batching · paging