The KV Cache: Sizing It, and the Attention Variants That Shrink It

Abstract

Autoregressive decoding is cheap only because it does not recompute the past: every generated token reuses the keys and values of every token before it. Those tensors are the KV cache, and its size follows from six numbers with no slack in the derivation. This article works out that formula, shows why the cache — not the weights — sets the batch-size ceiling at long context, and derives the arithmetic intensity that makes decode memory-bound. It then examines the architectural response: multi-query and grouped-query attention are two points on the $H_{kv}$ axis, and multi-head latent attention is a different move that leaves the head count alone and compresses what is stored instead. Finally it measures what those variants actually buy, and finds that an eight-times reduction in bytes returns about 2.4× in latency. Two candidate explanations are tested and rejected — access granularity and parallelism — leaving the work done per byte as the cause. It assumes the decoding loop from LLM Inference Introduction.

1. Why the Cache Exists

Decoding generates one token at a time. Producing token $t+1$ requires attention over all previous positions, which needs their keys and values — and those depend only on tokens already fixed, so they never change once computed.

Without a cache, generating the token at step $t$ re-runs attention over all $t$ previous positions, doing $O(t^2 d)$ attention work and re-reading every weight matrix each step. Over a run of $T$ tokens that sums to $O(T^3 d)$ — cubic, not merely quadratic. With a cache, each step computes one new key and value, appends them, and attends over stored state, which brings the run down to $O(T^2 d)$. The cache does not shave a constant factor off a recomputation; it changes the complexity class, trading a linear amount of storage for it.

That is the trade the rest of this article is about: the storage is not small.

2. Sizing the Cache

Each layer stores one key vector and one value vector per position per KV head. For a model with $L$ layers, $H_{kv}$ key-value heads, head dimension $d_h$, and $b$ bytes per element, a sequence of $T$ tokens costs

Everything except $T$ is fixed by the architecture, so the useful form is bytes per token:

and total cost is that constant times the number of live tokens, summed over every sequence in the batch.

Worked examples. In BF16 ($b = 2$):

The middle row is the reason grouped-query attention exists. A 70B model in BF16 has 140 GB of weights, which already needs two 80 GB GPUs; with multi-head attention, one 8K-token sequence would add 21.5 GB, and a batch of 32 would need 687 GB of cache — nearly five times the weights. The bottom row is the same model as actually shipped: the cache falls to an eighth, so that batch of 32 needs 86 GB, which one additional GPU covers.

Three properties of the formula are worth stating plainly.

It is linear in context length. Doubling the context doubles the cache. There is no reuse to exploit; every position must be retained.

It is linear in batch size, with no sharing. Two sequences share weights but never share cache — except for a common prefix, which is exactly the case PagedAttention makes exploitable.

It does not depend on the number of query heads. Only $H_{kv}$ appears. That is the loophole the next section exploits.

The 512 KB figure in the first row is not just arithmetic. Allocating one layer’s cache for one token as a (2, 1, 32, 128) BF16 tensor gives 16,384 bytes, and multiplying by 32 layers gives 524,288 bytes — exactly 512 KB. Every capacity number later in this series is built from that constant, so it is worth confirming once that the formula and the tensor layout agree.

3. Why Decode Is Memory-Bound

The cache is not just a capacity problem. It is read in its entirety on every decode step, and that read is what sets the speed.

Consider one layer of one decode step. Attention over $T$ cached positions performs $H_q T d_h$ multiply-accumulates for the scores and the same again for the value-weighted sum, so

Their ratio is the arithmetic intensity of attention at decode time:

The context length cancels. What survives is the ratio of query heads to KV heads — the group size — divided by the element width. In BF16 with multi-head attention, where $H_q = H_{kv}$, the intensity is exactly 1 FLOP per byte, against a ridge point of roughly 295 on an H100 (2,000 TFLOP/s BF16 over 3.35 TB/s HBM). Attention at decode time runs at about a three-hundredth of the machine’s arithmetic capability, and no amount of kernel tuning changes that; the operator is reading memory, not computing. The ratio $2g/b$ is hardware-independent; only the ridge point it is compared against changes from one GPU to the next.

