Multi-Query Attention: How Single Key-Value Head Sharing Slashed Transformer Serving Bottlenecks

Multi-query attention (MQA) is an architectural modification to the Transformer attention mechanism designed to resolve the memory bandwidth bottleneck during autoregressive token generation. First proposed by Noam Shazeer in the 2019 paper Fast Transformer Decoding: One Write-Head is All You Need, MQA alters the ratio of query, key, and value heads by sharing a single key head and a single value head across all query heads in each Transformer layer.

While standard multi-head attention (MHA) projects separate key and value tensors for every attention head, MQA reduces key-value (KV) tensor sizes by a factor proportional to the number of query heads. This structural reduction shrinks the runtime KV cache memory footprint, improves arithmetic intensity during decoding, and increases serving throughput across large language models.

The Memory Bandwidth Bottleneck in Transformer Serving

Transformer inference operates in two distinct execution phases: the prefill phase and the decoding phase.

Prefill Phase: When processing an input prompt, the model ingests all prompt tokens simultaneously. The attention computation evaluates matrix multiplications across sequence length $S$ , enabling compute-bound tensor core utilization via General Matrix Multiply (GEMM) operations.
Decoding Phase: During autoregressive generation, the model generates output tokens sequentially, one token at a time. Generating token $t$ requires computing attention between the new token's query vector and the key-value representations of all prior $t-1$ tokens.

To avoid recomputing keys and values for past tokens at every step, inference engines cache previous key and value activations in GPU High Bandwidth Memory (HBM). At each decoding step, the entire accumulated KV cache must be transferred from HBM into on-chip SRAM to compute attention.

+-------------------------------------------------------------------------+
|                  Autoregressive Decoding Memory Cycle                   |
|                                                                         |
|   +-------------------+      Read KV Cache       +------------------+   |
|   |  GPU Global HBM   | -----------------------> |   On-Chip SRAM   |   |
|   |  (Past Keys/Vals) |                          | (Compute Cores)  |   |
|   +-------------------+ <----------------------- +------------------+   |
|                              Append New KV                              |
+-------------------------------------------------------------------------+

For a single generated token, the number of arithmetic operations required for the attention layer is $O(S \cdot d_k)$ , where $S$ is sequence length and $d_k$ is head dimension. However, the data transferred across the memory bus scales as $O(S \cdot H \cdot d_k)$ , where $H$ is the number of attention heads.

Because arithmetic intensity (the ratio of floating-point operations to memory access bytes) falls below the hardware ridge point on modern accelerators such as the NVIDIA H100 GPU, decoding becomes memory-bandwidth-bound. Compute units remain underutilized while waiting for KV cache tensors to stream across memory channels.

Mathematical Formulation: MHA vs. MQA

In standard Multi-Head Attention (Vaswani et al., 2017), an input hidden state $X \in \mathbb{R}^{B \times S \times d_{\text{model}}}$ is linearly projected into $H$ query, key, and value heads:

$Q_i = X W_Q^i, \quad K_i = X W_K^i, \quad V_i = X W_V^i \quad (i \in \{1, \dots, H\})$

where projection weight matrices are defined as:

$W_Q^i \in \mathbb{R}^{d_{\text{model}} \times d_k}$
$W_K^i \in \mathbb{R}^{d_{\text{model}} \times d_k}$
$W_V^i \in \mathbb{R}^{d_{\text{model}} \times d_v}$

For head $i$ , the scaled dot-product attention computes:

$\text{head}_i = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i$

$\text{MHA}(X) = [\text{head}_1, \text{head}_2, \dots, \text{head}_H] W_O$

where $W_O \in \mathbb{R}^{(H \cdot d_v) \times d_{\text{model}}}$ .

Multi-Query Attention modifies this setup by maintaining $H$ distinct query heads while collapsing the key and value projections into a single shared head ( $H_{KV} = 1$ ):

$Q_i = X W_Q^i \quad (i \in \{1, \dots, H\})$ $K = X W_K$ $V = X W_V$

The corresponding parameter dimensions are:

$W_Q \in \mathbb{R}^{d_{\text{model}} \times (H \cdot d_k)}$ (split into $H$ heads of dimension $d_k$ )
$W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ (single projection)
$W_V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ (single projection)
$W_O \in \mathbb{R}^{(H \cdot d_v) \times d_{\text{model}}}$

Each query head $Q_i$ interacts with the identical shared key tensor $K$ and aggregates from the shared value tensor $V$ :

$\text{head}_i = \text{softmax}\left(\frac{Q_i K^T}{\sqrt{d_k}}\right) V$

$\text{MQA}(X) = [\text{head}_1, \text{head}_2, \dots, \text{head}_H] W_O$

During matrix operations in deep learning frameworks, $K$ and $V$ are broadcast along the head dimension to match the $H$ query heads during attention score computation.

KV Cache Memory Scaling and Bandwidth Arithmetic

The total KV cache memory consumption for a Transformer model during inference is governed by sequence length, batch size, number of layers, head count, head dimension, and numerical precision.

For standard Multi-Head Attention, the memory footprint is:

$\text{Memory}_{\text{MHA}} = 2 \times B \times S \times L \times H \times d_k \times P$

where:

$B$ = Batch size
$S$ = Context sequence length (tokens)
$L$ = Number of Transformer layers
$H$ = Number of attention heads
$d_k$ = Dimension per attention head
$P$ = Bytes per element (e.g., 2 for FP16/BF16, 1 for FP8)
The factor of 2 accounts for storing both keys and values.

For Multi-Query Attention, the memory footprint formula replaces $H$ with $1$ :

$\text{Memory}_{\text{MQA}} = 2 \times B \times S \times L \times 1 \times d_k \times P$

