Optimizing Attention: GQA & MQA Explained

The Problem: KV-Cache Memory Consumption in Multi-Head Attention

Let's start with a quick refresher on standard Multi-Head Attention (MHA). In the original Transformer architecture, each attention "head" has its own set of Query (Q), Key (K), and Value (V) projection matrices. This means the number of Q heads is equal to the number of K/V heads.

Recall the formula for the KV-cache size:
$\text{Cache Size}=2\times L\times N_h\times T_s\times d_h\times P$

The key variable for today's lesson is $N_h$, the number of attention heads. In MHA, this term represents the number of Q heads, K heads, and V heads, as they are all equal. Because the cache only stores Keys and Values, its size is directly proportional to the number of K and V heads.

The core idea behind MQA and GQA is simple: what if we could reduce the number of K and V heads without reducing the number of Q heads? This would shrink the KV-cache without crippling the model's ability to "ask" diverse questions via its multiple query heads.

Let's watch a short video that provides a high-level overview of MHA and introduces its more efficient variants, MQA and GQA.

A Spectrum of Attention Mechanisms

As the video explained, MHA, GQA, and MQA exist on a spectrum of efficiency versus representational power. The key difference lies in how many Key and Value heads are shared across the Query heads.

Let's use a diagram to make this crystal clear.

Let's break down the two main optimizations based on this visual:

1. Multi-Query Attention (MQA): The Extreme Optimization

   • Architecture: All Query heads share a single Key head and a single Value head.
   • Impact on KV-Cache: This is the most aggressive memory-saving technique. If a model has 32 attention heads, MQA reduces the size of the K and V tensors in the cache by a factor of 32 compared to MHA. This drastically lowers memory bandwidth requirements during inference.
   • Trade-off: The cost is a potential degradation in model quality. Since all queries must retrieve information from the same, single representation of keys and values, the model's capacity can be limited. It's like having 32 experts (queries) who are all forced to consult the same single, summarized reference book (the single K/V head).

2. Grouped-Query Attention (GQA): The Balanced Approach

   • Architecture: This is the pragmatic middle ground. The Query heads are divided into a smaller number of groups. Within each group, the Query heads share a common Key and Value head.
   • Example: A model like Llama 3 8B has 32 query heads but only 8 KV heads. This means the query heads are organized into 8 groups, with 4 query heads per group sharing one K/V head.
   • Impact on KV-Cache: The cache size is reduced by a factor of (Number of Q heads) / (Number of KV heads). In the Llama 3 8B example, this is a 4x reduction (32 / 8).
   • Trade-off: GQA has been empirically shown to offer a fantastic balance. It achieves nearly the same performance as MHA while capturing most of the speed and memory benefits of MQA. This is why it has become the de-facto standard in many modern LLMs.

From Architecture to PyTorch: The Engineer's View

Given your background, it's valuable to see how this architectural change translates into code. The implementation difference is subtle but powerful.

The following reading walks through the motivation and provides a full PyTorch implementation of GQA, comparing it directly to MHA.

The key implementation takeaway is the repeat_interleave operation. This is where the "sharing" happens. Instead of computing num_query_heads unique K/V tensors, we compute fewer (num_kv_heads) and then efficiently copy them to allow all num_query_heads to perform their attention calculation. This small change in the forward pass leads to massive savings in memory and bandwidth during inference.

Conclusion

You have now learned about two critical optimizations to the standard attention mechanism that are essential for deploying modern LLMs efficiently. While they may seem like small architectural tweaks, their impact on inference performance—especially latency and memory usage—is enormous.

Key Takeaways:

• The Motivation: Standard Multi-Head Attention (MHA) creates a large KV-cache during inference, which becomes a memory and bandwidth bottleneck.
• Multi-Query Attention (MQA): An aggressive optimization where all query heads share a single K/V head. This provides maximum memory savings but can degrade model quality.
• Grouped-Query Attention (GQA): A balanced approach where groups of query heads share K/V heads. It offers a "best of both worlds" compromise, significantly reducing the KV-cache size with minimal impact on performance.
• The Mechanism: The core implementation detail involves projecting to fewer K/V dimensions and then using an operation like repeat_interleave to align the K/V tensors with the query tensors for the attention computation.
• Practical Relevance: GQA is a key feature in leading models like Llama 3 and Mistral, enabling them to handle long contexts efficiently.

What's Next?

We have now examined several key optimizations related to the attention sub-layer of the Transformer block (RoPE, KV-caching, GQA/MQA). In the next lesson, we will shift our focus to the other major sub-layer: the Position-wise Feed-Forward Network (FFN). We will explore the different activation functions used in modern LLMs, such as SwiGLU and GeLU, and understand how they contribute to model performance.