How GQA and MQA Reduce Memory Bandwidth Bottlenecks in AI

This article explores how Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) alleviate memory bandwidth bottlenecks during the autoregressive decoding phase of large language models (LLMs). During AI text generation, models are typically slowed down by the speed at which data can be moved from memory to processors, rather than the speed of the processors themselves. By altering how attention mechanisms store and share state data, MQA and GQA significantly shrink the size of the Key-Value (KV) cache. This reduction minimizes data transfer overhead, directly resulting in faster inference speeds and higher throughput for modern AI applications.

The Memory Bandwidth Bottleneck in Autoregressive Decoding

Autoregressive decoding is the process by which an AI model generates text one token at a time. Each newly generated token depends on all the tokens that came before it. To avoid recalculating the mathematical representations of previous tokens at every step, models store this historical data in high-bandwidth memory (HBM). This stored data is known as the Key-Value (KV) cache.

During the generation of each new token, the GPU must read the entire KV cache from the memory and transfer it to the compute cores. Because the model only processes one token at a time during this phase, the actual mathematical computation takes microseconds. However, transferring gigabytes of KV cache data across the memory bus takes much longer. Consequently, the GPU compute cores spend most of their time sitting idle, waiting for data to arrive. This state is known as being "memory-bandwidth bound."

The Baseline Problem: Multi-Head Attention (MHA)

Standard Transformer models use Multi-Head Attention (MHA). In MHA, the model divides its attention mechanism into multiple "heads" so it can focus on different parts of the input sequence simultaneously.

Each Query head is paired with its own unique Key head and Value head. If a model has 64 Query heads, it also maintains 64 Key heads and 64 Value heads. For every single token generated, the KV cache must store and subsequently load the parameters for all 64 Key and Value heads. As sequences grow longer and batch sizes increase, the MHA KV cache balloons in size, completely saturating the memory bandwidth.

Multi-Query Attention (MQA): Maximum Compression

Multi-Query Attention (MQA) takes a drastic approach to shrinking the KV cache. Instead of giving every Query head its own Key and Value head, MQA forces all Query heads to share a single Key head and a single Value head.

If a model has 64 Query heads, MQA reduces the number of Key and Value heads from 64 down to 1. This architectural change yields immediate hardware benefits:

  • Drastic KV Cache Reduction: The size of the KV cache is reduced by a factor equal to the number of attention heads (e.g., a 64x reduction).
  • Unclogged Memory Bandwidth: For each generated token, the GPU only needs to fetch one set of Keys and Values from memory, cutting memory transfer time to a fraction of what MHA requires.
  • Higher Throughput: With freed-up memory capacity and bandwidth, systems can process much larger batch sizes simultaneously.

The primary drawback of MQA is that forcing all Queries to share a single Key-Value pair can lead to a slight degradation in the model's reasoning quality and output accuracy.

Grouped-Query Attention (GQA): The Optimal Balance

Grouped-Query Attention (GQA) was designed as a middle ground between the high quality of MHA and the high speed of MQA. Instead of all Query heads sharing just one Key-Value head, GQA divides the Query heads into a smaller number of groups. All the Query heads within a specific group share a single Key head and a single Value head.

For example, a model with 64 Query heads might be divided into 8 groups. Each group of 8 Query heads shares 1 Key-Value head, resulting in a total of 8 Key and Value heads.

GQA reduces the memory bandwidth bottleneck through the following mechanisms:

  • Proportional Bandwidth Savings: By reducing 64 KV heads to 8, the KV cache size and the required memory bandwidth per token are reduced by 8x compared to standard MHA.
  • Preserved Model Quality: By allowing multiple KV heads rather than just one, the model retains the capacity to capture complex, multi-faceted relationships in the text, achieving accuracy nearly identical to MHA.
  • Hardware Efficiency: The reduced memory transfers bring the ratio of computation-to-memory-access closer to an optimal level, keeping the GPU cores actively computing rather than waiting on data transfers.

Summary of Architectural Impact

Both MQA and GQA directly target the root cause of slow autoregressive decoding: the volume of data moved across the memory bus. By decoupling the number of Query heads from the number of Key and Value heads, these mechanisms slash the size of the KV cache. Fetching fewer Keys and Values per token directly translates to less strain on memory bandwidth, transforming a slow, memory-bound operation into a fast, efficient generation process.