KV cache memory math

The KV cache stores one key vector and one value vector per token, per layer, per attention head. For a model with L layers, H KV heads (the number of unique key-value heads — equal to the number of attention heads in standard MHA, fewer with GQA or MQA), d_head dimensions per head, and b bytes per element:

The factor of 2 accounts for storing both K and V. This is per token — multiply by the sequence length to get the total cache size for one request.

Concrete model calculations

Llama 3 8B

Architecture: 32 layers, 32 attention heads, 8 KV heads (grouped-query attention with 4 query heads per KV group), 128 dimensions per head, float16 (2 bytes per element).

At various context lengths:

• 2K context: $128 \text{KB} \cdot 2048 = 256 \text{MB}$
• 8K context: $128 \text{KB} \cdot 8192 = 1.0 \text{GB}$
• 32K context: $128 \text{KB} \cdot 32768 = 4.0 \text{GB}$
• 128K context: $128 \text{KB} \cdot 131072 = 16.0 \text{GB}$

The model weights (in float16) are approximately 16 GB. At 128K context, the KV cache equals the model's weight memory. On a single A100 (80 GB), this leaves $80 - 16 - 16 = 48 \text{GB}$ for activations, optimizer states, and overhead — feasible for a single request but problematic for batched serving.

Llama 3 70B

Architecture: 80 layers, 64 attention heads, 8 KV heads (GQA), 128 dimensions per head, float16.

At various context lengths:

• 2K context: $320 \text{KB} \cdot 2048 = 640 \text{MB}$
• 8K context: $320 \text{KB} \cdot 8192 = 2.5 \text{GB}$
• 32K context: $320 \text{KB} \cdot 32768 = 10.0 \text{GB}$
• 128K context: $320 \text{KB} \cdot 131072 = 40.0 \text{GB}$

The model weights in float16 are approximately 140 GB (typically sharded across multiple GPUs). In int4 quantization, weights compress to ~35 GB. At 128K context, the KV cache (40 GB) exceeds the quantized model weights. The cache is the dominant memory consumer — and unlike weights, it cannot be easily quantized without degradation (though recent work has made int8 KV cache practical).

Llama 3 405B

Architecture: 126 layers, 128 attention heads, 8 KV heads (GQA), 128 dimensions per head, float16.

At 128K context: 504 KB * 131072 ≈ 63 GB. This model requires multiple nodes for inference regardless — but the KV cache is a significant fraction of the per-node memory budget.

DeepSeek-V2

Architecture: 60 layers, 128 attention heads, but uses Multi-head Latent Attention (MLA) — a compressed KV representation where the KV head dimension is projected down to a latent dimension of 512 (shared across all heads) instead of storing separate KV per head group. Effective cache per token: $2 \cdot 60 \cdot 512 \cdot 2 = 122, 880 \text{bytes} ≈ 120 \text{KB}$. This is comparable to Llama 3 8B despite DeepSeek-V2 being a 236B-parameter model. MLA is an alternative to GQA that achieves even more aggressive cache compression at the cost of additional projection computation.

Batch serving: why the cache is the binding constraint

In production, an inference server handles multiple concurrent requests. Each request maintains its own KV cache — the cache cannot be shared between requests with different prompts (though prefix sharing is possible, covered in the next lesson).

For Llama 3 70B with 8K context, serving B concurrent requests requires:

• Model weights: ~35 GB (int4 quantized, loaded once, shared)
• KV cache: 2.5 GB * B
• Activations and overhead: ~2-4 GB

On a system with 8x A100 80 GB GPUs (640 GB total):

• Available for KV cache: 640 - 35 - 4 ≈ 601 GB
• Maximum concurrent requests at 8K context: 601 / 2.5 ≈ 240
• Maximum concurrent requests at 32K context: 601 / 10 ≈ 60
• Maximum concurrent requests at 128K context: 601 / 40 ≈ 15

Context length determines serving capacity more than any other factor. Doubling the context length halves the number of concurrent requests the system can serve. This is the fundamental reason why long-context inference is expensive — it is not the computation that costs, but the memory.

A KV cache calculator

def kv_cache_memory(
    num_layers: int,
    num_kv_heads: int,
    head_dim: int,
    seq_len: int,
    batch_size: int = 1,
    dtype_bytes: int = 2,  # 2 for float16/bfloat16, 1 for int8
) -> dict:
    """Calculate KV cache memory requirements."""
    per_token = 2 * num_layers * num_kv_heads * head_dim * dtype_bytes
    per_request = per_token * seq_len
    total = per_request * batch_size

    return {
        "per_token_bytes": per_token,
        "per_token_kb": per_token / 1024,
        "per_request_bytes": per_request,
        "per_request_gb": per_request / (1024 ** 3),
        "total_bytes": total,
        "total_gb": total / (1024 ** 3),
    }

