PidokuInfra

KV Cache Math

Basic Intermediate 1h 30m Difficulty 3/5 Topic 06 of 15

Prerequisites 05

Pure arithmetic. Work every example with a calculator. These calculations are what you’ll be asked to do in an interview and what you’ll do weekly on the job.


1. The formulas#

(1)  KV_per_token = 2 · L · n_kv_heads · head_dim · bytes_per_elem

(2)  KV_per_sequence = KV_per_token · context_length

(3)  KV_total = KV_per_sequence · batch_size

(4)  KV_budget = GPU_memory − weights − activation_headroom − framework_overhead

(5)  max_concurrency = KV_budget / KV_per_sequence

(6)  decode_bytes_per_step = weight_bytes + KV_per_token · context · batch

(7)  crossover: weight_bytes = KV_per_token · context · batch
                → context · batch = weight_bytes / KV_per_token

The 2 in (1) is for K and V. Everything else is read from config.json.

Diagram — Every term in the KV size formula#

flowchart LR
  TWO["2 - K and V"] --> X(("multiply"))
  L["layers"] --> X
  H["KV heads x head_dim"] --> X
  S["sequence length"] --> X
  B["concurrent sequences"] --> X
  D["bytes per element<br/>FP16 = 2"] --> X
  X --> R["KV cache bytes"]
  R --> C{"Fits in VRAM<br/>minus weights and overhead?"}
  C -->|"yes"| OK["Admit"]
  C -->|"no"| NO["Cut context or concurrency,<br/>quantize KV, or add GPUs"]

  class R,X memory
  class C queue
  class OK compute
  class NO warn
  class TWO,L,H,S,B,D neutral

2. Reading a config#

JSON
{
  "num_hidden_layers": 32,          →  L
  "hidden_size": 4096,              →  d
  "num_attention_heads": 32,        →  h
  "num_key_value_heads": 8,         →  n_kv_heads   (if absent, = num_attention_heads)
  "intermediate_size": 14336,       →  d_ff
  "vocab_size": 128256,             →  V
  "max_position_embeddings": 8192,  →  max context
  "torch_dtype": "bfloat16"         →  2 bytes
}

head_dim = hidden_size / num_attention_heads = 4096/32 = 128 (some configs specify head_dim explicitly — use it if present).


3. Worked example 1 — Llama-3-8B#

L=32, n_kv=8, head_dim=128, BF16

(1)  KV_per_token = 2 × 32 × 8 × 128 × 2
                  = 131,072 bytes = 128 KiB

(2)  at 8k context: 131,072 × 8,192 = 1.074 GB per sequence

Weights: 8.03e9 × 2 = 16.06 GB

On one 80 GB H100:
(4)  KV_budget = 80 − 16.06 − 2 (activations) − 2 (overhead) = 59.9 GB
(5)  max_concurrency at 8k = 59.9 / 1.074 = 55 sequences
                    at 4k  = 59.9 / 0.537 = 111
                    at 32k = 59.9 / 4.295 = 13
                    at 128k= 59.9 / 17.18 = 3

At 128k context you can serve three users on an H100. That is the entire economics of long-context serving in one line.


4. Worked example 2 — Llama-3-70B on 8×H100#

L=80, n_kv=8, head_dim=128, BF16, TP=8

(1)  KV_per_token = 2 × 80 × 8 × 128 × 2 = 327,680 bytes = 320 KiB
     Per GPU with TP=8: KV heads are split → 320/8 = 40 KiB per GPU per token

Weights: 70.6e9 × 2 = 141.2 GB total → 17.65 GB per GPU

Per GPU (80 GB):
(4)  KV_budget = 80 − 17.65 − 3 − 2 = 57.35 GB
(2)  per sequence at 8k, per GPU = 40 KiB × 8192 = 327.7 MB
(5)  max_concurrency = 57.35 GB / 0.328 GB = 175 sequences

     at 4k:   350
     at 32k:   43
     at 128k:  10

Note: with TP, both weights and KV are split across GPUs, so concurrency is computed per-GPU and the answer is the same for all of them (they hold the same sequences, different head slices).


5. Worked example 3 — the MHA penalty#

Same model shape but MHA (n_kv = 32) instead of GQA (n_kv = 8):

GQA-8:  KV_per_token = 2 × 32 × 8  × 128 × 2 = 131,072 B = 128 KiB
MHA:    KV_per_token = 2 × 32 × 32 × 128 × 2 = 524,288 B = 512 KiB

At 8k context, 80 GB H100, 16 GB weights, 60 GB budget:
  GQA-8: 60 GB / 1.074 GB = 55 sequences
  MHA:   60 GB / 4.295 GB = 13 sequences

→ GQA gives 4x the concurrency for this configuration.

For Llama-2-70B (MHA, L=80, h=64, head_dim=128):

KV_per_token = 2 × 80 × 64 × 128 × 2 = 2,621,440 B = 2.5 MiB per token
At 4k context: 10.2 GB PER SEQUENCE
On 8×A100 (640 GB) with 140 GB weights: 500 GB / 10.2 = 49 sequences

Compare to Llama-3-70B (GQA-8) in the same setup: ~390 sequences. 8x, from one architectural choice. This is why every serious model released after 2023 uses GQA or better.


6. Worked example 4 — the bandwidth crossover#

When does reading KV cost more than reading weights?

(7)  context × batch = weight_bytes / KV_per_token

