KV Cache Optimization: Efficient Context Storage

The GPU that ran out of room, not compute

A Bengaluru startup builds a Hindi-English customer support chatbot for a bank, running on a single rented A100 GPU with 40 GB of memory. The model is a 7-billion-parameter transformer, stored in fp16 (16-bit floating point), so its weights occupy about 13 GB — comfortably under 40 GB, with 27 GB to spare. The team assumes they can serve dozens of simultaneous chats. In production the server crashes with an out-of-memory error at the twelfth concurrent conversation, long before the GPU's compute is anywhere near saturated. The model never got bigger. What grew, silently, with every message exchanged in every open chat, was the KV cache — and by the time this chapter is done, you will be able to calculate exactly why 12 was the ceiling, not 11 or 20.

This is not a niche implementation detail. It is the single biggest reason large language model inference is expensive, and it is why an entire generation of serving systems — vLLM, TensorRT-LLM, Hugging Face's text-generation-inference — exist primarily to manage one data structure well.

What a decoder actually repeats, one token at a time

A transformer generates text autoregressively: it produces one token, appends it to the sequence, and runs the whole model again to produce the next token. Inside every self-attention layer, each token's embedding x is projected through three learned weight matrices to produce a query, a key, and a value:

q = x · W_Q
k = x · W_K
v = x · W_V

To compute the output for token t, its query q_t is compared against the keys of every token up to and including t (causal masking forbids looking ahead), producing attention scores; those scores are turned into weights via softmax, and the weights combine the corresponding values:

score(t, j) = (q_t · k_j) / sqrt(d)      for all j ≤ t
weights     = softmax(scores)
output_t    = Σ_j weights[j] · v_j

Here d is the per-head dimension, and the division by sqrt(d) prevents the dot products from growing large enough to saturate the softmax. Nothing in this formula is specific to KV caching yet — it is just what self-attention computes, at every layer, for every head.

The redundant work a naive decoder repeats

The wasteful part shows up when you generate a whole sentence. A naive implementation, given a growing sequence, simply reruns the full forward pass on the entire sequence-so-far at every step. To produce token t+1, it recomputes k_j and v_j for every token j from 1 to t — including tokens whose keys and values were already computed identically in the previous step, and the one before that, and so on back to the start. Those matrices never change once a token's embedding is fixed (causal masking means a token's own representation at a given layer does not depend on tokens generated after it), so this is pure waste.

The cost of that waste is not linear — it compounds. Let d be the model's hidden width. At decode step t, projecting all t tokens through W_Q, W_K, W_V costs roughly 3·t·d² floating-point operations, and computing the full t × t attention score matrix costs roughly t²·d. Summing over a generation of N tokens:

Σ(t=1..N) 3·t·d²  ≈ (3/2)·N²·d²        (projection cost, naive)
Σ(t=1..N) t²·d    ≈ (1/3)·N³·d         (attention cost, naive)

The attention term grows as N³ — not because attention itself is cubic, but because a quadratic computation is being repeated, wastefully, N times over. A cached decoder, by contrast, only ever computes the projections for the one new token at each step (cost O(d²), independent of t) and compares its single new query against the growing list of stored keys (cost O(t·d)):

Σ(t=1..N) d²    = N·d²                  (projection cost, cached)
Σ(t=1..N) t·d   ≈ (1/2)·N²·d            (attention cost, cached)

Caching does not make attention free or constant-time — total cost across a full generation is still O(N²), because each new token still has to look at every previous token's key. What caching removes is the redundant re-derivation of those keys and values, cutting the dominant term from O(N³) down to the true theoretical floor of O(N²). For a 500-token reply, that is roughly a 500-fold reduction in the projection-recomputation work alone.

A fully traced example: three tokens, by hand

To see that caching changes nothing about the answer, only the work, trace a tiny single-head attention with hidden dimension d = 2 across three tokens. Use these (deliberately simple, hand-computable) weights and inputs:

W_Q = [[1,0],[0,1]]     W_K = [[0,1],[1,0]]     W_V = [[1,1],[0,1]]
x1 = [1,0]   x2 = [0,1]   x3 = [1,1]

Projecting each token (row-vector times matrix, y_j = Σ_i x_i·W[i,j]) gives:

K1 = x1·W_K = [0,1]    V1 = x1·W_V = [1,1]
K2 = x2·W_K = [1,0]    V2 = x2·W_V = [0,1]
K3 = x3·W_K = [1,1]    V3 = x3·W_V = [1,2]
Q3 = x3·W_Q = [1,1]

Now compute the output at position 3 two ways. First, the naive path, which recomputes every key and value from scratch:

import numpy as np

W_Q = np.array([[1,0],[0,1]])
W_K = np.array([[0,1],[1,0]])
W_V = np.array([[1,1],[0,1]])
x1, x2, x3 = np.array([1,0]), np.array([0,1]), np.array([1,1])
scale = np.sqrt(2)

# naive: recompute K,V for ALL three tokens this step
K_naive = np.array([x1 @ W_K, x2 @ W_K, x3 @ W_K])
V_naive = np.array([x1 @ W_V, x2 @ W_V, x3 @ W_V])
q3 = x3 @ W_Q
scores  = (q3 @ K_naive.T) / scale
weights = np.exp(scores) / np.exp(scores).sum()
out3_naive = weights @ V_naive

By hand: scores = [Q3·K1, Q3·K2, Q3·K3]/√2 = [1, 1, 2]/1.4142 ≈ [0.7071, 0.7071, 1.4142]. Exponentiating gives [2.0281, 2.0281, 4.1133], summing to 8.1695, so weights ≈ [0.2482, 0.2482, 0.5035]. Combining with the values: 0.2482·[1,1] + 0.2482·[0,1] + 0.5035·[1,2] = [0.7517, 1.5034].

Now the cached path, where K1, V1, K2, V2 are already sitting in cache from the previous two steps, and only token 3's key and value are new:

# cached: K1,V1,K2,V2 already stored; only compute K3,V3
cache_K = [x1 @ W_K, x2 @ W_K]
cache_V = [x1 @ W_V, x2 @ W_V]
k3, v3 = x3 @ W_K, x3 @ W_V
cache_K.append(k3); cache_V.append(v3)
K_cached, V_cached = np.array(cache_K), np.array(cache_V)

scores2  = (q3 @ K_cached.T) / scale
weights2 = np.exp(scores2) / np.exp(scores2).sum()
out3_cached = weights2 @ V_cached

K_cached and V_cached are element-for-element identical to K_naive and V_naive, so out3_cached is exactly [0.7517, 1.5034] too — same answer, but this step only did one projection instead of three. That is the entire trick: a cache is a correctness-preserving memoization of K and V, nothing more.

Sizing the cache: a Llama-2-class worked example

Now scale that idea up to real numbers, and the "GPU ran out of room" scenario becomes computable. The memory a KV cache occupies, per token, per sequence, summed across every layer, is:

bytes_per_token = 2 × L × H_kv × d_head × bytes_per_value

where the leading 2 accounts for storing both K and V, L is the number of transformer layers, H_kv is the number of key/value heads, d_head is the dimension of each head, and bytes_per_value is 2 for fp16. Llama-2-7B's published configuration is L = 32, hidden size 4096, 32 attention heads with standard multi-head attention (so H_kv = 32, d_head = 128):

bytes_per_token = 2 × 32 × 32 × 128 × 2 = 524,288 bytes = 512 KiB

