Shared KV Cache

The axis nobody was using: layers

Attention builds three projections of every token — queries, keys, and values — and during autoregressive decoding the keys and values of all past tokens must be kept around so each new token can attend to them. That store is the KV cache, and its size is the sum of many independent copies: one per layer, per KV head, per token. The classic optimizations all attacked the head dimension: Multi-Query Attention (MQA) collapses all heads in a layer to one shared K and V, and Grouped-Query Attention (GQA) is the middle ground, a small number of KV heads each serving a group of query heads.

But there is a second dimension sitting untouched: the layer count L. In a standard transformer every one of the L layers computes and caches its own keys and values independently. Empirically, the keys and values adjacent layers compute are highly correlated — nearby layers read the context in similar ways. Cross-layer sharing exploits that redundancy: compute one KV set and let several layers use it. Because L is a factor in the cache size just like the head count is, shrinking the effective number of cached layers shrinks the cache proportionally.

The KV-cache formula and its four knobs

Write the cache size in bytes for a single sequence and it factors into exactly the knobs you can turn:

KV_bytes = 2 · b · N · L_kv · H_kv · d_h · p

  2     = one copy for K, one for V
  b     = batch size (concurrent sequences)
  N     = context length (tokens cached)
  L_kv  = number of DISTINCT cached KV layer-sets  ← layer axis
  H_kv  = number of KV heads per layer             ← head axis (GQA/MQA)
  d_h   = head dimension
  p     = bytes per element                        ← precision/quantization

Plain multi-head attention (MHA) sets L_kv = L (every layer distinct), H_kv = H (every head distinct), and p = 2 (fp16/bf16). Each optimization simply lowers one factor. GQA/MQA lower H_kv. KV quantization lowers p to 1 byte (int8) or 0.5 (int4). Cross-layer sharing lowers L_kv — and because the factors multiply, it stacks on the others rather than competing with them. That independence is the whole reason it is worth adding: it opens a factor no other technique touches.

GQA and MQA: sharing across heads (what this is NOT)

It is worth pinning down the contrast, because the two ideas are easy to blur. GQA and MQA share within a single layer, along the head axis. A layer with H = 32 query heads under MQA keeps just one K and one V broadcast to all 32 heads — a 32× cut on that layer’s KV; GQA with 8 KV groups keeps 8 K/V sets, a 4× cut that trades some of MQA’s quality loss back for capacity. Crucially, every layer still has its own KV store; you have shrunk each layer’s footprint but still pay for L of them.

Cross-layer sharing works on the perpendicular axis. It leaves the per-layer head structure alone — a shared KV set can itself be an MQA or GQA set — and instead reduces how many layers own a distinct copy. Where GQA asks “how many KV heads does this layer need?”, cross-layer sharing asks “how many layers need their own KV at all?” The questions are independent, which is why you apply both.

Cross-Layer Attention (CLA)

Cross-Layer Attention (Brandon et al., 2024) is the most direct expression of the idea. Partition the stack into groups of s adjacent layers. Within a group, one layer computes K and V; the other s − 1 layers skip their key/value projections and attend to the shared KV set. Queries stay per-layer — every layer still forms its own queries and output projection, so the layers are not identical — but they read a common set of keys and values.

The effect on the formula is immediate: L_kv = L / s. The common choice CLA2 (s = 2, sharing across pairs) halves the cached layer count and so halves the KV cache; CLA3 thirds it. Because CLA only removes KV projections — it does not touch the attention computation — it drops straight into an existing MHA/GQA implementation with no special kernels. The paper’s central result is a Pareto improvement: at fixed accuracy, CLA2 with MQA reaches roughly half the KV memory of MQA alone, freeing memory for longer context or larger batches.

YOCO: You Only Cache Once

YOCO (Sun et al., 2024) pushes sharing to its limit with a decoder–decoder design. The bottom half — the self-decoder — uses an efficient attention whose state is bounded regardless of context length (sliding-window or a retention/linear-attention variant), and emits a single global KV cache. The top half — the cross-decoder — does no self-attention; every one of its layers cross-attends to that one global KV set. Keys and values are, in effect, cached once for the whole upper stack rather than once per layer.

