LLM Basics 6 — Attention and the KV Cache

Series: 1 · What a model is on disk → 2 · Tokens and embeddings → 3 · A layer is matrix multiplication → 4 · The forward pass → 5 · Generating text → 6 · Attention and the KV cache → 7 · The hardware fit → Capstone · Why-LLM-Decode-Is-Memory-Bound

Every earlier lesson postponed one thing. Each of the 32 blocks has two halves: attention, then the feed-forward network that lesson 3 walked through. The feed-forward half processes each token’s vector on its own. So without attention, the vector at the end of “The capital of France is” could never learn that “France” was in the sentence.

Attention is the only step where one position takes in information from the others. This lesson covers how it works and what it looked like inside the real model, then turns to the KV cache, the trick that makes generation affordable and whose size grows with every token.

The specimen is still deepseek-coder-6.7b-instruct. For this lesson, the whole model — all 32 blocks, attention included — was rebuilt in about 80 lines of NumPy from the GGUF file’s weights. Its prediction for “The capital of France is” matches llama.cpp’s (the check is below), so the attention weights shown here are the model’s real ones.

The idea: a soft lookup

A hash table takes a key and returns exactly one value. Attention is a soft version of that lookup: instead of one exact match, every entry gets a relevance score, and the result is a blend of all the values, weighted by those scores.

At every position, the model computes three vectors from that token’s vector:

VectorRoleQuestion it answers
Query (Q)what this position is looking for“What do I need?”
Key (K)what this position offers to others“What do I contain?”
Value (V)what this position hands over if chosen“What do I give if you pick me?”

Each is one matvec — matrix-vector multiplication, from lesson 3 — using the block’s attn_q, attn_k, and attn_v matrices.

Then, for one position looking back over the sequence:

// q: this position's query.  k[j], v[j]: key and value of every position j <= this one
for (j = 0; j <= me; j++)
    score[j] = dot(q, k[j]) / sqrtf(128);      // relevance of position j
softmax(score, me + 1);                         // scores -> weights adding to 1
for (j = 0; j <= me; j++)
    out += score[j] * v[j];                     // blend of values, by weight
  • The dot product (multiply pairwise, then add — lesson 3) measures how well a query and a key line up.
  • Softmax (lesson 4) turns those scores into weights between 0 and 1 that add to 1.
  • The output is a weighted average of the values. A position with weight 0.93 contributes 93% of the blend.
  • out then passes through one more matvec, attn_output, and is added to the token’s vector — the same residual add as the feed-forward half (lesson 3).

Dividing by √128 keeps the dot products from growing so large that the softmax hands everything to a single position.

Three details that make it work

Heads: 32 lookups at once

The 4,096-number vectors are split into 32 heads of 128 numbers each, and each head runs its own lookup independently. That lets the model follow several relationships at once — one head can track the subject of the sentence while another tracks the previous word. The 32 outputs are put back side by side into 4,096 numbers before attn_output.

The causal mask: no peeking ahead

A position may attend only to itself and to earlier positions. During generation, later tokens don’t exist yet, and the model is trained under the same rule, so it never learns to depend on them. In the code above, that is the j <= me bound. Laid out as a grid of positions, the allowed entries form a triangle, as the real patterns below show.

Position: RoPE

A dot product doesn’t know where the two vectors sit in the sentence, so “dog bites man” and “man bites dog” would look alike. The model fixes this with RoPE (rotary position embedding): before the dot product, each query and key is rotated by an angle that depends on its position.

The rotation works on pairs of numbers, treating each pair as a point (x, y) in a plane:

x' = x * cos(angle) - y * sin(angle);    // rotate the point (x, y) by angle
y' = x * sin(angle) + y * cos(angle);

Rotating both vectors makes their dot product depend on how far apart the two positions are, not on where they are in absolute terms. Each of a head’s 64 pairs rotates at a different speed: fast pairs track nearby words, and slow pairs track distant ones.

Observed from this model’s metadata: rope.freq_base = 100,000, which sets that range of speeds, and rope.scale_linear = 4. The second value divides every position by 4. That stretches a model trained on shorter text to cover a longer context.

What attention looked like inside the real model

First, the check

