Attention with a KV Cache

Write one decoding step of attention that uses a key-value cache.

When a language model generates text, it adds one token at a time. Without a cache, each new token would recompute the keys and values of every earlier token. The KV cache stores them. At each step you only compute the key and value of the new token, append them to the cache, and attend from the new query over the whole cache.

Given the new token's q, k_new and v_new, and the caches k_cache and v_cache from earlier steps:

K = append k_new to k_cache        # shape (t + 1, d_k)
V = append v_new to v_cache        # shape (t + 1, d_v)
weights = softmax(K @ q / sqrt(d_k))
out     = weights @ V

Return (out, K, V) so that the caller can pass K and V to the next step. There is no mask, because the cache only holds the past and the present.

Function signature

def attend_with_cache(q, k_new, v_new, k_cache, v_cache):

Example

Input

k_cache = np.zeros((0, 2))
v_cache = np.zeros((0, 2))
out, k_cache, v_cache = attend_with_cache(np.array([1.0, 0.0]), np.array([1.0, 0.0]),
                                          np.array([2.0, 3.0]), k_cache, v_cache)
out, k_cache, v_cache = attend_with_cache(np.array([0.0, 1.0]), np.array([0.0, 1.0]),
                                          np.array([4.0, 5.0]), k_cache, v_cache)
out

Output

array([3.3395231, 4.3395231])

Reasoning

The first step has one cached token, so the output is its own value [2, 3]. In the second step the cache holds two keys. The scores are 0 and 1 / sqrt(2), so the weights are 0.330 and 0.670, and the output is 0.330 * [2, 3] + 0.670 * [4, 5]. Running the steps one by one gives exactly what causal attention over the full sequence gives.

Constraints and notes

• q and k_new have shape (d_k,) and v_new has shape (d_v,).
• The caches can start empty, with zero rows.
• Do not modify the input caches in place.