E2.2 · LESSON 2 OF 6

Growth with context

Calculate KV-cache memory and explain how layers, KV heads, context length and concurrent sequences affect it.

THE IDEA

The cache grows with every token

Every retained token adds a fixed amount of state: a key and a value for each KV head in each layer. In bytes:

bytes per token = 2 × layers × KV heads × head size × bytes per value

For L = 32, H_KV = 8, head size 128 and 2 bytes per value that is 2·32·8·128·2 = 131,072 bytes, or 128 KiB. A 4096-token sequence therefore needs 4096 × 128 KiB = 512 MiB, and doubling the context doubles it.

layers L32KV heads8head size128bytes per value2context tokens4,096sliding window (optional)1,024

per token = 2 (K and V) × 32 × 8 × 128 × 2 = 131,072 bytes (128 KiB) full cache for 4,096 tokens = 512 MiB with a 1,024-token window = 128 MiB

context	cache
1,024	128 MiB
2,048	256 MiB
4,096	512 MiB
8,192	1 GiB
16,384	2 GiB

Doubling the context doubles the cache, if every retained token contributes the same state. Windowed or hybrid layers cap what they retain and need their own accounting.

Change the model shape and the context. Add a sliding window to see a capped cache.

Predict first. If every retained token contributes the same state, what happens when the context doubles?

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

Work out the per-token and per-sequence size by hand, then write the cache-size functions.

layers L32KV heads8head size128bytes per value2context tokens4,096sliding window (optional)1,024

per token = 2 (K and V) × 32 × 8 × 128 × 2 = 131,072 bytes (128 KiB) full cache for 4,096 tokens = 512 MiB with a 1,024-token window = 128 MiB

context	cache
1,024	128 MiB
2,048	256 MiB
4,096	512 MiB
8,192	1 GiB
16,384	2 GiB

Doubling the context doubles the cache, if every retained token contributes the same state. Windowed or hybrid layers cap what they retain and need their own accounting.

KiB per token (reference model) MiB for 4096 tokens MiB with a 1024-token window

To do: L = 32 layers, H_KV = 8, head size 128, 2 bytes per value: bytes per token = 2·32·8·128·2. In KiB (1 KiB = 1024 bytes)?
To do: One sequence of 4096 tokens needs how many MiB (1 MiB = 1024 KiB)?
To do: A sliding-window layer keeps only the last 1024 tokens. For a 4096-token sequence, how many MiB?

Part 2. Do it in code

bytes_per_token(layers, kv_heads, head_dim, bytes_per_value): 2 · layers · kv_heads · head_dim · bytes_per_value.
cache_bytes(layers, kv_heads, head_dim, bytes_per_value, tokens): bytes_per_token · tokens.
to_mib(n_bytes): n_bytes / 1024².
windowed_cache_bytes(layers, kv_heads, head_dim, bytes_per_value, tokens, window): like cache_bytes but only min(tokens, window) tokens are retained.
tokens_that_fit(budget_bytes, per_token): the number of whole tokens that fit.

def bytes_per_token(layers, kv_heads, head_dim, bytes_per_value):
    pass

def cache_bytes(layers, kv_heads, head_dim, bytes_per_value, tokens):
    pass

def to_mib(n_bytes):
    pass

def windowed_cache_bytes(layers, kv_heads, head_dim, bytes_per_value, tokens, window):
    pass

def tokens_that_fit(budget_bytes, per_token):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

If every retained token contributes the same state, what happens when context doubles?

AThe KV state approximately doubles, because each retained token adds the same amount of state.BIt quadruples, because attention is quadratic.CIt stays the same, because the cache stores a summary.

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:

M_{KV}=2L H_{KV}d_h T B b.

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.