Build GPT · 9 of 10 · Python

Maintain cached keys and values

While a GPT writes, each new token attends to every token before it. The keys and values of those earlier tokens do not change from one step to the next, so working them out again every time is wasted effort. A KV cache stores them. Each step computes the key and value of the new token only, adds them to the store, and reuses the rest. Memory is limited, so in this exercise the oldest entries are dropped when the store is full.

Do the exercise in the app

Learn it

Words to know

KV cache: The stored keys and values of earlier tokens, kept so they are not recomputed at every step of generation.
query, key, value: The three roles in attention: a query is what a token looks for, a key is what it can be matched on, a value is what it passes on.
context window: The number of tokens the model reads at one time.

How it works

Add the new keys to the end of the stored keys, and the new values to the end of the stored values.
If there are now more than capacity rows, keep only the last capacity of each.
Keys and values must stay lined up: row i of one belongs to the same token as row i of the other. So both lists are cut in the same way.

Worked example

Input
keys = [[1]]
values = [[10]]
new_keys = [[2], [3]]
new_values = [[20], [30]]
capacity = 2

Output
[[[2], [3]], [[20], [30]]]

Step	Keys	Values
Stored	[1]	[10]
After adding	[1], [2], [3]	[10], [20], [30]
Keep the last 2	[2], [3]	[20], [30]

The oldest token, with key 1 and value 10, falls out of both lists together. Key 2 still sits beside value 20.

Why it matters

A cache is why a chat model keeps a steady pace as its reply grows. The workshop's GPT recomputes instead: its positions are learned, so when old tokens are dropped every kept token's key and value change.

Your turn

Add new rows to a key cache and a value cache, and keep only the most recent rows.

keys: the stored keys, one row per earlier token; may be empty
values: the stored values, lined up with keys
new_keys: the keys of the new tokens
new_values: the values of the new tokens, lined up with new_keys
capacity: the most rows to keep; positive
Returns: The list [keys, values] after adding the new rows and keeping the last capacity rows of each. Build new lists and leave the four inputs as they are; the checks cannot see this, but the caller can.

Example 1

Input
keys = [[1]]
values = [[10]]
new_keys = [[2], [3]]
new_values = [[20], [30]]
capacity = 2

Output
[[[2], [3]], [[20], [30]]]

Three rows would be stored, but the capacity is 2. The oldest row is dropped from both lists.

Example 2

Input
keys = []
values = []
new_keys = []
new_values = []
capacity = 2

Output
[[], []]

Nothing stored and nothing new: both caches stay empty.

Example 3

Input
keys = [[1]]
values = [[10]]
new_keys = [[2]]
new_values = [[20]]
capacity = 5

Output
[[[1], [2]], [[10], [20]]]

Two rows fit within a capacity of 5, so nothing is dropped.

Hints

Each hint gives away a little more. Open one only when you need it.

Hint 1: Where to start
Do the same two things to the keys and to the values: join, then trim.
Hint 2: The key step
keys + new_keys makes a new joined list without changing either input. A negative slice takes from the end.
Hint 3: The outline
Return [(keys + new_keys)[-capacity:], (values + new_values)[-capacity:]].

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

Join old and new into fresh lists; the inputs stay as they were. + on lists builds a new list. extend or += would change the caller's cache in place.
Keep the newest rows, cutting both lists at the same place. [-capacity:] takes the last capacity items, or all of them when there are fewer. The same cut on both keeps every key beside its value.

The same numbered steps are the comments in the solution.

Python

def solve(keys, values, new_keys, new_values, capacity):
    # 1. Join old and new into fresh lists; the inputs stay as they were
    all_keys = keys + new_keys
    all_values = values + new_values
    # 2. Keep the newest rows, cutting both lists at the same place
    return [all_keys[-capacity:], all_values[-capacity:]]

Watch out

Do not trim before adding. That keeps capacity old rows and then goes over the limit with the new ones.
Never cut the two lists differently. A key paired with another token's value gives wrong attention with no error to warn you.

Check yourself

What does a KV cache save, and why is that allowed?

The keys and values of earlier tokens. With a causal mask they do not depend on later tokens, so they never need to be recomputed.

Why must keys and values be trimmed in the same way?

Row i of each belongs to the same token. Cutting them differently would pair a key with the wrong value.