A Java engineer reading model.generate(prompt) sees one call return a string and
trusts that whatever happens token-by-token underneath is just "the same forward pass, run 100 times".
It is not: a correct decode loop never recomputes an old token's Key or Value, and the position it feeds
into position-aware attention has to stay
absolute even though each step only ever sees one new token.
Inside an LLM showed the KV cache saving work; this exercise builds the two pieces that make that
savings possible — rotary position embeddings (RoPE) and the cache itself — by hand, on top of the
attention you already wrote.
rope_freqs(d, base=10000) -> np.ndarray — shape (d/2,): one rotation
frequency per pair of dimensions, freq_i = 1 / base ** (2i / d). i=0 is the
fastest-turning pair (1 radian per token, at any base); i=d/2-1 is the slowest.apply_rope(x, positions, base=10000) -> np.ndarray — x is
(T, d), positions is (T,), one integer position per row. Rotate
each adjacent pair (x[..., 2i], x[..., 2i+1]) by angle
position * freq_i. This is the convention the published numbers below were computed with —
the other common one splits x into two halves instead of interleaving, and the trap
pairs-adjacent-vs-halves.py is that exact mix-up.decode_incremental(decoder, x_embedded) -> np.ndarray — x_embedded is
(T, D), one already-embedded token per row. decoder is provided: fixed random
Wq/Wk/Wv/Wo, and
decoder.project_q(x) / project_k(x) / project_v(x) /
project_out(x) that just apply those matrices — plus decoder.forward_full(x_embedded),
a complete non-incremental causal-attention forward pass, given so you have a ground truth to match.
Feed x_embedded one row at a time: at step t, project only that row
to get q, k, v; RoPE q and k at
position t (the absolute position, not the step index within some shorter window);
append k/v to a running cache; attend q against the whole
cache so far; project the result through decoder.project_out. Return the stacked logits,
(T, D).Two things worth reading before you start. rope_freqs(8, 10000) comes out
to a clean [1.0, 0.1, 0.01, 0.001] — each pair turns exactly 10× slower than the
last, because the exponent steps in units of 2/8. And the whole point of the cache: check 7
inspects the actual shape decoder.project_k was called with at every step. A
decode_incremental that is numerically correct but re-projects the growing prefix from
scratch each time (recompute-all.py) still passes the "matches forward_full" check —
it is only wasteful, not wrong — so a separate check exists purely to catch the person who "fixed" the
slow version by making it fast without checking it still returns the right numbers
(cache-off-by-one.py is that failure, the other direction).
i = np.arange(d // 2); return 1.0 / (base ** (2 * i / d)). No loop
needed; it is one vectorised expression over the pair index.i, angle
= position * freq_i. The rotated pair is the standard 2-D rotation:
x1' = x1*cos(angle) - x2*sin(angle), x2' = x1*sin(angle) + x2*cos(angle). Pull
the "even" and "odd" columns out with slicing — x[:, 0::2] is every x1,
x[:, 1::2] is every x2 — rotate them as whole arrays against the
(T, d/2) angle grid, then interleave the results back with the same
[:, 0::2] / [:, 1::2] assignment. A rotation cannot change a vector's length —
if your norm check fails, you added something instead of rotating it.for t in range(T): xt = x_embedded[t:t+1] (keep the leading axis, shape
(1, D), not (D,)). q = apply_rope(decoder.project_q(xt), [t], decoder.base),
same for k; v = decoder.project_v(xt) needs no RoPE (V carries no position
information — only Q and K do, because position only matters for where attention looks, not
what it reads). Append k, v to Python lists, and
np.concatenate them fresh each step to get the full K, V — that
concatenation is O(t) but the projection work per step is O(1), which
is the whole saving.q shape (1, d) attending over K shape
(t+1, d): scores = q @ K.T / sqrt(D) (no mask needed — the cache only
ever contains tokens up to and including t, so causality is automatic, not
enforced by zeroing anything), softmax over the last axis, @ V, then
decoder.project_out(...).t. They are equal in this exercise (nothing is ever
evicted from the cache), so that shortcut happens to pass every check here — but it silently breaks the
moment a real system starts dropping old entries (a sliding window, a compacted cache): the position a
token was generated at and its slot in a possibly-shrunk cache stop being the same
number. Use t, the generation step, never len(K_cache).