EasyCompilers

Vocabulary for this course

The terms the chapters assume, in plain words. You can open this from any chapter.

attention
A weighted blend of value rows, with weights decided by how well each query matches each key.
attention map
The table of scores showing how much each position attended to every other position.
causal mask
A rule that blocks each position from seeing later positions, so predictions can only use the past.
dimension
One axis of a list or grid of numbers; its length is how many values fit along that axis.
distribution
The pattern of how often each possible value appears, or how likely each one is.
dot product
Multiply matching numbers in two lists and add the results; a large total means the lists point the same way.
embedding
A learned list of numbers that stands for a token, with similar tokens placed near each other.
GPU
A processor built for many small calculations at once, which is why matrix work runs well on it.
gradient
A list of slopes, one per number, saying how much each nudge changes the result.
hard pointer
A choice that sends all the weight to one position, instead of spreading it across several.
head
One parallel slice of attention, with its own projections; a layer runs many so it can track different relations at once.
key
What a position advertises it can offer, compared against every query to decide who receives attention.
key axis
The direction in the key vectors whose width decides how far the scores are divided down.
layer
One stage of a model that reads the running vector, transforms it and passes it on.
layer norm
A per-position rescale that keeps the residual stream at a steady size before the next block.
logits
The raw scores a model produces before softmax turns them into probabilities.
matrix
A grid of numbers arranged in rows and columns, used to transform a list into another list.
negative infinity
A value smaller than every real number, used to mark a score that must never be chosen.
one-hot
A vector with a single one and zeros elsewhere, meaning all the weight sits on one choice.
output projection
The learned matrix applied after attention, which turns the mixed values back into the model's width.
parameter
A number inside a model that training adjusts, rather than an input it reads.
positional encoding
Extra numbers added to token embeddings so the model can tell which position each one came from.
projection
A learned matrix that turns a list of numbers into a different list, often shorter.
query
What a position is looking for, compared against every key to decide where to take information from.
residual stream
The running vector each block reads from and adds its output back into.
retrieval layer
A layer whose job is to pull specific stored information into the current position.
saturating
A score so large that the softmax gives almost all the weight to one position.
scale
The step that divides the scores by the square root of the key width, so they stay in a workable range.
score matrix
The grid holding every query's score against every key, one row per position.
soft alignment
Matching by degree rather than by a single choice, so several positions can contribute a little.
softmax
A rule that turns a list of scores into positive weights that sum to one.
token
A chunk of text — a word, part of a word, or a character — that the model treats as one unit.
value
The content a position passes on when attention selects it.
variance
The average squared distance of values from their average, so a spread-out set scores higher.
vector
An ordered list of numbers, where each position is one measured property.
weighted sum
Adding values together with different shares, where the shares come from the attention scores.
weights
The numbers a model multiplies inputs by before adding the results together.

Transformers from the ground up

KV cache and inference

Generation is a loop of one-token forward passes. The cache keeps the keys and values that do not change, turning a cubic decode into a quadratic one and a memory problem.

Part 9 of 10

Advanced17 min readBack to Courses

Decoding is a loop

A decoder-only model produces one token per forward pass. To produce the next token, it conditions on the prompt and on everything generated so far. The obvious implementation therefore passes the whole sequence through every layer again for each new token.

The cost of that naive approach is easy to state. Attention compares every position with every other. That is n squared comparisons for a prefix of n tokens inside each layer. Repeating the whole pass for each of the n generated tokens gives n cubed. Almost all of that work recomputes keys and values for tokens that have not changed.

The causal mask is what makes the past stable. Position i never attends to a position j greater than i. Appending a token therefore cannot change any earlier position’s key or value. Once a prefix has been processed, its projections are fixed for the rest of the generation. That is the property the cache uses.

What the cache holds

The cache stores a key vector and a value vector. It does so for every layer, every key-value head and every past position. Queries are not cached, because only the new position needs one. At step t the model projects the new token. It appends the token’s key and value and attends over the t cached pairs. Work per step becomes linear in t. The whole generation is quadratic.

The memory is easy to estimate. Multiply two tensors by layers, key-value heads, head width, positions, batch and bytes per value. Take a 32-layer model with 32 heads of width 128. At 8,192 positions in fp16 it needs 2 × 32 × 32 × 128 × 8192 × 2 bytes. That is exactly 4 GiB for one sequence.

Grouped-query attention is the usual way to shrink that number. Several query heads share one key-value head. The cache then keeps 8 heads instead of 32, and the same context costs 1 GiB. The queries still have their own projections. The model keeps most of its capacity while the cache falls by a factor of four.

Passing the cache back in

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

inputs = tokenizer("The cache stores", return_tensors="pt")
output = model(**inputs, use_cache=True)
past = output.past_key_values

next_token = output.logits[:, -1, :].argmax(dim=-1, keepdim=True)
output = model(next_token, past_key_values=past, use_cache=True)
print(output.logits.shape, len(output.past_key_values))

Cached and uncached agree

Run both paths over five tokens, confirm the outputs match, then append a sixth token and compare the projection counts again.

Output
Click "Run" to execute your code

Prefill and decode are different jobs

The prompt is processed in one parallel pass called prefill. Every position is known at once. The matrix products are large, and the work is compute-bound, meaning the maths takes longer than the memory reads. Prefill fills the cache and produces the first token. It therefore dominates time to first token.

Generation after that is decode: one token per sequence per step. Each step reads all the weights to do a fairly small multiply. Decode is therefore memory-bandwidth-bound, not compute-bound. Batching helps, because the weights are read once and used for every sequence in the batch. So serving throughput is quoted in tokens per second across a batch.

The split explains the metrics a serving system reports. Time to first token and the gap between tokens are separate numbers. Different resources limit them. A long prompt can be cheap per token while a long generation is expensive per token. A batch that improves throughput can make each request slower.

A cache is only valid for the exact prefix and the exact weights that produced it. Reusing it after editing the prompt, changing the model or applying a different chat template gives plausible text built on the wrong keys.

Reuse, eviction and sampling

Prefix caching reuses the keys and values of a shared prompt. Two requests that begin with the same system prompt can share the cache for those positions. The projections depend only on the prefix and the weights. Multi-turn chat and repeated instructions are the cases where this saves the most.

Long contexts eventually exceed memory. A sliding-window model keeps only the last W positions. Some serving stacks evict or quantise older entries. Eviction changes what the model can attend to. It is a quality decision as well as a memory one. A cache that silently drops the start of a document produces answers that ignore it.

Sampling sits on top of all of this. The logits for the new position are divided by a temperature. They are filtered by top-k or top-p, then sampled. The cache changes none of that distribution. It only makes computing it affordable. A cache bug usually shows up as repetition or drift, not as an exception. The shapes stay valid even when the contents are wrong.

Budget the cache

Compare the three configurations, then change the context to 32,768 and see how many sequences fit in 80 GiB.

Output
Click "Run" to execute your code

Check what you learned

One question at a time. A wrong answer explains itself — take another go.

Part 1 of 5

What does the key-value cache store?