M5.7 - MQA and GQA: KV-Head Sharing, Cache Bytes, and Grouped Decode Kernels

M5.7 - MQA and GQA: KV-Head Sharing, Cache Bytes, and Grouped Decode Kernels

0x00 Overview And Background

The first six M5 articles treat K/V as persistent serving state. They ask when that state is created, how requests are scheduled around it, where exact blocks live, which history remains resident, and whether prefill and decode should own the state on different workers.

M5.7 begins one step earlier. The cache allocator does not decide how many K/V heads exist. That count was fixed by the model architecture, encoded in the checkpoint, and inherited by every later serving component. Understanding why architectures began reducing that count requires following attention across two very different operating regimes.

From multi-head representation to incremental-decode pressure

The 2017 Transformer made multi-head attention a representation mechanism. Instead of computing one attention distribution in the full model width, it projects the hidden row into several lower-dimensional query, key, and value subspaces. Each query head can then learn a different score distribution before the head outputs are concatenated and projected back to the residual width. That history belongs to M3; the serving consequence begins with the projection ownership:

ordinary MHA checkpoint:
    H_q independently parameterized query heads
    H_q independently parameterized key heads
    H_q independently parameterized value heads

autoregressive cache consequence:
    every committed token contributes H_q key rows and H_q value rows
    at every cache-producing layer

During training or prompt prefill, many query rows exist at once. Matrix-matrix work can reuse projection weights and tiled attention can expose substantial parallelism. Incremental decoding changes the shape without changing the learned operation. A request normally contributes one new query row, yet that row must inspect all retained historical K/V rows. The model repeatedly reads an object whose size grows with context length while doing comparatively little work per loaded scalar. Batching helps projection and MLP weight reuse, but different requests still have independent, usually ragged histories. The persistent K/V read therefore remains a distinct bandwidth and capacity pressure.

This is the historical hinge. MQA and GQA were not introduced because the softmax equation had become incorrect. They change how many independently learned K/V write streams feed that equation so incremental inference stores and rereads less state.

2019: MQA moves diversity to the query side

Shazeer’s Fast Transformer Decoding: One Write-Head is All You Need keeps the many query heads but gives them one shared key head and one shared value head. The title’s “write-head” matters: each query head still computes its own scores, but all of them read the same stored K/V sequence. The model therefore retains query-side diversity while replacing H_q persistent K/V streams with one.

The paper’s measured table, reused in A27, shows why the proposal was about the incremental regime rather than a universal reduction in all attention cost:

Source: Noam Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, source-reported timing excerpt previously used in A27. The table is historical evidence from that paper’s workload and hardware, not a reproduced Transformer System benchmark.

MQA makes a strong modeling wager. Independent query projections alone must extract the distinctions that ordinary MHA could also encode through independent K/V projections. The byte saving follows algebraically once H_kv = 1; the quality effect does not. It depends on training, scale, data, objective, and the tasks used to evaluate the resulting checkpoint.

2023: GQA makes head ownership a tunable middle point

The GQA paper starts from an empirical problem with the two endpoints. MHA keeps full K/V capacity and full cache cost. MQA obtains the largest head-axis saving but can lose quality. Grouped-query attention chooses an intermediate number 1 < H_kv < H_q and assigns several query heads to each K/V head. The model now has a discrete architecture parameter: the number and membership of K/V groups.

GQA also changes the migration question. The paper does not merely train a new architecture from scratch. It studies converting an MHA checkpoint by pooling K/V heads within each group and then uptraining the converted model. That sequence separates three claims that are easy to collapse:

conversion:
    constructs tensors with the target GQA checkpoint shape

uptraining:
    adapts the model after independent K/V heads have been tied

evaluation:
    decides whether the resulting quality/latency point is acceptable

The original paper reports that mean pooling is a better conversion initialization than selecting the first head or random initialization in its experiments, and it studies a continued-training budget of roughly 5% of the original pretraining compute. Those are source-reported experimental findings, not mathematical consequences and not a universal migration budget. Loop 6 of this article will reconstruct the exact tensor conversion and the training contract before using those results.

One family, several independent design axes

The progression MHA -> GQA -> MQA changes head ownership inside a layer. It does not by itself choose which tokens remain visible, which physical pages hold them, which worker owns them, or how the rows are represented numerically. It also does not cover two adjacent architecture changes:

Later papers also ask which MHA heads should be merged. Mean pooling uses a fixed contiguous partition. Head alignment, trainable weighting, activation-informed groups, evolutionary search, key-driven dynamic grouping, and generalized Q/K/V grouping make different choices about the conversion or training problem. They can lead to the same final H_kv while producing different checkpoints. Loop 7 compares those methods by their actual paper contract, including the important fact that some reported experiments concern vision transformers rather than decoder-only language-model serving.

This article changes the question:

Before the allocator, scheduler, or transfer layer sees K/V,
how many distinct K/V heads did the model architecture create?

Classic multi-head attention, multi-query attention, and grouped-query attention differ at one local object:

H_q:
    number of query heads

H_kv:
    number of persistent key/value heads

In ordinary multi-head attention, H_kv = H_q. MQA sets H_kv = 1. GQA keeps an intermediate number of K/V heads and maps multiple query heads to each one. That sounds like a small notation change, but it crosses four contracts:

1. checkpoint contract: K/V projection weights have a different shape;
2. cache contract: every token stores fewer K/V rows;
3. kernel contract: several query heads must reuse one stored K/V head;
4. placement contract: sharding must preserve the query-to-KV grouping.

Source: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, Figure 2. The dashed edges identify query heads that share a key/value head.

The thesis of M5.7 is stricter than “GQA saves memory”:

Reducing H_kv reduces persistent bytes by H_q / H_kv
only if the checkpoint, cache layout, head mapping, and decode kernel
all preserve the reduced representation.

This article derives that statement and shows where it fails.

Reading contract

The article follows the state from model design to production:

1. derive the MHA projection and persistent-state baseline;
2. reduce ownership to MQA and then generalize it to GQA;
3. trace logical broadcast, grouped kernels, and backward accumulation;
4. reconstruct conversion, uptraining, and alternative grouping choices;
5. place K/V groups across ranks and distinguish logical from physical bytes;
6. validate checkpoint quality, cache state, HBM traffic, SLOs, and rollback.

Each stage can fail while tensor shapes still look plausible. For that reason, num_key_value_heads is treated as one field in a larger immutable architecture and runtime contract rather than as a sufficient description by itself.

What this article owns

M5.7 owns:

• MHA, MQA, and GQA K/V shapes;
• the per-token and per-request cache ledger;
• MHA-to-GQA conversion and uptraining boundaries;
• logical K/V broadcast versus physical expansion;
• grouped decode reuse and arithmetic intensity;
• K/V-head sharding and tensor-parallel constraints;
• a parity, byte, and performance validation protocol.

It does not own latent K/V representations. M5.8 will replace full K/V head rows with DeepSeek-style latent state and then open the absorbed decode kernel.

0x01 Local Object And Notation

Use:

The grouping factor is:

$$g = \frac{H_q}{H_{kv}}.$$

The required divisibility condition is:

H_q % H_kv == 0

The edge cases become:

The head mapping for contiguous groups is:

kv_head = q_head // g

That one integer mapping is model semantics. A runtime that uses a different group order can produce shape-correct tensors and wrong logits.

Immutable architecture-to-runtime manifest

H_q, H_kv, and d_h are necessary but not sufficient. A deployable GQA artifact needs one immutable manifest whose identity is consumed by checkpoint loading, cache creation, kernel planning, and parallel placement. Otherwise, each layer can independently infer a plausible but incompatible interpretation.

GQAManifest
    schema_version
    model_id
    checkpoint_sha256
    architecture_generation

    num_layers
    hidden_width d
    num_query_heads H_q
    num_kv_heads H_kv
    query_key_head_width d_qk
    value_head_width d_v
    query_to_kv_head[H_q]

    q_projection_shape
    k_projection_shape
    v_projection_shape
    output_projection_shape
    qkv_fusion_order
    rope_coordinate_order

    cache_layout
    cache_dtype_or_quant_schema
    block_size
    cache_generation_namespace

    backend_name
    backend_version_or_commit
    execution_mode
    materializes_repeated_kv

    tensor_parallel_degree
    query_heads_by_rank
    kv_heads_by_rank
    kv_replicas_by_rank

The projection shapes are derived from the architecture, not copied from a framework class name. Allowing different query/key and value widths for the general contract gives:

$$W_Q\in\mathbb{R}^{(H_qd_{qk})\times d},\qquad W_K\in\mathbb{R}^{(H_{kv}d_{qk})\times d},$$ $$W_V\in\mathbb{R}^{(H_{kv}d_v)\times d},\qquad W_O\in\mathbb{R}^{d\times(H_qd_v)}.$$

The common case uses d_qk = d_v = d_h, but the equality must be declared. Inferring it silently would make this article’s manifest unusable at the M5.8 boundary and in runtimes that support different value widths.

The query-to-K/V map must contain exactly H_q entries, each in [0, H_kv). Contiguous GQA uses:

$$\operatorname{kv}(i)=\left\lfloor\frac{i}{g}\right\rfloor, \qquad g=\frac{H_q}{H_{kv}}.$$

That equation is only the default. A searched or asymmetric grouping can use a non-contiguous map. In that case, the map, not integer division, is the checkpoint semantic. The K and V paths must consume the same map unless the architecture explicitly declares a different generalized attention design.

Original Transformer System reconstruction. The manifest is not a new model format; it is the minimum identity needed to reason about a model artifact and its serving path without silently reinterpreting grouped heads.

Canonical identity and publication order

Serialize the semantic fields in a canonical order and hash them:

$$I_{gqa}=\operatorname{SHA256} \left(\operatorname{canonical\_json}(\text{GQAManifest})\right).$$

Runtime-derived fields such as a temporary workspace pointer do not belong in that identity. Fields that alter tensor meaning do. A cache generation may be published only after all four receipts agree on I_gqa:

1. checkpoint receipt
       projection tensors and fused split order match the manifest

2. cache receipt
       stored H_kv, head widths, dtype/schema, layout, and block size match

3. kernel-plan receipt
       backend accepts H_q/H_kv, mode, layout, mask, and position contract

