Eight investigators, one archive

Imagine eight investigators reading the same case file. One looks for dates, another for amounts, a third for identities, and the others for contradictions, obligations, causal order, exceptions and missing evidence. Their questions differ. Now decide how the archive is organised.

In the first room, every investigator has a private catalogue and a private set of extracts. This is expensive, but each catalogue can classify the file differently. In the second room, pairs of investigators share a catalogue and a set of extracts. In the third, all eight share one. The investigators still ask eight questions, yet the archive may offer eight, four or one learned view of what can be found and what content is returned.

That is the useful intuition behind multi-head attention, grouped-query attention and multi-query attention. The choice is often described as a speed trick. It is more precise to call it a constraint on how many distinct memory projections the attention layer may maintain. Fewer projections shrink the key-value cache and reduce memory traffic during generation. They also remove degrees of freedom that may, or may not, matter to the task.

The central answer is therefore simple: keep query diversity separate from key-value diversity. A system can ask through many query heads while storing the context through fewer key and value heads. GQA is the adjustable middle, not an unrelated attention mechanism. Its engineering value comes from exposing a continuous design choice that ordinary “head count” language hides.

Part I

One attention block, two head counts

Scaled dot-product attention begins with three projections. A query represents what the current position is seeking. A key represents how a past position can be matched. A value represents what that past position contributes if selected. The original Transformer projected all three several times and ran the resulting attention heads in parallel. The stated purpose was to let the model attend through different representation subspaces and at different positions (Vaswani et al.).

The common shorthand says that a layer has H heads. That phrase collapses two quantities which modern decoder architectures often separate:

For query head i, let g(i) identify its key-value group. One attention output is then:

This equation exposes the exact constraint. Queries within a group remain different because each has its own query projection. They can therefore produce different attention weights over the same key matrix. What they cannot do is ask the context to be indexed through different key projections or to return content through different value projections. GQA shares the coordinate system of memory, not the question itself.

Keys choose; values carry

Vary one component at a time. Give two query heads the same values but different keys. The content available to both heads is identical, yet they can rank past positions differently because their matching spaces differ. Now give them the same keys but different values. Their attention maps can be identical, while the representation returned from each selected position differs. Keys govern selection; values govern transfer. Standard GQA shares both, so a quality change can arise from either lost routing freedom or lost payload freedom.

This thought experiment prevents a common diagnostic error. Inspecting only attention weights tests the key side of the mechanism. Two models can display similar maps while passing different value content into the residual stream. Conversely, different maps can still yield similar post-attention outputs when the selected values are redundant. A useful comparison therefore records attention logits, value-weighted head outputs and the combined output after the output projection. No single view is a complete causal account.

The parameter shapes reinforce the distinction. Let D be model width. Query projections contain approximately D H_qd_h parameters, while the key and value projections together contain approximately 2D H_kvd_h. The output projection remains tied to the concatenated query-head width. Reducing H_kv therefore shrinks K and V parameters and per-token K and V activations, but leaves the Q and output pathways largely unchanged. During generation, the decisive difference is that every earlier token leaves a K and V record at every layer. Projection weights are fixed model state; cached activations accumulate with the live workload.

Sharing keys and values as a pair is the convention compared in this article, not a mathematical necessity. One could share keys more aggressively than values, or compress one side differently. Such designs ask a broader architecture question. MHA, GQA and MQA keep the comparison controlled by moving one paired KV-head count while preserving the query-head count.

Many heads are not eight little experts

The investigator analogy is useful only if its failure is named. Human investigators arrive with stable roles and explicit goals. Attention heads are learned linear projections inside a coupled network. Some heads acquire interpretable patterns, but others are redundant, task-dependent or important only through interaction with later layers. Voita and colleagues found specialised heads in neural machine translation, yet also pruned most encoder heads with a small BLEU reduction. Michel and colleagues similarly showed that many trained heads could be removed, while a few layers remained sharply dependent on multiple heads (Voita et al.; Michel et al.).

