The Exact VRAM Mathematics of LLM Inference: KV-Cache, MQA/GQA, and PagedAttention

1. The Two Phases of LLM Generation

When an autoregressive Large Language Model (such as LLaMA-3, DeepSeek, or GPT-4) serves a request, inference proceeds in two fundamentally distinct operational regimes:

1. The Prefill Phase (Prompt Processing):
   • All $S$ input prompt tokens are ingested simultaneously.
   • Matrix multiplications are large and highly parallel: $\mathcal{O}(S^2)$ operations.
   • Compute-bound: Utilizes high Tensor Core compute throughput (TFLOPs).
2. The Decode Phase (Token Generation):
   • Generates tokens sequentially, one token at a time: $t_1, t_2, \dots, t_N$.
   • Matrix-vector products: batch size is small, arithmetic intensity is near zero.
   • Memory-bandwidth bound: The GPU spends $90\%$ of its time waiting for weights and previous keys/values to stream across the HBM memory bus.

To avoid recomputing keys and values for all past tokens at every step, we store them in GPU VRAM as the Key-Value (KV) Cache.

2. Deriving the Exact KV-Cache Memory Equation

Let us calculate the exact number of bytes required to cache a single token across an entire Transformer stack.

For each Transformer layer $l$:

• Key vector dimension: $d_{\text{key}} = n_{\text{kv\_heads}} \times d_{\text{head}}$
• Value vector dimension: $d_{\text{val}} = n_{\text{kv\_heads}} \times d_{\text{head}}$
• Precision byte size: $b$ bytes per parameter ($b = 2$ for FP16/BF16, $b = 1$ for FP8, $b = 0.5$ for INT4).

The total memory required per token across both Keys and Values is:

$\text{Bytes per Token} = 2 \times n_{\text{layers}} \times n_{\text{kv\_heads}} \times d_{\text{head}} \times b$

To find total KV-cache size for a batch of $B$ concurrent sequences of sequence length $L$:

$\text{Total KV Cache (Bytes)} = B \times L \times \left( 2 \times n_{\text{layers}} \times n_{\text{kv\_heads}} \times d_{\text{head}} \times b \right)$

3. Real-World Case Study: LLaMA-3 70B

Let us evaluate the architectural parameters of LLaMA-3 70B:

• Layers ($n_{\text{layers}}$): $80$
• Query Heads ($n_{\text{q\_heads}}$): $64$
• KV Heads ($n_{\text{kv\_heads}}$): $8$ (Grouped-Query Attention)
• Head Dimension ($d_{\text{head}}$): $128$
• Precision ($b$): $2$ bytes (FP16)

Plugging these numbers into our first-principles equation:

$\text{Bytes per Token} = 2 \times 80 \times 8 \times 128 \times 2 = 327,680 \text{ bytes} \approx 0.328 \text{ MB / token}$

Now compute the KV-cache requirement for a standard serving setup:

• Batch Size $B = 16$ concurrent users.
• Context Length $L = 8,192$ tokens.

$\text{Total Cache} = 16 \times 8,192 \times 327,680 \text{ bytes} \approx 42.95 \text{ GB}$

[!WARNING] The weights of LLaMA-3 70B quantized to 4-bit take $\sim 38 \text{ GB}$. The KV-cache for just 16 concurrent users at 8k context ($43 \text{ GB}$) exceeds the weight size of the model itself!

4. Architectural Solutions: MHA vs. MQA vs. GQA

How have AI architects solved this exponential memory bottleneck?

1. Multi-Head Attention (MHA) ($n_{\text{kv\_heads}} = n_{\text{q\_heads}}$):
   • Cache size is huge: $2.62 \text{ MB/token}$.
2. Multi-Query Attention (MQA) ($n_{\text{kv\_heads}} = 1$):
   • Compresses cache by $64\times$, but loses representation quality and expressive nuance.
3. Grouped-Query Attention (GQA) ($n_{\text{kv\_heads}} = 8$):
   • The golden mean adopted by modern foundation models. Yields an $8\times$ cache compression with negligible degradation in benchmark scores.

5. PagedAttention & Eliminating Virtual Fragmentation

In early serving frameworks (TensorRT-LLM v1, HuggingFace), memory was pre-allocated contiguously for the maximum sequence length (e.g. 8k).

Because actual response lengths vary widely, $60\% - 80\%$ of GPU VRAM was completely wasted due to internal and external fragmentation!

In 2023, Kwon et al. introduced PagedAttention (the core algorithm behind vLLM):

• Inspired by virtual memory paging in operating systems (OS).
• Divides the KV-cache of each request into fixed-size physical memory blocks (typically 16 or 32 tokens).
• Dynamic logical-to-physical block mapping tables allow non-contiguous physical storage.

This single algorithmic breakthrough nearly tripled serving throughput without altering a single weight of the model.