KV Cache Memory › Cache Size Calculation

With a clear understanding of what the cache stores, we can now derive the formula for total memory consumption. The calculation follows directly from multiplying the dimensions of the cached tensors by the number of bytes each element requires. There is no approximation involved: this is exact arithmetic about the shape of the tensors sitting in your GPU's memory. This derivation provides the basis for all deployment planning decisions, and once you work through it once, you will be able to estimate cache sizes in your head for any model architecture.

Consider the structure we described above. For each layer, we store both keys and values. Each of these has shape [batch_size, num_kv_heads, seq_len, head_dim]. The total number of elements in one such tensor is the product of these dimensions. Since we store two tensors (keys and values) per layer across all layers, the total element count becomes:

$$\text{Total Elements} = 2 \times L \times B \times T \times H_{\text{kv}} \times D_h$$

Converting elements to bytes requires multiplying by the bytes per element, giving us the complete formula:

$$\text{KV Cache Memory} = 2 \times L \times B \times T \times H_{\text{kv}} \times D_h \times \text{bytes}$$

where:

• $2$ accounts for storing both key and value matrices
• $L$: the number of transformer layers
• $B$: the batch size (number of concurrent sequences)
• $T$: the sequence length (number of tokens cached)
• $H_{\text{kv}}$: the number of key-value heads
• $D_h$: the head dimension (dimensionality of each attention head)
• $\text{bytes}$: the number of bytes per element (2 for FP16/BF16, 4 for FP32)

The key insight is that every term in this product is linear: there are no squares, no logarithms, no nonlinear interactions. Doubling any single factor doubles the memory. This linearity makes the formula both simple to use and relentlessly demanding in practice. You cannot cleverly reorganize the computation to escape the cost. Every factor must be paid in full.

This formula reveals the linear dependencies that govern cache memory. Doubling any of the multiplicative factors doubles the total memory. This linearity is both a blessing and a curse. It makes calculations simple and predictable, but it also means there are no economies of scale. Processing twice as many tokens always requires twice as much cache memory, with no possibility of compression through clever algorithms unless we explicitly introduce approximation techniques, which we cover in later chapters on cache quantization and pruning.

For standard Multi-Head Attention (MHA), $H_{\text{kv}}$ equals the number of attention heads $H$, and the relationship $H \times D_h = D_{\text{model}}$ holds. This relationship exists because each head operates on a slice of the model dimension, and together all heads span the full representation space. This architectural constraint allows us to substitute and simplify our formula:

$$\text{KV Cache Memory (MHA)} = 2 \times L \times B \times T \times D_{\text{model}} \times \text{bytes}$$

where:

• $D_{\text{model}}$: the model dimension (equal to $H \times D_h$, the full embedding size)
• $L$: the number of transformer layers
• $B$: the batch size
• $T$: the sequence length
• $\text{bytes}$: the number of bytes per element

This simplified form is particularly useful because $D_{\text{model}}$ is the first number reported in any model card or architecture description. Knowing that a model has dimension 4096, 32 layers, and you want to cache 4096 tokens in FP16, you can quickly compute:

$$2 \times 32 \times 1 \times 4096 \times 4096 \times 2 = 2{,}147{,}483{,}648 \text{ bytes} \approx 2 \text{ GB}$$

This mental arithmetic takes seconds and immediately tells you whether your hardware can support a given configuration.

For models using Grouped Query Attention (GQA), which we covered in the Modern Decoder Models section, the number of KV heads is smaller than the number of query heads. This architectural innovation allows multiple query heads to share the same key-value pairs, dramatically reducing the cache footprint. GQA achieves this reduction without proportionally harming model quality, making it one of the most impactful architectural changes for inference efficiency introduced in recent years.

To understand why GQA reduces cache size, consider what the cache stores. In standard MHA, every query head has its own dedicated key head and value head. A model with 64 query heads caches 64 sets of keys and 64 sets of values. In GQA with 8 KV heads, 8 groups of 8 query heads each share one set of keys and one set of values. The cache stores only 8 sets of keys and 8 sets of values, but all 64 query heads can still attend to their relevant context. The queries are still diverse and specialized; only the cached targets are shared.

If a model has $H$ query heads but only $H_{\text{kv}}$ key-value heads, the cache size reduces proportionally:

$$\text{Reduction Factor} = \frac{H}{H_{\text{kv}}}$$

where:

• $H$: the number of query attention heads (used for computation but not cached)
• $H_{\text{kv}}$: the number of key-value heads (the only heads whose projections are cached)

This reduction factor quantifies the memory savings from GQA. A model with 64 query heads and 8 KV heads achieves an 8-fold reduction in cache size compared to standard MHA. This 8-fold reduction can mean the difference between a model that requires multiple GPUs to deploy and one that fits on a single high-end card. The practical impact of GQA on deployment feasibility cannot be overstated.

Think of GQA as a shared note-taking system in a large meeting. In standard MHA, each participant (query head) takes their own complete notes (keys and values) about what everyone said. In GQA, one designated note-taker serves each small group of participants, and everyone in the group reads from the same set of notes. The quality of the discussion suffers only marginally because the notes are still detailed. But the storage cost falls by the number of participants per group.

Now let's examine concrete examples with popular model architectures. First, we define our model specifications and a calculation function: