Zenaique

Predict the KV cache size in GB for given model dimensions

Predict output·Hard·4.0 · 0·~2 min·Asked atCoinbaseModal LabsNVIDIA
Attempt it
Given a transformer with the following dimensions, predict the KV cache size in GB (to one decimal place) for a single request at full context. Use the standard formula `bytes = 2 * n_layers * n_kv_heads * head_dim * seq_len * dtype_bytes`.

Inputs:
- n_layers = 80
- n_kv_heads = 8   (this model uses GQA with group size 8)
- head_dim = 128
- seq_len = 131072 (128k tokens)
- dtype = FP16 (2 bytes per element)
- batch = 1

Report the result in GB (1 GB = 1024^3 bytes). Round to one decimal place.
TL;DR

Multiply 2 (K and V) by layers, KV heads, head dim, sequence length, and bytes per element. The result is about 20.0 GB at 128k context.

Memory aid
Sign in to see the mnemonic that makes this stick.
Easy to grasp

Picture a hotel where every guest who has ever checked in keeps their room forever, and you have to walk past all those rooms each time someone new arrives. The KV cache is that hallway of kept rooms. Its size is just a chain of multiplications: how many floors (layers), how many rooms per floor (KV heads), how big each room is (head dimension), how many guests so far (sequence length), and how many bytes each room costs (the number format). Multiply them all, then double it because every guest needs two rooms, one for keys and one for values. Do that for a long-context request and you discover the hallway can grow larger than the whole hotel building, which is the model weights.

Concept explanation~2 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

KV cache memory math is one of the highest-yield napkin calculations in LLM serving. An interviewer hands you a model's dimensions and asks how much memory one request's cache consumes at full context. The numbers themselves are easy multiplications; the difficulty lives entirely in remembering the formula correctly and executing the arithmetic without a slip.

This question gives a model that looks a lot like Llama 3.1 70B: 80 layers, 128 head dimension, 8 KV heads under grouped query attention, fp16 storage, and a 128k token context. The task is to predict the per-request cache size in gibibytes to one decimal place. The answer is a clean 20.0 GB, and the cleanliness is itself a clue worth understanding.

The deep dive below derives the formula factor by factor, runs the arithmetic explicitly, shows the two conversion and counting traps that catch most candidates, and then connects the number to the architectural reason it matters: grouped query attention is exactly the lever that turns an unshippable 160 GB into a manageable 20 GB.

Deriving the formula factor by factor

The cache holds the key and value vectors for every token already processed, at every layer, for every KV head. Walk the dimensions one at a time and the formula assembles itself.

Start with a single token at a single layer and a single KV head. It contributes one key vector of size head dim and one value vector of size head dim. That is the factor of 2 times head dim. Multiply by the number of KV heads to cover all heads in that layer, then by the number of layers to cover the whole stack. Multiply by the sequence length because every token contributes its own row. Finally multiply by bytes per element to turn element counts into raw bytes.

The assembled formula is the canonical one:

bytes=2LHkvdhTb\text{bytes} = 2 \cdot L \cdot H_{kv} \cdot d_h \cdot T \cdot b

Here L is the layer count, the KV head term is the GQA head count rather than the query head count, head dim is the per-head width, T is the token count, and b is bytes per element. Every factor is linear, which is why the cache grows so aggressively with context length.

Notice what is deliberately absent. There is no factor for the query heads, no model dimension, and no batch term in the per-request form. The query count never appears because queries are computed fresh each step and never stored. The model dimension is already decomposed into head count times head dim, so including it would double-count. Batch is a separate outer multiplier you apply only when sizing a whole server, not a single request. Getting the formula right is largely a matter of including exactly these six factors and resisting the urge to add a seventh.

Running the arithmetic for this model
The two arithmetic traps that sink candidates
Why GQA makes this number shippable
Extending the estimate to a real deployment
Sign in to unlock the full deep dive.

Situations where this technique stops working.

Sign in to see when this approach fails.

2–4 min · Everything important, quickly.

Sign in to see the quick scan of the deep dive.

Real products, models, and research that use this idea.

  • Llama 3.1 70B uses GQA with 8 KV heads, the exact configuration in this problem, so this 20 GB figure mirrors its real per-request long-context footprint.
  • vLLM sizes its paged KV blocks using precisely this byte formula to decide how many concurrent requests fit in the remaining HBM after weights.
Sign in to see more production examples.

What an interviewer would ask next. Try answering before peeking at the approach.

QWhat would this same model's cache cost if it used full MHA with 64 KV heads instead of GQA?
A

Scale the result linearly with KV heads. Multiply 20.0 GB by 64 over 8, which is 8, giving roughly 160 GB. That figure exceeds the bf16 weights of a 70B model, which is the whole motivation for GQA.

2 more follow-ups an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

Dropping the layer count, or forgetting the leading factor of 2 for K and V. Either error throws the answer off by an order of magnitude or by half.

Sign in to see all red flags and common mistakes.

60 second bullets to scan on the way to the call.

  • The six factors in the formula and the order to multiply them

  • Why the leading factor of 2 exists and what it counts

Sign in to unlock the revision sheet.

Primary sources. Browse if you want the original framing.

Similar questions

Same topic, related formats. Practice these next.

4 curated
Next question
What is the KV cache in transformer inference?
Flashcard·Easy