Memory budgeting
Calculate KV-cache memory and explain how layers, KV heads, context length and concurrent sequences affect it.
THE IDEA
Budgeting the memory
To plan capacity, start from the GPU's total memory, subtract the weights and other reserved memory, and what remains is the KV budget. Divide by the per-sequence size to see how many sequences fit at a given context length. Then allow for fragmentation and scheduling, because a real system cannot use every byte.
The weights fitting is only the first test. Long contexts times many users can exhaust the rest, and real systems lose some memory to fragmentation and scheduling.
Predict first. Can a model whose weights fit still fail under long concurrent contexts?
Chapter context
Causal attention can reuse keys and values for earlier tokens rather than recomputing them at every generated step. This saves work but creates state that grows with the number of retained token positions.
TRY IT YOURSELF
Part 1. Do it hands-on
Compute a KV budget and concurrency by hand, then write the budgeting functions with fragmentation.
The weights fitting is only the first test. Long contexts times many users can exhaust the rest, and real systems lose some memory to fragmentation and scheduling.
- To do: A 24 GiB GPU. The weights take 14 GiB and 2 GiB are reserved for workspace and other needs. How many GiB can the KV cache use?
- To do: Each 4096-token sequence needs 512 MiB. How many fit in 8 GiB?
- To do: With 32,768-token contexts (4 GiB each), how many fit?
Part 2. Do it in code
- kv_budget(total_bytes, weight_bytes, reserved_bytes): what is left for the KV cache. ValueError if it would be negative.
- usable_budget(budget, fragmentation): budget · (1 − fragmentation); fragmentation must be in [0, 1).
- max_concurrent(budget, seq_tokens, per_token): how many sequences of seq_tokens fit.
- fits(budget, lengths, per_token): True if all the sequences (a list of lengths) fit together.
- longest_context(budget, n_sequences, per_token): the longest equal context length that n_sequences sequences can have.
Stuck, or something not working? You can come back to it any time.
CHECK YOUR UNDERSTANDING
Can a model whose weights fit still fail under long concurrent contexts?
Chapter reference: math, worked example & code
This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.
Math and notation
For a uniform full-attention decoder:
L is layers, H_KV KV heads, d_h head dimension, T retained tokens per sequence, B equal-length unshared sequences and b bytes per value. The factor 2 counts keys and values. For unequal lengths, replace TB by the sum of retained lengths.
Worked example
For L=32, H_KV=8, d_h=128 and b=2, each retained token needs 2×32×8×128×2=131,072 bytes, or 128 KiB. At T=4096, one sequence needs 512 MiB. Four equally long independent sequences need 2 GiB, before allocator overhead and other memory.
Reference code
The following example uses the Python standard library and does not call a remote service.
layers,heads,dim,bytes_per_value=32,8,128,2
per_token=2*layers*heads*dim*bytes_per_value
assert per_token == 131072
assert per_token/1024 == 128
assert per_token*4096/1024**2 == 512
assert per_token*4096*4/1024**3 == 2
Python examples are provided for study and local execution; this page does not execute Python.
Common mistakes
Use KV heads rather than query heads. Distinguish decimal GB from binary GiB. Hybrid, sliding-window and quantized-cache architectures require their own detailed accounting.