query_heads = 8
tokens = 6
head_size = 8

for name, kv_heads in [("MHA", 8), ("GQA", 2), ("MQA", 1)]:
    queries_per_kv_head = query_heads // kv_heads
    mapping = [head // queries_per_kv_head for head in range(query_heads)]
    cached_scalars = 2 * kv_heads * tokens * head_size

    print(name, "mapping:", mapping)
    print(name, "cached scalars:", cached_scalars)

# Output:
"""
MHA mapping: [0, 1, 2, 3, 4, 5, 6, 7]
MHA cached scalars: 768
GQA mapping: [0, 0, 0, 0, 1, 1, 1, 1]
GQA cached scalars: 192
MQA mapping: [0, 0, 0, 0, 0, 0, 0, 0]
MQA cached scalars: 96
"""