Multi-Query Attention (MQA)

In standard Multi-Head Attention (MHA), each of the N query heads has its own dedicated key and value head, resulting in N distinct key and value projections. Multi-Query Attention (MQA) changes this by having all N query heads share a single key head and a single value head. This structural change is the source of its efficiency gains, particularly during the autoregressive decoding phase where the key-value cache must be stored in memory for all previous tokens in the sequence.

The most significant advantage of MQA is the dramatic reduction in the size of the key-value (KV) cache. For a model with h attention heads, d_k key dimension, and a sequence length L:

• MHA KV Cache Size: 2 * L * h * d_k
• MQA KV Cache Size: 2 * L * 1 * d_k This represents a reduction by a factor of h (e.g., 32x for a model with 32 heads). This is critical for long-context inference and high-throughput serving, where the KV cache can become the dominant memory consumer, often limiting batch sizes or sequence lengths.

The efficiency of MQA comes with a trade-off. By sharing a single key/value representation across all query heads, the model loses the representational capacity and specialization that individual key/value heads provide in MHA. This can lead to a measurable drop in model quality or perplexity on some tasks compared to a full MHA baseline. The shared representation acts as a bottleneck, forcing all queries to attend to a common, compressed set of features. This trade-off makes MQA most suitable for inference-constrained scenarios where the memory savings outweigh a modest performance cost.

Beyond memory savings, MQA accelerates inference. The smaller KV cache reduces memory bandwidth pressure when reading cached states for each new token generation. This leads to:

• Lower latency per generated token.
• Ability to support larger batch sizes within the same memory budget, increasing overall tokens/second throughput.
• More efficient use of GPU tensor cores as memory operations become less of a bottleneck. For large-scale deployment, these throughput gains directly translate to lower serving costs and improved user experience.

MQA exists on a spectrum of efficient attention variants:

• Multi-Head Attention (MHA): Full capacity, high memory cost.
• Multi-Query Attention (MQA): Minimal memory, shared KV heads.
• Grouped-Query Attention (GQA): A hybrid where N query heads are divided into G groups, each sharing a key/value head. GQA (e.g., with 8 groups) often recovers most of MHA's quality while retaining significant memory savings over MHA, making it a popular compromise. MQA is frequently used in models where extreme inference efficiency is paramount, such as those deployed on edge devices or for real-time applications.

MQA is implemented by modifying the projection layers in a transformer block. Instead of h separate linear projections for keys and values, only one is used. Its outputs are then broadcast to all attention heads. Key use cases include:

• On-device large language models where DRAM is severely limited.
• High-throughput inference servers for chatbots and code generation.
• Long-context processing (e.g., document summarization) where the KV cache for a 128K context would be prohibitive with MHA. Models like Falcon and certain configurations of LLaMA 2 have employed MQA to achieve practical deployment scales.

A hybrid attention mechanism that groups multiple query heads to share a single key head and a single value head. It is the direct architectural predecessor and successor to MQA.

• Key Innovation: Strikes a balance between the high quality of standard Multi-Head Attention (MHA) and the memory efficiency of Multi-Query Attention (MQA).
• Mechanism: For example, an 8-query-head model might use 4 groups (G=4), where each group of 2 query heads shares one key/value head.
• Trade-off: Reduces the size of the KV cache compared to MHA, but not as drastically as MQA. It often recovers most of the quality drop observed in pure MQA, making it a preferred choice in models like Llama 2 and Gemma.

The critical memory bottleneck that MQA and GQA are designed to optimize. During autoregressive generation (token-by-token), the Key and Value tensors for all previous tokens in the sequence are stored to avoid recomputation.

• Memory Cost: In standard MHA, the cache size scales with: batch_size * sequence_length * num_layers * num_heads * head_dim * 2 (for K and V).
• MQA/GQA Impact: By sharing key/value heads, MQA reduces this term from num_heads to 1, and GQA reduces it to num_groups. This is a direct, linear reduction in memory for long sequences and large batches.
• System Effect: A smaller KV cache enables longer context windows, larger batch sizes, and higher throughput on memory-constrained hardware.

An I/O-aware, exact attention algorithm that optimizes the computation of attention, which works in tandem with MQA/GQA to accelerate inference.

• Core Problem: Standard attention implementations make frequent reads/writes between fast GPU SRAM and slow High Bandwidth Memory (HBM), creating a bottleneck.
• Solution: FlashAttention recomputes attention scores on-the-fly within SRAM, avoiding the materialization of the large N x N attention matrix in HBM.
• Synergy with MQA: While MQA reduces the size of the KV cache, FlashAttention optimizes the speed of the attention operation itself. Using both techniques together yields the highest gains in end-to-end inference latency.

The standard attention mechanism used in the original Transformer architecture, which MQA modifies for efficiency.

• Mechanism: Each attention "head" maintains its own independent set of Query (Q), Key (K), and Value (V) projection matrices. This allows the model to jointly attend to information from different representation subspaces.
• Computational Cost: The forward pass computation for MHA and MQA is similar. The primary difference is in the autoregressive inference memory footprint due to the KV cache.
• Quality Benchmark: MHA is the quality baseline. MQA and GQA are evaluated based on how closely they can approximate MHA's performance while using significantly less memory.

A complementary efficiency technique that reduces the computational complexity of attention, whereas MQA reduces its memory footprint.

• Principle: Each token only attends to a fixed-size window of preceding tokens (e.g., 4096), rather than the entire sequence history.
• Complexity: Reduces attention complexity from O(n²) to O(n * window_size), enabling processing of extremely long sequences.
• Combined Use: Models like Mistral 7B and Llama 3 often employ both Sliding Window Attention (to handle long contexts computationally) and GQA (to manage the memory of the KV cache for that window).

A different architectural paradigm for efficiency that scales model capacity without a proportional increase in computation per token.

• Contrast with MQA: While MQA is an attention-layer optimization, MoE is a model-wide architecture. MoE uses a routing network to activate only a small subset of neural network "experts" for each input.
• Efficiency Goal: MQA targets memory-bound inference latency. MoE targets compute-bound training and inference costs, enabling models with hundreds of billions of parameters that are cheaper to run.
• System Design: Advanced models may incorporate both techniques: using MoE for massive parameter counts and MQA/GQA to ensure the active experts run with minimal memory overhead.