KV Cache State Managermedium

KV Cache State Manager

medium Inside an LLM 20–25 min

NumPy allowed · plain Python earns 10 bonus points

Logits become a next-token distribution

Softmax turns a model score vector into probabilities.

Open lesson →

01prompt

02tokens

03logits

04probabilities

attention→next token

mechanism65%

layer20%

pattern10%

model5%

Softmax(logits) → probability distribution

Maintain autoregressive key and value states during incremental decoding with optional sliding window eviction.

Task

Implement update_kv_cache(cached_keys, cached_values, new_key, new_value, window_size=None). Verify that len(new_key) == len(new_value) and that their dimensions match existing cached entries (if any). Append new_key to cached_keys and new_value to cached_values. If window_size is specified, retain only the most recent window_size vectors, evicting older entries from the front. If window_size <= 0, raise ValueError. Return a dict containing keys, values, seq_len (int), and evicted_count (int).

Function Signature

update_kv_cache(cached_keys: list[list[float]], cached_values: list[list[float]], new_key: list[float], new_value: list[float], window_size: int | None = None) -> dict

Examples

Example 1: Append without eviction
Input: cached_keys = [[0.1, 0.2]]
cached_values = [[0.5, 0.6]]
new_key = [0.3, 0.4]
new_value = [0.7, 0.8]
window_size = None
Output: {'keys': [[0.1, 0.2], [0.3, 0.4]], 'values': [[0.5, 0.6], [0.7, 0.8]], 'seq_len': 2, 'evicted_count': 0}

Example 2: Sliding window eviction
Input: cached_keys = [[1], [2]]
cached_values = [[10], [20]]
new_key = [3]
new_value = [30]
window_size = 2
Output: {'keys': [[2], [3]], 'values': [[20], [30]], 'seq_len': 2, 'evicted_count': 1}

Example 3: Dimension mismatch error
Input: cached_keys = [[1, 2]]
cached_values = [[1, 2]]
new_key = [1]
new_value = [1]
Output: Raises ValueError

Constraints

Vector dimension mismatches must raise ValueError.
window_size <= 0 must raise ValueError.
Must support empty initial cache (cold start).
A plain-Python solution earns 10 bonus points.

Python3Saved ✓

def update_kv_cache(cached_keys: list[list[float]], cached_values: list[list[float]], new_key: list[float], new_value: list[float], window_size: int | None = None) -> dict:

"""Append new key and value vectors to the inference KV cache.

Args:

cached_keys: Existing sequence of key vectors.

cached_values: Existing sequence of value vectors.

new_key: New key vector to append.

new_value: New value vector to append.

window_size: Optional maximum sliding window length.

Returns:

dict with 'keys', 'values', 'seq_len', and 'evicted_count'.

"""

# Your implementation here

pass

Case 1: Append without eviction

Input: cached_keys = [[0.1, 0.2]]
cached_values = [[0.5, 0.6]]
new_key = [0.3, 0.4]
new_value = [0.7, 0.8]
window_size = None

Expected: {'keys': [[0.1, 0.2], [0.3, 0.4]], 'values': [[0.5, 0.6], [0.7, 0.8]], 'seq_len': 2, 'evicted_count': 0}

Case 2: Sliding window eviction

Input: cached_keys = [[1], [2]]
cached_values = [[10], [20]]
new_key = [3]
new_value = [30]
window_size = 2

Expected: {'keys': [[2], [3]], 'values': [[20], [30]], 'seq_len': 2, 'evicted_count': 1}

Case 3: Dimension mismatch error

Input: cached_keys = [[1, 2]]
cached_values = [[1, 2]]
new_key = [1]
new_value = [1]

Expected: Raises ValueError