Grouped Query Attention: Memory Efficiency Through K/V Head Sharing

1. The Attention Efficiency Spectrum: MHA, MQA, and GQA

The challenge of optimizing the attention mechanism can be seen as a trade-off between two competing goals:

• Model Expressiveness (Performance): The ability of the model to learn complex patterns. In attention, this is maximized when each query head can learn its own unique key and value projections, allowing it to "look" for different types of information in the sequence.
• Memory Efficiency: The amount of GPU memory required, dominated by the size of the KV cache during inference.

Let's place three key architectures on this spectrum.

1. Multi-Head Attention (MHA): This is the original architecture we've studied. Each query head has its own dedicated key (K) and value (V) head. With N query heads, there are N corresponding K/V heads. This provides maximum expressiveness but also incurs the maximum KV cache memory cost.

2. Multi-Query Attention (MQA): At the opposite extreme lies MQA. Here, all query heads share a single K/V head. This is maximally memory-efficient, as you only need to cache one set of K and V vectors. However, it can degrade model quality because all query heads are forced to use the same key and value representations, limiting the diversity of information they can extract.

3. Grouped-Query Attention (GQA): GQA provides a practical and effective compromise. It organizes the query heads into groups, and all heads within a single group share one K/V head. Instead of having N K/V heads (like MHA) or just 1 (like MQA), GQA has g K/V heads, where 1 < g < N.

This diagram provides a clear visual summary of the three approaches.

To solidify this high-level understanding, please review the following text, which provides a concise summary of this trade-off.

2. The GQA Mechanism Explained

Now, let's dive into the mechanics of GQA. The core idea is simple: if you have N query heads and g key-value heads, you create g groups. The size of each group is N / g. All query heads within one group will share the same key and value projections.

The following video explains this grouping strategy and its implications.

As the video explains, GQA maintains more diversity than MQA because different groups still have different K/V heads, allowing them to capture different perspectives. However, it achieves significant memory savings over MHA by not requiring a unique K/V head for every single query head.

From Concept to Code

To see how this is implemented in practice using PyTorch, the key operation is to project the input to a smaller K and V dimension and then "repeat" those K and V tensors to match the number of query heads.

Let's look at the implementation details in the article we viewed earlier.

This implementation makes the mechanism concrete:

1. Project the input x into Q, K, and V. The q_proj is to the full dimension (d_model), but k_proj and v_proj are to a smaller dimension (kv_dim).
2. Reshape K and V to have num_kv_heads.
3. Use K.repeat_interleave(self.num_queries_per_kv, dim=1) to duplicate the K heads for each query head in the group. The same is done for V.
4. The rest of the attention calculation proceeds identically to MHA.

3. Quantifying the Memory Reduction

Now we can update our KV cache formula from the previous lesson to account for GQA.

Recall the MHA formula:
$\text{Cache Size (MHA)}=B\times S\times L\times N\times H\times 2\times \text{precision\_bytes}$
Here, N is the number of query heads, and since it's MHA, it's also the number of K/V heads.

The GQA formula:
In GQA, we are no longer storing N key-value heads. We are storing g key-value heads (where g = num_kv_heads). The formula becomes:
$\text{Cache Size (GQA)}=B\times S\times L\times g\times H\times 2\times \text{precision\_bytes}$

The memory reduction is therefore a factor of N / g.

The video you just watched provides a concise explanation of this exact point.

Worked Example: Llama 3

Many modern models, like Meta's Llama 3, use GQA. The Llama 3 8B model uses:

• N (num_query_heads) = 32
• g (num_kv_heads) = 8

The memory reduction factor is N / g = 32 / 8 = 4. This means GQA reduces the KV cache size by 4x compared to a standard MHA implementation, while maintaining much of the model's performance. This significant saving is what allows models like Llama 3 to handle longer context lengths more efficiently.

Conclusion

In this lesson, we dissected Grouped-Query Attention, a crucial optimization for modern Transformers.

Key Takeaways:

• GQA is an efficient attention mechanism that provides a balance between the expressiveness of Multi-Head Attention (MHA) and the extreme memory efficiency of Multi-Query Attention (MQA).
• The core mechanism of GQA involves organizing query heads into groups, where each group shares a single set of Key (K) and Value (V) heads.
• This sharing mechanism reduces the size of the KV cache by a factor of N / g, where N is the number of query heads and g is the number of K/V heads (groups).
• GQA has been widely adopted in state-of-the-art LLMs (like Llama and Mistral) to enable faster inference and support for longer context windows.

Preview of the Next Lesson:
We have now seen how to make the KV cache more efficient by reducing the number of key and value heads we need to store. But what if we could reduce the size of the K and V vectors themselves? In our next lesson, we will explore Multi-head Latent Attention (MLA), the key innovation in the DeepSeek models. MLA takes a different approach by using low-rank compression to shrink the K and V vectors, pushing the boundaries of memory efficiency even further.