Visual Explainer

The Transformer & KV Cache

A complete visual walkthrough — from embedding to generation — with interactive simulations and your full source code annotated.

What Is a Transformer?

The big picture

A Transformer is a neural network architecture that processes sequences by letting every token "look at" every other token through a mechanism called self-attention. Unlike RNNs that read tokens one-by-one left-to-right, a Transformer sees the entire sequence in parallel during training — which makes it massively faster to train on GPUs.

The core insight is simple: to understand any word in a sentence, you need context from other words. Self-attention provides a learned, weighted way to gather that context. The word "bank" means different things in "river bank" vs "bank account" — attention figures out which other tokens to focus on.

A language model (LM) transformer is trained to predict the next token given all previous tokens. At inference time, it generates one token at a time in a loop — each new token feeds back in as input. This is called autoregressive generation, and it's where KV Cache becomes essential.

Your code implements a decoder-only transformer — the same family as GPT. It has one layer with multi-head self-attention + an MLP, wrapped in Pre-LayerNorm residual blocks. Minimal, but architecturally identical to the building block of models like LLaMA.

Source Code

Your complete implementation — syntax highlighted

This is the full code we're explaining. Everything below — architecture diagrams, attention walkthrough, KV cache simulation — maps directly back to the classes and methods here.

Model Architecture

DummyLM — your model dissected

Your DummyLM stacks four components into a clean pipeline. A token ID enters at the top and a probability distribution over the vocabulary exits at the bottom. Here's each piece:

🔤

nn.Embedding

Lookup table mapping each of the 1000 token IDs to a learnable 256-dimensional vector.

1000 × 256

👁️

MultiHeadAttention

8 parallel attention heads (each 32-dim). Projects Q, K, V, computes scaled dot-product, merges heads, and projects output. This is where the KV Cache lives.

8 heads × 32d

⚡

MLP (Feed-Forward)

Two linear layers with GELU activation in between. Expands to 1024-dim hidden, then compresses back to 256. Adds non-linear representational power.

256→1024→256

🎯

LM Head

Final linear projection back to vocabulary size. Output logits are passed through argmax to select the most likely next token.

256 → 1000

Pre-LayerNorm: Your code applies LayerNorm before attention and MLP (not after). This is the modern convention used by GPT-2+, LLaMA, etc. It stabilizes training by normalizing inputs before each sub-layer, then adding the residual connection around the raw (un-normalized) output.

Attention Mechanics

The 8 steps inside MultiHeadAttention.forward()

Self-attention answers: "for each token, how much should I attend to every other token?" It does this by projecting the input into three roles — Query (what am I looking for?), Key (what do I contain?), and Value (what information do I carry?) — then computing a compatibility score between each Query and all Keys.

Attention(Q, K, V) = softmax( Q · K^T / √d_k ) · V

① Linear Projections

Input x is projected through learned weight matrices W_Q, W_K, W_V to produce Q, K, V. Each is shape (batch, seq, 256).

② Split into 8 Heads

Each matrix is reshaped: (B, seq, 256) → (B, 8, seq, 32). Each head independently attends in its own 32-dim subspace, learning different relationship patterns.

③ KV Cache Concat ★

If cached past keys/values exist, the new K and V are concatenated with the stored ones along the sequence dimension. Then the updated K, V are saved for next step. This is the key optimization — see Section 05.

④ Scaled Dot-Product

scores = Q · K^T / √32. Each query token gets a score against every key token. Division by √d_k prevents dot products from growing too large, which would push softmax into saturated regions.

⑤ Causal Mask

Lower-triangular mask sets future positions to −∞ before softmax. This ensures token i can only attend to tokens ≤ i — preserving the autoregressive property. During cached decoding with a single Q token, no mask is needed.

⑥ Softmax → Weighted Sum

Softmax normalizes scores into attention weights (probabilities). The output is the weighted sum of Values: tokens with higher scores contribute more information.

⑦ Merge Heads

Transpose back and reshape (B, 8, seq, 32) → (B, seq, 256). The 8 independent perspectives are concatenated into a single representation.

⑧ Output Projection

Final W_O linear layer mixes information across heads and produces the attention block output, which gets added back via the residual connection.

The KV Cache

Trading memory for speed

During autoregressive generation, the model generates one token at a time. Without cache, every new step feeds the entire growing sequence back through the model — recomputing Q, K, V for all previous tokens. This is wasteful because the Keys and Values for past tokens never change (they only depend on the fixed past input, not the future).

The KV Cache exploits this by storing the K and V tensors from every previous step. On each new decode step, only the single new token is projected to get new Q, K, V. The new K, V are appended to the cache, and attention is computed with Q (1 token) against the full cached K, V (all tokens). The savings are enormous:

Without Cache

At step n, project all n tokens through W_Q, W_K, W_V. Compute full n×n attention matrix. Total projection cost over N steps: O(N²).

With Cache

At step n, project only 1 token. Compute 1×n attention row. Append to cache. Total projection cost: O(N). Memory cost: store K, V for all past tokens.

Two-phase generation: Your code splits generation into a Prefill phase (processes the full prompt in one pass, filling the cache) and a Decode phase (generates tokens one-by-one, reading from cache). This is the standard pattern used by all production LLM inference engines.

Interactive Simulation

Step through generation

Click through each decode step to see what happens inside attention. Without cache: every token is recomputed (red). With cache: past K/V are stored (blue) and only the new token (orange) is computed.

Recomputed

New token

Prefill

Cached K/V

New (computed)

Performance Comparison

Generating 200 tokens from a 10-token prompt

Your test script generates 200 tokens with and without KV cache, then asserts identical outputs. The table below shows why caching is faster, and the bars show how much.

Metric	Without Cache	With Cache
Q/K/V projections per step	All n tokens	1 token only
Attention matrix size	n × n (full square)	1 × n (single row)
Input to model per step	Full sequence (growing)	Last token only
Extra memory	None	O(n · d) per layer
Total projection ops (N tokens)	Σ(P+i) × 3 = O(N²)	P×3 + N×3 = O(N)

Deriving the numbers — Without Cache (65,700):

Your prompt is P=10 tokens and you generate N=200 tokens. Each step projects Q, K, and V (×3) for the entire sequence so far:

Step 0: seq_len = 10 → 10 × 3 = 30 projections
Step 1: seq_len = 11 → 11 × 3 = 33 projections
Step 2: seq_len = 12 → 12 × 3 = 36 projections
⋮
Step 199: seq_len = 209 → 209 × 3 = 627 projections

Total = 3 × Σ(i=0 to 199) (10 + i) = 3 × (10 + 11 + 12 + … + 209)
= 3 × ^{200 × (10 + 209)}⁄₂ = 3 × ^{200 × 219}⁄₂ = 3 × 21,900 = 65,700

Deriving the numbers — With Cache (627):

Prefill: process all 10 prompt tokens once → 10 × 3 = 30 projections
Decode: 199 steps × 1 new token each × 3 = 199 × 3 = 597 projections

Total = 30 + 597 = 627

That's it — the cache means you never re-project old tokens. The ratio is 65,700 / 627 ≈ 104.8× fewer projection operations.

Without Cache

total token-projections

With Cache

total token-projections

Without KV Cache—

With KV Cache—

assert torch.equal(...): Your test confirms both paths produce identical output sequences. The KV cache is a pure optimization — it changes how computation is scheduled, not what is computed.