Llama-3-8B:  16.06e9 / 131,072 = 122,543 token-slots
Llama-3-70B: 141.2e9 / 327,680 = 430,908 token-slots

For the 8B model:

batch 32  → crossover at context 3,830
batch 64  → crossover at context 1,915
batch 128 → crossover at context   957

At batch 128 and 4k context, KV reads are 4x the weight reads. Your decode step is dominated by KV bandwidth, and weight quantization gives you almost nothing while KV quantization gives you 2x.

Do this calculation for your workload. It tells you which optimization to invest in, and it’s frequently the opposite of what people assume.


7. Worked example 5 — decode step time#

Llama-3-8B, one H100 (3.35 TB/s), batch 32, context 4096:

(6)  bytes = 16.06 GB (weights) + 131,072 × 4,096 × 32 (KV)
           = 16.06 GB + 17.18 GB
           = 33.24 GB

     T_step = 33.24 / 3350 = 9.92 ms

     throughput = 32 / 0.00992 = 3,226 tokens/sec
     ITL per user = 9.92 ms

Real systems achieve 60-75% → expect ~2,000-2,400 tok/s and ITL ~13-16 ms.

Same model, batch 32, context 32768:

     bytes = 16.06 + 137.4 = 153.5 GB
     T_step = 45.8 ms
     throughput = 699 tok/s        ← 4.6x worse throughput at 8x context
     ITL = 45.8 ms                 ← 4.6x worse latency too

Long context hurts both throughput and latency, not just memory.


8. Worked example 6 — full capacity plan#

Requirement: serve Llama-3-70B, 2,000 concurrent users, average context 6k, p95 ITL < 60 ms.

Step 1 — KV per sequence
  KV_per_token = 320 KiB
  at 6k context = 1.966 GB per sequence

Step 2 — total KV needed
  2,000 × 1.966 GB = 3,932 GB

Step 3 — how many 8×H100 nodes?
  Per node: 640 GB − 141 GB (weights, TP=8) − 24 GB (overhead+activations) = 475 GB KV
  Nodes needed for KV: 3,932 / 475 = 8.3 → 9 nodes

Step 4 — check the latency constraint
  Per node, sequences = 475 / 1.966 = 241
  Decode bytes per step per GPU = 17.65 GB (weights) + 241 × 6144 × 40 KiB
                                = 17.65 + 59.2 = 76.9 GB
  T_step = 76.9 / 3350 = 23 ms theoretical
  With TP AllReduce overhead (~30%) and 70% efficiency: ~42 ms
  → ITL ≈ 42 ms  ✓ meets the 60 ms SLO

Step 5 — throughput check
  Per node: 241 sequences / 0.042 s = 5,738 tokens/sec
  9 nodes: 51,600 tokens/sec aggregate

Step 6 — headroom
  Add N+1 redundancy and 30% spike headroom: 9 × 1.3 + 1 = 13 nodes

ANSWER: 13 nodes of 8×H100 (104 GPUs).

Sanity check this against intuition: 2,000 users at ~24 tokens/sec each = 48,000 tokens/sec. Our 9-node capacity is 51,600. Consistent. Good.


9. Quick reference table#

KV bytes per token (BF16) for common models:

ModelLn_kvhead_dimKV/tokenat 8kat 128k
Llama-3-8B328128128 KiB1.0 GB16 GB
Llama-3-70B808128320 KiB2.6 GB41 GB
Llama-2-7B (MHA)3232128512 KiB4.2 GB67 GB
Llama-2-70B (MHA)80641282.5 MiB20.5 GB328 GB
Mistral-7B (GQA-8)328128128 KiB1.0 GB16 GB
Qwen2.5-7B (GQA-4)28412856 KiB0.46 GB7.3 GB
Mixtral-8x7B (GQA-8)328128128 KiB1.0 GB16 GB

(Verify these against current configs — architectures change.)


10. Hands-on exercise#

A. Build a calculator. Write a script that takes a HuggingFace model id, downloads its config.json, and prints: parameters, weight bytes at each precision, KV per token, KV per sequence at several context lengths, and max concurrency on a GPU you specify. You will use this constantly.

B. Verify against reality. Run vLLM with --max-model-len set and read the startup log — it prints the number of KV blocks allocated. Convert to sequences and compare to your calculator’s prediction. Explain any discrepancy (hint: block granularity, and vLLM reserves some memory).

C. The crossover. For three models, compute and plot the (batch, context) crossover curve. For each, name the optimization you’d prioritize in three different traffic regimes.

D. Capacity plan. Do a full plan like section 8 for: Qwen2.5-32B, 500 concurrent users, average context 16k, p95 ITL < 80 ms, on A100-80GB. Show every step.

E. KV quantization. Recompute example 5 with FP8 KV cache (1 byte instead of 2). What’s the throughput improvement at context 32768? At context 1024? Explain the difference.


11. Interview questions#

  1. Write the KV cache formula and explain every term.
  2. Compute KV per token for Llama-3-70B. Now for Llama-2-70B. Why the 8x difference?
  3. How many concurrent 8k-context sequences fit on one H100 running Llama-3-8B? Show your work.
  4. At what (batch, context) do KV reads exceed weight reads for an 8B model?
  5. A customer wants 128k context. What does that do to your capacity and cost?
  6. Size a cluster for 2,000 concurrent users on a 70B model. State your assumptions.
  7. Your ITL doubled when you enabled 32k context support, even at the same batch size. Why?

12. Further reading#

↑↓ navigate↵ openesc close