Next tokenNumPy rebuildllama.cpp (lesson 4)
Paris92.6%92.7%
"1.5%1.5%
<0.9%0.8%
?0.8%0.8%
:0.7%0.7%
{0.6%0.6%

The two implementations agree to within a few tenths of a percentage point; the small differences come from rounding inside each one’s arithmetic. A mistake in the attention code — a wrong rotation, a missing mask — would not produce " Paris" at 92.6%.

A head that finds “France”

Observed: layer 23, head 12. Each row is a position looking back; each number is how much weight it gives to each earlier position. · marks entries the causal mask forbids.

Looking from ↓ / at →ThecapitalofFranceis
The1.000····
capital0.9800.020···
of0.9670.0090.025··
France0.9330.0120.0180.037·
is0.0240.0010.0010.9300.045
  • From is, this head puts 93% of its weight on France. Through this head, the vector at the last position takes in what France carries — the fact the next token depends on.
  • The other positions have nothing to fetch, so they park their weight on The.
  • Observed: 26 of the model’s 1,024 heads (32 layers × 32) give France more than half their weight from is. The strongest are layer 23 head 12 (0.93), layer 17 head 11 (0.89), and layer 21 head 0 (0.87).
  • A small number of specialized heads did the important work for this prompt.

Most attention goes to the first token

Observed: weight from is to each position, averaged over all 32 heads, for selected layers:

LayerThecapitalofFranceis
00.1860.1290.2780.1400.267
10.2630.0750.1560.2010.305
20.7730.0470.0600.0610.059
80.7760.0570.0460.0700.052
160.7200.0720.0380.1180.052
240.8570.0430.0140.0530.033
310.6590.0640.0620.1130.102
  • Layers 0 and 1 spread their attention around, favoring nearby words.
  • From layer 2 on, most heads put most of their weight on the first token, The. Across every layer and head, the average is 77%.
  • The standard explanation is the attention sink. Softmax forces a head’s weights to add to 1, so a head with nothing useful to fetch still has to put its weight somewhere. Models learn to park it on the first position, which acts as a “no-op” slot.
    • That explanation comes from published research on this behavior, not from this experiment.
    • What this experiment shows is the pattern itself.
  • So “attention” mostly means “attend to nothing in particular”, with a few heads doing pointed work, like the France head above.

The KV cache

Why earlier work can be saved

Lesson 5’s loop appends one token and runs the model again. Run naively, every trip would recompute every earlier position. But the causal mask means earlier positions never see later ones, so appending a token cannot change anything already computed for the positions before it. In particular, their keys and values stay exactly the same.

So engines compute each position’s K and V once, store them, and reuse them on every later step. That store is the KV cache. Each new token then needs only:

  1. its own Q, K, and V — one position’s worth of matvecs, in every block
  2. its K and V appended to the cache
  3. its query compared against every cached key, and a blend of every cached value

Observed: the rebuild appended Paris to the prompt both ways.

Next-token predictionPositions processed
Recompute everything. 39.75%, ." 8.56%, < 7.96%6
Use the KV cache. 39.75%, ." 8.56%, < 7.96%1
Largest difference in any of the 32,256 logits0.000016

The results are the same. The logits run up to about 25, so a largest difference of 0.000016 is under one part in a million: floating-point rounding, from adding the same numbers in a different order. The cached step simply skips work whose answer was already known.

What the cache costs

The cache trades compute for memory, and the memory grows with every token.

Computed from this model’s dimensions, storing each number as FP16, a 16-bit (2-byte) floating-point format:

per token:  2 (K and V) × 32 layers × 4,096 numbers = 262,144 numbers × 2 bytes ≈ 0.5 MB
Context lengthKV cache
2,048 tokens~1.1 GB
16,384 tokens (this model’s maximum)~8.6 GB

At full context, the cache is larger than the 7.2 GB model itself.

  • Newer models shrink it with grouped-query attention (GQA). Several query heads share one set of keys and values.
    • This model has 32 key/value heads for 32 query heads, so there is no sharing (attention.head_count_kv = 32 in its metadata).
    • A model with 8 key/value heads for 32 query heads needs a cache a quarter the size.
  • The cache is also read on every generated token, because step 3 above touches every cached key and value.

That second point explains a measurement from lesson 4. Decode slowed from 33.2 to 29.6 tokens/s after a 1,921-token prompt.

  • The weights didn’t change. About 7.0 GB is read per token: the whole file, except the embedding table, which is only indexed.
  • But the cache now held about 1.0 GB, so each token read about 8.0 GB instead of 7.0.
  • If speed is proportional to bytes read, that predicts 33.2 × 7.0 / 8.0 ≈ 29 tokens/s, against 29.6 measured.

The prediction assumes llama.cpp’s default FP16 cache, and it ignores smaller costs, so treat it as a sanity check rather than a model. But it shows the mechanism: long contexts slow generation because the cache is more bytes to read.

And reading the prompt gets expensive too

During prefill (lesson 4), every prompt position attends to every earlier one, so the number of query-key scores grows with the square of the prompt length. That is 1,921² ≈ 3.7 million per head per layer for lesson 4’s prompt. Doubling a prompt quadruples that part of the work, which is why very long contexts need tricks of their own.

At a million tokens the same count is about 10¹² — a trillion scores per head per layer, or about half that once the causal mask removes the future positions. Long-context models cannot afford that, and the usual way out is sparse attention.

One fix: sparse attention

The idea comes straight from what the real model showed above: most of a head’s weight lands on a handful of positions. The France head put 93% on one token, and across the model 77% of all weight went to the first token. Computing a score for every earlier position, only to find most of them near zero, is mostly wasted work. Sparse attention tries to skip that work.

It runs in two passes:

  1. A cheap index pass picks where to look. The earlier positions are grouped into blocks — runs of consecutive positions. Small, low-cost queries and keys, much cheaper than the full ones, score every block, and the model keeps the top-k: the k highest-scoring blocks.
  2. Full attention runs only on those blocks. The real query is compared against the keys of the selected blocks, and blends only their values. Everything outside them is skipped.

The expensive part now grows with the number of blocks kept, not with the length of the context.

  • It is an approximation, unlike the KV cache. The cache gives exactly the same answer as recomputing (tested above). Sparse attention changes what the model can see: if the index pass misses a block that mattered, the token never sees it. It works because models are designed and trained to attend this way.
  • The course’s model does not use it. The NumPy rebuild above computes full attention — every head scores every earlier position — and it reproduces llama.cpp’s prediction, so this model attends densely.
  • One published design is MiniMax’s sparse attention. It shares one cheap index query per group of query heads (the GQA groups from the KV-cache section), scores blocks by the highest score inside each, and lets different groups pick different blocks in the same pass.

Where this leaves us

Every piece of the model has now been opened: the file, the tokens, the matvecs, the forward pass, the loop, and attention. Two of the lessons so far — matvec versus matmul in lesson 3, and the cache read here — kept ending at the same place: memory bandwidth. Lesson 7 looks at the hardware itself — what a GPU is, where its memory sits, and why that decides speed — which sets up the capstone.

On a DGX Spark

Reproduce it: every number in this lesson is printed by lessons/06_attention_rebuild.py in the course’s companion code repository, llm-bytes-up-code. Run it with --server to repeat the comparison against llama.cpp. The run this lesson quotes is recorded in that repository’s outputs/2026-10-12/, produced from a fresh clone on the Spark.

Everything here was run on the course’s DGX Spark on 2026-10-12.

  • The rebuild is the Rebuild class in lessons/06_attention_rebuild.py, about 80 lines of NumPy.
    • It reads the GGUF with the gguf Python package and dequantizes each block’s Q8_0 weights to FP32 as it reaches them.
    • It runs embedding → 32 × (RMSNorm, Q/K/V, RoPE, causal attention, output projection, residual, RMSNorm, SwiGLU feed-forward, residual) → final norm → output.weight → softmax.
    • The model’s settings were read from the file’s metadata: 32 heads, 32 key/value heads, 128 dimensions per head, RoPE base 100,000, linear RoPE scale 4, and an RMSNorm epsilon of 1e-6.
    • The RoPE pairing follows llama.cpp’s “normal” layout for this architecture, which rotates adjacent pairs (0,1), (2,3), and so on.
    • The final check against llama.cpp is the agreement table above.
  • It ran on the Spark’s CPU, at about 5 seconds per forward pass of 5–6 tokens, and about 9 GB of memory. Nearly all of that time is dequantizing 7 GB of weights; the arithmetic itself is small.
  • The attention weights were captured from inside the rebuild, for every layer and head.
  • The KV-cache test ran the 5-token prompt once while keeping each layer’s K and V (with RoPE applied to K). It then processed Paris alone at position 5 against that cache, and compared the result with a fresh 6-token run.
  • The cache-size figures are computed, from the model’s dimensions and FP16 storage. They were not read from llama.cpp, whose log at default verbosity doesn’t report them.