4. placement receipt
       every query head has a local, replicated, or communicated K/V owner

5. publish
       attach I_gqa and committed_generation to the visible cache state

This ordering prevents a particularly dangerous failure: a checkpoint can load successfully because the flattened K/V row count is correct while a different head permutation is assumed by the runtime. The output shape remains valid; only numerical parity reveals that query heads are reading the wrong learned K/V features.

def validate_gqa_manifest(m):
    assert m.num_query_heads > 0
    assert m.num_kv_heads > 0
    assert m.num_query_heads % m.num_kv_heads == 0
    assert len(m.query_to_kv_head) == m.num_query_heads
    assert set(m.query_to_kv_head) <= set(range(m.num_kv_heads))

    assert m.q_projection_shape == [m.num_query_heads * m.d_qk, m.hidden_width]
    assert m.k_projection_shape == [m.num_kv_heads * m.d_qk, m.hidden_width]
    assert m.v_projection_shape == [m.num_kv_heads * m.d_v, m.hidden_width]
    assert m.output_projection_shape == [m.hidden_width,
                                         m.num_query_heads * m.d_v]

    assert m.execution_mode in {"prefill", "decode", "append", "verify"}
    assert m.backend_version_or_commit is not None
    assert every_query_head_has_a_kv_owner(m)
    return canonical_hash(m)

The manifest remains immutable. A conversion from MHA to GQA creates a new checkpoint and a new manifest identity; an uptraining checkpoint creates a new checkpoint hash; a backend re-plan creates a new plan receipt but does not rewrite architecture semantics. Reusing old cache pages across any incompatible identity is a correctness failure, not a cache hit.

0x02 MHA Baseline: One Persistent Write Stream Per Query Head

For a decoder layer using ordinary MHA:

Q: [B, H_q, S_q, d_h]
K: [B, H_q, T,   d_h]
V: [B, H_q, T,   d_h]

Projection rows and parameter count

Start with the general widths declared by the manifest. Biases are omitted so the matrix terms remain visible:

$$P_Q=dH_qd_{qk},\qquad P_K=dH_qd_{qk},$$ $$P_V=dH_qd_v,\qquad P_O=dH_qd_v.$$

Therefore the four attention projections contain:

$$P_{\text{proj,MHA}} =P_Q+P_K+P_V+P_O =2dH_q(d_{qk}+d_v).$$

In the common equal-width case, d_qk = d_v = d_h and d = H_q d_h, so:

$$P_{\text{proj,MHA}}=4d^2.$$

This is a checkpoint ledger, not the parameter count of a complete Transformer layer. Norms, biases, MLP parameters, and any gated branch remain outside it. The identity check is row-oriented:

Q rows: H_q * d_qk
K rows: H_q * d_qk
V rows: H_q * d_v
O input columns: H_q * d_v

A fused W_qkv stores the same rows in one allocation. Fusion changes layout and launch structure, not the semantic split or parameter count.

Source-to-tensor trace for one decode step

Let x be the current hidden rows with shape [B, S_q, d]. An explicit MHA path is:

def mha_attention_step(x, cache, layer, mask):
    # x: [B, S_q, d]
    q_flat = x @ layer.W_Q.T        # [B, S_q, H_q * d_qk]
    k_flat = x @ layer.W_K.T        # [B, S_q, H_q * d_qk]
    v_flat = x @ layer.W_V.T        # [B, S_q, H_q * d_v]

    q = q_flat.view(B, S_q, H_q, d_qk).transpose(1, 2)
    k_new = k_flat.view(B, S_q, H_q, d_qk).transpose(1, 2)
    v_new = v_flat.view(B, S_q, H_q, d_v).transpose(1, 2)

    # M5.1 owns publication. This view may include the transaction's new rows,
    # but an aborted step must not leave them committed.
    k_hist, v_hist = cache.view_with_candidate(k_new, v_new)

    scores = (q @ k_hist.transpose(-2, -1)) / sqrt(d_qk)
    probs = softmax(apply_mask(scores, mask), dim=-1)
    heads = probs @ v_hist

    merged = heads.transpose(1, 2).reshape(B, S_q, H_q * d_v)
    return merged @ layer.W_O.T, (k_new, v_new)

The trace exposes three tensor identities that abbreviated equations often hide:

1. score scaling uses the query/key head width d_qk, not model width d;
2. the cache stores projected K/V rows before any query-head broadcast;
3. the output projection consumes H_q value results even if a later architecture shares the stored value heads.

Source: Full Stack Optimization of Transformer Inference: a Survey, figure excerpt previously used in A27. It supplies conceptual operator context; the tensor shapes and accounting equations in this article are independently reconstructed.

Forward work, persistent writes, and historical reads

At one decode step, S_q = 1. The current query is small, but the operation must read historical K and V for all T tokens and all H_q heads.

Let N = B S_q current rows. Under the M4 convention of two FLOPs per multiply-add, the four projections perform approximately:

$$F_{\text{proj,MHA}}\approx2N P_{\text{proj,MHA}}.$$

The QK and probability-value products perform:

$$F_{QK}\approx2BH_qS_qTd_{qk},$$ $$F_{PV}\approx2BH_qS_qTd_v.$$

Softmax, masking, reductions, and position transforms are separate operation terms. In the equal-width case, the two matrix products reduce to:

$$F_{\text{attn,MHA}}\approx4BH_qS_qTd_h.$$

The newly produced K/V payload is:

$$D_{\text{append,MHA}} =BS_qH_q(d_{qk}+d_v)b.$$

The dominant historical K/V payload requested by one attention layer is:

$$D_{\text{history,MHA}} =BTH_q(d_{qk}+d_v)b.$$

This is a logical unique-payload ledger. Real HBM traffic can be higher because of rereads, padding, poor partitioning, or cache misses, and lower than a naive sum when a tile remains in a useful cache level. The full logical score tensor would contain B H_q S_q T elements, but an IO-aware kernel need not materialize it in HBM. That temporary schedule remains M4.3’s concern.

The persistent payload per token per layer is:

$$M_{token,layer}^{MHA}=2H_qd_hb.$$

For B requests of equal retained length T across L layers:

$$M_{KV}^{MHA}=2BTLH_qd_hb.$$

Example:

L   = 32
H_q = 32
d_h = 128
b   = 2 bytes for fp16/bf16

Then one token contributes:

per layer:
    2 * 32 * 128 * 2 = 16,384 bytes = 16 KiB

across 32 layers:
    16 KiB * 32 = 512 KiB

A single 4096-token request therefore carries about:

512 KiB/token * 4096 tokens = 2 GiB

This counts only K/V payload. Page padding, allocator metadata, scale tensors, workspace, replication, and fragmentation are separate terms.

The distinction from M4 matters. FlashAttention can avoid writing the full score and probability matrices to HBM, but autoregressive serving still retains K/V between calls. IO-aware attention changes a temporary execution schedule; MQA/GQA change the persistent representation handed to that schedule.

Reconciled decode example

Use the article’s running dimensions and one decode row:

B = 1, S_q = 1, T = 4096
d = 4096, H_q = 32, d_qk = d_v = 128
b = 2 bytes

The per-layer baseline is:

Across 32 layers, this idealized MHA request reads 2 GiB of historical K/V for one generated token and appends 512 KiB of new K/V payload. It also performs projection, MLP, norm, position, sampling, and runtime work not counted in that 2 GiB. The comparison is deliberately narrow: later loops may divide the K/V-head terms by g, but they must not divide unrelated work by the same factor.

0x03 MQA: Many Readers, One K/V Write Head

Multi-query attention keeps many query heads but emits one K and one V head:

Q: [B, H_q, S_q, d_h]
K: [B, 1,   T,   d_h]
V: [B, 1,   T,   d_h]

What changes in the checkpoint

Only the K/V projection row counts collapse to one head. Q and the input width of O remain query-head sized:

$$P_{\text{proj,MQA}} =dH_qd_{qk}+dd_{qk}+dd_v+dH_qd_v =d(H_q+1)(d_{qk}+d_v).$$

For equal head widths and d = H_q d_h:

$$P_{\text{proj,MQA}} =2d^2\left(1+\frac{1}{H_q}\right).$$

Relative to the MHA attention-projection baseline:

$$\frac{P_{\text{proj,MQA}}}{P_{\text{proj,MHA}}} =\frac{H_q+1}{2H_q}.$$

That ratio approaches one half, not 1/H_q, because Q and O are still full width. The 1/H_q reduction applies to K/V projection rows and cache state, not to the entire attention block.

A fused projection therefore changes from an MHA output width of H_q*d_qk + H_q*d_qk + H_q*d_v to:

$$d_{\text{fused,MQA}}=H_qd_{qk}+d_{qk}+d_v.$$

Source-local A27 implementation figure derived from the pinned llm-foundry attention example cited by A27. It is code context, not evidence that this historical implementation is the current runtime API.

Full tensor trace and implicit broadcast

The shared K/V sequence does not need a query-head axis:

def mqa_attention_step(x, cache, layer, mask):
    # x: [B, S_q, d]
    fused = x @ layer.W_qkv.T
    q_flat, k_flat, v_flat = split(
        fused,
        [H_q * d_qk, d_qk, d_v],
        dim=-1,
    )

    q = q_flat.view(B, S_q, H_q, d_qk).transpose(1, 2)
    k_new = k_flat.view(B, S_q, d_qk)
    v_new = v_flat.view(B, S_q, d_v)
    k_hist, v_hist = cache.view_with_candidate(k_new, v_new)

    # q: [B,H_q,S_q,d_qk], k_hist: [B,T,d_qk]
    scores = einsum("bhsd,btd->bhst", q, k_hist) / sqrt(d_qk)
    probs = softmax(apply_mask(scores, mask), dim=-1)
    heads = einsum("bhst,btc->bhsc", probs, v_hist)

    merged = heads.transpose(1, 2).reshape(B, S_q, H_q * d_v)
    return merged @ layer.W_O.T, (k_new, v_new)

Written elementwise, query head h reads the same stored key row:

$$S_{b,h,s,t} =\frac{1}{\sqrt{d_{qk}}} \sum_{r=0}^{d_{qk}-1}Q_{b,h,s,r}K_{b,t,r},$$