Those results do not prove that one head is universally enough. They show why head count is not a direct count of independent cognitive functions. A head can be redundant for one evaluation distribution and necessary after shift. Two heads can encode the same function in rotated coordinates. A head that looks linguistically clear may be causally dispensable, while a diffuse head may matter through its output projection. The relevant quantity for GQA is functional incompatibility inside a proposed group, not the story attached to each attention map.

A decisive negative control

Suppose an MHA layer has four query heads. Its first two heads use identical key and value projection matrices, and its next two use another identical pair. Replacing those four stored pairs with two shared pairs changes no attention score and no returned value. The layer is exactly a two-group GQA layer written wastefully.

This control matters because it blocks a simplistic explanation. If GQA matches MHA, the result may not mean that aggressive sharing is always harmless. It may mean that the trained MHA had already learned duplicate or near-duplicate KV subspaces. Conversely, a quality drop after grouping indicates that some removed distinction mattered, but it does not reveal whether keys, values, grouping choice, conversion method or insufficient retraining caused the loss.

Part II

Why decoding makes shared KV valuable

A decoder processes a prompt in two different regimes. During prefill, many prompt tokens are projected and attended in parallel. During decode, the model generates one new token at a time. The new token produces a query, key and value at every layer. Its key and value are appended to the cache. Its query attends over all cached positions.

The cache prevents the model from recomputing keys and values for every earlier token. It replaces repeated arithmetic with stored state. The price is that each decoding step must read a growing collection of K and V vectors. Shazeer introduced MQA specifically because this incremental regime can become limited by memory bandwidth rather than arithmetic throughput (Shazeer). GQA preserves the same causal lever while allowing more than one KV view (Ainslie et al.).

The cache equation

Under the usual assumption that keys and values both use head width d_h, the logical cache payload is:

The equation explains the attraction of sharing. At fixed layer count, batch, context, head width and precision, changing from MHA to GQA reduces logical cache in the ratio H_kv/H_q. MQA reduces it to 1/H_q. The same factor applies to the bytes that a decode kernel must fetch for cached keys and values, subject to tiling, locality, sharding and reuse.

It does not follow that all attention arithmetic falls by the same factor. Every query head still forms scores against the sequence and mixes values. A well-designed kernel stores a KV head once and reuses it across its query group, but the number of query-specific score vectors remains H_q. K and V projection work shrinks, and cache traffic shrinks sharply. Query projection, output projection, feed-forward work and much of the attention arithmetic remain. GQA is primarily a memory-capacity and memory-movement intervention, not a blanket division of Transformer FLOPs.

Worked example: a serving batch at 8K context

Consider a synthetic decoder with 32 layers, 32 query heads, head width 128 and BF16 cache elements. Eight requests each carry 8,192 cached tokens. The logical cache is 32 GiB with MHA, 8 GiB with eight KV heads, and 1 GiB with MQA.

The saved 24 GiB between MHA and GQA can become a larger batch, a longer context, fewer devices or room for other runtime state. Which benefit appears depends on the serving policy. A scheduler that cannot admit more work gains no throughput merely because memory is free. A kernel that materialises repeated KV tensors can surrender part of the theoretical advantage. A tensor-parallel layout can also replicate a single MQA head across partitions; Ainslie and colleagues note that several GQA groups can align more naturally with model partitions than one shared head.

Capacity, bandwidth and latency are different claims

The 32, 8 and 1 GiB figures establish a logical-capacity result. They do not state total device memory, because weights, activations, temporary workspaces, allocator pages, quantisation metadata and replicated state also consume space. Nor do they state observed memory traffic. A kernel may reuse a KV tile across several query heads, read some data from a nearer cache, or perform extra layout transformations. The architecture determines how much distinct KV state exists; the runtime determines how that state is placed and moved.

