The cache data structure·5 min

What the cache holds

0 / 1

Before any attention math, get the container right. The KV cache is a per-layer store of keys and values: one cache per attention layer, each holding a growable K and V. The logical shape of each is [n_heads, seq, head_dim] = [12, seq, 64] — twelve heads, one row per cached position, sixty-four dimensions per head. append(k, v) adds exactly one position (the K and V for the newest token); length is how many positions are stored; clone() makes an independent copy. GPT-2 has 12 attention layers, so the model holds 12 of these caches — one per layer, each minding its own history.

Why `[n_heads, seq, head_dim]`

Attention runs per head. Head h only ever dots its own 64-dimensional slice of the query against its own 64-dimensional slice of every cached key — it never sees the other heads. Laying the buffer out as [n_heads, seq, head_dim] keeps each head's history in one contiguous plane: to score the new token, head h walks down its [seq, 64] slice and touches nothing else. The seq axis is the only one that grows; n_heads (12) and head_dim (64) are fixed by the architecture.

In practice you can store each position as one flat 768-vector (12 × 64, all heads concatenated) and slice out head h's 64 values with an offset — the logical shape is still [12, seq, 64]. What matters is that the seq axis is the growable one and each head reads only its own band.

One layer's K buffer, one head, growing per decode step

head h's view of K  —  shape [seq, head_dim] = [seq, 64]

decode step:        t=0        t=1        t=2        t=3
                  ┌──────┐  ┌──────┐  ┌──────┐  ┌──────┐
       pos 0      │ k₀…  │  │ k₀…  │  │ k₀…  │  │ k₀…  │   ← 64 dims
                  └──────┘  ├──────┤  ├──────┤  ├──────┤
       pos 1                │ k₁…  │  │ k₁…  │  │ k₁…  │
                            └──────┘  ├──────┤  ├──────┤
       pos 2                          │ k₂…  │  │ k₂…  │
                            length=2  └──────┘  ├──────┤
       pos 3                                    │ k₃…  │
       length=1            length=3   └──────┘
                                                 length=4

append(k,v) adds ONE row (one position) along the seq axis.
After k appends:  length == k  and the buffer is [12, k, 64].

Three operations, nothing more

append(k, v) grows seq by one — never by zero, never by two. length returns the current seq. clone() returns a deep, independent copy (you'll need it in Chapter 6 so reusing a saved prefix can't mutate the original). That's the whole interface for this chapter; attention comes later and reads this store, it doesn't change its shape.

The two killer bugs live in the layout

Almost every KV-cache failure in this course traces back to one of two structural mistakes:

Transposed layout — storing [seq, n_heads, head_dim] instead of [n_heads, seq, head_dim]. Heads and positions get interleaved, and head h ends up dotting against the wrong slice. It often runs and produces plausible-looking garbage.
Off-by-one growth — append adding zero or two positions, so length drifts from the true seq. One wrong position quietly shifts every downstream index.

That's why this chapter grades the structure in isolation, with atol=0. If the buffer's shape and contents are exact here, neither bug can hide inside attention later.

References

Why cache at all?Build the KV cache