and uses its own probability distribution with the same value row:

$$O_{b,h,s,c} =\sum_{t=0}^{T-1} \operatorname{softmax}(S_{b,h,s,:})_tV_{b,t,c}.$$

Broadcast is an indexing rule in those equations. It does not require storing H_q copies of K or V. An explicit repeated reference is useful for numerical testing, but it is not the desired cache representation:

k_reference = k_hist[:, None, :, :].expand(B, H_q, T, d_qk)
v_reference = v_hist[:, None, :, :].expand(B, H_q, T, d_v)

assert_close(
    direct_mqa(q, k_hist, v_hist),
    ordinary_mha(q, k_reference, v_reference),
)

Source-local A27 code illustration based on the llm-foundry attention path. It shows semantic broadcast. Whether a concrete backend materializes the expanded view is a separate allocation and traffic question.

MQA work and traffic ledger

Projection work follows the smaller checkpoint:

$$F_{\text{proj,MQA}}\approx2BS_qP_{\text{proj,MQA}}.$$

The attention matrix products remain query-head sized:

$$F_{\text{attn,MQA}} \approx2BH_qS_qT(d_{qk}+d_v).$$

The unique K/V write and historical-read terms become:

$$D_{\text{append,MQA}} =BS_q(d_{qk}+d_v)b,$$ $$D_{\text{history,MQA}} =BT(d_{qk}+d_v)b.$$

The model still produces B*H_q*S_q*T logical scores and B*H_q*S_q*d_v head outputs. MQA therefore increases the reuse opportunity per loaded K/V row; it does not remove the query-head work that consumes that row.

Each query head still computes a separate score distribution:

for q_head in range(H_q):
    scores = q[:, q_head] @ K_shared.transpose(-2, -1)
    out[:, q_head] = softmax(scores) @ V_shared

The stored payload becomes:

$$M_{KV}^{MQA}=2BTLd_hb.$$

Relative to MHA:

$$\frac{M_{KV}^{MQA}}{M_{KV}^{MHA}}=\frac{1}{H_q}.$$

For the preceding 32-head example:

per token per layer:
    2 * 1 * 128 * 2 = 512 bytes

per token across 32 layers:
    16 KiB

one 4096-token request:
    64 MiB

The payload ratio is exactly 32x for this ledger. The end-to-end speedup is not 32x: projection work, output projection, MLPs, launch overhead, synchronization, and imperfect K/V reuse remain.

Noam Shazeer’s MQA paper identifies incremental decoding as a bandwidth problem: the current row is cheap compared with repeatedly loading historical K/V. MQA attacks that read term. Its tradeoff is on the learned write side: all query heads now inspect the same stored K/V features. Whether quality remains acceptable is a trained-model result, not a consequence of the byte formula.

Reconciled MQA example

For d=4096, H_q=32, d_qk=d_v=128, T=4096, and bf16/fp16 K/V:

Across 32 layers, one 4096-token MQA request stores 64 MiB of K/V payload instead of 2 GiB in the MHA baseline. Whether one generated token actually causes only 64 MiB of aggregate historical K/V traffic depends on work partitioning, cache hierarchy, sharding, and the selected grouped kernel.

The training and serving contracts are also distinct. A model trained natively with MQA has learned under shared K/V projections. A shape-converted MHA model has not. Conversion and continued training are treated in Loop 6 rather than being smuggled into this byte comparison. Likewise, the 2022 Efficiently Scaling Transformer Inference paper analyzes partitioning choices for MQA; Loop 8 owns the resulting replication and communication ledger.

0x04 GQA: A Trained Middle Point

GQA chooses 1 < H_kv < H_q:

Q: [B, H_q,  S_q, d_h]
K: [B, H_kv, T,   d_h]
V: [B, H_kv, T,   d_h]

Projection and fused-layout derivation

GQA keeps full-width Q/O paths and scales K/V rows with H_kv:

$$P_{\text{proj,GQA}} =d(H_q+H_{kv})(d_{qk}+d_v).$$

Relative to MHA:

$$\frac{P_{\text{proj,GQA}}}{P_{\text{proj,MHA}}} =\frac{H_q+H_{kv}}{2H_q} =\frac{g+1}{2g}.$$

The fused QKV output width is:

$$d_{\text{fused,GQA}} =H_qd_{qk}+H_{kv}d_{qk}+H_{kv}d_v.$$

These formulas place MQA and MHA at the endpoints. Setting H_kv=1 recovers the MQA ledger; setting H_kv=H_q recovers MHA. This endpoint reconciliation is a useful formula test because it catches accidentally shrinking Q or O.

The head map is the executable architecture

Let m(h) map query head h to one stored K/V head. For contiguous balanced groups, m(h)=floor(h/g). The general grouped equations are:

$$S_{b,h,s,t} =\frac{1}{\sqrt{d_{qk}}} \sum_r Q_{b,h,s,r}K_{b,m(h),t,r},$$ $$O_{b,h,s,c} =\sum_t\operatorname{softmax}(S_{b,h,s,:})_t V_{b,m(h),t,c}.$$

The map can be non-contiguous after a searched or asymmetric conversion:

H_q = 8, H_kv = 2

contiguous balanced map:
    query head: 0 1 2 3 4 5 6 7
    K/V owner:  0 0 0 0 1 1 1 1

valid non-contiguous map:
    query head: 0 1 2 3 4 5 6 7
    K/V owner:  0 1 0 1 0 1 0 1

Both have four query heads per K/V head and identical cache bytes. They are not the same checkpoint. A runtime that replaces the second map with integer division silently changes the model.

def gqa_attention_step(x, cache, layer, head_map, mask):
    fused = x @ layer.W_qkv.T
    q_flat, k_flat, v_flat = split(
        fused,
        [H_q * d_qk, H_kv * d_qk, H_kv * d_v],
        dim=-1,
    )
    q = q_flat.view(B, S_q, H_q, d_qk).transpose(1, 2)
    k_new = k_flat.view(B, S_q, H_kv, d_qk).transpose(1, 2)
    v_new = v_flat.view(B, S_q, H_kv, d_v).transpose(1, 2)
    k_hist, v_hist = cache.view_with_candidate(k_new, v_new)

    # Reference gather. A production grouped kernel should consume head_map
    # directly rather than allocate these query-head-sized tensors.
    k_by_query_head = k_hist[:, head_map, :, :]
    v_by_query_head = v_hist[:, head_map, :, :]
    scores = einsum("bhsd,bhtd->bhst", q, k_by_query_head) / sqrt(d_qk)
    probs = softmax(apply_mask(scores, mask), dim=-1)
    heads = einsum("bhst,bhtc->bhsc", probs, v_by_query_head)
    merged = heads.transpose(1, 2).reshape(B, S_q, H_q * d_v)
    return merged @ layer.W_O.T, (k_new, v_new)

The code is deliberately a correctness reference. Advanced indexing often materializes the gathered result. A direct grouped kernel computes the same index relation while retaining [B,H_kv,T,*] storage and arranging work so K/V tiles are reused within the group.

Backward is a grouped reduction

Logical broadcast has a backward consequence. Every query head in group j contributes to the one learned K/V head:

$$\frac{\partial\mathcal{L}}{\partial K_j} =\sum_{h:m(h)=j} \frac{\partial\mathcal{L}}{\partial K^{\text{logical}}_h},$$ $$\frac{\partial\mathcal{L}}{\partial V_j} =\sum_{h:m(h)=j} \frac{\partial\mathcal{L}}{\partial V^{\text{logical}}_h}.$$

FlashAttention-2 describes MQA/GQA as index manipulation in the forward pass and summation of dK and dV across logically duplicated heads in the backward pass. The direct kernel and an explicit-repeat reference are equivalent only if that reduction is preserved.

Source: FlashAttention-2, source excerpt previously used in A27. It supports the forward-index/backward-reduction contract; it does not imply every framework fallback avoids a physical repeat.

GQA work, state, and endpoint ledger

The persistent payload is:

$$M_{KV}^{GQA}=2BTLH_{kv}d_hb.$$

Relative to MHA:

$$\frac{M_{KV}^{GQA}}{M_{KV}^{MHA}} =\frac{H_{kv}}{H_q} =\frac{1}{g}.$$

For H_q = 32, H_kv = 8, and g = 4:

per token per layer:
    2 * 8 * 128 * 2 = 4096 bytes = 4 KiB

per token across 32 layers:
    128 KiB

one 4096-token request:
    512 MiB

This is 4x smaller than MHA and 8x larger than MQA for the same layer/head dimensions. H_kv is therefore both a modeling choice and a serving capacity parameter.

Projection work is:

$$F_{\text{proj,GQA}}\approx2BS_qP_{\text{proj,GQA}}.$$

Attention matmul work remains:

$$F_{\text{attn,GQA}} \approx2BH_qS_qT(d_{qk}+d_v).$$

Unique append and historical payloads scale with H_kv:

$$D_{\text{append,GQA}} =BS_qH_{kv}(d_{qk}+d_v)b,$$ $$D_{\text{history,GQA}} =BTH_{kv}(d_{qk}+d_v)b.$$

For the running H_q=32, H_kv=8, d=4096, d_h=128, T=4096 example:

The 5/8 projection ratio and 1/4 K/V ratio answer different questions. Using the latter as an end-to-end attention FLOP or latency prediction is a category error.

Checkpoint conversion is not a reshape

An MHA checkpoint contains distinct K/V projections for each head. Converting that checkpoint changes learned parameter ownership, not merely a tensor view:

W_Q_old: [H_q, d_qk, d]       retained
W_K_old: [H_q, d_qk, d]       pooled by the declared head groups
W_V_old: [H_q, d_v, d]        pooled by the declared head groups
W_O_old: [d, H_q * d_v]       retained

The same rule applies to K/V bias rows when those parameters exist. A fused checkpoint first has to be split according to its declared qkv_order, head axis, transposition convention, and any RoPE permutation already applied by the checkpoint loader. Pooling arbitrary byte ranges from a fused matrix is not a valid conversion algorithm.

Let G_j be the explicit set of old K/V heads assigned to new K/V head j. Arithmetic mean pooling is the solution of a small, precise optimization problem:

