Author(s): Naimish Baranwal

Originally published on Towards AI.

From Vanilla Transformers to Modern LLMs: The Evolution of Transformer Architecture — Part 2: Inference efficiency (KV Cache)

Intuition Behind KV Cache

In the previous part, we discussed Query (Q) and Key (K) while understanding positional information. A simple way to think about attention is:

• Query (Q) → What am I looking for?
• Key (K) → What information does each token represent, so I can decide whether it is relevant to the query?
• Value (V) → What information should I actually retrieve if that token is relevant?

You can think of it like a search system:

Query = the question
Key = the information used to determine whether a token is relevant to the Query
Value = the actual information retrieved

This is also why positional information in RoPE is applied primarily to Q and K: attention decides which tokens should interact through the Q-K relationship.

But KV Cache solves a different problem.

During autoregressive generation (LLM inference), an LLM generates one token at a time. At every new step, the newly generated token needs to attend to all the previous tokens.

Suppose the model has already generated:

The→ name→ is → Naimish

When generating the next token, the model needs the Key and Value representations of all previous tokens.

Without KV caching, the growing sequence would cause the model to recompute the K and V representations of previously processed tokens at each decoding step.

For example:

Step 1:
The
→ calculate K₁, V₁
Step 2:
The Name
→ calculate K₁, V₁ again
→ calculate K₂, V₂
Step 3:
The name is
→ calculate K₁, V₁ again
→ calculate K₂, V₂ again
→ calculate K₃, V₃

Notice the repetition.

The Keys and Values for "The" and "name" have already been calculated, but we keep calculating them again.

KV Cache: The Basic Idea

Instead of throwing away the previously calculated K and V, we store them in memory.

Previous tokens
 ↓
Calculate K and V
 ↓
 Store them
 ↓
 KV Cache

When the next token arrives, we only calculate the new K and V and append them to the existing cache.

Conceptually:

Step 1:
K₁, V₁
↓
Cache = [K₁, V₁]

Step 2:
New K₂, V₂
↓
Cache = [K₁, K₂], [V₁, V₂]

Step 3:
New K₃, V₃
↓
Cache = [K₁, K₂, K₃], [V₁, V₂, V₃]

Therefore, instead of repeatedly recalculating the past, the model reuses the previously computed Key and Value representations.

That is the basic idea behind KV Cache:

Compute the past once, store it, and reuse it during future generation.

This significantly reduces redundant computation during autoregressive inference and is one of the most important techniques used to make LLM generation efficient.

However, this optimization introduces another problem:

If we keep storing K and V for every previous token, how much GPU memory will the KV Cache consume as the context grows?

Before getting into that memory problem, let’s first understand exactly how Q, K, and V are calculated.

Mathematics Analogy

Now that we understand the intuition behind KV Cache, let’s look at the mathematics. We are discussing inference time here.

During training, the model learns its parameters, including the projection matrices used to generate Q, K, and V. Once training is complete, these weights are fixed during inference.So during inference, we use these already-trained weights to calculate Q, K, and V for the input tokens.

1. How Are Q, K, and V Created?

Suppose our input contains four tokens:

The
name
is
naimish

For simplicity, assume each token is represented using a 3-dimensional embedding.

Our input matrix (X) is:

X =
[
 [ 0.1, 2.3, -3.4 ], → The
 [ 0.4, 3.7, -7.9 ], → name
 [-1.2, -0.4, 12.7 ], → is
 [ 0.8, 9.5, -2.2 ] → naimish
]

So the shape of (X) is:
X : 4 × 3
where
4 = number of tokens
3 = embedding dimension

The Transformer contains three learned projection matrices:

Wq → Query projection matrix
Wk → Key projection matrix
Wv → Value projection matrix
For our simplified example, suppose:
Wq =
[
 [0.12, 1.3, 8.4],
 [2.4, 5.7, 2.9],
 [1.2, 7.4, 1.7]
]

and similarly we have (W_k) and (W_v).

These matrices were learned during training and are fixed during inference.We then calculate:

Q = XWq

K = XWk

V = XWv
Therefore:

Input X
 ↓
 ┌───────┬───────┬───────┐
 ↓ ↓ ↓
Wq Wk Wv
 ↓ ↓ ↓
 Q K V

The result is three matrices containing the Query, Key, and Value representations for our tokens.

For a simplified single-head example:

X : 4 × 3

Wq : 3 × 3
Wk : 3 × 3
Wv : 3 × 3

Therefore:
Q = XWq → 4 × 3

K = XWk → 4 × 3

V = XWv → 4 × 3

In a real multi-head Transformer, these projections are divided across attention heads, but the basic idea remains the same.

2. From Q, K, V to Attention

Now we have: Q, K and V

The next question is:

How does the model decide which tokens should receive more attention?

We first calculate the similarity between every Query and every Key using the dot product.

This is written as:

**Attention Scores** = QKᵀ

Then we scale the scores by the square root of the Key dimension:
Scaled Scores = QKᵀ / √dₖ

Then we apply Softmax:
**Attention Weights** = softmax(QKᵀ / √dₖ)

Finally, these attention weights are multiplied by the Value matrix:

**Context vector** = softmax(QKᵀ / √dₖ)V

