Fixed-capacity KV cache

Problem

A decoder makes one token at a time. Without a cache, each new token computes the keys and values of the whole prefix again. A KV cache stores the key and the value of each step once. Then each new token needs only one query against the stored keys and values.

Task

Write a class KVCache and a function decode_step.

KVCache(max_len, d_k, d_v) makes an empty cache with room for max_len steps. It has these methods:

1. append(k, v) adds one or more steps after the steps that are already in the cache. A 1D k of shape (d_k,) is one step. A 2D k of shape (n, d_k) is n steps. The same rule applies to v.
2. keys() returns the stored keys, in the order of the appends, with shape (len, d_k).
3. values() returns the stored values with shape (len, d_v).
4. len(cache) returns the number of stored steps.

append raises ValueError in two cases:

• the new steps do not fit, that is, len(cache) + n > max_len;
• k and v have different numbers of steps.

When append raises, the cache does not change.

decode_step(q, cache) computes the attention of one query over the cache. Let K be cache.keys() and V be cache.values(). The scores are s = K q / \sqrt{d_k}. The weights are the softmax of s. The output is the weights times V.

Input and output

• append returns None.
• decode_step returns a tuple (output, weights).
• output is a float array of shape (d_v,).
• weights is a float array of shape (len(cache),).

Constraints

• 1 <= max_len <= 256 and 1 <= d_k, d_v <= 64.
• Scores can be larger than 1000. The output must stay finite.
• The cache stores a copy of the data. A later change to the caller's k or v does not change the cache.
• Do not change the input arrays.
• Use numpy only.

Examples

pythonCopycache = KVCache(4, 2, 2)
cache.append(np.array([[1.0, 0.0], [0.0, 1.0]]), np.array([[1.0, 2.0], [3.0, 4.0]]))
len(cache)
cache.append(np.array([1.0, 1.0]), np.array([5.0, 6.0]))
len(cache)
cache.values()

textCopy2
3
[[1. 2.]
 [3. 4.]
 [5. 6.]]

pythonCopyoutput, weights = decode_step(np.array([1.0, 0.0]), cache)

textCopyoutput
[3. 4.]
weights
[0.40111209 0.19777581 0.40111209]

pythonCopycache.append(np.zeros((2, 2)), np.zeros((2, 2)))

textCopyValueError: the cache is full

After the error, len(cache) is still 3.