$$J_j(W)=\sum_{i\in G_j}\lVert W-W_i\rVert_F^2.$$

Differentiating elementwise gives

$$\frac{\partial J_j}{\partial W} =2|G_j|W-2\sum_{i\in G_j}W_i.$$

The stationary point is therefore

$$W_j^*=\frac{1}{|G_j|}\sum_{i\in G_j}W_i,$$

and the positive quadratic Hessian makes it the unique minimizer. Applied to the K/V weights and biases:

$$W^{K,new}_j=\frac{1}{|G_j|}\sum_{i\in G_j}W^{K,old}_i, \qquad W^{V,new}_j=\frac{1}{|G_j|}\sum_{i\in G_j}W^{V,old}_i.$$

This derivation says exactly what the mean preserves: the unweighted centroid under squared parameter-space distance. It does not prove that the mean minimizes activation error, language-model loss, or downstream quality. First head, random initialization, weighted means, activation-aware alignment, and searched grouping answer different optimization questions. The GQA paper reports mean pooling outperforming first-head selection and random initialization in its T5-Large ablation; that is source-reported evidence for that experiment, not a theorem that the mean is best for every checkpoint.

Source: A27’s reuse of Figure 1 from Ainslie et al., GQA. The figure explains mean-pooled initialization. It does not show the continued-pretraining, evaluation, cache-identity, or release transactions required after conversion.

The implementation must use the declared group membership rather than infer contiguous groups from a reshape:

def pool_heads(tensor, groups):
    # tensor: [H_q, rows_per_head, d] or [H_q, rows_per_head]
    assert sorted(h for group in groups for h in group) == list(range(len(tensor)))
    return stack([tensor[group].mean(dim=0) for group in groups])


def convert_fused_mha_to_gqa(checkpoint, manifest, groups):
    q, k, v = split_fused_projection(
        checkpoint.qkv_weight,
        order=manifest.qkv_order,
        head_axis=manifest.head_axis,
        transposed=manifest.projection_transposed,
    )
    q_bias, k_bias, v_bias = split_optional_bias(checkpoint.qkv_bias, manifest)

    # RoPE coordinate permutations belong to checkpoint reconciliation. They
    # must be settled before heads are compared or pooled.
    assert manifest.rope_coordinates_reconciled

    k_new = pool_heads(k, groups)
    v_new = pool_heads(v, groups)
    k_bias_new = None if k_bias is None else pool_heads(k_bias, groups)
    v_bias_new = None if v_bias is None else pool_heads(v_bias, groups)

    candidate = checkpoint.copy()
    candidate.qkv_weight = fuse_projection(q, k_new, v_new, manifest.qkv_order)
    candidate.qkv_bias = fuse_optional_bias(q_bias, k_bias_new, v_bias_new, manifest)
    candidate.output_weight = checkpoint.output_weight  # unchanged
    candidate.architecture_generation = checkpoint.architecture_generation + 1
    candidate.checkpoint_sha256 = hash_tensor_tree(candidate)
    return candidate

For the non-contiguous map from Loop 5, G_0={0,2} and G_1={1,3}. reshape(H_kv, group_size, ...) would silently pool {0,1} and {2,3} and produce a shape-correct but semantically different checkpoint. The conversion receipt must therefore record the exact old-head sets, source tensor names, split convention, whether biases existed, and hashes before and after pooling.

A pinned MHA checkpoint schema, not an imagined tensor

As a concrete source-schema trace, consider openlm-research/open_llama_7b_v2 at revision e5961def23172a2384543940e773ab676033c963. Its pinned config.json declares hidden_size=4096, num_attention_heads=32, num_hidden_layers=32, fp16 weights, and Transformers 4.30.1. It has no separate K/V-head field. The matching pinned Transformers implementation constructs Q, K, V, and O as full-width linear maps and reshapes each projected K/V tensor to [B,32,T,128]; this is an MHA source schema.

The checkpoint index names layer-zero tensors model.layers.0.self_attn.{q,k,v,o}_proj.weight, places all four in pytorch_model-00001-of-00002.bin, and reports 13,476,839,424 bytes of indexed tensor payload. The immutable evidence receipt records the config, index, and implementation hashes. That trace verifies tensor names, shard ownership, and the shapes implied jointly by configuration and constructor code. It does not claim to have parsed the binary tensor headers, and it does not claim that this public checkpoint was the GQA paper’s T5 checkpoint or that the toy conversion below is a released OpenLLaMA derivative.

For this separate-projection checkpoint, conversion would pool K and V rows before writing new K/V tensors. A fused-QKV checkpoint adds a mandatory split step first. These are two storage encodings of the same semantic operation, not permission to guess an axis from a matching element count.

Uptraining is a new training run

Averaging ties degrees of freedom that the MHA checkpoint previously used independently. The GQA paper therefore defines a two-step process: convert the checkpoint, then continue pretraining so the rest of the network can adapt to the new K/V ownership. In the paper’s T5.1.1 experiments, the converted model uses the original pretraining setup and dataset; the main result uses alpha=0.05 of the original pretraining steps and reports roughly 600 TPUv3 chip-days. The paper also reports diminishing returns at 10% and explicitly notes that its main evidence is encoder-decoder T5 rather than a decoder-only comparison against training from scratch.

Those coordinates must travel with the claim. 5% is neither a universal data fraction nor a guaranteed recovery budget. A production migration freezes a new training manifest:

source checkpoint hash and architecture generation
conversion map, tensor layout, bias policy, and candidate hash
dataset snapshot, tokenizer, objective, sequence-length mixture
optimizer-state policy, optimizer hyperparameters, LR schedule, warmup
global batch, token/step budget, precision, parallel layout, RNG seeds
validation tasks, serving workloads, quality and latency thresholds

The original optimizer slots no longer have a one-to-one shape correspondence for pooled K/V parameters. They must be reinitialized or transformed by an explicit, tested rule; silently attaching stale MHA slots to the GQA candidate is invalid. The paper states that it reuses the original pretraining recipe, but that statement does not by itself specify a portable optimizer-slot migration algorithm for another checkpoint format.

def convert_uptrain_release(source, manifest, groups, training, gates):
    candidate = convert_fused_mha_to_gqa(source, manifest, groups)
    assert candidate.checkpoint_sha256 != source.checkpoint_sha256
    assert training.optimizer_state_policy in {"reinitialize", "explicit_transform"}

    validate_shapes_and_finite_values(candidate)
    compare_initialization_logits(source, candidate)  # diagnostic, not equality
    trained = continue_pretraining(candidate, training)

    receipts = evaluate_quality_and_serving(trained, gates)
    assert receipts.quality_pass and receipts.runtime_pass
    return release_with_new_cache_namespace(trained, receipts)

The transaction is deliberately longer than a tensor edit:

inspect -> split -> map -> pool -> reassemble -> hash -> validate
        -> uptrain -> evaluate -> release

The converted initialization, the uptrained checkpoint, and the released artifact are three different identities. None may reuse K/V pages generated by the source MHA artifact: head count, row ownership, and cache shape have all changed. Rollback means routing new requests to the original MHA release and its cache namespace, not relabeling GQA pages as MHA state.

The durable engineering lesson is:

mean pooling creates a least-squares parameter-space initialization
continued training adapts the changed architecture
task and runtime evaluation decide whether the tradeoff is releasable
new architecture identity invalidates old cache state

The GQA paper’s quality-versus-latency plot is the empirical counterpart to that rule:

Source: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, Figure 3. The points average the paper’s T5 task suite and report TPUv4 time per sample. GQA-XXL approaches the reported MHA-XXL quality at MQA-like latency in this experiment; the figure is evidence for a trained tradeoff, not a hardware-independent law.

The byte derivation predicts the direction of K/V traffic. It cannot predict the y-axis. That requires conversion, uptraining, and task evaluation. It also cannot transfer the x-axis unchanged to a GPU backend whose grouped kernel, batch shape, or memory hierarchy differs from the paper’s TPU setup.

The same paper sweeps the number of GQA groups instead of comparing only the three endpoints. In its T5-XXL measurement, moving from one MQA group to eight groups adds little time per sample, while the curve rises more sharply toward 32 groups and reaches the MHA endpoint at 64. This is useful evidence against a linear latency model: halving H_kv halves the cache-byte term, but kernel and system time need not move proportionally.

Source: Ainslie et al., GQA, Figure 6. The experiment uses GQA-XXL with input length 2048 and output length 512 on the paper’s TPU setup. The shape of this curve supports a measured middle ground in that system; GPU kernels, batching, and memory hierarchies can produce different absolute times and crossover points.

Alternative conversion methods do not share one contract

Mean pooling fixes the group map and optimizes an unweighted parameter-space objective. Later methods modify at least one of those choices: coordinate alignment, merge weights, grouping objective, group cardinality, per-layer head count, or whether the map is fixed at inference. Calling all of them “better GQA” hides the part a checkpoint converter and serving kernel need.

The table is a contract comparison, not a leaderboard. Each paper uses a different base model, data, recovery budget, metric, and hardware coordinate.

Align Heads: change coordinates before merging

Two heads can encode related features in rotated coordinates. Directly averaging their matrices can cancel those features even when a change of basis would align them. Align Heads collects calibration K/V activations and solves an orthogonal Procrustes problem:

$$R^*=\arg\min_{R^TR=I}\lVert XR-Y\rVert_F^2.$$

If X^T Y = U Sigma V^T, one orientation of the solution is R*=U V^T. The exact left/right convention must match the tensor layout. Orthogonality is valuable because a compensating transform can preserve a linear computation; for a value/output pair, for example,

$$(V R)(R^T W_O)=V(RR^T)W_O=VW_O.$$

For Q/K, applying a common orthogonal basis preserves the unpositioned dot product, but RoPE constrains which transforms commute with position rotations. The paper therefore uses RoPE-compatible two-dimensional subspaces rather than an arbitrary dense rotation. That is a checkpoint-algebra requirement, not a post-hoc runtime flag.

Source: A27 excerpts from Align Attention Heads Before Merging Them. These figures describe calibration and coordinate transformation; reported similarities and downstream scores remain paper results.

The paper then scores candidate groups from pairwise aligned-head similarity and searches for a high-similarity partition. This changes G_j before any weights are tied:

$$S(\mathcal{G})=\sum_{G_j\in\mathcal{G}} \sum_{a<b,\;a,b\in G_j}\operatorname{sim}(a,b).$$

Finally, newly pooled group projections coexist with original heads during a pruning schedule. A schematic interpolation is

$$\widetilde W_i=z_i W_i^{old}+(1-z_i)W_{m(i)}^{group},$$

with hard-concrete/L_0 control driving old-head masks toward zero and distillation terms constraining the student. The paper uses teacher SFT, vanilla KL and BiLD losses; its method is therefore not “Procrustes plus one average.” After pruning, the old heads are removed and the target is a standard GQA checkpoint.

Source: A27 excerpt from the Align Heads pruning pipeline. The final fixed checkpoint can use ordinary GQA only after transformed weights, group map, mask schedule, and removal of old heads have all reconciled.

WGQA: learn the merge, then fold it

Weighted GQA replaces the uniform 1/|G_j| coefficients with learned scalar, row-wise, or column-wise coefficients during fine-tuning:

$$W^{K,new}_j=\sum_{i\in G_j}a^K_i\odot W^K_i, \qquad W^{V,new}_j=\sum_{i\in G_j}a^V_i\odot W^V_i.$$

Scalar coefficients add 2H_q trainable values per attention layer in the paper’s parameterization; row/column variants add broader vectors. Initializing the coefficients to the mean-pooling value recovers the ordinary GQA starting point. Once training ends, the coefficients can be multiplied into and summed with the projection matrices, so they need not remain an extra decode operation. “No inference overhead” is therefore a folding claim for the paper’s fixed trained weights, not evidence that training-time aggregation is free or that every exporter performs the fold.

Source: A27 excerpt from Weighted Grouped Query Attention. Its reported improvements use T5-small/base downstream fine-tuning; they are not decoder-only serving measurements.

AsymGQA: activation-informed and unequal groups

AsymGQA uses calibration activations to compare K heads and V heads, searches for similar members, and allows unequal cardinalities. K and V groupings are searched independently in the paper. An activation objective can distinguish heads whose weights look different but whose outputs are similar on the calibration distribution; it also makes the result dataset-sensitive.

Source: A27 excerpt from Optimised Grouped-Query Attention Mechanism for Transformers. The paper evaluates LLaMA2 and reports source-local gains. Unequal groups do not automatically fit kernels that require a constant H_q/H_kv ratio.

For a fixed irregular map, semantics remain simple:

for q_head in range(H_q):
    k_owner = key_map[layer][q_head]
    v_owner = value_map[layer][q_head]
    out[q_head] = attention(q[q_head], k_cache[k_owner], v_cache[v_owner])

But a production backend must either consume those maps, reorder query heads and output-projection columns into a supported layout, or fall back. If K and V maps differ, the cache ABI must store both. Padding an unequal partition into equal groups can erase the intended capacity result.

QCQA: search a layer-specific quality/capacity frontier

QCQA makes the grouping search itself multi-objective. Its inexpensive proxy is the weight-sharing error (WSE):

$$L_{WSE}=\sum_j\sum_{i\in G_j} \left( \mathbb{E}[(W^K_i-W^K_{G_j})^2] +\mathbb{E}[(W^V_i-W^V_{G_j})^2] \right).$$

This is closely related to Loop 6’s centroid objective, but QCQA uses it as a fitness proxy across candidate partitions and combines it with cache capacity in NSGA-II search. It can use equal-cardinality or arbitrary-cardinality groups and can leave sensitive layers as MHA, so H_kv becomes a per-layer vector, not one model-wide scalar.

Source: A27 excerpt from QCQA. The paper evaluates LLaMA2/OPT-derived checkpoints and optional Alpaca fine-tuning. WSE is a paper-validated proxy in those experiments, not a proof of downstream quality for every workload.

Search output must be serialized as architecture data:

layer 0: H_kv=32, identity map       # retained MHA
layer 1: H_kv=6, irregular groups
layer 2: H_kv=4, equal groups
...

That affects checkpoint row counts, per-layer cache bytes, tensor-parallel placement, and kernel dispatch. A single global num_key_value_heads field is insufficient for this architecture.

KDGQA and DGQA: dynamic allocation is a different runtime question

KDGQA allocates query heads in proportion to normalized key-head norms during each forward pass. DGQA updates an allocation signal over training windows, using norm differences or an exponential moving average. A simplified allocation rule is

$$n_j=\lVert K_j\rVert_2,\qquad q_j\approx\operatorname{round} \left(H_q\frac{\widehat n_j}{\sum_r\widehat n_r}\right),$$

followed by an integer reconciliation so sum_j q_j=H_q. Zero range, rounding, minimum occupancy, and deterministic tie-breaking are part of the algorithm, not implementation trivia.

Source: A27 excerpt from Beyond Uniform Query Distribution. Its experiments uptrain and fine-tune Vision Transformers for image classification. They do not verify autoregressive K/V-cache traffic, decode kernels, or serving stability.

If a map changes during serving, a kernel must consume the map for that forward and batching may combine requests with different maps. The number of stored K/V heads can stay fixed while scheduling and reuse change. The current paper does not establish that this tradeoff benefits decoder-only LLM serving, so M5.7 treats it as a design boundary rather than a recommendation.

GQKVA changes both sides of attention

GQKVA partitions Q and K/V projections and evaluates combinations of query and K/V groups. Its family includes MHA, MQA, GQA, MKVA, GKVA, and more general grouped-Q/grouped-KV combinations.

Source: A27 excerpt from GQKVA. The paper studies Vision Transformer pretraining and model-size/accuracy tradeoffs. A GQKVA result cannot be inserted into a decoder GQA cache by changing only H_kv.

This variant belongs in the comparison because it reveals a boundary: once query projection ownership and cross-group output composition change, the Loop 5 fixed map query_head -> stored_kv_head no longer fully specifies the operator. It needs a different architecture manifest, parameter derivation, kernel, and training contract.

0x05 Logical Broadcast Is Not Physical Expansion

Many reference implementations expose GQA by logically repeating K/V heads:

def repeat_kv(x, n_rep):
    # x: [B, H_kv, T, d_h]
    b, h_kv, t, d_h = x.shape
    x = x[:, :, None, :, :]                # [B, H_kv, 1, T, d_h]
    x = x.expand(b, h_kv, n_rep, t, d_h)   # stride-0 group axis
    return x.reshape(b, h_kv * n_rep, t, d_h)

For H_kv = 2 and g = 4:

stored object:
    [B, 2, T, d_h]

logical grouped view:
    [B, 2, 4, T, d_h]

query-head view:
    [B, 8, T, d_h]

expand itself can use a zero stride instead of allocating repeated storage. The following reshape, a layout conversion, a backend fallback, or a contiguous-input requirement may still materialize an H_q-sized tensor.

The correctness oracle and the performance oracle are different:

correctness:
    repeated and grouped implementations produce the same attention result

performance:
    persistent storage and HBM reads remain proportional to H_kv, not H_q

For contiguous groups, repeat_kv encodes the same map as integer division. For a non-contiguous head_map, a gather or a kernel that consumes the map is required. The allocation states are distinct:

The backward reference must retain the gather/repeat lineage until gradients are reduced into [H_kv,...] K/V projection rows. Detaching or independently optimizing the repeated copies would train MHA-shaped state while claiming GQA.

PyTorch’s current scaled_dot_product_attention documentation exposes enable_gqa, requires the query-head count to be divisible by the K/V-head count, and requires equal K and V head counts. The docs also describe GQA as experimental and backend-constrained. A math fallback can be numerically valid while missing the intended traffic reduction. Backend/version verification is therefore part of the claim.

0x06 Grouped Decode Kernels: Reuse Must Be In The Schedule

The efficient head mapping is direct:

for q_head in range(H_q):
    kv_head = q_head // g
    attend(q[:, q_head], k_cache[:, kv_head], v_cache[:, kv_head])

But assigning one independent thread block to every query head can make blocks on different SMs reload the same K/V head. The cache is small in storage yet read redundantly in execution.

FlashInfer describes head-group fusion for short-query GQA: merge the query head-within-group axis into an effective query-row axis so one work unit can load a K/V tile and use it for all query heads in that group.

Source: FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving, Figure 11, as previously imported in Transformer Series A27. The effective row count becomes g * S_q; this is a work-partitioning transformation, not a cache copy.

For decode with S_q = 1:

ordinary row count per request/KV head:
    1

head-group-fused row count:
    g

This both increases row-level work and exposes K/V reuse within the same work unit. At longer query lengths, S_q itself may provide enough rows, so the same fusion is less critical. That is why prefill and decode can select different attention schedules for the same GQA checkpoint.

A first-order arithmetic-intensity derivation

For one layer, one decode query row, and cached length T, ignore the small current Q/output traffic and count the dominant attention operations.

QK and PV together perform approximately:

$$F_{attn}\approx4TH_qd_h\quad\text{FLOPs}.$$

The factor four is two FLOPs per multiply-add for QK plus two for PV. If each stored K and V scalar is read once from HBM, grouped K/V traffic is:

$$D_{KV}\approx2TH_{kv}d_hb\quad\text{bytes}.$$

The idealized arithmetic intensity is therefore:

$$I_{GQA} =\frac{F_{attn}}{D_{KV}} \approx\frac{4TH_qd_h}{2TH_{kv}d_hb} =\frac{2g}{b}\quad\text{FLOPs/byte}.$$

For bf16/fp16, b = 2, so:

idealized I_GQA ~= g FLOPs/byte

Examples:

This is not a complete roofline model. It excludes score/softmax state, metadata, partial reductions, output writes, cache misses, padding, and launch overhead. It isolates the architectural reuse opportunity.

If a backend physically expands or independently reloads K/V as though H_kv = H_q, the dominant traffic becomes:

$$D_{bad}\approx2TH_qd_hb,$$

and the same intensity falls back to:

$$I_{bad}\approx\frac{2}{b},$$

or about 1 FLOP/byte at bf16. The model config still says GQA, the logits can still be correct, and the intended bandwidth gain has disappeared.

Bytes saved are not automatically latency saved

A proper benchmark must separate:

cache allocation bytes:
    does storage scale with H_kv?