That ratio is measurable rather than merely derivable. On a B300 (2,250 TFLOP/s BF16, 8 TB/s HBM, so a ridge point of 281 FLOP/byte), a paged decode kernel with $H_q = H_{kv} = 8$ and $d_h = 128$ computes at an arithmetic intensity of 1.0 FLOP/byte at every batch size and context length tested, and reaches at most 75% of HBM peak — while prefill at the same head dimension reaches 78% of compute peak.

One consequence is easy to get wrong. Increasing the batch size does not change the intensity. Batching amortizes weight reads across sequences, which is why it works so well for the feed-forward layers, but every sequence has its own KV cache, so the attention read scales with the batch exactly as the arithmetic does. Attention stays at the same intensity no matter how many requests are in flight.

What batching does change is whether the memory system is actually saturated. Those are different statements, and the distinction matters. Measured on the same kernel, at a context length of 16,384:

Intensity never leaves 1.0 FLOP/byte, yet throughput improves 51× per request. The intensity fixes the ceiling — you can never escape the memory roof — while the batch size decides how close to that ceiling you get. A single sequence leaves the HBM channels mostly idle; thirty-two concurrent sequences keep them busy. This is why serving throughput is a scheduling problem, not just a kernel problem.

4. The $H_{kv}$ axis: MHA, MQA, GQA

The three variants in this section are one idea at three strengths: reduce how many key-value heads are stored. Nothing else in the formula is available — $L$ and $d_h$ are the model, $T$ is the request, $b$ is the storage format — so $H_{kv}$ is the only quantity left to move. A fourth variant that moves something else follows in the next section.

4.1 Multi-head attention (MHA)

Every query head gets its own key and value. Eight query heads, eight KV heads, eight caches per token:

  q0   q1   q2   q3   q4   q5   q6   q7      query heads (H_q = 8)
   |    |    |    |    |    |    |    |
  k0   k1   k2   k3   k4   k5   k6   k7      KV heads  (H_kv = 8)
   |    |    |    |    |    |    |    |
  c0   c1   c2   c3   c4   c5   c6   c7      stored per token

  group size g = 1        KV heads stored = 8        intensity = 2/b

The original: $H_{kv} = H_q$, every head with its own projection. Maximum expressiveness, maximum cache. For the 70B model above it is 2.5 MB per token, which is unaffordable at any serious batch size or context length.

4.2 Multi-query attention (MQA)

Keep all eight query heads, share a single key and value among them:

  q0   q1   q2   q3   q4   q5   q6   q7      query heads (H_q = 8)
   |    |    |    |    |    |    |    |
   +----+----+----+----+----+----+----+---->
                        |
                       k0                    KV heads  (H_kv = 1)
                        |
                       c0                    stored per token

  group size g = 8        KV heads stored = 1        intensity = 16/b

Proposed in 2019 for exactly this reason: keep all $H_q$ query heads but project a single shared key and value. The cache falls by a factor of $H_q$ — 64× for the model above — and decode intensity rises by the same factor.

The cost is quality. Collapsing to one KV head is a real reduction in what attention can represent, and models trained this way show measurable degradation and, in some reports, training instability. MQA is the right answer when memory is the binding constraint and the quality loss is acceptable.

4.3 Grouped-query attention (GQA)

The interpolation, and now the default. Partition the query heads into groups; every head within a group shares one key-value pair. With two groups:

  q0   q1   q2   q3   q4   q5   q6   q7      query heads (H_q = 8)
   |    |    |    |    |    |    |    |
   +----+----+    |    |    +----+----+
        |         |    |         |
       k0         |    |        k1           KV heads  (H_kv = 2)
        |         |    |         |
       c0                     c1             stored per token

  group size g = 4        KV heads stored = 2        intensity = 8/b

  g = 1  <-----------------+-----------------> g = H_q
     MHA                  GQA                   MQA
                        (g = 8 typical)

At $g = 8$, the cache is one eighth of MHA and the intensity is eight times better, with quality close to MHA.

Two practical properties explain its adoption. First, the natural choice is $H_{kv}$ equal to the tensor-parallel degree, so each GPU holds exactly one KV head and no cache is replicated across devices. Second, an existing MHA checkpoint can be converted rather than retrained: mean-pool the key and value projections within each group and continue training for a small fraction of the original budget. Llama 2’s 70B model, Mistral, and most models released since use GQA.

4.4 The variants side by side

MLA is deliberately absent from that table; it is not on this axis and gets its own section next.

The formulas are worth applying to one concrete model rather than three separate ones. The 70B shape from §2 ($L = 80$, $H_q = 64$, $d_h = 128$, BF16) at 8K context:

The last column is the point of the whole article. Going from MHA to GQA makes the cache eight times smaller and the intensity eight times higher, so the same kernel reads an eighth as many bytes while the ratio of arithmetic to bytes rises by the same factor. The bytes are what cost time, so capacity and speed improve together.

But only up to a point, and the limit is worth stating here rather than leaving to the measurement section: an eight-times reduction in bytes does not buy an eight-times reduction in latency. Measured, it buys about 2.4×.

4.5 What the variants do to the weights

Everything so far has been about the cache. The variants also change the model’s parameter shapes, and they do so asymmetrically — which is the part that is easy to get wrong. Only the key and value projections shrink. The query and output projections keep their full width in every variant.

With $D = H_q d_h$ the model width:

The reason is structural rather than incidental. There are always $H_q$ query heads, so $W_Q$ must emit $H_q d_h = D$ columns, and $W_O$ must consume $D$ and return $D$. Only the K and V projections emit one column block per KV head, and there are fewer of those.

It is worth being precise here because the phrasing “each head has its own $W_Q$, $W_K$, $W_V$” invites a wrong mental picture. In code there is one matrix, and the heads are slices of its output:

# MHA: Q, K, V all project to D columns
self.W_Q = nn.Linear(D, D)
self.W_K = nn.Linear(D, D)
Q = self.W_Q(X).view(B, L, H_q, d_h).transpose(1, 2)   # heads = column blocks

# GQA: only K and V get narrower
self.W_K = nn.Linear(D, H_kv * d_h)          # D/g columns worth
K = self.W_K(X).view(B, L, H_kv, d_h).transpose(1, 2)

Splitting one $(D, D)$ weight into $H_q$ column blocks is exactly equivalent to having $H_q$ separate $(D, d_h)$ matrices, because a linear layer’s output columns are independent. It is also why production stacks fuse the three projections into one qkv_proj: one GEMM beats three.

Applied to the 70B shape, the asymmetry shows up as a real parameter difference:

GQA saves 9.4 billion parameters against multi-head — roughly 13% of a 70B model. That is a larger effect than it first appears, because attention is being compared against itself; the FFN in the same layer is 704 M parameters, so GQA attention (151 M) is about a fifth of the block.

One last consequence is practical, and it explains why the checkpoint-conversion trick in §4.3 works at all. Moving down this table is a reshape: going from MHA to GQA discards columns within each group and averages what remains, which is exactly the mean-pooling operation used to initialise from an existing checkpoint. Moving up, or going to MQA, is not a reshape — 87% of $W_K$’s columns have no counterpart — so those models are trained from scratch. Shape compatibility runs one way only.

5. Compressing instead of dropping heads: MLA

The three variants above sit on a single axis: how many KV heads get stored. Multi-head latent attention is not on that axis. It leaves $H_{kv}$ alone and shrinks what is stored per head, so the parameter-shape argument above does not apply to it and neither does the group-size machinery.

Instead of reducing the number of KV heads, MLA compresses keys and values jointly into a low-rank latent vector and caches that, reconstructing per-head keys and values on the fly. The cached object is one vector of dimension $d_c$ per token per layer, with $d_c$ far smaller than $H_{kv} d_h$, plus a small decoupled component $d_r$ carrying rotary position information that cannot be compressed away.

DeepSeek-V2 makes those dimensions concrete: $d_c = 512$ compressed plus $d_r = 64$ decoupled RoPE. That is 576 elements per token per layer, against $2 H_{kv} d_h = 2 \cdot 128 \cdot 128 = 32{,}768$ for an equivalent multi-head layer — about 1.8% of the naive per-layer cache. The paper reports roughly a 93% end-to-end reduction; the two figures measure different baselines, since the end-to-end number compares against a sibling model with its own layer count and head configuration rather than against a like-for-like MHA layer. The trade is arithmetic: the reconstruction is extra work on every step, and the kernel is more complex than either MHA or GQA.

  stored per token, per layer:

     MHA                     MLA
  +---------+            +---------+
  |   k0    |            |         |
  |   k1    |            |  latent |  d_c = 512   compressed K and V
  |   ...   |   ---->    |         |
  |  k_{63} |            +---------+
  |   v0    |            |  rope   |  d_r = 64    cannot be compressed
  |   ...   |            +---------+
  |  v_{63} |
  +---------+            576 elements
  32,768 elements

  Heads stay at 64. The bytes are compressed, not the head count.