At a 4096-token context (a full conversation with system prompt), one sequence's cache is 512 KiB × 4096 = 2,097,152 KiB, which is exactly 2 GiB — note this equals the model's own hidden size arithmetic doubling back on itself cleanly, since H_kv × d_head = 4096, the hidden size, for plain multi-head attention. Compare that to the model weights: 7 billion parameters at 2 bytes each is 14,000,000,000 bytes, or ≈13.04 GiB. A single fully-loaded conversation therefore costs about one-seventh as much memory as the entire model. The A100's advertised 40 GB is a decimal figure equal to ≈37.25 GiB of physical memory; after loading weights (13.04 GiB), roughly 24.2 GiB remains. Divide by 2 GiB per full-context conversation: 24.2 / 2 ≈ 12. That is the ceiling the Bengaluru team hit — not a bug, an inevitable consequence of the arithmetic above, and it ignores activation memory, which would push the real number lower still.

Where the memory actually goes

The diagram below contrasts the two decoding strategies at the moment a fourth token is being generated after three tokens already exist.

Shrinking the cache without shrinking the model

Once the cache is understood as the true bottleneck, four independent optimization strategies follow directly from the bytes-per-token formula and from how that storage is managed.

Fewer KV heads — Multi-Query and Grouped-Query Attention. The formula's H_kv term need not equal the number of query heads. Multi-Query Attention (MQA, Shazeer, 2019) uses a single shared key/value head for all query heads; Grouped-Query Attention (GQA, Ainslie et al., 2023 — used in Llama-2's 34B and 70B variants) uses a small group, e.g. 8, instead of 1 or 32. Cache size scales linearly with H_kv, so cutting 32 KV heads to 8 cuts cache memory by exactly 4×, independent of everything else in the model. Applied hypothetically to the 7B example: bytes_per_token drops from 512 KiB to 128 KiB, and a 4096-token conversation drops from 2 GiB to 0.5 GiB — turning the GPU's 12-conversation ceiling into roughly 48. This is precisely why Llama-2-70B, despite having 80 layers and a wider 8192 hidden size, uses only 8 KV heads and ends up with a smaller per-token cache (320 KiB) than the 32-KV-head 7B model (512 KiB): 2 × 80 × 8 × 128 × 2 = 327,680 bytes, versus a hypothetical full-MHA 70B (64 KV heads) at 2,621,440 bytes — GQA buys the 70B model an 8× reduction over what plain MHA would have cost it.

Paging instead of pre-allocating — PagedAttention. Before vLLM (Kwon et al., 2023, UC Berkeley), serving frameworks typically reserved one contiguous memory block sized for the maximum possible sequence length, per request, up front. Suppose a service reserves a full 4096-token buffer for every request but replies average 800 tokens: (4096 − 800) / 4096 ≈ 80% of every reservation sits idle for the whole conversation — classic internal fragmentation. PagedAttention borrows the operating system's virtual-memory trick: the cache is split into small fixed-size blocks (e.g. 16 tokens each) drawn on demand from a shared pool, and each sequence keeps a block table mapping logical positions to physical blocks, exactly as a page table maps virtual to physical addresses. Waste drops to a fraction of one block per sequence, and identical prefixes (e.g. a shared system prompt across many users) can have their blocks referenced by multiple sequences instead of duplicated.

Fewer bits per value — quantized cache. Storing K and V in int8 instead of fp16 halves cache memory again, at a small, measurable accuracy cost; some production systems push to int4 for a further halving, usually reserved for the least-recently-used portion of very long contexts.

Bounding the cache instead of growing it — sliding windows and attention sinks. For streaming or very long contexts, StreamingLLM (Xiao et al., 2023) observed that the first few tokens of a sequence absorb a disproportionate share of attention regardless of their content (an "attention sink"), so evicting them destabilizes the model even when they're semantically irrelevant. The fix keeps a small fixed set of sink tokens' K/V permanently, plus a sliding window of the most recent tokens, and evicts everything in between — capping cache size at a constant regardless of how long the conversation runs, instead of letting it grow to O(N).

The misconception: caching does not defeat quadratic attention

