Multi-Query & Grouped-Query Attention

Eight query heads can attend to the same conversation and produce eight different outputs. Do they need eight different key/value histories? Multi-query attention says no: keep the queries distinct and share one K/V representation. Grouped-query attention shares within smaller groups. The change reduces stored history while preserving separate attention weights for each query head.

Standard Multi-Head Attention (MHA) gives each query head a corresponding K/V head.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} During full-attention decode, growing cached history consumes capacity and contributes memory traffic.^{[2]Reference 2Efficiently Scaling Transformer Inferencehttps://proceedings.mlsys.org/paper_files/paper/2023/hash/c4be71ab8d24cdfb45e3d06dbfca2780-Abstract-mlsys2023.html} Whether that traffic dominates the whole model depends on context, batch size, active weights, kernels, and hardware. Head sharing targets the K/V term, not every source of latency.

For equal head widths and precision, an 8:1 query/KV ratio gives eightfold smaller raw K/V payload.^{[3]Reference 3Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150[4]Reference 4GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} Real concurrency and speed also depend on other state, physical replication, and latency constraints. DeepSeek-V2's Multi-Head Latent Attention (MLA) instead caches a learned compact representation and a positional component.^{[5]Reference 5DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelhttps://arxiv.org/abs/2405.04434} These are trained architecture choices; selecting a smaller cache shape alone cannot convert an arbitrary checkpoint.

Why KV cache memory becomes the bottleneck

Understanding why MQA, GQA, and MLA exist starts with the physical hardware tension of autoregressive token generation. Inference runs in two distinct phases: prefill and decode.

Prefill offers multiple uncached token rows to process together, possibly in chunks. For one head, $QK^\top$ has shape $(S \times d_k)(d_k \times S)$. Larger row counts can improve reuse, but shape alone does not establish measured HBM intensity or saturated Tensor Cores. Attention tiling, activation traffic, chunk size, and kernel overhead matter.

Ordinary decode advances each request by one consumed token per step, while a batch can contain many requests. Full-attention queries use retained K/V from the prefix and the current token. Caching avoids repeatedly projecting old positions; it does not remove the growing cost of attending to their history.

The kernel processes K/V in tiles rather than requiring the whole cache to fit on chip. Head sharing can reduce distinct cached bytes that need reading, but actual HBM traffic depends on reuse between query groups and kernel implementation.

The autoregressive memory wall and operational intensity

Quantify this memory traffic with the roofline model. Consider a single attention layer decoding one token for a batch of $B$ sequences with context length $S$, $H_q$ query heads, $H_{kv}$ key-value heads, head dimension $d_k$, and numerical precision $P$ bytes per element (for example, $P = 2$ for FP16 or BF16).

For equal key/value widths, count the two main matrix products (omitting softmax, projections, and other operations):

1. Query-key dot products ($Q K^\top$): $2 \cdot B \cdot H_q \cdot S \cdot d_k$ FLOPs.
2. Softmax-value weighting ($A V$): $2 \cdot B \cdot H_q \cdot S \cdot d_k$ FLOPs.

Total attention decode compute: $\text{Attention FLOPs} = 4 B H_q S d_k$

Assume each distinct cached K/V value is read from HBM once and reused across its query group. The ideal cache-read payload is:

1. Key cache read: $B \cdot H_{kv} \cdot S \cdot d_k \cdot P$ bytes.
2. Value cache read: $B \cdot H_{kv} \cdot S \cdot d_k \cdot P$ bytes.

Total attention HBM read traffic: $\text{Attention Memory Traffic} = 2 B H_{kv} S d_k P \text{ bytes}$

Under those assumptions, the matrix-product FLOPs divided by ideal cache-read bytes are:

$\text{Operational Intensity} = \frac{4 B H_q S d_k \text{ FLOPs}}{2 B H_{kv} S d_k P \text{ bytes}} = \frac{2}{P} \times \frac{H_q}{H_{kv}} \quad [\text{FLOPs/byte}]$

Look at what happens in standard Multi-Head Attention (MHA), where $H_q = H_{kv}$. With 16-bit floats ($P = 2$ bytes per element):

$\text{Operational Intensity}_{\text{MHA}} = \frac{2}{2} \times 1 = \mathbf{1 \text{ FLOP/byte}}$

H100 SXM's rounded shipping specification implies about 989.5 TFLOPS of dense BF16 Tensor Core throughput and 3.35 TB/s HBM bandwidth; the advertised 1,979 TFLOPS includes structured sparsity.^{[6]Reference 6H100 GPUhttps://www.nvidia.com/en-us/data-center/h100/[7]Reference 7NVIDIA Hopper Architecture In-Depthhttps://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/} Its simple dense-BF16/HBM ridge is approximately:

$\text{Ridge Point} = \frac{989.5 \times 10^{12} \text{ FLOP/s}}{3.35 \times 10^{12} \text{ bytes/s}} \approx \mathbf{295.4 \text{ FLOPs/byte}}$

At a measured 1 FLOP/byte, this roofline bounds arithmetic rate at 3.35 TFLOPS, about 0.34% of dense peak. That percentage does not measure idle cycles. The ideal calculation also omits query/output traffic and repeated loads; tiny kernels, communication, or other limits can keep actual performance below either roof.

Count the stored numbers first

Before looking at full serving configurations, trace how many numbers one token adds to the cache in a miniature layer: 2 attention heads and 4 values per head.

Multiply by 8 requests and 2,048 cached positions: $16 \times 2048 \times 8 = 262,144$ numbers for one layer. Across 80 uniform layers, that is 20,971,520 numbers before counting bytes. This is a capacity calculation, not a measurement of which resource limits serving.

def kv_elements(layers: int, tokens: int, batch: int, kv_heads: int, head_dim: int) -> int:
    dimensions = (layers, tokens, batch, kv_heads, head_dim)
    if any(type(value) is not int or value < 0 for value in dimensions):
        raise ValueError("Dimensions must be nonnegative integers.")
    if layers == 0 or kv_heads == 0 or head_dim == 0:
        raise ValueError("Layers, KV heads, and head dimension must be positive.")
    return 2 * layers * tokens * batch * kv_heads * head_dim

tiny_one_layer = kv_elements(layers=1, tokens=2048, batch=8, kv_heads=2, head_dim=4)
scaled_layers = kv_elements(layers=80, tokens=2048, batch=8, kv_heads=2, head_dim=4)
print("tiny one-layer elements:", tiny_one_layer)
print("with 80 layers:", scaled_layers)

tiny one-layer elements: 262144
with 80 layers: 20971520

In standard MHA, each of the $h$ attention heads maintains its own Key ($K$) and Value ($V$) projections. The attention computation for a single layer uses softmax over scaled query-key scores:

$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

Here batch is omitted, softmax runs over key positions, and causal attention masks future positions before softmax. During full-prefix processing, $Q$, $K$, and $V$ have shape $h \times n \times d_k$ in this equal-width example. During one-token decode, $Q$ has shape $h \times 1 \times d_k$, while K/V span the cached prefix including the current token.

KV cache memory with standard MHA

For standard Multi-Head Attention (MHA) with $h$ heads of dimension $d_k$, attach a byte count to that picture. The leading factor of 2 accounts for separate Key and Value tensors. $h$ counts heads, $d_k$ counts values inside each head, and bytes per element comes from the storage format: 2 bytes for FP16 or BF16, 1 byte for INT8 or FP8.

Multiplying across all dimensions yields the complete serving estimate:

$\text{KV cache per token per layer} = 2 \times h \times d_k \times P \text{ bytes}$

$\text{Total KV cache} = 2 \times L \times S \times B \times h \times d_k \times P \text{ bytes}$

Here $L$ is layer count, $S$ is stored tokens per request, $B$ is concurrent batch size, and $P$ is bytes per precision format. This formula represents the logical dense payload: equal K/V widths, uniform full-attention layers, no shared prefixes, and no allocator overhead.

Throughout the tables, GB means $10^9$ bytes and TB means $10^{12}$ bytes. KiB means $2^{10}$ bytes and GiB means $2^{30}$ bytes. Context labels use 4K = 4,096, 32K = 32,768, 128K = 131,072, and 1M = 1,048,576 positions. Read actual configured head dimensions; they need not equal model width divided by query heads.

[!WARNING] $2 \times h \times d_k \times P$ is only the cache for one token in one layer. For a full serving estimate, you must multiply by layers, sequence length, and batch size.

Concrete example (72B-style decoder dimensions)

Put concrete numbers into the formula. Consider dimensions matching a large 70B- to 72B-parameter decoder: 80 layers, 64 query heads, and a head dimension of 128 (model width $d_{\text{model}} = 64 \times 128 = 8192$). If this model used standard MHA ($h_{kv} = 64$):

$\text{KV cache}_{\text{MHA}} = 2 \times 80 \times 4096 \times 32 \times 64 \times 128 \times 2 \text{ bytes} \approx \mathbf{343.6 \text{ GB}}$

For these hypothetical MHA dimensions, about 344 GB of raw K/V exceeds a 72-billion-parameter BF16 weight estimate of 144 GB. Parameter count alone does not specify head layout. Real checkpoints with this size can use GQA, so this is not their measured deployment footprint.

The same total of 131,072 token positions can arrive as 32 requests at 4K or one long-context document request at 128K.

That's the difference between an MHA cache that can't fit on an 80 GB GPU and a GQA-8 cache around 43 GB.

def kv_cache_gb(kv_heads: int, tokens: int = 4096, batch: int = 32) -> float:
    bytes_used = 2 * 80 * tokens * batch * kv_heads * 128 * 2
    return bytes_used / 1e9

for name, kv_heads in (("MHA", 64), ("GQA-8", 8), ("MQA", 1)):
    print(f"{name}: {kv_cache_gb(kv_heads):.2f} GB")

MHA: 343.60 GB
GQA-8: 42.95 GB
MQA: 5.37 GB

Multi-query attention (MQA)

One shared K/V pair

Multi-Query Attention (MQA), introduced by Noam Shazeer in 2019, takes the most aggressive architectural cut possible.^{[3]Reference 3Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150} All $H_q$ query projections remain intact, but a single Key projection and a single Value projection serve every query head in the layer.

Storage drops by the full query-head count:

$\text{Raw Payload Reduction Factor}_{\text{MQA}} = \frac{H_q}{H_{kv}} = \frac{H_q}{1} = H_q$

$\text{MQA KV cache per token per layer} = 2 \times 1 \times d_k \times P \text{ bytes}$

Where $1$ represents the single shared KV head. Cache size per token per layer scales with $d_k$ rather than with $H_q \cdot d_k$.

For the 72B-style dimensions with batch 32 and sequence length 4096, MHA requires 343.6 GB of KV cache. MQA slashes that by 64x down to ~5.37 GB.

The operational intensity of attention decode jumps from 1 FLOP/byte to:

$\text{Operational Intensity}_{\text{MQA}} = \frac{2}{P} \times \frac{H_q}{1} = \frac{2}{2} \times 64 = \mathbf{64 \text{ FLOPs/byte}}$

This raises the ideal matrix-product/cache-read ratio and reduces retained payload. The measured kernel need not attain that ratio or improve latency proportionally.

Performance vs. quality tradeoffs in MQA

That massive memory savings comes with a representational price. Every query head still computes its own query vector $q_i$, meaning heads can weight cached tokens differently through softmax. But all query heads project onto the identical Key subspace and retrieve from the identical Value subspace.

Shared K/V restricts the layer's parameterization while queries still choose different token weights. MQA can reduce quality, but the size of any loss depends on training and task; it is not evidence that a particular head necessarily owns one linguistic function.^{[3]Reference 3Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150[4]Reference 4GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} Evaluate trained checkpoints instead of inferring quality from cache size alone.

A minimal MQA decode step

This minimal implementation illustrates how query heads interact with shared K/V storage. Q has one vector per query head, while K and V are single shared sequences. Each query scores the same cached K and blends the shared V:

import math

def dot(a: list[float], b: list[float]) -> float:
    return sum(x * y for x, y in zip(a, b, strict=True))

def softmax(xs: list[float]) -> list[float]:
    peak = max(xs)
    exps = [math.exp(x - peak) for x in xs]
    total = sum(exps)
    return [e / total for e in exps]

def mqa_decode_step(
    query_heads: list[list[float]],
    keys: list[list[float]],
    values: list[list[float]],
) -> list[list[float]]:
    """Attend h query heads over one shared, entirely visible K/V prefix."""
    if not query_heads or not keys or not values or len(keys) != len(values):
        raise ValueError("Need queries and a nonempty aligned K/V prefix.")
    d_k, d_v = len(query_heads[0]), len(values[0])
    if not d_k or not d_v or any(len(row) != d_k for row in query_heads + keys):
        raise ValueError("Queries and keys must have a common positive width.")
    if any(len(row) != d_v for row in values):
        raise ValueError("Values must have a common positive width.")
    if any(not math.isfinite(x) for row in query_heads + keys + values for x in row):
        raise ValueError("Inputs must be finite.")
    scale = math.sqrt(len(query_heads[0]))
    outputs: list[list[float]] = []
    for query in query_heads:
        weights = softmax([dot(query, key) / scale for key in keys])
        outputs.append([
            sum(weight * value[dim] for weight, value in zip(weights, values, strict=True))
            for dim in range(len(values[0]))
        ])
    return outputs

# Two query heads, one shared cache of two tokens, d_k = 2.
q_heads = [[1.0, 0.0], [0.0, 1.0]]
shared_k = [[1.0, 0.0], [0.0, 1.0]]
shared_v = [[1.0, 0.0], [0.0, 1.0]]
out = mqa_decode_step(q_heads, shared_k, shared_v)
assert len(out) == 2
assert len(out[0]) == 2
print("head0 output:", [round(x, 3) for x in out[0]])
print("head1 output:", [round(x, 3) for x in out[1]])
print("shared KV tokens:", len(shared_k))

head0 output: [0.67, 0.33]
head1 output: [0.33, 0.67]
shared KV tokens: 2

Grouped-query attention (GQA)

Groups of queries, fewer K/V heads

Grouped-Query Attention (GQA), proposed by Ainslie et al. in 2023, is the balanced compromise between MHA and MQA.^{[4]Reference 4GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} Instead of one KV head for all queries (MQA) or one KV head per query (MHA), GQA divides the $H_q$ query heads into $g$ groups, where $1 < g < H_q$ and $H_q$ is divisible by $g$.

$\text{Group Size } G = \frac{H_q}{g} = \text{queries per KV group}$

Each group of $G$ queries shares a single Key head and a single Value head:

$\text{Raw Payload Reduction Factor}_{\text{GQA}} = \frac{H_q}{g}$

$\text{Ideal Tokens in a Fixed Raw-KV Budget Multiplier}_{\text{GQA}} = \frac{H_q}{g}$

$\text{Operational Intensity}_{\text{GQA}} = \frac{2}{P} \times \frac{H_q}{g} \quad [\text{FLOPs/byte}]$

With $H_q = 64$ and $g = 8$, each cached K or V tensor has shape $8 \times n \times d_k$ instead of $64 \times n \times d_k$. Raw payload falls eightfold. If the kernel reads each distinct cached value once, the ideal ratio above rises from 1 to 8 FLOPs/byte. Neither total model traffic nor end-to-end throughput must change eightfold.

def kv_group_for_query(query_head: int, query_heads: int, kv_heads: int) -> int:
    if any(type(value) is not int for value in (query_head, query_heads, kv_heads)):
        raise ValueError("Head counts and index must be integers.")
    if query_heads <= 0 or kv_heads <= 0 or query_heads % kv_heads:
        raise ValueError("Positive query-head count must be divisible by KV heads.")
    if not 0 <= query_head < query_heads:
        raise ValueError("Query-head index is out of range.")
    return query_head // (query_heads // kv_heads)

assignments = [kv_group_for_query(head, query_heads=8, kv_heads=2) for head in range(8)]
print("query to KV group:", assignments)
print("cache reduction:", 8 // 2, "x")

query to KV group: [0, 0, 0, 0, 1, 1, 1, 1]
cache reduction: 4 x

GQA in practice

These published model versions illustrate different grouping ratios. They are architectural examples, not a current model ranking or an exhaustive list of frontier designs. Context limits and layer layouts require separate checks:

Llama 3.1 fixes eight KV heads across 8B, 70B, and 405B, while the query count and sharing ratio grow.^{[9]Reference 9The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783} The original April 2024 Llama 3 8B/70B releases had 8K context, not 128K. GQA itself does not establish a context limit or make context extension reliable.

[!NOTE] Gemma 2 alternates 4K sliding-window and 8K global attention; sum each layer's actual retained positions when estimating cache.^{[12]Reference 12Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf} The original Mistral 7B paper describes sliding-window attention,^{[14]Reference 14Mistral 7B.https://arxiv.org/abs/2310.06825} while the v0.3 configuration above has sliding_window: null.^{[11]Reference 11Mistral-7B-v0.3 configurationhttps://huggingface.co/mistralai/Mistral-7B-v0.3/blob/main/config.json} A family name alone does not specify retention policy. Qwen2.5's model card and default configuration also distinguish supported extended context from configured context.^{[8]Reference 8Qwen2.5-72B-Instruct model card and configurationhttps://huggingface.co/Qwen/Qwen2.5-72B-Instruct}

Tensor Parallelism interactions

Head ratios meet an engineering constraint when sharding across GPUs. Tensor parallelism splits layer computation; head sharding is one attention layout, not a universal definition of how every weight matrix is partitioned.

In serving engines like vLLM, simple head sharding requires query heads to divide evenly across TP ranks ($H_q \pmod{\text{TP}} = 0$).^{[15]Reference 15vLLM Llama attention implementationhttps://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama.py}

For KV heads, two distinct regimes exist:

1. $\text{TP} \le g$: In this vLLM Llama layout, $g$ must divide evenly across TP ranks. Each rank stores $g / \text{TP}$ KV heads without KV-head replication.
2. $\text{TP} > g$: This implementation requires $\text{TP}$ divisible by $g$ and gives each rank one KV head. Each original KV head appears on $\text{TP} / g$ ranks. Other attention or distributed layouts can differ.

Consider the Llama 3.1 70B head layout ($H_q = 64, g = 8$) at $\text{TP} = 16$:

• Query heads split cleanly: $64 / 16 = 4$ query heads per GPU.
• But there are only 8 KV heads. To ensure each GPU rank has KV state, each KV head is duplicated across $16 / 8 = 2$ GPUs.
• Each rank stores one KV head, so the aggregate is sixteen physical head copies rather than eight. With all other dimensions equal, this doubles GQA payload relative to the unreplicated layout. Compared with MHA's 64 physical heads, the raw reduction factor is now 4x instead of 8x; this is not a claim that half of all saved bytes disappeared.

def vllm_sharding_plan(query_heads: int, kv_heads: int, tensor_parallel: int) -> str:
    if any(type(value) is not int or value <= 0 for value in (query_heads, kv_heads, tensor_parallel)):
        return "invalid: head counts and TP must be positive integers"
    if query_heads % kv_heads != 0:
        return "invalid: query heads must divide into KV groups"
    if query_heads % tensor_parallel != 0:
        return "invalid: query heads must divide evenly across TP shards"
    if tensor_parallel <= kv_heads:
        if kv_heads % tensor_parallel != 0:
            return "invalid: KV heads must divide evenly across TP shards"
        return f"even: {kv_heads // tensor_parallel} KV heads per shard"
    if tensor_parallel % kv_heads != 0:
        return "invalid: TP must be divisible by the KV-head count for replication"
    replicas_per_kv_head = tensor_parallel // kv_heads
    return f"replicated: each KV head appears on {replicas_per_kv_head} shards"

print("TP=4:", vllm_sharding_plan(query_heads=64, kv_heads=8, tensor_parallel=4))
print("TP=16:", vllm_sharding_plan(query_heads=64, kv_heads=8, tensor_parallel=16))
print("24Q, 6KV, TP=4:", vllm_sharding_plan(query_heads=24, kv_heads=6, tensor_parallel=4))
for tp in (4, 16):  # both valid for this 64Q/8KV layout
    local_kv_heads = max(1, 8 // tp)
    physical_kv_heads = tp * local_kv_heads
    print(f"TP={tp} aggregate KV head copies: {physical_kv_heads}; "
          f"matching MHA/raw GQA ratio: {64 / physical_kv_heads:.1f}x")

TP=4: even: 2 KV heads per shard
TP=16: replicated: each KV head appears on 2 shards
24Q, 6KV, TP=4: invalid: KV heads must divide evenly across TP shards
TP=4 aggregate KV head copies: 8; matching MHA/raw GQA ratio: 8.0x
TP=16 aggregate KV head copies: 16; matching MHA/raw GQA ratio: 4.0x

Converting MHA to GQA via uptraining

An existing MHA checkpoint provides learned parameters, but reducing its KV-head count changes the function. Uptraining reuses those parameters as an initialization and then adapts the model to sharing.

Ainslie et al. introduced uptraining: initialize grouped projections from an MHA checkpoint, then continue pretraining to adapt to sharing.^{[4]Reference 4GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245}

The mean-pooling initialization

The conversion begins by grouping existing Key and Value projection matrices and averaging them within each group.

Suppose an MHA model has $H_q = 64$ heads and converts to GQA with $g = 8$ heads. Each group contains $G = 64 / 8 = 8$ original heads. Let $W_{K, i}^{\text{MHA}} \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W_{V, i}^{\text{MHA}} \in \mathbb{R}^{d_{\text{model}} \times d_k}$ denote the projection slices for head $i$.

The initial weights for group $j \in \{0, \dots, g-1\}$ are computed as the arithmetic mean:

$W_{K, j}^{\text{GQA}} = \frac{1}{G} \sum_{i \in \text{group } j} W_{K, i}^{\text{MHA}}, \quad W_{V, j}^{\text{GQA}} = \frac{1}{G} \sum_{i \in \text{group } j} W_{V, i}^{\text{MHA}}$

The original experiments favored mean pooling over selecting one head or random initialization.^{[4]Reference 4GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} That does not prove averaging preserves variance, feature meaning, or outputs. Opposing projections can cancel. Use the existing weights as a candidate initialization and measure the converted function.

If the attention layer uses compatible bias vectors, pool them with the same grouping. Keep query and output projections at initialization; subsequent training can update them.

def mean_pool_heads(heads: list[list[float]], group_size: int) -> list[list[float]]:
    if type(group_size) is not int or group_size <= 0 or not heads or len(heads) % group_size:
        raise ValueError("Need a positive group size that divides nonempty heads.")
    if not heads[0] or any(len(head) != len(heads[0]) for head in heads):
        raise ValueError("Head slices must have the same nonzero width.")
    pooled: list[list[float]] = []
    for start in range(0, len(heads), group_size):
        group = heads[start : start + group_size]
        pooled.append([
            sum(values) / len(group)
            for values in zip(*group, strict=True)
        ])
    return pooled

mha_k_heads = [[1.0, 3.0], [3.0, 5.0], [10.0, 12.0], [14.0, 16.0]]
print("GQA initial K heads:", mean_pool_heads(mha_k_heads, group_size=2))

GQA initial K heads: [[2.0, 4.0], [12.0, 14.0]]

The uptraining recipe and compute budget

Mean pooling need not preserve attention outputs. Predict what happens when two heads' keys point in opposite directions and their values emphasize different positions:

import math

def attend(keys: list[float], values: list[float]) -> float:
    # One scalar query of 1; two entirely visible positions; head width 1.
    peak = max(keys)
    weights = [math.exp(key - peak) for key in keys]
    return sum(w * v for w, v in zip(weights, values, strict=True)) / sum(weights)

old_keys = [[1.0, -1.0], [-1.0, 1.0]]
old_values = [[10.0, 0.0], [0.0, 10.0]]
pooled_keys = [sum(pair) / 2 for pair in zip(*old_keys, strict=True)]
pooled_values = [sum(pair) / 2 for pair in zip(*old_values, strict=True)]
before = [attend(k, v) for k, v in zip(old_keys, old_values, strict=True)]
after = [attend(pooled_keys, pooled_values)] * 2
assert all(not math.isclose(a, b) for a, b in zip(before, after, strict=True))
print("original head outputs:", [round(x, 3) for x in before])
print("pooled head outputs:", [round(x, 3) for x in after])

original head outputs: [8.808, 8.808]
pooled head outputs: [5.0, 5.0]

This counterexample proves a functional change, not a universal severity of quality loss. A practical conversion must adapt and evaluate the model:

1. Initialize: Pool K/V projection slices and compatible biases; retain query/output projections at conversion time.
2. Adapt: The paper continued T5.1.1's original pretraining setup and dataset for 5% of original steps, using T5's Adafactor recipe. It did not establish a universal cosine schedule, reduced peak learning rate, or synthetic causal-LM recipe.
3. Evaluate: Its T5-XXL comparison reported average task score 47.1 for GQA-8 versus 47.2 for MHA across summarization, translation, and TriviaQA. It explicitly excluded GLUE. This is a historical experiment, not a guaranteed recovery percentage for every model.^{[4]Reference 4GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245}

Beyond GQA: Multi-head latent attention (MLA)

GQA reduces the number of cached KV heads. MLA, introduced in DeepSeek-V2 and retained in DeepSeek-V3, instead uses a learned low-rank representation.^{[5]Reference 5DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelhttps://arxiv.org/abs/2405.04434[16]Reference 16DeepSeek-V3 Technical Report.https://arxiv.org/abs/2412.19437}

Those models have 128 query heads. An absorbed decode implementation computes content scores against compressed state, avoiding expansion of historical content keys. This does not prove equality to an arbitrary MHA model or impose one cache/kernel implementation on every serving engine.

Low-rank joint KV compression and decoupled RoPE

In standard attention, a hidden state $h_t \in \mathbb{R}^d$ projects into separate key and value tensors of size $h \times d_h$.

MLA replaces independent KV projections with a single low-rank down-projection matrix $W^{DKV} \in \mathbb{R}^{d_c \times d}$, where $d_c \ll h \cdot d_h$:

$c_t^{KV} = W^{DKV} h_t \in \mathbb{R}^{d_c}$

In DeepSeek-V2/V3, $d_c = 512$, while $h \cdot d_h = 128 \times 128 = 16,384$. That single 512-dimensional vector $c_t^{KV}$ represents the compressed content of both Key and Value for token $t$.

Why the positional component is separate

In modern Transformers, Rotary Position Embedding (RoPE) applies position-dependent 2D rotation matrices $\mathcal{R}_t$ to key and query vectors based on token index $t$.

RoPE is linear for a fixed position, and matrix multiplication remains associative. The obstacle is position dependence and noncommutation with the learned up-projection. If we rotated the expanded content keys, one head's score would be:

$(\mathcal{R}_t q_t)^\top \mathcal{R}_s W_i^{UK} c_s^{KV} = \left((W_i^{UK})^\top \mathcal{R}_s^\top \mathcal{R}_t q_t\right)^\top c_s^{KV}$

Here $q_t \in \mathbb{R}^{d_h}$ and the rotations are $d_h \times d_h$. The absorbed query on the right depends on the cached position $s$, so it cannot be computed once for the current token and reused across all history. Moving an arbitrary head-space rotation into the differently sized latent space would generally change the function.^{[5]Reference 5DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelhttps://arxiv.org/abs/2405.04434}

This two-dimensional fixture makes that dependence visible. Predict whether the two cached positions can use the same absorbed query:

import torch

torch.set_num_threads(1)
up = torch.tensor([[1.0, 2.0], [0.0, 1.0]], dtype=torch.float64)
query = torch.tensor([1.0, 0.0], dtype=torch.float64)  # current query already rotated
latent = torch.tensor([1.0, 2.0], dtype=torch.float64)
rotations = [torch.eye(2, dtype=torch.float64),
             torch.tensor([[0.0, -1.0], [1.0, 0.0]], dtype=torch.float64)]
absorbed_queries = []
for position, rotation in enumerate(rotations):
    absorbed = up.T @ rotation.T @ query
    direct = torch.dot(query, rotation @ up @ latent)
    regrouped = torch.dot(absorbed, latent)
    torch.testing.assert_close(direct, regrouped)
    absorbed_queries.append(absorbed)
    print(f"position {position}: direct={direct.item():.1f}, regrouped={regrouped.item():.1f}")
assert not torch.equal(absorbed_queries[0], absorbed_queries[1])
print("one shared absorbed query works here:", False)

position 0: direct=5.0, regrouped=5.0
position 1: direct=-2.0, regrouped=-2.0
one shared absorbed query works here: False

DeepSeek solves this with decoupled RoPE. Keys and queries are partitioned into separate content components and positional components:

1. Content Key (compressed): $k_{t,i}^C = W_i^{UK} c_t^{KV} \in \mathbb{R}^{d_h}$ Where $W_i^{UK} \in \mathbb{R}^{d_h \times d_c}$ is the up-projection matrix for head $i$.
2. RoPE Key (uncompressed, shared): $k_t^R = \text{RoPE}(W^{KR} h_t) \in \mathbb{R}^{d_r}$ Where $W^{KR} \in \mathbb{R}^{d_r \times d}$ and $d_r = 64$.
3. Content Query and RoPE Query: Queries undergo a similar low-rank compression into $c_t^Q \in \mathbb{R}^{d_c'}$, then up-project into $q_{t,i}^C \in \mathbb{R}^{d_h}$ and $q_{t,i}^R = \text{RoPE}(W_i^{QR} c_t^Q) \in \mathbb{R}^{d_r}$.

The complete attention score between query token $t$ and cached token $s$ for head $i$ is the sum of content and positional dot products:

$\text{Score}_{t,s,i} = \frac{(q_{t,i}^C)^\top k_{s,i}^C + (q_{t,i}^R)^\top k_s^R}{\sqrt{d_h + d_r}}$

Matrix absorption for content attention

Here is the central mathematical insight of MLA. During autoregressive decoding, does the engine ever need to compute $k_{s,i}^C$ and cache it?

Look at the content dot product:

$(q_{t,i}^C)^\top k_{s,i}^C = (q_{t,i}^C)^\top \left( W_i^{UK} c_s^{KV} \right)$

By associativity of matrix multiplication:

$(q_{t,i}^C)^\top \left( W_i^{UK} c_s^{KV} \right) = \left( (W_i^{UK})^\top q_{t,i}^C \right)^\top c_s^{KV}$

Notice that $(W_i^{UK})^\top q_{t,i}^C$ depends strictly on the current query token $t$ and static model weights $W_i^{UK}$. It has zero dependence on past token positions $s$.

We can define the absorbed query for head $i$:

$\tilde{q}_{t,i} = (W_i^{UK})^\top q_{t,i}^C \in \mathbb{R}^{d_c}$

Now the content attention score against any cached token $s$ simplifies to:

$(q_{t,i}^C)^\top k_{s,i}^C = \tilde{q}_{t,i}^\top c_s^{KV}$

The dot product is computed directly between the 512-dimensional absorbed query $\tilde{q}_{t,i}$ and the 512-dimensional cached latent vector $c_s^{KV}$!

The identical mathematical trick applies to the Value vectors. In standard attention, the layer computes $\sum_s A_{t,s,i} v_{s,i}$. In MLA:

$o_{t,i} = \sum_s A_{t,s,i} \left( W_i^{UV} c_s^{KV} \right) = W_i^{UV} \left( \sum_s A_{t,s,i} c_s^{KV} \right)$

The attention weights $A_{t,s,i}$ directly multiply the cached latent vectors $c_s^{KV} \in \mathbb{R}^{d_c}$. The up-projection $W_i^{UV}$ is applied once to the accumulated 512-dimensional output, and can even be fused directly into the final attention output projection matrix $W^O$.

import torch

torch.manual_seed(42)
d_c, d_h = 512, 128
w_uk = torch.randn(d_h, d_c, dtype=torch.float64)
c_kv = torch.randn(d_c, dtype=torch.float64)  # cached latent vector in KV cache
q_c = torch.randn(d_h, dtype=torch.float64)    # content query for current token

# Path 1: Naive decompression (up-project key in cache, then dot product)
k_decompressed = w_uk @ c_kv
score_naive = torch.dot(q_c, k_decompressed)

# Path 2: Matrix absorption (absorb W_UK into query, dot product with cached latent)
q_absorbed = w_uk.T @ q_c
score_absorbed = torch.dot(q_absorbed, c_kv)

print("naive score:", round(float(score_naive), 4))
print("absorbed score:", round(float(score_absorbed), 4))
torch.testing.assert_close(score_naive, score_absorbed, rtol=1e-12, atol=1e-10)
print("numerically close:", True)

naive score: 398.0423
absorbed score: 398.0423
numerically close: True

The regrouping is exact algebra within this MLA parameterization; floating-point paths need numerical tolerance. An absorbed implementation can retain compressed content plus positional keys rather than full historical K/V. Prefill kernels or fallback implementations can choose to materialize expanded tensors, so inspect the engine's actual allocations.

MLA vs. GQA: exact payload comparison

What does an MLA serving engine write into HBM for each token position?

1. The compressed content latent vector $c_t^{KV}$ ($d_c = 512$ numbers).
2. The decoupled RoPE key $k_t^R$ ($d_r = 64$ numbers).

Total numbers stored per token per layer: $\text{MLA KV Cache Payload} = d_c + d_r = 512 + 64 = \mathbf{576 \text{ numbers}}$

Compare chosen raw cache layouts with 128 queries. The MHA/GQA/MQA rows use 128-wide keys and values. The MLA row uses 128-wide content components plus a separate 64-wide positional component. These counts do not imply identical score dimensions, compute, trained functions, or quality:

For these layouts, $32{,}768 / 576 \approx 56.9$ and $2{,}048 / 576 \approx 3.56$ are raw element-count ratios. All retain separate query outputs, but their parameterizations and representational constraints differ. Low-rank compression does not preserve every possible full-MHA function.

def gqa_cached_values_per_token(kv_heads: int, head_dim: int) -> int:
    return 2 * kv_heads * head_dim

gqa_8_values = gqa_cached_values_per_token(kv_heads=8, head_dim=128)
mla_example_values = 512 + 64  # compact content latent plus positional component
equivalent_gqa_heads = mla_example_values / (2 * 128)
print("GQA-8 cached values:", gqa_8_values)
print("MLA-style cached values:", mla_example_values)
print("MLA-style width in GQA-head units:", equivalent_gqa_heads)

GQA-8 cached values: 2048
MLA-style cached values: 576
MLA-style width in GQA-head units: 2.25

Serving impact

Memory capacity vs. memory bandwidth

When evaluating serving performance, separate two distinct constraints:

1. Memory capacity: Available VRAM limits how much state can be resident. Cache, weights, workspaces, and reservations must fit on each device. Admission and scheduling policies can impose a lower concurrency limit.
2. Memory bandwidth: The rate at which bytes move between HBM and compute units. When those transfers are the limiting resource, reducing traffic can shorten decode. Kernel execution, communication, and queueing also contribute to ITL and TPOT.

For equal K/V widths and precision, GQA reduces the raw cache payload by $H_q/H_{kv}$. MLA's ratio depends on its latent and positional dimensions instead. Smaller representations create room for more state and can reduce attention traffic; the resulting latency and throughput depend on how the engine uses them.

Concurrency ceilings (72B-style dimensions, 4K context)

Set aside a fixed budget of 10 GiB (10.74 GB) of VRAM dedicated purely to logical KV cache payload. With 4,096 tokens per request, 80 layers, 128 head dimension, and FP16 precision:

• MHA requires ~10.74 GB per request. That 10 GiB budget can admit exactly 1 concurrent request.
• GQA-8 requires ~1.34 GB per request. The identical budget admits 8 concurrent requests.
• MQA requires ~0.17 GB per request. The budget admits 64 concurrent requests.

Now suppose arrivals supply eight requests. MHA's raw budget permits only one resident request; GQA-8 permits all eight. That creates a batching opportunity. Measure whether the larger batch improves completed tokens per second while meeting latency targets: more admitted requests can also wait longer, and a smaller cache doesn't guarantee compute saturation.

Long-context impact and serving ceilings

Longer sequences amplify the architectural difference. At 128K context, per-request state scales linearly:

The hypothetical 128K MHA payload is 343.6 GB, exceeding four 80 GB devices even before weights. Five devices is only a capacity lower bound, not a valid sharding plan: the head counts and engine may disallow TP=5. GQA-8 reduces this payload to 42.9 GB, but feasibility and cost still require per-rank weights, state, workspaces, and a supported parallel layout.

def kv_cache_gb(kv_heads: int, tokens: int) -> float:
    return 2 * 80 * tokens * kv_heads * 128 * 2 / 1e9

mha_128k = kv_cache_gb(kv_heads=64, tokens=131_072)
gqa_128k = kv_cache_gb(kv_heads=8, tokens=131_072)
print(f"MHA 128K cache: {mha_128k:.1f} GB")
print(f"GQA-8 128K cache: {gqa_128k:.1f} GB")
print("GQA cache plus 144 GB weights fits in 4x80 GB raw:", gqa_128k + 144 <= 320)

MHA 128K cache: 343.6 GB
GQA-8 128K cache: 42.9 GB
GQA cache plus 144 GB weights fits in 4x80 GB raw: True

How serving engines handle GQA head expansion

When computing attention scores, the query tensor has $H_q$ heads while the cached Key and Value tensors have only $H_{kv}$ heads. How does a production inference engine match them?

A reference implementation can use repeat_interleave to expand K/V to the query-head count. If it materializes those tensors, the temporary working set grows to the MHA shape. That can erase important runtime savings, even though the persistent cache remains compact. The parity example below deliberately uses this simple reference path.

FlashAttention's flash_attn_with_kvcache supports fewer K/V heads than query heads, and FlashInfer provides grouped-attention kernels.^{[17]Reference 17FlashAttentionhttps://github.com/Dao-AILab/flash-attention[18]Reference 18FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving.https://arxiv.org/abs/2501.01005} Inspect the selected backend and temporary allocations. A head-count API or mathematically equivalent reference expansion doesn't prove that a particular dispatch avoids materialization.

def kv_index_for_query(query_head: int, query_heads: int, kv_heads: int) -> int:
    if any(type(x) is not int for x in (query_head, query_heads, kv_heads)):
        raise ValueError("Head counts and index must be integers.")
    if query_heads <= 0 or kv_heads <= 0 or query_heads % kv_heads or not 0 <= query_head < query_heads:
        raise ValueError("Invalid head counts or query index.")
    return query_head // (query_heads // kv_heads)

def gqa_scores(
    queries: list[list[float]],
    kv_keys: list[list[float]],
) -> list[float]:
    """Raw dot products at one cached position, not full attention."""
    if not queries or not kv_keys or not queries[0]:
        raise ValueError("Queries and KV keys must be nonempty.")
    if any(len(row) != len(queries[0]) for row in queries + kv_keys):
        raise ValueError("All query and key widths must match.")
    scores: list[float] = []
    for i, query in enumerate(queries):
        kv_key = kv_keys[kv_index_for_query(i, len(queries), len(kv_keys))]
        scores.append(sum(q * k for q, k in zip(query, kv_key, strict=True)))
    return scores

queries = [[float(i + 1), 0.0] for i in range(8)]
kv_keys = [[1.0, 0.0], [10.0, 0.0]]
scores = gqa_scores(queries, kv_keys)
assert scores == [1.0, 2.0, 3.0, 4.0, 50.0, 60.0, 70.0, 80.0]
print("query to KV index:", [kv_index_for_query(i, 8, 2) for i in range(8)])
print("per-head scores:", scores)

query to KV index: [0, 0, 0, 0, 1, 1, 1, 1]
per-head scores: [1.0, 2.0, 3.0, 4.0, 50.0, 60.0, 70.0, 80.0]

KV cache update correctness

An attention formula can be correct while a cache cursor, mask, or position error corrupts its inputs. Test token-by-token state updates as well as the grouped computation.

For ordinary full attention with RoPE on Q/K, trace one decode step. Here $t$ is the next logical sequence position, not necessarily a physical buffer offset:

1. Projection: Project the hidden state of the new token into $Q$ ($H_q$ heads), $K$ ($H_{kv}$ heads), and $V$ ($H_{kv}$ heads).
2. RoPE application: Transform Q/K using the model's position IDs and RoPE convention. In this example, use position $t$, leave V unrotated, and retain already rotated keys.
3. Cache write: Store the new K/V at the physical locations assigned to position $t$. A dense example might use (batch, n_kv_heads, max_seq, d_k); paged, sliding-window, and MLA caches have different layouts.
4. Attention computation: Attend the newly projected query heads against all cached key and value positions from $0$ to $t$.
5. Cursor increment: Advance the sequence length cursor for this request.

def append_kv(cache: list[str], cursor: int, token_value: str) -> int:
    if type(cursor) is not int or cursor != len(cache):
        raise ValueError("cursor would overwrite or skip cache state")
    cache.append(token_value)
    return cursor + 1

cache: list[str] = []
cursor = append_kv(cache, cursor=0, token_value="token-0-kv")
cursor = append_kv(cache, cursor=cursor, token_value="token-1-kv")
print("cache length:", len(cache))
try:
    append_kv(cache, cursor=1, token_value="stale-write")
except ValueError as exc:
    print("blocked:", exc)

cache length: 2
blocked: cursor would overwrite or skip cache state

Common cache update bugs and diagnostic checks

Predict the visible keys

A causal mask must know where the query sits in the stored history. PyTorch SDPA's is_causal=True uses upper-left alignment for a non-square query/key matrix.^{[19]Reference 19torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html} With one query and three cached keys, it exposes only the first key. That is wrong when the query is the newest token and all three keys are its valid prefix. FlashAttention uses a different, bottom-right convention for non-square causal attention.^{[17]Reference 17FlashAttentionhttps://github.com/Dao-AILab/flash-attention}

Here every score is zero, so attention averages the visible values. Predict whether each call returns 1 or 4 before running it.

import torch
import torch.nn.functional as F

q = torch.zeros(1, 1, 1, 2, dtype=torch.float64)
k = torch.zeros(1, 1, 3, 2, dtype=torch.float64)
v = torch.tensor([1.0, 2.0, 9.0], dtype=torch.float64).reshape(1, 1, 3, 1)

upper_left = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=True)
all_prefix = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
torch.testing.assert_close(upper_left, torch.ones_like(upper_left))
torch.testing.assert_close(all_prefix, torch.full_like(all_prefix, 4.0))
print("upper-left causal mask:", upper_left.item())
print("latest query sees complete prefix:", all_prefix.item())

upper-left causal mask: 1.0
latest query sees complete prefix: 4.0

Disabling the mask is valid for this one latest query with no padding or future keys. Chunked queries and ragged batches need masks derived from their positions and valid lengths.

Compare cached and recomputed outputs

Test both execution paths within the same model. Choose absolute and relative tolerances for the dtype, scale, and backend; there is no universal FP32 or FP16 threshold. Close logits can still select different tokens near an argmax tie. Converting MHA weights to GQA changes the model, so exact agreement with the old MHA model is a separate, generally unmet condition.

The following CPU float64 fixture tests a small GQA attention layer and output projection, not a complete pretrained transformer. It uses explicit query positions for cached masking and expanded heads only as a reference.

import math
import torch
import torch.nn.functional as F

torch.manual_seed(7)
torch.set_num_threads(1)

def grouped_attention(q, k, v, query_start=0):
    # q: [B,H,L,D], k: [B,G,S,D], v: [B,G,S,Dv]
    if any(t.ndim != 4 or not t.is_floating_point() for t in (q, k, v)):
        raise ValueError("Need floating rank-four tensors.")
    if any(t.dtype != q.dtype or t.device != q.device for t in (k, v)):
        raise ValueError("Inputs must share dtype and device.")
    b, h, length, d = q.shape
    groups, stored = k.shape[1:3]
    if min(*q.shape, *k.shape, *v.shape) <= 0:
        raise ValueError("All dimensions must be positive.")
    if k.shape[0] != b or k.shape[:3] != v.shape[:3] or k.shape[-1] != d or h % groups:
        raise ValueError("Batch, KV positions, widths, or head groups don't align.")
    if type(query_start) is not int or query_start < 0 or query_start + length > stored:
        raise ValueError("Queries must correspond to positions inside the stored prefix.")
    if any(not bool(torch.isfinite(t).all()) for t in (q, k, v)):
        raise ValueError("Inputs must be finite.")
    queries = q.reshape(b, groups, h // groups, length, d)
    scores = torch.einsum("bgrld,bgsd->bgrls", queries, k) / math.sqrt(d)
    query_positions = torch.arange(query_start, query_start + length, device=q.device)
    visible = torch.arange(stored, device=q.device)[None, :] <= query_positions[:, None]
    weights = scores.masked_fill(~visible, -torch.inf).softmax(dim=-1)
    output = torch.einsum("bgrls,bgsd->bgrld", weights, v)
    return output.reshape(b, h, length, v.shape[-1])

def rope(tensor, start):
    d = tensor.shape[-1]
    if d % 2:
        raise ValueError("Paired-coordinate RoPE requires even head width.")
    positions = torch.arange(start, start + tensor.shape[-2], dtype=tensor.dtype)
    frequencies = 10000.0 ** (-torch.arange(0, d, 2, dtype=tensor.dtype) / d)
    angles = positions[:, None] * frequencies[None, :]
    even, odd = tensor[..., 0::2], tensor[..., 1::2]
    return torch.stack((even * angles.cos() - odd * angles.sin(),
                        even * angles.sin() + odd * angles.cos()), dim=-1).flatten(-2)

B, H, G, T, D = 2, 8, 2, 6, 4
model_dim = H * D
x = torch.randn(B, T, model_dim, dtype=torch.float64)
wq = torch.randn(model_dim, H * D, dtype=torch.float64) / math.sqrt(model_dim)
wk = torch.randn(model_dim, G * D, dtype=torch.float64) / math.sqrt(model_dim)
wv = torch.randn(model_dim, G * D, dtype=torch.float64) / math.sqrt(model_dim)
wo = torch.randn(H * D, 11, dtype=torch.float64) / math.sqrt(H * D)

def project(tokens, start=0):
    b, n, _ = tokens.shape
    q = (tokens @ wq).reshape(b, n, H, D).transpose(1, 2)
    k = (tokens @ wk).reshape(b, n, G, D).transpose(1, 2)
    v = (tokens @ wv).reshape(b, n, G, D).transpose(1, 2)
    return rope(q, start), rope(k, start), v

def to_logits(heads):
    return heads.transpose(1, 2).reshape(B, heads.shape[2], H * D) @ wo

q, k, v = project(x)
expanded_k = k.repeat_interleave(H // G, dim=1)
expanded_v = v.repeat_interleave(H // G, dim=1)
reference = F.scaled_dot_product_attention(q, expanded_k, expanded_v,
                                          dropout_p=0.0, is_causal=True)
full = grouped_attention(q, k, v)
torch.testing.assert_close(full, reference, rtol=1e-10, atol=1e-10)

k_cache = torch.empty(B, G, 0, D, dtype=x.dtype)
v_cache = torch.empty_like(k_cache)
cached_outputs = []
for position in range(T):
    q_new, k_new, v_new = project(x[:, position:position + 1], start=position)
    k_cache = torch.cat((k_cache, k_new), dim=2)
    v_cache = torch.cat((v_cache, v_new), dim=2)
    cached = grouped_attention(q_new, k_cache, v_cache, query_start=position)
    q_prefix, k_prefix, v_prefix = project(x[:, :position + 1])
    recomputed = grouped_attention(q_prefix, k_prefix, v_prefix)[:, :, -1:]
    torch.testing.assert_close(to_logits(cached), to_logits(recomputed), rtol=1e-10, atol=1e-10)
    cached_outputs.append(cached)

torch.testing.assert_close(torch.cat(cached_outputs, dim=2), full, rtol=1e-10, atol=1e-10)

# Two new queries after an existing prefix: each has a different causal boundary.
for start in (1, 3):
    stop = start + 2
    q_chunk, _, _ = project(x[:, start:stop], start=start)
    chunk = grouped_attention(q_chunk, k[:, :, :stop], v[:, :, :stop], query_start=start)
    torch.testing.assert_close(chunk, full[:, :, start:stop], rtol=1e-10, atol=1e-10)
    unmasked = F.scaled_dot_product_attention(
        q_chunk, expanded_k[:, :, :stop], expanded_v[:, :, :stop], dropout_p=0.0)
    # The first query can wrongly read the second query's future K/V.
    assert not torch.allclose(unmasked[:, :, :1], chunk[:, :, :1], rtol=1e-10, atol=1e-10)

print("grouped matches repeated-head SDPA:", True)
print("cached logits match full-prefix recompute at every step:", True)
print("two-query chunks match full-prefix slices:", True)
print("unmasked chunks leak future values:", True)
print("stored K shape:", tuple(k_cache.shape))
print("logical K/V element reduction:", (expanded_k.numel() + expanded_v.numel()) //
      (k_cache.numel() + v_cache.numel()))

grouped matches repeated-head SDPA: True
cached logits match full-prefix recompute at every step: True
two-query chunks match full-prefix slices: True
unmasked chunks leak future values: True
stored K shape: (2, 2, 6, 4)
logical K/V element reduction: 4

In a two-query chunk starting at position 3, the first query may see keys 0 through 3, while the second may also see key 4. The example compares both rows with full-prefix execution and confirms that an unmasked call changes the first row. The single-query shortcut from the earlier fixture therefore doesn't generalize to a chunk. This test uses equal valid lengths across the batch; ragged requests still need their own length and position metadata.

Common pitfalls

Understanding head sharing math prevents common deployment mistakes that lead to production outages or underwhelming benchmark results.

Conflating model weights with KV cache memory

• Symptom: An engineering team calculates that an unquantized 70B model requires 140 GB of VRAM, deploys it across two 80 GB GPUs (160 GB total VRAM), and watches the server crash with Out-Of-Memory errors after serving just four concurrent long-context requests.
• Cause: Model weights are static. The KV cache is dynamic per-request memory that grows with concurrent batch size and sequence length ($B \times S$).
• Remedy: Budget weights, cache, temporary tensors, runtime reservations, and allocation overhead separately on each device. Measure peak memory at the intended sequence lengths and concurrency; a fixed reserve percentage can't establish fit.

Assuming GQA accelerates pretraining throughput

• Symptom: A team switches their training pipeline from MHA to GQA and wonders why pretraining steps take almost the identical wall-clock time.
• Cause: Fewer K/V projections reduce some parameters and projection FLOPs, but every query head still computes attention scores and weighted values. Training time also depends on backward kernels, activations, communication, and the rest of the model.
• Remedy: Measure training and serving separately. GQA can affect both; a large decode-cache saving doesn't predict a matching training speedup.

Believing an 8x KV cache cut guarantees an 8x throughput speedup

• Symptom: Your measured raw cache shrinks by 8x, while completed tokens per second changes little.
• Cause: Cache reads are one part of execution. Weight traffic, feed-forward layers, other kernels, communication, and scheduling can dominate. Their relative costs change with context length, batch size, precision, and hardware.
• Remedy: Measure service latency and throughput, then profile the limiting work. Try batching within a latency budget if arrivals and memory allow it. Report the measured gain and workload instead of inferring a speedup from bytes saved.

Expanding KV heads with repeat_interleave in serving kernels

• Symptom: A custom serving implementation of GQA runs out of memory during inference despite using an 8x smaller KV cache.
• Cause: The implementation called repeat_interleave in PyTorch to expand KV heads before running attention, temporarily materializing the full MHA tensor in VRAM.
• Remedy: Use native grouped-attention kernels like FlashAttention or FlashInfer that take separate query and KV head counts directly.

Overlooking KV head replication in tensor parallelism

• Symptom: Deploying an 8-KV-head model across 16 GPUs uses significantly more KV cache memory than anticipated.
• Cause: When TP degree exceeds the KV head count ($\text{TP} > g$), engines like vLLM replicate KV heads so each GPU owns at least one KV head.
• Remedy: Account for the engine's physical copies. In the illustrated vLLM Llama layout, TP=8 avoids KV-head replication; TP=16 may still be necessary for other memory or latency constraints. Compare per-rank fit, communication, and serving results rather than choosing TP from the KV count alone.

Validate a GQA conversion before rollout

Before promoting a newly converted or trained GQA model to production, run this systematic validation checklist:

1. Parity testing: Compare cached and full-prefix execution within the converted model, using justified absolute and relative tolerances. Include one-token and multi-token cached queries.
2. Batch consistency: Run test requests with ragged sequence lengths (e.g., prompt lengths 128, 512, and 2,048 in the same batch) to ensure per-request cursors don't overwrite adjacent cache buffers.
3. Tensor parallel verification: Test the production layout against a trusted reference layout that can fit the model. Compare outputs and measured per-rank allocations, including replicated heads.
4. Context configuration: Verify the checkpoint's supported context, positional configuration, and retention rules. Test long-prefix quality and retrieval at the intended lengths; a larger allocation or RoPE scaling flag doesn't establish model capability.
5. Quality and serving evaluation: Set acceptable quality changes against the old model before testing held-out tasks. Separately measure latency, throughput, and peak memory at the intended workload. Cache parity doesn't certify the conversion's quality.