# Llama 3 8B: 32 layers, 8 KV heads, 128 head_dim
llama_8b = kv_cache_memory(32, 8, 128, seq_len=8192)
print(f"Llama 3 8B @ 8K: {llama_8b['per_request_gb']:.2f} GB per request")
print(f"  Per token: {llama_8b['per_token_kb']:.0f} KB")

# Llama 3 70B: 80 layers, 8 KV heads, 128 head_dim
llama_70b = kv_cache_memory(80, 8, 128, seq_len=8192)
print(f"Llama 3 70B @ 8K: {llama_70b['per_request_gb']:.2f} GB per request")

# Batch serving: 32 concurrent requests at 8K context
llama_70b_batch = kv_cache_memory(80, 8, 128, seq_len=8192, batch_size=32)
print(f"Llama 3 70B @ 8K, 32 requests: {llama_70b_batch['total_gb']:.1f} GB total")

# Effect of int8 KV cache quantization
llama_70b_int8 = kv_cache_memory(80, 8, 128, seq_len=8192, batch_size=32, dtype_bytes=1)
print(f"Llama 3 70B @ 8K, 32 requests (int8): {llama_70b_int8['total_gb']:.1f} GB total")

Output:

Llama 3 8B @ 8K: 1.00 GB per request
  Per token: 128 KB
Llama 3 70B @ 8K: 2.50 GB per request
Llama 3 70B @ 8K, 32 requests: 80.0 GB total
Llama 3 70B @ 8K, 32 requests (int8): 40.0 GB total

Why GQA and MQA exist

Grouped-Query Attention (GQA, Ainslie et al. 2023) and Multi-Query Attention (MQA, Shazeer 2019) reduce the number of KV heads. Standard multi-head attention (MHA) uses one KV head per query head. In Llama 3 8B's GQA with 32 query heads and 8 KV heads, four query heads share each KV group. This reduces KV cache by 4x compared to full MHA:

• MHA (32 KV heads): $$2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 524, 288 \frac{\text{bytes}}{\text{token}} = 512 \text{KB}$$
• GQA (8 KV heads): $$2 \cdot 32 \cdot 8 \cdot 128 \cdot 2 = 131, 072 \frac{\text{bytes}}{\text{token}} = 128 \text{KB}$$

The 4x reduction translates directly to 4x more concurrent requests at the same memory budget, or 4x longer context at the same batch size. MQA takes this further with a single KV head shared across all query heads — an 8x reduction for Llama 3 8B — but with more quality degradation. GQA is the standard compromise used by Llama 3, Mistral, and Gemma 2.

The impact on serving economics

Consider serving Llama 3 70B on 8x A100 80 GB GPUs (640 GB total, ~600 GB available after weights):

• With full MHA (64 KV heads): $$\text{cache}\_\text{per}\_\text{token} = 2 \cdot 80 \cdot 64 \cdot 128 \cdot 2 = 2, 621, 440 \text{bytes} ≈ 2.5 \frac{\text{MB}}{\text{token}}$$. At 8K context: 20 GB per request. Max concurrent: ~30.
• With GQA (8 KV heads): cache_per_token = 327,680 bytes ≈ 320 KB/token. At 8K context: 2.5 GB per request. Max concurrent: ~240.

GQA enables 8x more concurrent requests with minimal quality loss (Ainslie et al. reported <0.5% degradation on most benchmarks). This is why every major LLM released since 2023 uses GQA — the serving cost reduction is too large to ignore.

Why not always use MQA?

Multi-Query Attention (1 KV head) maximizes cache compression. For Llama 3 70B, MQA would reduce cache to $$2 \cdot 80 \cdot 1 \cdot 128 \cdot 2 = 40, 960 \frac{\text{bytes}}{\text{token}} = 40 \text{KB}$$ — an 8x reduction over GQA. But with only one KV head shared across 64 query heads, the key and value representations lack capacity to support the diversity of attention patterns the model needs. Shazeer (2019) originally proposed MQA and showed it worked well for inference-heavy translation models. For generative LLMs on diverse tasks, GQA with 8 KV heads has empirically proven to be the sweet spot — enough compression for practical serving, enough capacity for strong performance across benchmarks.