Multi-Query Attention (MQA)

The defining architectural change in MQA is that all attention heads share a single set of Key and Value projection weights. Unlike standard Multi-Head Attention where each head has distinct K and V linear layers, MQA collapses them into one. The Query projections remain independent per head, allowing the model to still attend to different representation subspaces while drastically reducing the number of parameters and the size of the KV cache.

The primary motivation for MQA is to slash the memory footprint of the Key-Value (KV) cache during inference. In a standard model with h heads, the cache stores h sets of keys and values per token. MQA stores only 1 set.

• Memory Savings: Reduces KV cache size by a factor of h (e.g., 96x reduction for a 96-head model).
• Bottleneck Relief: Shifts the inference bottleneck from memory bandwidth to compute, enabling much faster token generation.

MQA offers a significant speedup in autoregressive decoding by reducing the data that must be loaded from GPU memory at each step. However, this aggressive sharing can lead to a minor degradation in model quality compared to standard Multi-Head Attention.

• Throughput: Enables much higher batch sizes during inference.
• Quality: The capacity constraint from shared K and V can slightly reduce the model's ability to capture fine-grained syntactic or semantic nuances.

MQA represents one extreme of the attention spectrum, while Grouped-Query Attention (GQA) is a direct interpolation between MQA and standard Multi-Head Attention. GQA partitions Query heads into G groups, with each group sharing a single Key and Value head.

• MQA: G = 1 (all heads share one K/V).
• GQA: 1 < G < h (balances speed and quality).
• MHA: G = h (no sharing, maximum quality).

GQA was adopted by models like Llama 2 to recover the quality lost by MQA while retaining most of the speed benefits.

MQA was popularized by Google's PaLM model and the original Chinchilla paper, which demonstrated its effectiveness at scale. It is a critical technique for serving large language models in latency-sensitive applications.

• PaLM: Used MQA to achieve state-of-the-art few-shot results with efficient serving.
• Chinchilla: Showed MQA had minimal performance degradation compared to MHA when trained on an optimal compute budget.
• Legacy: While foundational, many newer open-source models have migrated to GQA for a better quality-speed Pareto frontier.

MQA requires custom attention kernels for maximum efficiency. Standard FlashAttention implementations must be adapted to handle the asymmetric shape where Query heads outnumber Key/Value heads.

• Training: MQA trains faster than MHA due to reduced FLOPs in the attention calculation.
• Kernel Design: The 'multi-query' pattern requires tiling strategies that broadcast the single K and V tensors across all Query heads to fully utilize GPU SRAM and avoid redundant HBM reads.

The standard attention mechanism where each head has its own Query, Key, and Value projection weights. This allows the model to jointly attend to information from different representation subspaces at different positions.

• KV Cache Size: Scales linearly with the number of heads (h × d_k per token)
• Memory Bandwidth: High, as the full per-head KV cache must be loaded during autoregressive decoding
• Quality: Generally highest expressivity, serving as the upper bound for attention fidelity

MHA is the baseline against which MQA and GQA trade-offs are measured.

An interpolation between MHA and MQA that partitions Query heads into groups, with each group sharing a single set of Key and Value heads.

• Groups: Typically 2-8 groups, balancing quality and speed
• KV Cache Reduction: Reduces KV cache size by a factor equal to the number of groups
• Adoption: Used in Llama 2 70B, Llama 3, and Mistral models

GQA recovers most of the quality lost by MQA while maintaining significant memory bandwidth savings, making it the practical sweet spot for large-scale deployment.

An inference optimization that stores previously computed Key and Value vectors for all generated tokens, preventing redundant re-computation during autoregressive decoding.

• Problem Solved: Without caching, each new token would require re-computing attention over the entire sequence
• Memory Footprint: Grows linearly with batch size, sequence length, and number of heads
• MQA Impact: MQA's single KV head reduces cache size by a factor of h compared to MHA

The KV cache is the primary bottleneck that MQA and GQA are designed to alleviate, making them critical for high-throughput serving.

An IO-aware exact-attention algorithm that minimizes high-bandwidth memory (HBM) reads and writes by fusing operations in a tiled manner within GPU SRAM.

• Tiling: Computes attention in blocks that fit in fast SRAM, avoiding materializing the full N×N attention matrix in HBM
• Recomputation: Recomputes attention statistics during the backward pass rather than storing them
• Compatibility: Works with MHA, MQA, and GQA, providing orthogonal speedups

FlashAttention addresses the memory bandwidth bottleneck from a different angle than MQA, and the two techniques are complementary in production inference stacks.

The inherent computational limitation of standard self-attention where both memory and time complexity scale as O(n²) with respect to input sequence length n.

• Cause: The attention score matrix requires computing pairwise interactions between all tokens
• Impact: Makes processing long documents (>8K tokens) prohibitively expensive
• Solutions: MQA addresses the memory bandwidth dimension of this problem, while sparse attention and linear attention variants tackle the compute dimension

MQA does not reduce the quadratic compute complexity but significantly reduces the memory bandwidth required to load the KV cache during decoding.

A model compression technique that identifies and removes redundant or low-importance attention heads from a trained Transformer.

• Motivation: Many attention heads learn overlapping or trivial patterns
• Methods: Magnitude-based pruning, Taylor expansion scoring, or learned masking
• Relationship to MQA: Pruning permanently removes heads, while MQA structurally shares KV projections without discarding Query diversity

Head pruning can be applied on top of MQA models to further reduce parameters, though MQA's already-shared KV heads leave fewer redundant heads to prune compared to MHA.