GQA Repeat Padding

Problem

Compute the key/value repetition factor for grouped-query attention and the number of dummy query heads needed to make the query-head count a multiple of the KV-head count: groups = ceil(num_query_heads / num_kv_heads).

The signature is gqa_repeat_padding(num_query_heads, num_kv_heads). Return [groups, padded_heads - num_query_heads].

Hint

Ceiling division tells you how many query heads each KV head must serve.

Test cases

case 1

in: [8,2]

expected: [4,0]

case 2

in: [12,5]

expected: [3,3]

case 3

in: [1,1]

expected: [1,0]

case 4

in: [10,4]

expected: [3,2]

Starter codePython

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

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