In the formula, the cross-decoder’s contribution collapses to L_kv ≈ 1 instead of L/2, and the self-decoder’s state is a small, context-independent window rather than a growing per-layer cache. Net KV memory becomes nearly independent of depth — roughly an L-fold reduction for a deep model — and prefill gets dramatically faster, because the global cache is built once instead of layer by layer. YOCO is a heavier commitment than CLA: an architecture you train from scratch, not a drop-in edit.

The memory math, unified

Both schemes are the same substitution — replace L with a smaller effective L_kv — so one expression covers them:

reduction_factor  =  L / L_kv

  MHA / GQA :  L_kv = L          →  factor 1   (no layer sharing)
  CLA, group s :  L_kv = L / s   →  factor s   (CLA2 → 2×)
  YOCO      :  L_kv ≈ 1     →  factor ~L  (near depth-independent)

The layer factor multiplies whatever the head and precision axes already gave you. Total cache relative to a plain MHA baseline is the product of three ratios:

cache / MHA_cache  =  (H_kv / H) × (L_kv / L) × (p / 2)
                        head axis      layer axis    precision axis

Nothing in one ratio constrains the others, which is why a well-tuned stack reaches for all three at once: GQA for heads, CLA or YOCO for layers, int8 for precision. Each is a modest cut alone; multiplied, they are the difference between a cache that fits in memory and one that does not. Note the ordering: GQA and CLA are architectural, baked in at training time, while quantization is a post-hoc deployment knob — so fix the sharing factors first, train, then treat KV quantization as the elastic final multiplier.

A worked example

Take a 32-layer model: hidden size 4096, H = 32 query heads, head dimension d_h = 128. Serve one sequence (b = 1) at N = 8192 tokens in fp16 (p = 2). Start from plain MHA:

MHA:  2 · 1 · 8192 · 32 · 32 · 128 · 2
      = 4,294,967,296 bytes  ≈  4.0 GB

Now turn the knobs one at a time. GQA with H_kv = 8 KV heads divides by 4. Add CLA2 (L_kv = 16) and divide by another 2. Add int8 KV quantization (p = 1) and divide by another 2:

GQA-8        : 4.0 GB / 4          =  1.0 GB
+ CLA2       : 1.0 GB / 2          =  512 MB
+ int8 KV    : 512 MB / 2          =  256 MB    (16× vs MHA)

YOCO instead : L_kv ≈ 1 not 16  →  ~32 MB from the GQA-8 point (~128×)

The 4× from GQA and the 2× from CLA2 do not overlap — they act on different factors — so they compound to 8× before quantization even enters. YOCO’s near-depth-independent cache is in a different league again, at the cost of a bespoke architecture.

The quality trade

None of this is free — you are removing representational capacity, and the question is only how much accuracy it costs per byte saved. The empirical pattern is consistent: mild sharing is nearly free, aggressive sharing bites. CLA2 lands close to the baseline on the memory–quality Pareto frontier, often recoverable by spending a fraction of the saved memory on more parameters. CLA3 and beyond, which force three or more layers to agree on one KV set, degrade more visibly, because distant layers genuinely want to attend differently.

YOCO makes a larger bet and, at scale, reportedly matches a standard transformer baseline while slashing memory and prefill cost — but it does so by changing the architecture, so its quality profile is its own, not a knob on a vanilla model. The honest framing: layer sharing trades a small, tunable amount of quality for a large, multiplicative memory win, and the sweet spot for most deployments is the gentle end.

Practical implications and pitfalls

For CPU and small-model serving, the layer axis is quietly the highest-leverage one. On a memory-bandwidth-bound decode — exactly the regime a CPU runs in — the KV cache is not just a capacity limit, it is traffic: every generated token re-reads the whole cache, so halving L_kv with CLA2 halves that per-token read and can directly speed up decode, not merely fit more context.

The pitfalls are mostly conceptual. First, do not confuse this with request-level prefix sharing — reusing a common system prompt’s KV across users — which shrinks the cache along the b and N axes and is fully complementary. Second, CLA and YOCO are training-time choices; you cannot bolt them onto released weights, so plan them up front. Third, watch the sharing factor: pushing CLA3/CLA4 for a bigger number usually costs more quality than the memory is worth. Start at CLA2, measure, and go deeper only if the accuracy budget allows.