Latency is a third claim. At one decode step, each active query must read information from an increasingly long history. Fewer distinct KV heads reduce the history width available for that read. The reduction improves inter-token latency only to the extent that moving this history lies on the critical path. If the device is waiting on feed-forward weights, cross-device communication, host scheduling or another request, a smaller cache can leave token time almost unchanged. Capacity can still improve even when latency does not, because more sequences may fit concurrently.

For the synthetic fixture, GQA-8 uses one quarter of MHA's logical KV payload. A fourfold batch increase is therefore the absolute KV-only capacity ceiling, not a throughput forecast. Model weights and other state reduce the attainable increase, while queue shape and scheduling determine whether admitted requests become useful throughput. A cache ratio is a mechanism measurement; an endpoint speed-up is a system measurement. Both belong in the release evidence, and neither should be inferred from the other.

Thought experiment: the same model in two queues

Place the same model behind two endpoints. Queue A summarises 20,000-token documents into 150 tokens, one request at a time. Queue B answers chats with 2,000-token prompts and 2,000 generated tokens, maintaining a large continuous batch. Keep model weights and hardware fixed.

Queue A spends much of its time in prefill. Its long prompt creates heavy attention and projection work, but only a short decode. Queue B repeatedly streams cached state for many sequences while weights are amortised across the batch. Reducing H_kv is likely to matter more to Queue B. This remains an empirical prediction, not a guarantee, because kernels, quantisation, memory hierarchy and scheduling change the balance.

This thought experiment supplies a serious systems negative control. If the measured endpoint is prefill-bound, network-bound, queue-bound or dominated by feed-forward weight movement, GQA may reduce reserved memory without materially reducing end-to-end latency. A failure to speed up does not falsify the cache equation. It shows that the removed traffic was not the active bottleneck.

Orthogonal mechanisms should not be credited to GQA

FlashAttention reduces reads and writes associated with attention computation by tiling exact attention around the GPU memory hierarchy (Dao et al.). PagedAttention manages the dynamic allocation and sharing of KV blocks to reduce fragmentation and duplication in serving (Kwon et al.). GQA reduces the width of the KV state created by the model architecture. These mechanisms can be composed, but they solve different problems.

A useful audit asks three separate questions: how many KV elements logically exist, how efficiently the runtime allocates them, and how many bytes the kernel transfers. A page manager cannot recover capacity removed by a poor architecture choice. An architecture with a small logical cache can still run badly through an unsuitable kernel. Model geometry, memory management and kernel IO belong to different evidence lines.

Part III

flowchart LR
  A["Eight investigators, one archive"] --> B["One attention block, two head<br/>counts"]
  B --> C["Why decoding makes shared KV<br/>valuable"]
  C --> D["The quality cost is a grouping<br/>problem"]
  D -. "re-test" .-> A
  classDef first fill:#0a1424,stroke:#087A98,stroke-width:2px,color:#e6f3ff
  classDef step fill:#0a1424,stroke:#3f7da0,stroke-width:1px,color:#e6f3ff
  class A first
  class B,C,D step

The quality cost is a grouping problem

Keys and values do not play interchangeable roles. Keys determine which past positions receive weight. Values determine what representation flows forward after those weights are chosen. Sharing keys asks several query heads to rank positions in one learned key space. Sharing values asks them to receive content through one learned value space. GQA normally shares both.

This removes capacity in a structured way. It does not force all query heads in a group to attend identically because their queries differ. It does force their attention logits to be computed against the same keys and their selected content to come from the same values. The quality risk rises when a proposed group contains heads whose useful K or V projections cannot be represented well by one shared pair.

What a useful grouping objective would preserve

The easiest conversion objective is weight proximity: place K and V matrices that look similar into the same group, then average them. It is cheap and can be useful, but weights are only an indirect description of behaviour. The same matrix difference can be irrelevant on the activations the model actually visits, while a small difference can be decisive near a sharp attention boundary. Rotated or permuted internal coordinates further weaken raw element-by-element comparisons.