The diagram marks the difference that matters for everything after this section: MHA, GQA, and MQA are all points on one axis — how many KV heads get stored — whereas MLA shrinks the representation instead, keeping the head count where it was. It is the only variant here that does not change $H_{kv}$.

That difference extends to the weights, and in the opposite direction from the previous section. GQA removes parameters by narrowing $W_K$ and $W_V$. MLA adds them: a down-projection $W_{DKV}$ from $D$ to the latent width $d_c$, plus up-projections $W_{UK}$ and $W_{UV}$ back out to per-head keys and values. Fewer bytes are stored but more matrices exist. In DeepSeek-V2 those up-projections are folded algebraically into $W_Q$ and $W_O$ at inference time, so the extra cost is mostly a training-time and kernel-complexity cost rather than a serving one — which is why MLA can be worth it despite having more parameters than the GQA model it replaces.

6. Measuring the Variants

The formulas above predict that cutting bytes per token cuts decode time in proportion. That prediction is testable, and testing it turns out to be more informative than confirming it.

Sweep $H_{kv} \in \{8, 4, 2, 1\}$ with $H_q = 8$ fixed, so the group size takes the values a real model would choose: $g = 1$ is MHA, $g = 8$ is MQA. Same kernel, same page size, same dtype; only the head axis of the cache changes. At a 16K context and batch 32:

Bytes fall by exactly the factor the formula predicts: eight. Latency falls by 2.38×. The prediction is wrong by a factor of 3.4, and it is not wrong because the machine was idle — repeating the sweep at batch 64 and 64K context, where the baseline already sustains 88.9% of peak, still yields only 2.12×.

The reason is visible in the bandwidth column. Achieved bandwidth falls as the group size rises: from 74.6% of HBM peak at $g=1$ to 22.1% at $g=8$. The tensor does shrink exactly as promised; what does not shrink proportionally is the memory system’s ability to stream it.

The obvious explanation — that narrower reads coalesce worse — turns out to be wrong, and it is worth showing the falsification rather than asserting the answer. In the standard NHD layout the contiguous run for one token is $H_{kv} d_h b$ bytes, so the two candidates are access granularity and parallelism. Both can be varied independently.

Hold the head count at one and widen the contiguous run by raising $d_h$, and bandwidth does not move at all:

An eight-times wider slice changed bandwidth by one percentage point, while latency tracked bytes exactly. Access granularity is not the cause. Nor is parallelism: at fixed $H_{kv}=1$, going from 16 to 512 concurrent requests — 32× more independent streams — moved bandwidth only from 19.7% to 24.0%.

What does track bandwidth is the work per byte. For each KV element a decode kernel reads, it performs $g$ dot products and $g$ exponentiations, so raising $g$ raises the arithmetic per byte in exact proportion to the byte saving. At $g=8$ the kernel is doing eight times the work per byte that it does at $g=1$, and the exponentiation path — not HBM — sets the pace. That is the same effect the later schedule article describes, where exponential throughput becomes limiting on hardware whose matrix units outgrew it, here showing up in a serving kernel rather than a prefill kernel.

Three consequences follow, and they are more useful than the prediction was.

This is why production GQA stops at $g = 8$. Beyond that point the per-byte arithmetic dominates and further sharing stops paying: the last doubling of $g$ from 4 to 8 bought 17% of latency while costing more than half the bandwidth efficiency. That is a stronger argument for GQA-as-default than “quality is close to MHA”, and one that no model-quality study would have surfaced.

It suggests where MLA sits, though this article does not measure it. MLA reduces bytes by compression rather than by dropping heads, so its per-token record is a single 576-element latent — and with $H_{kv}$ left alone, its work per byte should be much closer to MHA than to MQA. If that holds, MLA is the one variant that takes the byte saving without the arithmetic penalty, at the cost of the reconstruction the previous section describes. Whether it does is a question for the serving article.

It generalises past attention, and it is testable. Reducing bytes while keeping the work per byte constant — quantization, layer sharing, windowing — should not pay this penalty. Reducing bytes by increasing work per byte, as group size does, should. That distinction predicts the shape of the curve for any future variant, and it is a better guide than “fewer bytes is faster”.

7. The Other Levers

