← Lesson
The KV Cache
Joeven
Run
Reset
Python loads on first run
past_keys = [] def prefill(prompt_vecs): past_keys.clear() past_keys.extend(prompt_vecs) return len(past_keys) def decode_one(new_vec): scores = [sum(a * b for a, b in zip(new_vec, k)) for k in past_keys] past_keys.append(new_vec) return scores, len(past_keys) prompt = [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]] print("prefill tokens", prefill(prompt), "cache", len(past_keys)) print("cache rows", past_keys) scores, n = decode_one([0.9, 0.1]) print("decode scores vs cache", [round(s, 3) for s in scores], "cache now", n) print("last cache row", past_keys[-1]) print("we did not recompute the prompt keys") scores2, n2 = decode_one([0.0, 1.0]) print("second decode scores", [round(s, 3) for s in scores2], "cache now", n2)
Run to execute this in your browser. Nothing is sent to a server.