A stronger objective is functional. On representative inputs x, compare the original and grouped attention logits, the value-weighted head outputs, and the layer output after the output projection. In compact form, a conversion can seek groups that minimise E_x[||A_MHA(x) − A_GQA(x)||²], where A denotes the chosen internal or layer-level behaviour and the expectation is taken over a calibration distribution. The symbol does not define one universal loss. It makes the evidence choice explicit: the conversion should preserve the behaviour that the release test actually values.

The calibration distribution must contain the distinctions likely to be compressed. Random text can make every grouping look similar because it rarely activates specialised retrieval. A code model needs repeated identifiers, scope boundaries and distant definitions. A multilingual model needs scripts, morphology and word-order patterns across languages. A document model needs distractors, tables and long-range references. Grouping is an empirical clustering problem under a task distribution, not a tensor-indexing convenience.

Layer uniformity is another hidden assumption. A fixed H_kv at every layer simplifies kernels and sharding, yet pruning studies suggest that redundancy is not evenly distributed. Some layers may tolerate one KV view while others need several. A layer-specific design can spend capacity more selectively, but it increases architecture and serving complexity. The appropriate baseline is therefore not automatically a sophisticated irregular design. It is the simplest group schedule whose measured quality benefit exceeds its implementation cost.

Minimal worked example: group compatible functions, not neighbours

Take four query heads. Head 1 and Head 2 both need to find entity mentions, but ask for different entities. Head 3 needs temporal markers. Head 4 needs negation scope. Two KV heads are available.

Grouping Heads 1 and 2 is plausible: their queries can differ while sharing an entity-oriented key space and value representation. Grouping Head 2 with Head 3 merely because they are adjacent in the weight tensor is riskier. One shared key projection must now make entity identity and temporal order simultaneously easy to retrieve. The model may learn such a compromise during pretraining, but a post-hoc average can erase a distinction the original layer used.

Conversion makes the hidden assumption visible

A model trained as GQA can adapt its queries, keys, values and later layers to the sharing constraint from the start. Converting an existing MHA checkpoint is different. Ainslie and colleagues initialise each GQA group by mean-pooling the original K and V projection heads, then continue pretraining. They reported that five per cent of the original pretraining steps recovered a favourable trade-off, with GQA-8 close to MHA quality and near MQA speed in their T5 XXL study.

Mean pooling assumes that heads inside a group live in sufficiently compatible coordinates. Yet neural representations admit reparameterisations. Two heads can implement related functions with rotated value bases, making raw weight averages destructive even when a transformed alignment would merge well. A later peer-reviewed study by Jin and colleagues uses Procrustes transformations to increase within-group KV similarity while preserving the model function before pruning, then applies supervised fine-tuning (Jin et al.). Quality-aware and asymmetric grouping proposals make the same broader point through different methods: which heads share can matter as much as how many share.

The synthetic stress test

The companion lab starts with four MHA KV heads arranged as two exact duplicate pairs. Two-head GQA is therefore identical. It then perturbs each original K and V projection independently, mean-pools the heads back into two groups, and measures the relative change in the attention output. As within-group dispersion rises, the compressed function departs from the original.

The curve is deliberately not a model benchmark. It isolates one causal feature: disagreement among heads being merged. Real quality depends on whether output differences occur on task-relevant inputs, whether later layers compensate, and whether retraining adapts the whole network. The lab nevertheless supplies a useful conversion diagnostic. A layer whose KV activations remain close under aligned representations is a stronger sharing candidate than a layer with several incompatible clusters.

What published evidence establishes, and what it does not

Ainslie and colleagues provide the cleanest direct comparison. In their reported T5 XXL setting, MHA required 1.51 seconds per sample and scored 47.2 on their average of task metrics. Uptrained MQA required 0.24 seconds and scored 46.6. GQA with eight groups required 0.28 seconds and scored 47.1. This supports the existence of a useful middle point for that architecture, training recipe, hardware and task collection. It does not establish that eight groups are generally optimal.