Changing the architecture is not the only way to shrink the cache, and the alternatives compose with it.

Quantize it. The cache is a tensor like any other, and quantizing it to INT8 or INT4 halves or quarters it directly. The effect on both formulas is exact and worth stating: $b$ appears in the numerator of bytes/token and in the denominator of intensity, so moving BF16 ($b=2$) to FP8 ($b=1$) halves every capacity number in this article and doubles every intensity. The 512 KB constant becomes 256 KB, the 137 GB batch becomes 68 GB, and MHA’s intensity goes from 1.0 to 2.0 — still two orders of magnitude below the ridge point, which is why quantizing the cache helps but does not change the regime.

It has its own asymmetry worth knowing: keys carry channel-wise outliers and want per-channel scales, while values are well behaved under per-token scaling. At long context this is often the single largest win available, because it is orthogonal to everything else in this article.

Bound what is kept. Sliding-window attention caps the cache at the window size instead of the context length, making it constant in $T$ rather than linear. Attention-sink methods keep the first few positions permanently alongside a recent window, recovering most of the quality that a naive window loses. Eviction policies go further and drop positions judged unimportant.

Share across layers. Adjacent layers can be made to share one set of keys and values, dividing the cache by the sharing factor along $L$ — the same idea as GQA applied to the layer axis rather than the head axis.

Stop wasting what you allocate. PagedAttention does not make the cache smaller; it removes the fragmentation and over-reservation that made the allocated cache several times the used cache, and lets sequences with a common prefix share physical blocks.

8. What This Means for a Serving System

Three consequences follow from the formula and the intensity result.

The cache sets the batch-size ceiling at long context. Weights are a fixed cost paid once; cache is a per-sequence cost that grows with the context. Past a few thousand tokens the cache is the term that decides how many requests fit, which is why throughput at 32K context looks nothing like throughput at 512.

Decode speed tracks cache bytes — but sub-linearly, once work per byte rises. Since attention at decode is memory-bound with intensity $2g/b$, anything that reduces bytes read per token helps: a larger group size, a compressed latent, a quantized cache, a shorter window. Anything that only reduces arithmetic does not. But as §6 measures, an eight-times reduction in bytes obtained by raising the group size yields about 2.4× in latency, not eight, because the same change multiplies the arithmetic per byte and the kernel stops saturating memory. Reducing bytes without raising work per byte is the stronger move.

Prefill and decode want different things. Prefill processes many tokens at once and is compute-bound; decode processes one and is bandwidth-bound. The KV cache is written in the first phase and read in the second, which is why serving systems increasingly schedule the two separately, and sometimes on different hardware.

9. Summary

• KV cache size is $2 L H_{kv} d_h T b$ bytes, linear in context and batch, and independent of the number of query heads.
• Decode attention has arithmetic intensity $2g/b$ where $g = H_q/H_{kv}$ — one FLOP per byte for multi-head attention in BF16, roughly 280× below the hardware’s ridge point. Batching does not change the intensity, though it does decide how much of the memory system’s bandwidth is actually used.
• MQA, GQA, and MLA all reduce what has to be stored per token. The first two shrink $H_{kv}$; MLA keeps the head count and compresses the cached representation. GQA at 8 groups is the current default because it recovers most of the memory at a small and measurable quality cost, and converts cleanly from an existing MHA checkpoint.
• Cutting bytes per token helps decode less than proportionally when the saving comes from raising the group size. Measured over $g \in \{1,2,4,8\}$, an eight-times reduction in bytes returned 2.4× in latency, because the same change multiplies the arithmetic per byte and achieved bandwidth falls from 75% of peak to 22%. Tests that varied access granularity and parallelism independently ruled both out as the cause.
• Quantization, windowing, cross-layer sharing, and paged allocation are orthogonal levers that compose with the architectural choice — and they reduce bytes without narrowing the read, so they do not pay the same penalty.

References

1. Noam Shazeer. Fast Transformer Decoding: One Write-Head is All You Need. 2019.
2. Joshua Ainslie et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023.
3. DeepSeek-AI. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. 2024.
4. Reiner Pope et al. Efficiently Scaling Transformer Inference. MLSys 2023.
5. Woosuk Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023.
6. Guangxuan Xiao et al. Efficient Streaming Language Models with Attention Sinks. ICLR 2024.
7. Zirui Liu et al. KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache. ICML 2024.