One decode step with a KV cache
Hard25 min10 hidden tests
A chat model generates one token at a time. Recomputing keys and values for the whole conversation at every step would be wasteful, so it keeps a KV cache: the key and value vectors of every earlier token. Write decode_step(x, Wq, Wk, Wv, k_cache, v_cache) for one attention head. Project the new token's hidden state x to a query, a key and a value, append the key and value to the caches, then let the query attend over every cached token, itself included, with scaled dot-product attention. Return the attention output and the two grown caches. With an empty cache, the new token can only attend to itself, so the output is its own value vector.
Decode one token at a time
q = x Wq; append k, v; out = softmax(K q / √d) VToken 0: q = (2, 1), and its key (1, 2) and value (3, 2) start the cache. With nothing else cached it can only attend to itself, so the output is its own value, (3, 2).
Cache has 1 rows. Output (3, 2).
Examples
Empty cache: the output is the token's own value [3, 2]
decode_step([1.0, 0.0, 2.0], [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]], [[1.0, 0.0], [0.0, 1.0], [0.0, 1.0]], [[1.0, 2.0], [0.0, 1.0], [1.0, 0.0]], np.zeros((0, 2)), np.zeros((0, 2)))
returns ([3, 2], [[1, 2]], [[3, 2]])
The next token, with the first one already cached
decode_step([0.0, 1.0, 0.0], [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]], [[1.0, 0.0], [0.0, 1.0], [0.0, 1.0]], [[1.0, 2.0], [0.0, 1.0], [1.0, 0.0]], [[1.0, 2.0]], [[3.0, 2.0]])
returns ([2.00928, 1.66976], [[1, 2], [0, 1]], [[3, 2], [0, 1]])
Requirements
xhas shape (d_model,).Wq,WkandWvhave shape (d_model, d_head).k_cacheandv_cacheare arrays of shape (t, d_head), and t can be 0 (np.zeros((0, d_head))).q = x @ Wq,k = x @ Wk,v = x @ Wv. Appendkandvas the last rows of the caches, giving shape (t + 1, d_head).- Weights are the softmax of
K @ q / sqrt(d_head)over all t + 1 rows. Stay finite when scores are in the thousands. The output isweights @ V. - Return a tuple
(out, K, V): a (d_head,) float array and the two (t + 1, d_head) float caches, oldest token first.
The idea behind it
When a transformer generates text, each new token needs attention over every earlier token. The keys and values of earlier tokens never change, because a causal model never lets them see later tokens. A KV cache stores them, so each step computes the projections for one new token, appends one row to each cache and runs one query against the cache. That turns the work per step from reprocessing the whole sequence into a single row of attention. No mask is needed at decode time: the cache holds only the past and the token itself, which is exactly what causal attention allows. The cache grows with every token and with every layer and head, which is why long contexts are limited by memory. Grouped-query attention, in which several query heads share one set of keys and values, exists to shrink it.
Tested in interview case Fix the attention layer, then add a KV cache.