Shazeer’s earlier experiments likewise found large incremental-decoding gains with only small average quality changes, and sometimes no loss on a selected beam-search metric. Head-pruning studies show that substantial redundancy can exist. On the other side, the GQA paper reports that MQA can degrade quality, especially before sufficient uptraining, and later work continues to improve grouping and conversion. The responsible synthesis is therefore conditional: sharing often removes redundant capacity, but the safe compression ratio is learned by the model and revealed by task evidence.

Part IV

Choose the KV count through evidence

The wrong decision rule is “GQA is modern, therefore use eight KV heads”. Eight is a configuration, not a law. Llama 3 reports eight KV heads to improve inference speed and reduce cache size, which demonstrates a serious implementation choice rather than an optimum for other models (The Llama 3 Herd of Models). The architecture decision should begin with the workload and end with a release test.

Recent evidence makes the workload dependence more explicit. Chen and colleagues jointly varied model size, total attention-head count and KV-head count for long-context modelling. In experiments up to 1.2 billion parameters, they reported a configuration for 128K context with more than 50 per cent lower memory use and FLOPs than the Llama 3 GQA configuration at matched loss (Chen et al.). This result is bounded to their models and search procedure. It does not prescribe one group count for larger checkpoints. It does show that context length belongs inside the optimisation objective, and that a familiar configuration can cease to be economical when the workload changes.

1. Describe the workload before the model

Record prompt-length and output-length distributions, active batch, concurrency, beam or sampling strategy, prefix reuse, latency objectives and memory limits. Report percentiles, not one representative request. H_kv matters most when cached tokens and active sequences make KV capacity or bandwidth consequential. A single short completion can hide a problem that appears only under sustained batching.

2. Locate the actual bottleneck

Profile prefill and decode separately. Measure time to first token, inter-token latency, throughput, cache occupancy, memory-bandwidth utilisation, compute utilisation, queueing and communication. Repeat at several context lengths and batch sizes. If the decode slope grows with cached tokens and memory traffic rises, KV sharing has a plausible mechanism. If latency remains flat while feed-forward or interconnect time dominates, another intervention deserves priority.

3. Sweep H_kv as a capacity parameter

Compare MHA, several GQA group counts and MQA under matched model width, training budget, data and evaluation. For checkpoint conversion, compare contiguous mean pooling with at least one grouping based on activation or function similarity. Evaluate immediately after conversion and after bounded adaptation, so the evidence separates destructive initialisation from an architecture that cannot recover.

4. Test where shared memory should hurt

Average benchmark scores can conceal the exact distinctions that KV sharing removes. Include long-range retrieval with distractors, entity binding, order and negation, code symbol resolution, multilingual morphology and any domain-specific evidence task that depends on precise selection. Add the duplicated-head negative control and a deliberately incompatible grouping. A credible evaluation should show both a case where sharing is free and a case where it is not.

5. Treat kernel and sharding support as part of the architecture

A model configuration is not deployable merely because its tensor shapes are legal. Confirm that the serving kernel consumes compact KV tensors without expanding them, that H_kv divides H_q, and that KV groups align with tensor-parallel partitions. Official implementations such as NVIDIA Transformer Engine expose the number of GQA groups separately from query heads, reflecting this operational distinction. Benchmark the exact precision, cache layout and attention kernel intended for release.

6. Release by non-inferiority, then watch the workload

Define the quality boundary before seeing results. For each critical task and slice, set an allowed non-inferiority margin, confidence procedure and minimum sample size. Separately set the serving improvement that would justify the change, such as a cache-capacity reduction, a tail inter-token-latency improvement or a throughput increase at a stated service level. A smaller H_kv passes only when it meets both sides. A statistically uncertain quality loss is not cancelled by an attractive memory ratio.

Retain the MHA or higher-group checkpoint and its serving manifest as the rollback unit. H_kv is embedded in projection shapes, so it is not usually a harmless runtime knob on one checkpoint. After release, monitor prompt and output lengths, active batch, cache occupancy, tail token latency and the quality slices most sensitive to selection and value transfer. Workload drift can remove the economic benefit, while domain drift can activate distinctions absent from the conversion corpus. The release decision has an expiry condition: re-evaluate when the workload, kernel, parallel layout or task distribution changes materially.

