KV Cache Append Step

Problem

Append one new key vector and one new value vector to the key/value caches. Return [new_k_cache, new_v_cache] with the new vector at the end of each cache.

The signature is kv_cache_append(k_cache, v_cache, new_k, new_v). Each cache is a list of vectors.

Hint

Decoding appends one position per step instead of re-running the whole prefix.

Test cases

case 1

in: [[[1,0]],[[0,1]],[2,2],[3,3]]

expected: [[[1,0],[2,2]],[[0,1],[3,3]]]

case 2

in: [[],[],[1],[2]]

expected: [[[1]],[[2]]]

case 3

in: [[[1,2]],[[3,4]],[5,6],[7,8]]

expected: [[[1,2],[5,6]],[[3,4],[7,8]]]

case 4

in: [[[1]],[[1]],[0],[0]]

expected: [[[1],[0]],[[1],[0]]]

Starter codePython

def kv_cache_append(k_cache, v_cache, new_k, new_v):
    # Your code here
    pass

4 test cases · run your code against them in the editor