HBM bytes read:
    does the kernel reuse each group tile?

kernel duration:
    did lower traffic expose another bottleneck?

end-to-end TPOT:
    did attention matter enough in the full decoder step?

For small T, launch and fixed work can dominate. For large T, K/V traffic usually becomes more visible. For large batches, weight reuse and scheduler behavior change again. One throughput number cannot validate the mechanism.

Allocation oracle before the pinned runtime trace

The full pinned framework and kernel trace appears in 0x0A. Here the local question is narrower: can one semantic owner map produce identical results through explicit repetition and direct grouped indexing while exposing different temporary-allocation receipts? An API flag alone cannot answer that question, and a parity test alone cannot establish GPU traffic.

The reference expansion and direct grouped paths can consume the same owner map:

def materialized_reference(q, k, v, owner):
    # Explicit oracle allocation: [H_q,T,d_h] for both K and V.
    k_rep = [k[owner[h]] for h in range(len(q))]
    v_rep = [v[owner[h]] for h in range(len(q))]
    return ordinary_attention(q, k_rep, v_rep)


def direct_grouped(q, k, v, owner):
    # No repeated K/V object is part of this function's contract.
    return [attention_one_head(q[h], k[owner[h]], v[owner[h]])
            for h in range(len(q))]

For the Loop 10 fixture H_q=4, H_kv=2, T=3, d_h=2, and two bytes per scalar:

persistent grouped K/V payload:
    2 * H_kv * T * d_h * b = 48 bytes

explicit reference-repeat temporary:
    2 * H_q * T * d_h * b = 96 bytes

direct grouped repeat temporary:
    0 bytes by the oracle contract

The zero is narrow: the pure grouped oracle creates no repeated K/V object. It is not a claim about hidden workspaces, compiler copies, page-table traffic, or HBM transactions in a production backend.

Serving adds plan state that an operator call does not expose

A serving backend must bind the architecture to a cache ABI and a concrete execution plan before run:

plan = backend.plan(
    num_qo_heads=manifest.H_q,
    num_kv_heads=manifest.H_kv,
    head_dim=manifest.d_h,
    kv_layout=manifest.kv_layout,
    page_size=cache.page_size,
    page_indptr=cache.indptr,
    page_indices=cache.indices,
    last_page_len=cache.last_page_len,
    q_dtype=manifest.q_dtype,
    kv_dtype=manifest.kv_dtype,
    workspace=workspace,
)

assert plan.architecture_hash == manifest.architecture_hash
assert plan.cache_generation == cache.generation
output = backend.run(plan, q, cache)

The requested backend name is not enough. A release receipt must record the selected backend/kernel, pinned revision, plan hash, workspace size, layout, page size, head counts, dtypes, mask/position mode, query length, and cache generation. Replanning after a shape or backend change creates a new receipt. Reinterpreting an existing cache under a different H_kv, owner map, layout, or page schema is forbidden.

This also separates lifecycle modes:

The executable evidence bundle compares three CPU paths: an explicit materialized repeat, direct owner-map indexing, and grouped-by-K/V-head execution. It proves output parity and reports created repeated scalars. It does not substitute for GPU allocation, traffic, or latency measurement.

0x07 Sharding, Ownership, And Communication

Tensor parallelism must preserve the local mapping between query and K/V heads. Suppose P ranks shard query heads. Define R_j as the set of ranks holding a physical replica of K/V head j, and r_j=|R_j|. Logical request-cache bytes remain

$$M_{logical}=2BTLd_hb\sum_{j=0}^{H_{kv}-1}1 =2BTLH_{kv}d_hb,$$

while aggregate resident bytes are

$$M_{physical}=2BTLd_hb\sum_{j=0}^{H_{kv}-1}r_j.$$

Only when every r_j=1 are the two ledgers equal. This distinction is about physical copies, not whether the model semantically has extra learned heads.

Clean divisible case

If both counts divide evenly:

local H_q  = H_q / P
local H_kv = H_kv / P
local g    = local H_q / local H_kv = global g

Each rank can own complete query-head groups and its corresponding K/V cache. No cross-rank K/V fetch is required inside attention.

The placement invariant is stronger than divisibility:

for every rank p and every local query head h:
    owner = query_to_kv_head[h]
    owner is local on p
    OR the execution plan declares how p obtains that owner's K/V

With arbitrary maps from Loop 7, balancing query-head counts is not enough; whole groups may have unequal sizes and K/V maps may differ.

Fewer K/V heads than ranks

If H_kv < P, every rank cannot receive a distinct K/V shard. Choices include:

MQA makes this conflict obvious: one K/V head cannot be evenly head-sharded over many tensor-parallel ranks without replication or a different partition. A serving ledger should therefore distinguish:

logical cache bytes per request
physical resident bytes after replication
bytes communicated per decode step

The lowest logical H_kv is not always the fastest topology for a specific model, GPU count, and collective fabric.

A 16-rank reconciliation

Use the running H_q=32, d_h=128, T=4096, L=32, bf16 example and place two contiguous query heads on each of 16 ranks.

The MQA row is deliberately surprising. A 32x semantic cache reduction becomes only 2x aggregate resident reduction under this query-head-sharded placement. That is not a defect in MQA; it is a mismatch between model ownership and the chosen parallel axis.

For GQA-8, using only eight attention-parallel ranks gives each rank one full four-query group and one K/V head. The physical cache then returns to 512 MiB. Other model operations may still use 16-way tensor parallelism, but the transition requires a redistribution or a different composite parallel plan.

Replication traffic versus remote history

Replicated caches remove remote history reads from the attention inner loop, but replicas have to be populated. The minimum payload for distributing one new decode row to existing replicas across all layers is

$$M_{append,replicate} =2BLd_hb\sum_j(r_j-1).$$

For the table’s GQA placement this is 128 KiB per generated token per request; for 16-way replicated MQA it is 240 KiB. These are logical payload lower bounds. Collective algorithms, packet headers, alignment, routing, and links traversed can increase fabric traffic.

Keeping one unique K/V owner per head avoids replica memory. If every non-owner consumes the full history remotely, a lower-bound demand is

$$M_{remote,history} \ge 2BTLd_hb\sum_j(c_j-1),$$

where c_j is the number of ranks whose query heads consume head j. Caching those remote rows locally turns the design back into replication. Fetching only the new row amortizes communication but also builds replicas over time. The system must choose explicitly; “unique storage with local attention” is not a free fourth option.

Projection sharding determines collectives

The Efficiently Scaling Transformer Inference paper contrasts MHA head sharding with an MQA layout that uses another partition axis. Its graph is useful because it shows that changing K/V ownership can introduce redistribution around attention even while the K/V cache itself shrinks:

Source: A27 excerpt from Pope et al., Efficiently Scaling Transformer Inference. The collectives shown belong to the paper’s tensor layout and partition choices. They are not a universal collective sequence for every MQA/GQA implementation.

The exact collective depends on how W_Q, W_K, W_V, and W_O are column/row sharded and whether batch, sequence, head, or hidden dimensions are distributed. A release manifest therefore records tensor ownership and the phase-specific collective plan, rather than deriving it from H_kv alone:

projection output layout -> optional redistribution -> grouped attention
-> optional redistribution -> output projection input layout

A payload formula is not a latency model. All-reduce, all-gather, reduce-scatter, all-to-all, broadcast, peer reads, and multicast have different topology sensitivity. The article can derive bytes; only a pinned runtime trace and measurement can establish the implementation’s route and time.

Cache schema needs more than a tensor shape

A robust cache descriptor includes:

model/checkpoint identity
layer and K/V-head count
group mapping or ordering
head width and dtype/quantization schema
tensor-parallel shard and owner
position/scaling convention
valid token range and commit state

Two caches with [B, 8, T, 128] can still be incompatible if their head order, checkpoint, quantization scales, or shard ownership differ.

0x08 Adjacent Architecture Axis: Cross-Layer Sharing

MQA/GQA reduce K/V heads within each layer. Cross-layer mechanisms reduce the number of distinct layer K/V sources. Let L_kv be the number of unique cache-producing layer groups:

$$M_{KV}=2BTH_{kv}d_hbL_{kv}.$$

Relative to an MHA baseline with L independent layers:

$$\text{payload ratio} =\frac{H_{kv}}{H_q}\frac{L_{kv}}{L}.$$

That product assumes every producer has the same K/V-head count and widths. The non-uniform form is the actual accounting identity:

$$M_{KV} =BTb\sum_{p\in\mathcal{P}} \left(H^K_{p}d^K_{p}+H^V_{p}d^V_{p}\right).$$

The corresponding K/V projection parameter ledger is

$$P_{KV} =d\sum_{p\in\mathcal{P}} \left(H^K_{p}d^K_{p}+H^V_{p}d^V_{p}\right),$$

before biases. These sums matter for irregular layer groups, layer-specific head counts, or different K and V widths. L_kv alone is not enough to reconstruct the checkpoint or its cache.

Persistent bytes and core attention reads are different sums. Every consumer layer still executes attention over its producer’s history:

$$D_{history,core} \approx BTb\sum_{l=0}^{L-1} \left(H^K_{source(l)}d^K_{source(l)} +H^V_{source(l)}d^V_{source(l)}\right).$$

Uniform cross-layer sharing divides the stored and projected K/V state by the sharing factor, but it does not remove the consumer layers from this read sum. A cache-capacity claim and a core-attention-latency claim therefore need different evidence.

The head-sharing sequence follows the MHA/MQA/GQA taxonomy in the book’s Figure 8.8 and the GQA paper. The layer-sharing sequence follows Cross-Layer Attention, MLKV, and YOCO. Tensor shapes, ownership maps, and the combined byte ledger are an independent Transformer System reconstruction.

The producer map is architecture state

Let source(l) return the layer whose K/V projections serve query layer l. The set of distinct cache producers is:

$$\mathcal{P}=\{\operatorname{source}(l):0\le l<L\}, \qquad L_{kv}=|\mathcal{P}|.$$

For the four-layer visual:

query layer l:    0  1  2  3
source(l):        0  0  2  2
unique producers: {0, 2}
L_kv:             2