Executable artefact

The MHA, GQA and MQA lab

The page includes two executable forms. The calculator below makes the cache equation interactive. The Python artefact computes attention directly, proves the duplicated-head equivalence, estimates a production-shaped cache and stresses mean-pooled grouping as KV heads diverge.

Assumptions and expected result

The lab uses synthetic Gaussian activations and weights, four query heads and two prospective GQA groups. It requires NumPy. The positive case duplicates MHA KV heads within each group, so GQA must match exactly. The negative case perturbs each head before mean pooling, so sufficiently incompatible groups must produce a large output difference. The deployment fixture must report 32, 8 and 1 GiB for MHA, GQA-8 and MQA.

"""MHA, GQA and MQA laboratory for Paper 58.

Python 3.10+ and NumPy 1.24+.
All inputs and outputs are synthetic. The lab demonstrates three claims:
1. GQA is exactly equivalent to MHA when the MHA key/value heads inside each
   group are duplicates.
2. Logical KV-cache size scales with the number of distinct KV heads, not the
   number of query heads.
3. Compressing incompatible KV heads into one group can change the attention
   output substantially.

Run:
    python3 many_heads_gqa_mqa_lab.py
"""

from __future__ import annotations

import json
import math
from dataclasses import asdict, dataclass

import numpy as np


@dataclass(frozen=True)
class StressPoint:
    perturbation: float
    within_group_dispersion: float
    relative_output_error: float


def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    """Numerically stable softmax."""
    shifted = x - np.max(x, axis=axis, keepdims=True)
    exp = np.exp(shifted)
    return exp / np.sum(exp, axis=axis, keepdims=True)


def grouped_attention(
    x_history: np.ndarray,
    x_now: np.ndarray,
    w_query: np.ndarray,
    w_key: np.ndarray,
    w_value: np.ndarray,
) -> np.ndarray:
    """Return one causal attention output per query head.

    Shapes:
      x_history: [tokens, model_width]
      x_now:      [model_width]
      w_query:    [query_heads, model_width, head_width]
      w_key:      [kv_heads, model_width, head_width]
      w_value:    [kv_heads, model_width, head_width]
      output:     [query_heads, head_width]

    MHA uses kv_heads == query_heads.
    GQA uses 1 < kv_heads < query_heads.
    MQA uses kv_heads == 1.
    """
    query_heads = w_query.shape[0]
    kv_heads = w_key.shape[0]
    if query_heads % kv_heads != 0:
        raise ValueError("query_heads must be divisible by kv_heads")
    if w_key.shape != w_value.shape:
        raise ValueError("key and value projection tensors must have equal shape")

    query = np.einsum("m,hmd->hd", x_now, w_query)
    key = np.einsum("tm,gmd->gtd", x_history, w_key)
    value = np.einsum("tm,gmd->gtd", x_history, w_value)

    queries_per_group = query_heads // kv_heads
    key_for_query = np.repeat(key, queries_per_group, axis=0)
    value_for_query = np.repeat(value, queries_per_group, axis=0)

    scores = np.einsum("hd,htd->ht", query, key_for_query)
    scores /= math.sqrt(query.shape[-1])
    weights = softmax(scores, axis=-1)
    return np.einsum("ht,htd->hd", weights, value_for_query)


def mean_pool_heads(weights: np.ndarray, target_kv_heads: int) -> np.ndarray:
    """Mean-pool contiguous MHA heads into a smaller number of KV heads."""
    source_heads = weights.shape[0]
    if source_heads % target_kv_heads != 0:
        raise ValueError("source heads must be divisible by target KV heads")
    group_size = source_heads // target_kv_heads
    return weights.reshape(target_kv_heads, group_size, *weights.shape[1:]).mean(axis=1)


