KV Cache Update

Task

Implement kv_cache_update(key_cache: list[list[float]], value_cache: list[list[float]], new_keys: list[list[float]], new_values: list[list[float]], start_position: int) -> tuple[list[list[float]], list[list[float]]]. Return updated key and value caches by writing a contiguous block of new rows beginning at start_position.

Requirements

• key_cache and value_cache contain the same number of cached positions; new_keys and new_values contain the same number of new positions.
• start_position is between zero and the current cache length, inclusive.
• Preserve every cached row before start_position.
• Replace existing rows covered by the new block.
• Append rows when the new block extends beyond the current cache length.
• Preserve existing suffix rows after the written range.
• Update key and value caches at exactly the same positions.
• Do not mutate any input; return the updated key cache followed by the updated value cache.
• An empty update returns independent copies of both caches.

Example

key_cache = [[1.0, 0.0], [0.0, 1.0]]
value_cache = [[10.0, 0.0], [0.0, 20.0]]

kv_cache_update(
    key_cache,
    value_cache,
    [[1.0, 1.0]],
    [[30.0, 40.0]],
    start_position=2,
)
# (
#   [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]],
#   [[10.0, 0.0], [0.0, 20.0], [30.0, 40.0]],
# )