A uniform sharing factor s can use source(l) = s * floor(l / s), giving L_kv = ceil(L / s). That is one partition, not a universal rule. Non-divisible and non-uniform designs need an explicit source table rather than inferring ownership from tensor shape.

The homogeneous formula is only a shorthand. If producer p owns H_kv(p) heads, the general persistent payload is

$$M_{KV}=2BTd_hb\sum_{p\in\mathcal P}H_{kv}(p).$$

It reduces to 2 B T d_h b L_kv H_kv only when every producer has the same head count and width. This matters for QCQA-like per-layer head budgets or MLKV configurations whose producer groups are not uniform.

For a six-layer example with source=[0,0,2,2,2,5], the producer groups have sizes two, three, and one. With B=1, T=1024, H_kv=2, d_h=64, and bf16:

unique producers:                         {0, 2, 5}
persistent K/V payload:                   1.5 MiB
same-head-count GQA without layer share:  3.0 MiB
MHA baseline with H_q=8:                 12.0 MiB
core K/V read demand across six layers:   3.0 MiB per decode step

The persistent payload is one eighth of the MHA baseline because the head-axis factor is 2/8 and the layer-axis factor is 3/6. The core read term is twice the stored payload because six consumers read histories produced by only three writers. Group size affects reuse opportunities in a schedule, but no source paper establishes that those reads disappear from HBM on an arbitrary backend.

At one decode step, a producer layer creates the new K/V row. A consumer layer still creates its own query, but reads the producer’s persistent history:

producer layer p:
    Q_p:       [B, H_q,  1, d_h]
    new K_p:   [B, H_kv, 1, d_h]
    new V_p:   [B, H_kv, 1, d_h]
    cache[p]:  [B, H_kv, T + 1, d_h] for K and V

consumer layer l, source(l) = p:
    Q_l:       [B, H_q,  1, d_h]
    K_l, V_l:  not independently projected or cached
    attention: Q_l reads cache[p]

The execution contract can be written as:

def cross_layer_decode(hidden, layer, source, cache):
    q = layer.q_proj(layer.q_norm(hidden))
    producer = source[layer.index]

    if producer == layer.index:
        kv_input = layer.kv_norm(hidden)
        new_k = layer.k_proj(kv_input)
        new_v = layer.v_proj(kv_input)
        cache[producer].append(new_k, new_v)

    k_history, v_history = cache[producer].view_committed()
    attended = grouped_attention(q, k_history, v_history)
    return layer.finish(hidden, attended)

The pseudocode is a dataflow sketch, not a drop-in implementation. The actual normalization placement and residual path belong to the checkpoint. The important invariant is that cache identity follows producer, while query and residual computation still follow layer.

Parameters, projection work, and history reads diverge

For equal Q/K/V head widths and homogeneous H_kv, every query layer retains its Q projection and output-projection columns, while only producers retain K/V projections. The attention projection parameter count becomes

$$P_{proj}=2LdH_qd_h+2L_{kv}dH_{kv}d_h.$$

The first term is Q plus O across all L query layers; the second is K plus V across the producer layers. One-token projection FLOPs under the article’s multiply-add convention are 2 P_proj. By contrast, every consumer still forms H_q attention rows, so core QK/PV work remains

$$F_{QK+PV}=2BH_qS_qT(d_{qk}+d_v)L.$$

Likewise, absent cross-layer SRAM residency or explicit fusion, logical historical K/V read demand is

$$M_{history,reads} =2BTd_hb\sum_{l=0}^{L-1}H_{kv}(\operatorname{source}(l)).$$

With a uniform head count this scales with L, not L_kv: each query layer reads the shared history for its own attention operation. Persistent storage, K/V projection work, and attention-history reads are therefore three separate ledgers.

For the running L=32, H_q=32, H_kv=8, d=4096, d_h=128, T=4096 example, adding a uniform CLA2 producer map gives:

The final two rows are why a 2x smaller resident cache is not a 2x core attention speedup claim.

This multiplicative ledger is useful, but the mechanisms are not runtime deduplication. Cross-Layer Attention, MLKV, and YOCO train explicit mappings or architectures in which multiple query layers consume shared K/V state.

The boundary warning is the important part: copying one ordinary layer’s K/V pointer into another layer is not an exact serving optimization. It changes the model unless the checkpoint was trained or adapted for that sharing rule.

Training and runtime consequences

The papers do not reach layer sharing through one interchangeable recipe:

CLA is orthogonal to MHA/MQA/GQA, but its paper’s strongest experiments combine CLA2 with MQA. It also reports that factors above two and tested non-uniform patterns had different perplexity tradeoffs. Those observations are design evidence for the paper’s 1B/3B training coordinates, not a universal optimal sharing factor.

MLKV writes the cache element count as 2 b s m g d_k, where m<L is the number of K/V-producing layers and g<H_q is the number of K/V heads per producer. That is the same two-axis product derived here, but the paper’s conversion and 5% data result belongs to Pythia-160M and its recorded hyperparameters. It does not inherit the GQA paper’s T5 quality guarantee.

YOCO must not be squeezed into source(l) as if all ordinary decoder layers shared a pointer. Its self-decoder first computes contextual memory M; one learned pair of projections creates global K_hat,V_hat; every cross-decoder layer forms its own query against that memory. Efficient self-attention still keeps bounded state, so “only cache once” means one sequence-length-scaled global cache plus self-decoder state, not literally one tensor allocation.

CLA: native adjacent-layer sharing

Cross-Layer Attention removes K/V projections from selected layers and makes those layers consume activations from an earlier producer. The paper trains 1B- and 3B-parameter models from scratch and searches attention type, head width, sharing factor, and learning rate. Its reported recommendation is empirical: sharing consecutive pairs (CLA2) was the most robust point in its experiments, especially with MQA. It is not a proof that adjacent pairs are optimal for another scale or data distribution.

The paper’s systems section makes three boundaries explicit. CLA slightly reduces parameters and forward/backward FLOPs because fewer K/V projections exist. It can reduce training K/V activations and inference cache residency. It does not directly reduce core decode-attention bandwidth, because each consumer layer still rereads the shared history. Across pipeline stages, a sharing group must be co-located or its K/V activations must cross the stage boundary.

MLKV: checkpoint conversion across both axes

MLKV studies a different migration. It starts from Pythia-160M and Pythia-410M checkpoints, mean-pools K/V weights within head groups and across source layers, increases MLP width to compensate for removed K/V parameters, then continues pretraining. Its Pythia-160M setup uses six million documents from deduplicated The Pile, packed into 2.46 million length-2048 rows, with fixed order across variants; the paper reports two A100 80GB GPUs and about 22 hours per uptraining run. Those coordinates are source-reported evidence from small models, not a portable recovery recipe.

MLKV’s notation permits m cache-producing layer groups and g K/V heads per producer, giving 2*b*s*m*g*d_k cache elements under its equal-width assumptions. It reports useful quality/memory points around sharing factors two and three, but also reports collapse at its most extreme one-head/one-producer configuration. Its generalized implementation did not show a significant speedup from layer sharing; the paper itself notes that every layer still fetches the history. That negative runtime result belongs beside the memory claim.

YOCO: one global cache is not a CLA source map

YOCO is a native decoder-decoder architecture. Its first half is a self-decoder using bounded-state efficient self-attention such as gated retention or sliding-window attention. The output of that stack produces one global K/V object:

$$\hat K=\operatorname{LN}(X^{L/2})W_K, \qquad \hat V=\operatorname{LN}(X^{L/2})W_V.$$

Every cross-decoder layer owns a distinct query projection but attends to that same global object. The paper’s “cache once” shorthand therefore means one global full-context K/V cache plus bounded self-decoder state, not literally zero state in the first half. Its prefill can skip the cross-decoder because the global cache depends on the self-decoder output, whereas generation still runs the cross-decoder. That dataflow cannot be reproduced by pointing ordinary Transformer layers at layer zero’s cache.

These three contracts lead to different release artifacts:

CLA:   natively trained producer map
MLKV:  converted + uptrained producer/head map
YOCO:  natively trained self-decoder/cross-decoder graph

Only CLA and MLKV fit the source(l) abstraction directly. YOCO remains in the comparison because it reduces the layer axis, but its cache producer and prefill graph require a different runtime ABI.

Layer sharing also changes placement. If a producer and one of its consumer layers land in different pipeline stages, the system must either co-locate the sharing group or transfer the producer’s K/V activations across the stage boundary. The smaller persistent cache does not erase that dependency.

For a producer p whose consumers span c_p additional stages, distributing one new decode row has payload lower bound

$$M_{stage,append}=2BH_{kv}d_hb\sum_{p\in\mathcal P}c_p.$$

Sending the full history instead scales the same expression by T; retaining received rows builds a stage-local replica and must be counted as physical cache. For the worked CLA2 map, placing each producer/consumer pair together gives zero cross-stage K/V payload. Splitting every pair across stages gives 16 crossings and 64 KiB of new-row payload per request per token before protocol overhead. Pipeline placement is thus part of the architecture deployment receipt, not a scheduler afterthought.

Finally, cache bytes and core attention reads are different ledgers. The CLA paper notes that every consumer attention layer still reads the shared history for its own query computation. Reducing L_kv therefore reduces persistent storage and K/V projection work, but does not automatically divide the attention kernel’s historical K/V traffic by the same factor.

0x09 Shape And Byte Ledger

Tensor shapes

Concrete fp16/bf16 cache ledger

Assume L=32, H_q=32, d_h=128, b=2:

The table excludes block padding, metadata, workspace, scale tensors, and parallel replicas. Those must be added to a physical-residency measurement.

0x0A Pinned Runtime Source Trace

The architecture equations establish what a correct grouped implementation must preserve. They do not establish which path a framework dispatches, whether that path expands K/V physically, or how many HBM transactions a kernel issues. Those are separate source, allocation, and profiler questions. An API flag is not an allocation receipt. The executable Loop 10 evidence therefore follows a five-rung ladder:

Original Transformer System reconstruction. Higher rungs depend on the lower ones, but none can be inferred from an API name alone.

The trace pins the exact files, revisions, hashes, line ranges, and markers used by the local audit:

The full machine-readable receipt is runtime_source_trace.json. Its distinction is deliberate:

API acceptance:
    the shape and metadata are legal at the public boundary