def cache_bytes(
    *,
    layers: int,
    batch: int,
    tokens: int,
    kv_heads: int,
    head_width: int,
    bytes_per_element: int,
) -> int:
    """Logical K plus V cache bytes, excluding allocator and page overhead."""
    values = (layers, batch, tokens, kv_heads, head_width, bytes_per_element)
    if any(value <= 0 for value in values):
        raise ValueError("all cache dimensions must be positive")
    return layers * batch * tokens * kv_heads * head_width * 2 * bytes_per_element


def gibibytes(byte_count: int) -> float:
    return byte_count / (1024**3)


def run_lab() -> dict[str, object]:
    rng = np.random.default_rng(58)

    query_heads = 4
    grouped_kv_heads = 2
    model_width = 12
    head_width = 3
    history_tokens = 7

    x_history = rng.normal(size=(history_tokens, model_width))
    x_now = rng.normal(size=(model_width,))
    w_query = rng.normal(size=(query_heads, model_width, head_width)) / math.sqrt(model_width)

    base_key = rng.normal(size=(grouped_kv_heads, model_width, head_width)) / math.sqrt(model_width)
    base_value = rng.normal(size=(grouped_kv_heads, model_width, head_width)) / math.sqrt(model_width)

    # Negative control: an MHA model whose KV heads are exact duplicates inside
    # each prospective GQA group has the same function as that GQA model.
    duplicated_mha_key = np.repeat(base_key, query_heads // grouped_kv_heads, axis=0)
    duplicated_mha_value = np.repeat(base_value, query_heads // grouped_kv_heads, axis=0)
    mha_output = grouped_attention(
        x_history, x_now, w_query, duplicated_mha_key, duplicated_mha_value
    )
    gqa_output = grouped_attention(x_history, x_now, w_query, base_key, base_value)
    exact_equivalence_error = float(np.max(np.abs(mha_output - gqa_output)))

    noise_key = rng.normal(size=duplicated_mha_key.shape) / math.sqrt(model_width)
    noise_value = rng.normal(size=duplicated_mha_value.shape) / math.sqrt(model_width)

    stress_points: list[StressPoint] = []
    for perturbation in (0.0, 0.025, 0.05, 0.1, 0.2, 0.4, 0.8):
        stressed_key = duplicated_mha_key + perturbation * noise_key
        stressed_value = duplicated_mha_value + perturbation * noise_value
        pooled_key = mean_pool_heads(stressed_key, grouped_kv_heads)
        pooled_value = mean_pool_heads(stressed_value, grouped_kv_heads)

        full_output = grouped_attention(
            x_history, x_now, w_query, stressed_key, stressed_value
        )
        compressed_output = grouped_attention(
            x_history, x_now, w_query, pooled_key, pooled_value
        )
        relative_error = float(
            np.linalg.norm(full_output - compressed_output) / np.linalg.norm(full_output)
        )

        key_grouped = stressed_key.reshape(grouped_kv_heads, 2, -1)
        value_grouped = stressed_value.reshape(grouped_kv_heads, 2, -1)
        pooled_key_flat = pooled_key.reshape(grouped_kv_heads, 1, -1)
        pooled_value_flat = pooled_value.reshape(grouped_kv_heads, 1, -1)
        key_dispersion = np.linalg.norm(key_grouped - pooled_key_flat) / np.linalg.norm(
            key_grouped
        )
        value_dispersion = np.linalg.norm(
            value_grouped - pooled_value_flat
        ) / np.linalg.norm(value_grouped)
        dispersion = float((key_dispersion + value_dispersion) / 2)

        stress_points.append(
            StressPoint(
                perturbation=perturbation,
                within_group_dispersion=dispersion,
                relative_output_error=relative_error,
            )
        )

    deployment = {
        "layers": 32,
        "batch": 8,
        "tokens": 8192,
        "query_heads": 32,
        "gqa_kv_heads": 8,
        "head_width": 128,
        "bytes_per_element": 2,
    }
    mha_cache = cache_bytes(
        layers=deployment["layers"],
        batch=deployment["batch"],
        tokens=deployment["tokens"],
        kv_heads=deployment["query_heads"],
        head_width=deployment["head_width"],
        bytes_per_element=deployment["bytes_per_element"],
    )
    gqa_cache = cache_bytes(
        layers=deployment["layers"],
        batch=deployment["batch"],
        tokens=deployment["tokens"],
        kv_heads=deployment["gqa_kv_heads"],
        head_width=deployment["head_width"],
        bytes_per_element=deployment["bytes_per_element"],
    )
    mqa_cache = cache_bytes(
        layers=deployment["layers"],
        batch=deployment["batch"],
        tokens=deployment["tokens"],
        kv_heads=1,
        head_width=deployment["head_width"],
        bytes_per_element=deployment["bytes_per_element"],
    )

    # Positive and negative assertions make the teaching claims executable.
    assert exact_equivalence_error < 1e-12
    assert mha_cache // gqa_cache == 4
    assert mha_cache // mqa_cache == 32
    assert stress_points[-1].relative_output_error > 0.8
    assert stress_points[-1].relative_output_error > stress_points[2].relative_output_error

    return {
        "assumptions": {
            "data": "synthetic Gaussian activations and projection weights",
            "grouping": "contiguous query heads, equal group size",
            "cache": "logical K and V payload only, no allocator or page overhead",
        },
        "exact_equivalence_max_abs_error": exact_equivalence_error,
        "deployment_cache_gib": {
            "MHA_32_KV_heads": gibibytes(mha_cache),
            "GQA_8_KV_heads": gibibytes(gqa_cache),
            "MQA_1_KV_head": gibibytes(mqa_cache),
        },
        "stress_curve": [asdict(point) for point in stress_points],
    }


if __name__ == "__main__":
    print(json.dumps(run_lab(), indent=2))

{
  "exact_equivalence_max_abs_error": 0.0,
  "deployment_cache_gib": {
    "MHA_32_KV_heads": 32.0,
    "GQA_8_KV_heads": 8.0,
    "MQA_1_KV_head": 1.0
  },
  "stress_curve_last_point": {
    "within_group_dispersion": 0.4801368281,
    "relative_output_error": 0.9104757097
  }
}

stateDiagram-v2
  [*] --> Claim
  state "Keys choose, values carry" as Claim
  state "Many heads are not eight<br/>little experts" as Test
  state "A decisive negative control" as Revision
  Claim --> Test: expose mechanism
  Test --> Decision: evidence survives
  Test --> Revision: rival survives
  Revision --> Claim: narrow or rebuild
  Decision --> [*]

Compact glossary

Query head

One learned projection of the current representation used to score past positions.

KV head

One learned key projection and one value projection whose activations are cached across tokens.

KV cache

Per-layer stored keys and values for earlier tokens, reused during autoregressive decoding.

Group ratio

H_q/H_kv, the number of query heads sharing each KV pair.

Prefill

Parallel processing of the supplied prompt before the first generated token.

Decode

Sequential generation in which each new query reads the accumulated cache.

Spend distinct memory where it earns its keep

MHA, GQA and MQA are best understood as points on one controlled axis. Query heads determine how many learned questions can be posed in parallel. KV heads determine how many learned memory projections those questions can use, how much state must be cached, and how much historical state may need to move during decoding.

The architecture decision changes when these counts are separated. Do not ask whether many heads are good in the abstract. Ask whether an additional KV head preserves a distinction that improves the intended task enough to justify its cache, bandwidth and sharding cost. Do not claim a speed-up from a formula alone. Show that the target workload is decode-bound and that the runtime keeps shared KV compact.

The defensible default is the smallest H_kv that survives discriminating quality tests on the real serving path. GQA is valuable because it turns that principle into a tunable architecture, leaving MHA and MQA as its two endpoints rather than rival doctrines.