MHA GQA MQA Parameter Comparison

Problem

Count attention projection parameters for three variants with head_dim = d_model / num_heads: MHA (Q, K, V and output each d_model x d_model), GQA (K and V use num_kv_heads * head_dim columns) and MQA (a single KV head).

The signature is attention_params_by_variant(d_model, num_heads, num_kv_heads). Return [mha_params, gqa_params, mqa_params] as integers, ignoring biases.

Hint

Q and the output projection are always d_model x d_model; only K and V shrink with the KV head count.

Test cases

case 1

in: [512,8,8]

expected: [1048576,1048576,589824]

case 2

in: [512,8,2]

expected: [1048576,655360,589824]

case 3

in: [64,4,1]

expected: [16384,10240,10240]

case 4

in: [1024,16,4]

expected: [4194304,2621440,2228224]

Starter codePython

def attention_params_by_variant(d_model, num_heads, num_kv_heads):
    # Your code here
    pass

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