Lesson 22 / 30

Memory Math: Weights plus KV Cache

Estimate total memory for a model, quantization and context length before you download anything.

By · 15 years in web development

The full memory bill

Total memory needed is roughly weights + KV cache + compute buffers + overhead. Weights you already know how to estimate from parameters and bits per weight. The KV cache is the part people forget. For every token in the context, every layer stores a key and a value vector for each KV head, so its size is 2 x layers x kv_heads x head_dim x context_length x bytes_per_value. Modern models use grouped-query attention (GQA), where several query heads share one KV head, which shrinks this a lot, but it still grows linearly with context. A model that fits comfortably at 4K context can fail to fit at 32K. Compute buffers for intermediate activations add a further amount that grows with batch size. When you plan a deployment, pick the context length you actually need, not the maximum the model supports, and leave headroom for the operating system and other apps.

KV cache for a Llama-3-8B-like shape

32 layers, 8 KV heads, head dimension 128. The q8_0 figure approximates 8.5 bits per value.

def kv_cache_gb(layers, kv_heads, head_dim, ctx, bytes_per_val=2):
    # 2 = one K and one V tensor per layer
    return 2 * layers * kv_heads * head_dim * ctx * bytes_per_val / 1e9

# Llama-3-8B-like shape: 32 layers, 8 KV heads (GQA), head_dim 128
for ctx in (2048, 8192, 32768):
    print(f"ctx={ctx:6}: f16 {kv_cache_gb(32, 8, 128, ctx):.2f} GB, "
          f"q8_0 ~{kv_cache_gb(32, 8, 128, ctx, 1.0625):.2f} GB")

Output:

ctx=  2048: f16 0.27 GB, q8_0 ~0.14 GB
ctx=  8192: f16 1.07 GB, q8_0 ~0.57 GB
ctx= 32768: f16 4.29 GB, q8_0 ~2.28 GB

What fills your VRAM

Each block adds to the total; only the first is fixed by the model file. The stages, in order: model weights -> KV cache (grows with context) -> compute buffers -> headroom for the OS and other apps.

4 connected stages, left to right: model weights, KV cache (grows with context), compute buffers, headroom for the OS and other apps.
Figure 8.1: model weights -> KV cache (grows with context) -> compute buffers -> headroom for the OS and other apps

Quantize the KV cache when context is the problem

llama.cpp can store the cache at lower precision with -ctk q8_0 -ctv q8_0 (V-cache quantization needs flash attention: -fa on in recent builds, which may also enable it automatically, or plain -fa in older ones). Ollama offers OLLAMA_KV_CACHE_TYPE=q8_0 together with OLLAMA_FLASH_ATTENTION=1. q8_0 roughly halves cache memory with little quality loss; q4_0 saves more but hurts more.

Quick check: A model fits at 4K context but runs out of memory at 32K. Which component grew?

  • The KV cache, which scales linearly with context length
  • The model weights
  • The tokenizer vocabulary
  • The GGUF metadata
Answer

The KV cache, which scales linearly with context length — Weights are fixed by the file. The KV cache stores keys and values for every token of context, so 8 times the context means about 8 times the cache.