Skip to content
algo150Open the app
Build GPT · 10 of 10 · Python

Share keys across query heads

The KV cache holds keys and values for every head of every layer, and for a long text that is a great deal of memory. Grouped-query attention shrinks it. The model keeps all of its query heads but builds fewer key and value heads, and each key/value head is shared by a group of query heads. With 8 query heads and 2 key/value heads, the cache is a quarter of the size.

Do the exercise in the app

Learn it

Words to know

grouped-query attention
Attention in which several query heads share one key and value head, to save memory.
attention head
One attention calculation. A layer runs several side by side, each with its own queries, keys and values.
KV cache
The stored keys and values of earlier tokens, kept so they are not recomputed at every step of generation.

How it works

  1. The group size is q_heads // kv_heads: how many query heads share one key/value head.
  2. Query heads are grouped in order. The first group uses key/value head 0, the next group uses head 1, and so on.
  3. So query head i uses key/value head i // group_size.

Worked example

Input
q_heads = 8
kv_heads = 2

Output
[0, 0, 0, 0, 1, 1, 1, 1]
Query headsCalculationKey/value head
0, 1, 2, 30 // 4 up to 3 // 40
4, 5, 6, 74 // 4 up to 7 // 41
query01234567
key/value00001111
Eight query heads, in the top row, in two groups of four. Under each is the key/value head it reads. A group shares one, so two sets of keys and values are stored, not eight.

The group size is 8 // 2 = 4. Whole-number division drops the remainder, so the first four query heads all land on 0.

Why it matters

Llama 3 and Mistral 7B use grouped-query attention, with 32 query heads and 8 key/value heads. With equal head counts it is ordinary multi-head attention, and with one key/value head it is called multi-query attention.

Your turn

Work out which key/value head each query head uses in grouped-query attention.

q_heads
the number of query heads; positive
kv_heads
the number of key/value heads; positive, and it divides q_heads exactly
Returns
A list of length q_heads: for each query head in order, the index of its key/value head. Groups are equal in size and follow one another.

Example 1

Input
q_heads = 8
kv_heads = 2

Output
[0, 0, 0, 0, 1, 1, 1, 1]

Groups of 4: query heads 0 to 3 share head 0, and 4 to 7 share head 1.

Example 2

Input
q_heads = 4
kv_heads = 4

Output
[0, 1, 2, 3]

Equal counts give groups of 1. Every query head has its own key/value head, as in ordinary multi-head attention.

Example 3

Input
q_heads = 3
kv_heads = 1

Output
[0, 0, 0]

One key/value head is shared by all three query heads.

Hints

Each hint gives away a little more. Open one only when you need it.

  1. Hint 1: Where to start

    First work out how many query heads share one key/value head.

  2. Hint 2: The key step

    With a group size of 4, query heads 0 to 3 map to 0 and 4 to 7 map to 1. Whole-number division by the group size does that.

  3. Hint 3: The outline

    Set group_size = q_heads // kv_heads. Return [i // group_size for i in range(q_heads)].

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

  1. How many query heads share one key/value head. The task promises that kv_heads divides q_heads exactly, so every group is the same size.
  2. Whole-number division sends each group to its own head. // drops the remainder: indexes 0 to group_size - 1 give 0, the next group_size give 1, and so on.

The same numbered steps are the comments in the solution.

Python
def solve(q_heads, kv_heads):
    # 1. How many query heads share one key/value head
    group_size = q_heads // kv_heads
    # 2. Whole-number division sends each group to its own head
    return [i // group_size for i in range(q_heads)]

Watch out

  • i % kv_heads also gives valid head numbers, but it deals them out in turn, 0, 1, 0, 1, instead of in groups that follow one another.
  • Divide by the group size, not by kv_heads. With 8 and 2 those are 4 and 2, and the wrong one gives head numbers up to 3.

Check yourself

What does grouped-query attention save?

Memory in the KV cache. Fewer key and value heads are stored, while the number of query heads stays the same.

What are the two extreme cases of grouped-query attention?

Equal head counts give ordinary multi-head attention. One key/value head for all queries gives multi-query attention.

← Maintain cached keys and values