Zenaique

Complete the per token KV cache memory formula

Fill in blank·Medium·4.0 · 0·~1 min·Asked atLtimindtreeModal LabsUniphore
Attempt it
bytes_per_token = x n_layers x n_kv_heads x head_dim x bytes_per_value. The leading factor counts the and value tensors stored for every past token at every layer.
TL;DR

bytes_per_token = 2 x n_layers x n_kv_heads x head_dim x bytes_per_value. The 2 counts the key and value tensors stored for every past token at every layer.

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

The cache holds two tensors per past token at each layer: one for the key and one for the value. The key is like an index card telling the next token what kind of information this past token offers; the value is the actual information. Each is a slab of shape n_kv_heads by head_dim. Two slabs per layer per token, multiplied across all layers, times the bytes each number occupies (2 for fp16, 1 for fp8), gives the per-token cache footprint. The 2 is small but load-bearing: if you mistake it for 1, you halve your memory estimate and your serving capacity planning is wrong by a factor of two.

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.

This formula is the single most-used piece of arithmetic in LLM serving. Capacity planning, GPU sizing, batch-size choice, context-length feasibility, and architecture decisions like GQA vs MLA all read off this one line. If you can write it from memory and explain every term, you can hold an end-to-end LLM serving discussion.

The full formula is bytes_per_token = 2 x n_layers x n_kv_heads x head_dim x bytes_per_value, with the leading 2 counting the key and value tensors. This walkthrough derives each term, plugs in real configs, and ends with the modern adjustments (MLA, sliding window, fp8 cache) that bend the formula without breaking its structure.

Derive the formula from one attention layer

Consider one attention layer at one position. The layer needs K and V vectors to compute attention later. Each is shaped [n_kv_heads, head_dim] because each KV head holds its own per-head K and V vector. The K tensor has n_kv_heads x head_dim values; so does V.

Total values stored at this layer for this token: 2 x n_kv_heads x head_dim. The 2 covers K and V together. Each value occupies bytes_per_value bytes (2 for fp16/bf16, 1 for fp8, 0.5 for int4).

Multiply by the number of layers because every layer has its own attention sublayer with its own K and V to store:

bytes per token=2LHkvdheadb\text{bytes per token} = 2 \cdot L \cdot H_\text{kv} \cdot d_\text{head} \cdot b

That is the whole story for vanilla MHA and GQA models. Multiply by context length T for per-sequence cache. Multiply by batch size for total serving footprint.

Plug in real configs
Why the leading 2 matters and why the second blank is 'key'
Modern variants: MLA, sliding window, fp8 cache
Putting numbers to it: 2026 frontier-model context
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.
python
# Per-token per-layer KV cache size in bytes
def kv_bytes_per_token_per_layer(n_kv_heads, head_dim, dtype_bytes=2):
    # factor 2 covers K and V together
    return 2 * n_kv_heads * head_dim * dtype_bytes

# Llama 3 8B: 8 KV heads, head_dim 128, fp16
print(kv_bytes_per_token_per_layer(8, 128, 2))  # 4096 bytes
# Full cache for 32 layers at 8192 context, batch 1:
# 32 * 8192 * 4096 = 1.07 GB

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

  • Llama 3 8B in fp16: 2 x 32 x 8 x 128 x 2 = 128 KB per token; 1 GiB for an 8192-token sequence.
  • Llama 3 70B in fp16: 2 x 80 x 8 x 128 x 2 = 320 KB per token; about 2.5 GiB for an 8192-token sequence.
Sign in to see more production examples.

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

QHow does the formula change for fp8 cache quantisation, and what is the typical quality impact?
A

bytes_per_value drops from 2 to 1, halving the per-token cache. Quality impact is small if you use per-head or per-channel scaling and avoid quantising the first and last few layers. Production deployments on H100 and B200 ship fp8 cache by default in vLLM and TRT-LLM.

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 leading 2 because the K and V tensors look like one combined object, or treating bytes_per_value as a constant rather than a precision-dependent factor (2 for fp16, 1 for fp8, 0.5 for int4).

Sign in to see all red flags and common mistakes.

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

  • Every term in the formula and what it represents architecturally

  • Why the leading factor is 2 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
Explain scaled dot product attention.
Short answer·Medium