Build your own KV cache
Single-token cached attention·6 min

One token

0 / 1

In decode you have a populated cache of length P and exactly one new token. You compute Q, K, V for that token only — not for the whole sequence. You append its K and V to the cache (it now holds P+1 rows). Then the new token's Q attends over all P+1 cached K/V — every past position plus itself — softmaxes the scores, and takes the weighted sum of the cached V. That single 768-vector is the attention output for the step.

That asymmetry is the whole point of the cache: you project one row, but you attend over the entire history, because the history's K and V are already sitting in the cache from earlier steps. Nothing past gets recomputed.

The causal mask is trivial here. In a full forward pass, position i may only attend to positions ≤ i, so you mask out the upper triangle. In decode, the new token is the last position — everything in the cache is in its past, including the row you just appended. There is nothing to mask. The new token legitimately attends to all P+1 positions, full stop.

Where positions come from

GPT-2 has no rotary or relative positions. It uses learned absolute position embeddings — a table wpe indexed directly by the token's absolute index. The embedding for the new token is wte[id] + wpe[t], where t is its absolute position in the sequence.

The rule you must internalize: the new token's position is t = cache.length before you append. A cache holding P rows occupies positions 0..P-1, so the incoming token is position P — which is exactly the current cache.length. Read the length first, embed at that position, then append.

Off-by-one on the position corrupts every logit

This is the single most common KV-cache bug. wpe[t] is added to the token before it ever reaches attention, so getting t wrong shifts the input to the entire stack. Use P+1 (the length after append) and every position is one slot too far; reset to 0 and the model thinks every generated token is the start of the sequence. Either way the output silently degrades — no crash, just wrong numbers. The fix is one line: const t = cache.length read before cache.append(...).

One new token over a cache of 4
cache before append:  K[0] K[1] K[2] K[3]      (length P = 4, positions 0..3)
                       V[0] V[1] V[2] V[3]

new token's position:  t = cache.length = 4       <- read BEFORE append
embed:                 x = wte[id] + wpe[4]

project the new token (only):   Q_new   K_new   V_new
append K_new, V_new:   K[0] K[1] K[2] K[3] K[4]    (length is now P+1 = 5)
                       V[0] V[1] V[2] V[3] V[4]

per head (head_dim 64, 12 heads):
  scores = [ Q_new·K[0], Q_new·K[1], Q_new·K[2], Q_new·K[3], Q_new·K[4] ] / sqrt(64)
                                                          \__ includes itself (position 4)
  w      = softmax(scores)            <- length 5, no mask needed
  out    = w[0]V[0] + w[1]V[1] + w[2]V[2] + w[3]V[3] + w[4]V[4]   -> [64]

concat 12 heads -> [768] -> c_proj -> attention output [768]