So the complete scaled dot-product attention formula is:

Attention(Q, K, V)=softmax(QKᵀ / √dₖ)V

The important distinction is:

Q × Kᵀ
 ↓
How relevant is each token?
 ↓
Softmax
 ↓
Attention weights
 ↓
Multiply by V
 ↓
Context / attended representation

And this is exactly where KV Cache becomes useful during autoregressive generation. The Query for the newly generated token must interact with the Keys of all previous tokens:

New Q
 ↓
Q × [K₁, K₂, K₃, ..., Kₙ]ᵀ

But those previous Keys have already been calculated. So instead of recalculating them:

Previous K + V
 ↓
 KV Cache
 ↓
 Reuse them

Only the new token's K and V need to be calculated and appended to the cache.

That is the fundamental mathematical reason KV Cache reduces redundant computation during autoregressive inference.

Visual Representation:

Limitations of KV Cache

KV Cache solves one major problem: repeatedly recomputing the Keys and Values of previous tokens.

But it creates a new problem:

We are now storing all those Keys and Values in GPU memory.

As the context grows, the cache grows as well.

For example, suppose we have generated:

10 tokens
 ↓
10 sets of K and V

1,000 tokens
 ↓
1,000 sets of K and V100,000 tokens
 ↓
100,000 sets of K and V

For a **7B-parameter model** with **32 layers, 32 attention heads, head dimension \(d = 128\)**, running in **FP16** with a **batch size of 1**, the KV Cache can consume a surprising amount of GPU memory:

Context Length Approx. KV Cache Memory
-------------------------------------------
4K tokens ≈ 1 GB
32K tokens ≈ 8 GB
128K tokens ≈ 32 GB

This means that at **128K tokens**, the KV Cache alone can require around **32 GB of VRAM**—before we even account for the memory needed to store the model weights and other runtime requirements.

This is why KV Cache becomes a serious memory bottleneck as context length increases.

So KV Cache trades computation for memory.

Why does the cache become so large?

Consider a Transformer with:

Number of layers = L
Number of KV heads = H_KV
Head dimension = d
Context length = T

For every layer and every token, we need to store both K and V.
Therefore,Ignoring batch size, data type, and implementation-specific details,
 the approximate number of cached values is proportional to:

KV Cache elements ∝ 2 × L × T × H_KV × d_head

The 2 comes from storing both: K + V

The important part is:

KV Cache ∝ Context Length × Number of KV Heads × Head Dimension × Number of Layers

This means that as the context length increases, the KV Cache can become extremely large.

Notice something important: KV cache memory is directly proportional to the number of Key-Value heads. This means that reducing the number of K/V heads can directly reduce inference memory.

MHA → many KV heads
GQA → fewer KV heads
MQA → one KV head

For a long-context model, this memory requirement can become one of the major bottlenecks during inference.

The Core Problem with Multi-Head Attention

Now let’s connect this to Multi-Head Attention (MHA).

In standard MHA, every attention head has its own Query, Key, and Value projections:

Head 1 → Q₁, K₁, V₁
Head 2 → Q₂, K₂, V₂
Head 3 → Q₃, K₃, V₃
...
Head H → Qₕ, Kₕ, Vₕ

During autoregressive inference, we therefore need to cache:

K₁, V₁
K₂, V₂
K₃, V₃
...
Kₕ, Vₕ

for every previous token and every Transformer layer.

The Query vectors are not the problem here because the Query of the new token is computed and used at each decoding step.The major memory cost comes from storing the historical:

K+V

for all tokens.

So we arrive at a new question:

Do we really need a separate Key and Value cache for every attention head?

This question led researchers to explore more memory-efficient attention architectures.

Conclusion

KV Cache is one of the simplest but most important optimizations introduced for efficient LLM inference.

Instead of recalculating the Keys and Values of all previous tokens at every generation step, we calculate them once, store them in memory, and reuse them.This significantly reduces redundant computation.

But every optimization comes with a trade-off. KV Cache replaces a computation problem with a memory problem. As the context length, number of layers, and number of Key-Value heads increase, the amount of data that must remain in memory also increases.

And this becomes especially important with Multi-Head Attention (MHA), where every attention head maintains its own Key and Value representations.That brings us to an important question:

Do we really need a separate K and V for every Query head?

This question led to the evolution of more memory-efficient attention architectures such as MQA (Multi-Query Attention), GQA (Grouped-Query Attention), and eventually more sophisticated approaches such as MLA (Multi-Head Latent Attention).

At the same time, reducing KV-cache memory is only one part of making attention efficient. The actual computation and memory movement involved in attention also became a major bottleneck, leading to techniques such as Flash Attention and other efficient-attention mechanisms.

TL;DR

KV Cache
 ↓
Avoid recomputing past K and V
 ↓
Faster autoregressive inference
 ↓
But...
 ↓
KV Cache consumes GPU memory
 ↓
Especially expensive with MHA
 ↓
Can we reduce the number of K/V representations?
 ↓
MQA → GQA → MLA

So the evolution is not about replacing the Transformer overnight.

Instead, researchers kept asking:

Can we keep the strengths of attention while making it faster, cheaper, and more memory-efficient?

That is exactly what we will explore in the next part.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI