Grouped-Query Head Assignment

Problem

Assign each query head to a KV head for grouped-query attention, spreading num_query_heads as evenly as possible across num_kv_heads groups: the first num_query_heads % num_kv_heads groups receive one extra head.

The signature is gqa_head_assignment(num_query_heads, num_kv_heads). Return a list of length num_query_heads holding the KV head index for each query head.

Hint

Assign contiguous blocks of query heads; the remainder heads go to the lowest-numbered KV groups.

Test cases

case 1

in: [8,2]

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

case 2

in: [12,5]

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

case 3

in: [4,4]

expected: [0,1,2,3]

case 4

in: [1,1]

expected: [0]

Starter codePython

def gqa_head_assignment(num_query_heads, num_kv_heads):
    # Your code here
    pass

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