Decoding with a KV Cache

The problem

While generating text, a model produces one token at a time, and each new token attends to all the previous ones. The keys and values of the past tokens never change (the model is causal: they depend only on tokens before them), so recomputing them at every step is wasted work. The KV (key-value) cache stores them: each step computes the query, key and value of the new token only, appends its key and value to the cache, and attends over the whole cache.

You will write one step of attention with a cache, for one head, and the formula for the size of the cache.

One step with a cache

Use one attention head with weights p = {"Wq", "Wk", "Wv"}, each of shape $(D, d)$. For the new token with embedding x_t of shape $(D,)$:

1. compute $q = x_t W_q$, $k = x_t W_k$, $v = x_t W_v$, each of $d$ numbers;
2. append $k$ and $v$ as a new row of the cached keys and values: the cache now holds $t + 1$ rows;
3. scores: the dot product of $q$ with every cached key, divided by $\sqrt{d}$;
4. softmax of the scores, and the output is the weighted average of the cached values.

No causal mask is needed: everything in the cache is in the new token's past (or is the token itself).

Example. $D = d = 2$, $W_q = W_k = I$ (identity), $W_v = \mathrm{diag}(1, 2)$ (doubles the second number). Three tokens arrive: $(1, 0)$, $(0, 1)$, $(1, 1)$.

• Step 1: $q = (1, 0)$, cache holds one key $(1, 0)$ and one value $(1, 0)$. The only weight is 1, so the output is $(1, 0)$.
• Step 2: $q = (0, 1)$; keys $(1, 0), (0, 1)$; values $(1, 0), (0, 2)$. Scores $(0, 1)/\sqrt{2}$, weights $(0.33, 0.67)$, output $0.33 \cdot (1, 0) + 0.67 \cdot (0, 2) = (0.33, 1.34)$.
• Step 3: $q = (1, 1)$; keys $(1, 0), (0, 1), (1, 1)$; values $(1, 0), (0, 2), (1, 2)$. Scores $(1, 1, 2)/\sqrt{2}$, weights $(0.25, 0.25, 0.50)$, output $(0.75, 1.50)$.

The size of the cache

For every layer, every key/value head and every cached token, the cache stores one key and one value of $d$ numbers each:

where $b$ is the number of bytes per value, $L$ is the number of layers, $G$ the number of key/value heads, $d$ the head size, $T$ the number of cached tokens, and the 2 is for keys and values.

Example. $L = 32$, $G = 32$, $d = 128$, $T = 8192$, 2 bytes per value (the BF16 (brain floating point) format): $2 \cdot 32 \cdot 32 \cdot 128 \cdot 8192 \cdot 2 = 4{,}294{,}967{,}296$ bytes, 4 GiB (gibibytes) for one sequence. With $G = 8$ it is 1 GiB.

What to implement

• attend_step(x_t, cache, p): cache is a dict with "k" and "v", arrays of shape $(t, d)$ for the $t$ tokens seen so far (it starts with shape $(0, d)$). Return (out, new_cache), with out of shape $(d,)$. Do not modify the input cache.
• kv_cache_bytes(n_layers, n_kv_heads, head_dim, n_tokens, bytes_per_value): the formula above.

Requirements

• Running attend_step on the tokens one by one must give exactly the rows of full causal attention over the whole sequence (the tests compare with a reference).
• Scale scores by $1/\sqrt{d}$ and use a stable softmax.

Hint

np.vstack([cache["k"], k[None]]) appends a row to a table; k[None] turns the vector of $d$ numbers into a table of one row.