Multi-Query Attention Repeat KV

Problem

Expand grouped-query / multi-query key-value heads. Each KV head is a matrix of shape (seq_len, head_dim); repeat each head consecutively num_repeats times so the result lines up with the query heads.

The signature is repeat_kv(kv_heads, num_repeats). Return a list of head matrices.

Hint

Each KV head serves num_repeats consecutive query heads, so copies stay adjacent.

Test cases

case 1

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

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

case 2

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

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

case 3

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

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

case 4

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

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

Starter codePython

def repeat_kv(kv_heads, num_repeats):
    # Your code here
    pass

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