The most common misreading of this topic is believing that a KV cache turns generation into a constant-time-per-token operation, or even that it eliminates the quadratic cost of attention altogether. It does neither. Look again at the cached-path formula derived earlier: Σ(t=1..N) t·d ≈ (1/2)·N²·d. Every new token's query still has to be compared against every key currently in the cache, and the cache keeps growing by one entry per step. Token 4,000 of a long generation is measurably slower to produce than token 40, because its attention step scans four thousand cached keys instead of forty — that per-step cost grows linearly with context length, and summed across a full generation the total is still O(N²), not O(N). What the cache eliminates is the redundant re-derivation of keys and values that a naive implementation would otherwise repeat at every step, which is what pushes total cost to O(N³). The cache achieves the best possible asymptotic complexity for causal attention; it does not escape that complexity. This is also why techniques like sliding-window caching exist at all — if caching already made each step free, nobody would need to bound the cache size for long contexts.

Active recall

Attempt each question before reading its answer.

1. Why does the KV cache store keys and values but never queries?
2. A model has hidden size 2048, 24 layers, 16 attention heads of dimension 128 each, no GQA, and stores its cache in fp16. Compute the cache memory, in MiB, for one sequence at a 2048-token context.
3. Without a KV cache, why does the dominant cost of generating N tokens scale as O(N³) rather than O(N²), even though a single self-attention pass over N tokens is only O(N²)?
4. If a model switches from 32 KV heads to 8 KV heads (GQA) with everything else unchanged, by what factor does its KV cache shrink, and why does this not require retraining the query-side attention heads to a smaller count?
5. Why does reserving one fixed-size contiguous memory block per request, sized for the maximum possible sequence length, waste memory even when the block is never over-filled?
6. A conversation grows past a model's designed context window using a sliding-window cache with attention sinks. What specifically gets evicted, what is deliberately kept even though it's old, and why?

Answers.

1. A query is used exactly once, the instant its token is generated, to attend over the keys and values available at that moment; it is then discarded and never looked up by any future token. Keys and values, by contrast, are looked up repeatedly — every later token's query attends over them — so persisting them avoids repeated recomputation. There is nothing to gain from caching something that is used once and never referenced again.
2. Per-token bytes = 2 × L × H × d_head × bytes_per_value = 2 × 24 × 16 × 128 × 2 = 196,608 bytes = 192 KiB. Across 2048 tokens: 192 KiB × 2048 = 393,216 KiB = 384 MiB.
3. A single attention pass over a fixed N-token sequence costs O(N²) once. But a naive autoregressive decoder without a cache reruns that full computation from scratch at every one of the N generation steps — effectively performing an O(t²) attention pass (and an O(t) set of projections) at every t from 1 to N. Summing t² over t = 1 to N grows as O(N³); it is N separate O(N²)-scale computations stacked on top of each other, not one.
4. Cache size scales linearly with the number of KV heads (H_kv in the formula), so dropping from 32 to 8 KV heads shrinks the cache by exactly 4×. This works without retraining query heads to match because GQA groups the (unchanged) query heads into clusters that each share one KV head — e.g. 32 query heads split into 8 groups of 4, each group's queries attending over one shared key/value pair. The number of query heads, and therefore the model's representational richness on the query side, is untouched.
5. Because actual reply lengths vary and are usually far shorter than the maximum, most of each pre-allocated block sits empty for the request's entire lifetime — the same internal-fragmentation problem that motivated demand paging in operating systems. The block isn't over-filled; it's under-used, and that unused space cannot be lent to any other request because it's reserved as one contiguous chunk.
6. The middle portion of the conversation — older tokens that have scrolled out of the sliding window — gets evicted. The first few "sink" tokens are deliberately kept even though they are the oldest, because attention heads route a disproportionate share of their attention mass onto early positions regardless of content; removing them destabilizes the softmax distribution for every later token, so a small fixed set is pinned in cache permanently while the window of recent tokens slides forward around it.

Think About It

Think about this: How would you explain kv cache optimization: efficient context storage to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where kv cache optimization: efficient context storage is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting kv cache optimization: efficient context storage to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind kv cache optimization: efficient context storage, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.