Numerical Comparison

Consider a representative 7-billion parameter language model architecture with $L = 32$ layers, $H = 32$ heads, $d_k = 128$ , and 16-bit precision ( $P = 2$ bytes):

MHA KV Cache per Token: $2 \times 32 \times 32 \times 128 \times 2 = 524,288 \text{ bytes} = 512 \text{ KB}$
MQA KV Cache per Token: $2 \times 32 \times 1 \times 128 \times 2 = 16,384 \text{ bytes} = 16 \text{ KB}$

| Batch Size ( $B$ ) | Context Length ( $S$ ) | MHA KV Cache Size | MQA KV Cache Size | Memory Reduction | | :--- | :--- | :--- | :--- | :--- | | 1 | 2,048 | 1.05 GB | 32.8 MB | 32x | | 1 | 8,192 | 4.19 GB | 131.1 MB | 32x | | 16 | 8,192 | 67.11 GB | 2.10 GB | 32x | | 64 | 8,192 | 268.44 GB | 8.39 GB | 32x | | 64 | 32,768 | 1,073.74 GB | 33.55 GB | 32x |

At a batch size of 64 and context length of 8,192 tokens, MHA demands 268 GB of memory purely for the KV cache, exceeding the total HBM capacity of three 80 GB GPUs. Under MQA, the entire KV cache requires 8.39 GB, allowing the full batch to execute on a single accelerator.

Representation Capacity and Trade-Offs

While MQA reduces memory bandwidth requirements, sharing key-value projections introduces architectural trade-offs in representational expressivity.

+-------------------------------------------------------------------------+
|                  Subspace Representation Comparison                     |
|                                                                         |
|   Multi-Head Attention (MHA):                                           |
|   Head 1: Q1 x K1^T ---> Scores 1 x V1 (Distinct Subspace 1)            |
|   Head 2: Q2 x K2^T ---> Scores 2 x V2 (Distinct Subspace 2)            |
|   Head H: QH x KH^T ---> Scores H x VH (Distinct Subspace H)            |
|                                                                         |
|   Multi-Query Attention (MQA):                                          |
|   Head 1: Q1 x K^T  ---> Scores 1 x V  (Identical Value Subspace)       |
|   Head 2: Q2 x K^T  ---> Scores 2 x V  (Identical Value Subspace)       |
|   Head H: QH x K^T  ---> Scores H x V  (Identical Value Subspace)       |
+-------------------------------------------------------------------------+

Shared Key Geometry: Because all query heads attend to the same key matrix $K$ , the metric space for similarity scoring is constrained. Distinct query vectors $Q_i$ can still attend to different positions along the sequence, but they evaluate dot products against a uniform key representation.
Identical Value Subspaces: In MHA, each head computes a linear combination of head-specific value vectors $V_i$ . In MQA, every head computes a linear combination of the exact same value vectors $V$ . The post-attention projections must rely entirely on the output projection matrix $W_O$ to differentiate head representations before residual addition.
Task-Specific Degradation: As shown in evaluations by Ainslie et al. (2023), MQA achieves competitive perplexity on broad web corpora but exhibits measurable performance drops in tasks requiring dense multi-hop associative retrieval, code parsing, and fine-grained entity tracking.

Landmark Deployments in Frontier Architectures

Despite capacity trade-offs, several major foundation models adopted MQA to maximize decoding speed and context scalability:

PaLM (Pathways Language Model): Google's 540-billion parameter PaLM (Chowdhery et al., 2022) used Multi-Query Attention across TPU v4 clusters, establishing MQA as a viable production standard for ultra-large models.
PaLM 2: Google's PaLM 2 (Anil et al., 2023) maintained MQA across all four model sizes (Gecko, Otter, Bison, and Cliff) to maximize throughput.
Falcon Series: The Technology Innovation Institute deployed MQA in Falcon-7B and Falcon-40B (Almazrouei et al., 2023), demonstrating high serving throughput in open-weight models.
StarCoder: BigCode's StarCoder (Li et al., 2023) utilized MQA to maintain an 8,192-token context window for code generation with minimal GPU memory overhead.
ChatGLM2-6B: Zhipu AI's ChatGLM2-6B transitioned from MHA to MQA, yielding a reported 4.2x increase in inference speed and support for 32K context windows.

The Evolution: From MQA to GQA and MLA

The development of Multi-Query Attention triggered subsequent innovations in KV cache compression:

+-------------------------------------------------------------------------+
|                  Attention Evolution and KV Compression                 |
|                                                                         |
|   MHA (2017)           MQA (2019)           GQA (2023)       MLA (2024) |
|   H Query Heads        H Query Heads        H Query Heads    Low-Rank   |
|   H KV Heads           1 KV Head            G KV Groups      Latent KV  |
|   [Max Capacity]       [Max Compression]    [Balanced]       [Decoupled]|
+-------------------------------------------------------------------------+

Grouped-Query Attention (GQA): Introduced by Ainslie et al. (2023), GQA partitions $H$ query heads into $G$ groups, assigning one key-value head per group ( $1 < G < H$ ). By selecting intermediate values (such as $G = 8$ for $H = 32$ or $H = 64$ ), models such as Llama 3 and Mistral achieve the quality of MHA while retaining the majority of MQA's memory savings.
Multi-Head Latent Attention (MLA): Introduced in DeepSeek-V2 and DeepSeek-V3 (Liu et al., 2024), MLA compresses keys and values into a low-rank latent vector $c_t^{KV}$ , decoupling positional embeddings via RoPE while reconstructing multi-head key and value matrices on the fly without expanding the stored KV cache.

Multi-Query Attention established the foundational principle that decoupling query head count from key-value head count is essential for high-throughput, low-latency LLM serving.