Multi-Query AttentionMQA

Self-attention lies at the heart of large language model training and inference. As context windows expand and parameter

Self-attention lies at the heart of large language model training and inference. As context windows expand and parameter counts grow, standard Multi-Head Attention (MHA) has exposed a critical bottleneck: excessive GPU memory consumption and sluggish inference speed. To address these limitations, researchers have introduced several architectural variants. Among them, Multi-Query Attention (MQA), proposed by Noam Shazeer in the 2019 paper "Fast Transformer Decoding: One Write-Head is All You Need," stands out as a widely adopted optimization that streamlines the attention mechanism without sacrificing the core benefits of multi-head computation.
In conventional MHA, every attention head maintains its own independent set of Query, Key, and Value projection matrices. With h heads, the model stores h separate copies of Keys and Values. This redundancy becomes increasingly expensive during inference, particularly when the model must retain a large KV cache for long sequences.

MQA simplifies this design by allowing all attention heads to share a single set of Key and Value matrices, while each head retains its own independent Query matrix. The result is an asymmetry: h Query heads, but only one shared Key head and one shared Value head. This structural change dramatically reduces the memory footprint of the attention layer while preserving the ability of each head to attend to different parts of the input through distinct queries.

Solving the Memory Bandwidth Bottleneck

To appreciate why MQA matters, it helps to distinguish the two phases of LLM inference. The prefill stage is compute-bound: the GPU processes the input prompt in parallel, and floating-point throughput is the limiting factor. The decoding stage, by contrast, is memory-bound. Each newly generated token requires loading the historical Key and Value cache from high-bandwidth memory (HBM) into the faster on-chip SRAM before attention scores can be computed.

Under MHA, the KV cache scales linearly with the number of heads. If a model has 32 attention heads, it must maintain 32 separate Key and Value tensors for every layer and every sequence position. This massive cache quickly exhausts memory bandwidth, forcing the decoder to wait on data movement rather than computation. MQA collapses the Key and Value dimensions from h to 1, shrinking the KV cache by a factor of h. In practice, this means:

  • Reduced memory bandwidth pressure: Fewer tensors need to be fetched from HBM during each decoding step.
  • Larger batch sizes: With less memory consumed per sequence, more sequences can run concurrently.
  • Longer context windows: The same hardware can support extended sequences without running out of memory.
  • Higher decoding throughput: Autoregressive generation runs significantly faster because the memory wall is lowered.

MQA, MHA, and GQA: An Architectural Comparison

The relationship between MHA, MQA, and Grouped-Query Attention (GQA) can be summarized by how they distribute Query, Key, and Value heads:

  • Multi-Head Attention (MHA): h Query heads, h Key heads, h Value heads. Offers the strongest representational capacity but carries the largest KV cache and the slowest inference speed.
  • Multi-Query Attention (MQA): h Query heads, 1 Key head, 1 Value head. Minimizes cache size and maximizes speed, though extreme cases may show a slight drop in accuracy compared with full MHA.
  • Grouped-Query Attention (GQA): h Query heads, g Key and Value heads where 1 < g < h. Query heads are partitioned into groups, each sharing a KV head. GQA serves as a middle ground, retaining much of MHA's accuracy while delivering most of MQA's efficiency. Modern models such as Llama 3 and Mistral commonly use GQA to balance performance and throughput.

Implementation Sketch

The following simplified PyTorch code illustrates the essential difference between MHA and MQA. In MQA, the Key and Value projections output a single head dimension rather than expanding to match the number of Query heads:

import torch
import torch.nn as nn

class MultiQueryAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        # Each query head has its own projection
        self.q_proj = nn.Linear(d_model, d_model)
        
        # MQA core: one shared Key and one shared Value projection
        self.k_proj = nn.Linear(d_model, self.head_dim)
        self.v_proj = nn.Linear(d_model, self.head_dim)
        
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch_size, seq_len, _ = x.shape

        # Project and reshape queries into multiple heads
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
        q = q.transpose(1, 2)

        # Keys and values have no head dimension; broadcasting handles the rest
        k = self.k_proj(x).unsqueeze(1)
        v = self.v_proj(x).unsqueeze(1)

        # Scaled dot-product attention
        scores = torch.matmul(q, k.transpose(-1, -2)) / (self.head_dim ** 0.5)
        attn_weights = torch.softmax(scores, dim=-1)
        context = torch.matmul(attn_weights, v)

        # Merge heads and apply output projection
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.out_proj(context)

Because k and v carry a singleton head dimension, PyTorch automatically broadcasts them across all query heads during the matrix multiplications. This avoids storing redundant copies while keeping the computation graph straightforward.

Training Considerations and Practical Trade-offs

MQA delivers substantial inference gains, but training an MQA model from scratch requires care. Naïvely converting a pretrained MHA checkpoint to MQA—for example, by averaging or pruning Key and Value weights—often leads to a measurable rise in perplexity and a degradation in generation quality. For best results, MQA models should be pretrained from initialization rather than retrofitted after the fact.

Even with proper training, some practitioners find that reducing KV heads all the way to one is too aggressive. GQA offers a practical compromise: by selecting a small number of shared KV heads (such as 4 or 8), developers can recover much of MHA's accuracy while still enjoying a dramatically smaller cache than full multi-head attention. In production deployments, the choice between MQA and GQA usually depends on whether maximum throughput or maximum fidelity is the priority.

Conclusion

Multi-Query Attention is a pivotal engineering innovation for deploying large language models at scale. By sharing a single Key and Value matrix across all attention heads, MQA directly attacks the memory bandwidth bottleneck that constrains autoregressive decoding. The resulting reduction in KV cache size enables larger batches, longer contexts, and faster token generation. Although it may involve minor accuracy trade-offs in certain settings, MQA has established a foundation for efficient, high-throughput inference and directly inspired subsequent architectures such as Grouped-Query Attention.