source trace:
    a pinned implementation exposes a particular expansion or direct path

numerical parity:
    reference expansion and direct grouping agree on outputs

allocation evidence:
    the measured execution did or did not create repeated storage

traffic evidence:
    profiler counters establish memory-system behavior

latency/SLO evidence:
    the full workload benefits after dispatch, launch, and scheduler costs

Loop 10’s CPU oracle closes only the numerical and explicit-allocation parts of that ladder. It compares materialized repetition, direct owner-map indexing, and grouped-by-K/V-head execution under one immutable manifest. Its exact parity and zero explicit repeat allocation for the direct path are local executable facts. They are not GPU traffic, kernel-selection, or production-latency facts.

0x0B Validation Protocol

An architecture-level cache claim should pass four gates.

1. Numerical parity

For the same grouped checkpoint and mask:

reference:
    explicitly repeat K/V to H_q and run ordinary attention

candidate:
    run grouped kernel directly on H_kv cache

compare:
    logits/output under stated dtype tolerances

Test S_q=1, short and long T, mixed sequence lengths, causal masks, padding, and each supported H_q/H_kv ratio.

2. Allocation parity

Inspect persistent allocation after prefill and decode:

expected payload proportional to:
    2 * B * T * L * H_kv * d_h * b

unexpected:
    persistent H_q-sized repeated cache

3. Traffic evidence

Profile bytes and cache behavior. A correct grouped implementation should not load the same K/V tile once independently for every query head when reuse can be kept in one work unit or cache level.

4. End-to-end quality and SLOs

For converted models, evaluate task quality before and after adaptation. For serving, sweep context length, batch/active sequences, TP degree, and backend. Report TPOT distributions and goodput, not only one isolated kernel maximum.

The evaluation manifest is immutable

A valid comparison changes the intended GQA treatment while freezing every unrelated coordinate. Record the baseline and candidate identities before any measurement:

EvaluationManifest
    dataset_manifest_hash
    prompt_and_request_order_hash
    tokenizer_hash
    quality_tasks_and_metrics

    baseline_checkpoint_hash
    candidate_checkpoint_hash
    architecture_hash
    query_to_kv_map_hash

    backend_revision
    kernel_plan_hash
    hardware_and_topology
    q_dtype, kv_dtype, accumulator_dtype

    context_lengths
    active_request_counts
    query_lengths
    page_sizes
    tensor_parallel_degrees
    warmup_and_repetition_policy
    SLO_thresholds

The converted initialization and the uptrained checkpoint are separate candidates. Otherwise recovery from continued training can be attributed to mean pooling, or a poor initialization can be misreported as the final GQA quality point.

Keep the complete lineage rather than only the two endpoints:

source checkpoint
    -> conversion recipe and converted initialization
    -> training run and uptrained checkpoint
    -> runtime plan
    -> released generation

Each arrow has an input identity and an output identity. Reusing one identifier for the converted initialization, trained checkpoint, and release would erase which stage produced a quality recovery or runtime regression.

Quality and numerical gates answer different questions

Numerical parity compares two implementations of the same grouped checkpoint. Quality evaluation compares different model artifacts:

numerical gate:
    expanded reference versus direct grouped backend
    fixed checkpoint, tokens, positions, mask, and cache history

quality gate:
    source MHA versus converted initialization versus uptrained GQA
    fixed tokenizer, dataset split, prompts, decoding policy, and metric code

For a loss or perplexity metric whose lower value is better, define relative regression against the approved baseline:

$$\Delta_{quality} =\frac{Q_{candidate}-Q_{baseline}}{Q_{baseline}}.$$

The gate is Delta_quality <= epsilon_quality, with the tolerance declared per task before evaluation. Averaging can hide a catastrophic slice, so the global gate and every protected slice must pass. A source-reported paper frontier can help choose candidate tolerances; it is not a locally reproduced result.

Reconcile logical payload, physical residency, and measured traffic

The expected physical allocation for one workload slice is:

$$M_{expected} =M_{logical}R_{replica} +M_{padding}+M_{metadata}+M_{scales}+M_{workspace}.$$

Compare an allocator snapshot to that full ledger:

$$\epsilon_{resident} =\frac{|M_{observed}-M_{expected}|}{M_{expected}}.$$

Passing the byte gate means the allocator agrees with the declared placement and cache ABI within a predeclared tolerance. It does not establish efficient reads. A profiler supplies a separate read-amplification coordinate:

$$A_{HBM} =\frac{D_{HBM,measured}}{D_{history,logical}}.$$

The denominator is the article’s logical unique-history demand for the exact request batch. Metadata, page-table loads, partial reductions, output traffic, and compulsory effects can make the ratio exceed one. The useful comparison is the ratio across pinned backends and shapes, accompanied by profiler counters, not a universal expectation that A_HBM = 1.

Tail latency and goodput close the serving claim

Measure prefill and decode separately, then measure the integrated service. At minimum retain:

For interval W, goodput is:

$$G_{SLO}=\frac{N_{completed\ within\ SLO}}{W}.$$

Promotion can require both:

candidate_TPOT_p99 <= baseline_TPOT_p99 * (1 + epsilon_tail)
candidate_goodput >= baseline_goodput * (1 - epsilon_goodput)

A mean kernel speedup cannot compensate for worse p99 or lower goodput unless the release policy explicitly accepts that tradeoff.

Original Transformer System reconstruction. No single green metric promotes a candidate; the publication transaction consumes all receipts and retains the prior generation for rollback.

Sensitivity matrix: the mechanism must survive coordinate changes

The minimum sweep crosses architecture and workload axes rather than varying one convenient prompt:

The release result is attached to this matrix. It must not be generalized to an untested topology or maximum context.

Publication, canary, and rollback are one transaction

A GQA candidate changes checkpoint rows, head ownership, cache shape, and often the kernel plan. Publish it under a fresh immutable identity:

CandidateGeneration
    checkpoint_hash
    gqa_manifest_hash
    conversion_and_uptraining_receipts
    backend_revision
    kernel_plan_hash
    evaluation_manifest_hash
    evaluation_receipt_hash
    cache_generation_namespace

The release sequence is:

def publish_gqa(candidate, baseline, gates):
    assert all(gate.status == "pass" for gate in gates)
    assert candidate.cache_namespace != baseline.cache_namespace
    assert candidate.plan_hash != ""

    warm_candidate_without_sharing_baseline_cache(candidate)
    canary = route_fraction(candidate, fraction=0.01)
    monitor(canary, quality=True, errors=True, bytes=True,
            tpot_tail=True, goodput=True)
    if any_rollback_trigger_fires(canary):
        stop_admission(candidate)
        drain_or_abort_candidate_requests_by_policy()
        route_new_requests(baseline)
        invalidate_candidate_cache_namespace()
        return "rolled_back"
    return staged_promotion(candidate)

Rollback restores the previous checkpoint and its matching architecture, backend plan, and cache namespace. Reusing candidate pages under the baseline would reinterpret H_kv, head order, layout, or generation state. Existing candidate requests may drain only while their complete generation remains loaded and isolated.

The rollback controller must fence candidate pages before the previous route becomes authoritative. Fencing means new lookups cannot resolve the candidate namespace; it does not rename those pages into baseline state.

The ordered rollback receipt is:

STOP_NEW_CANDIDATE_ADMISSION
PIN_PREVIOUS_ACTIVE_RELEASE
DRAIN_OR_ABORT_CANDIDATE_REQUESTS
FENCE_CANDIDATE_CACHE_NAMESPACE
RESTORE_PREVIOUS_ROUTE
RETAIN_FAILED_RECEIPTS_FOR_AUDIT
ROLLBACK_COMPLETE

Failure injection must exercise stale plans, incompatible cache namespaces, direct-kernel dispatch failure, fallback materialization, quality and traffic gate failures, and a canary SLO failure. Safe recovery is a compatible fallback, recomputation, blocked promotion, or rollback; it is never cache reinterpretation.

Predeclare triggers for numerical divergence, quality canaries, allocator-byte drift, backend fallback, error rate, TPOT p99, and goodput. A rollback path that exists only after an incident is not a release contract.

0x0C Failure Lab

0x0D Claims, Boundary, And Handoff

Externally grounded claims:

• MQA shares K/V heads while retaining multiple query heads and was proposed to reduce incremental-decoding memory bandwidth.
• GQA interpolates between MHA and MQA, and its paper studies conversion plus uptraining from MHA checkpoints.
• FlashAttention-2 describes index-based MQA/GQA support instead of requiring physical K/V duplication, and FlashInfer describes head-group fusion to reuse shared K/V for short queries.
• Current PyTorch SDPA exposes experimental GQA support with explicit head-count constraints and backend limitations.

Local derivations and synthesis:

• ideal persistent cache reduction relative to MHA is H_q / H_kv = g;
• ideal K/V-only decode arithmetic intensity is approximately 2g/b FLOPs per byte when each stored K/V scalar is read once;
• physical expansion or redundant per-query-head reads can erase that runtime benefit without changing model outputs;
• tensor-parallel replication and communication must be added to logical cache bytes before making a placement claim.

M5.7 stops at full K/V rows with fewer owners:

M5.7:
    cache [H_kv, T, d_h] K and V rows
    make several query heads reuse each stored row

M5.8:
    replace those rows with latent content plus a RoPE component
    absorb reconstruction into the query/output path
    execute the representation with MLA-specific kernels

0x0E Publish Gate

References

Local source articles:

• Transformer Series A27 - MQA and GQA
• M3.2 - Multi-Head Attention: Projection and Head Layout
• M4.5 - Decode Attention Kernels and Backend Boundaries
• M5.1 - Autoregressive Serving State
• M5.3 - KV Memory Layout and Exact Reuse
• M5.6 - Prefill/Decode Disaggregation and KV Placement
• M5.8 - MLA Latent Cache and Absorbed Kernels

Primary and official sources:

• Fast Transformer Decoding: One Write-Head is All You Need
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving
• PyTorch scaled dot product attention documentation
• PyTorch variable-length attention documentation
• FlashInfer single-decode API
• Transformers Llama source at v5.14.1
• Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
• MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding
• You Only Cache Once: Decoder-Decoder Architectures for Language Models