Inference and Serving, lesson 3
How big the cache gets
Work out KV cache memory from layers, heads, length and batch size.
The KV cache sits in GPU memory beside the weights, and it grows with every token of every request. This lesson works out its size by hand from a model's config file, and then how many requests fit on one GPU. You need what the cache holds, from "Caching keys and values".
By the end, you can:
Part 1. Bytes for one token
Work out a model's cache per token from its config.
Count what one token leaves behind in Llama 2 7B. Its config file lists 32 layers and 32 attention heads, and each head works with vectors of 128 numbers. The cache stores each number in 16 bits, which is 2 bytes, just like the weights.
Example 1.1
One token of Llama 2 7B
At every layer, every head stores one key and one value for each token. How many bytes is one token's share?
- One head, one layer128 + 128 = 256 numbers
A key and a value, 128 numbers each.
- Every head256 × 32 = 8,192
Each of the 32 heads keeps its own pair.
- Every layer8,192 × 32 = 262,144
Each layer runs its own attention, so it keeps its own cache.
- Into bytes262,144 × 2 = 524,288 bytes
That's 512 KiB (a KiB is 1,024 bytes), about half a megabyte for one token.
Answer524,288 bytes, or 512 KiB, per token.
The number of weights doesn't appear in that sum. What counts is how many keys and values each token leaves: one pair per head per layer. Two models with the same number of weights can have very different caches, as part 3 shows.
Head size
The length of each head's query, key and value vectors. It's the model's width divided by its heads: 4,096 ÷ 32 = 128 for Llama 2 7B. Config files call the width hidden_size and the heads num_attention_heads, and some list head_dim directly.
bytes per token = 2 × heads × head size × layers × bytes
Two 2s turn up in that formula, and they're easy to mix up. The first is the key and the value, two vectors per head. The last is the 2 bytes of a 16-bit number. Store the cache in 8 bits and the last one becomes 1, halving everything.
Your turn
Llama 2 13B has 40 layers and 40 heads of size 128, and keeps its cache in 16 bits. How big is one token's cache?
Part 2. Times length, times batch
Scale the per-token cache to whole requests and a full batch.
Every token in a request leaves the same amount behind, and every request being served keeps a cache of its own. So the per-token figure gets multiplied twice: by the tokens in each request, and by the number of requests.
Example 2.1
Eight full requests to Llama 2 7B
Llama 2 7B reads up to 4,096 tokens. How much cache do 8 requests at that length hold, and how does it compare with the model's 13.5 GB of 16-bit weights?
- One full request524,288 × 4,096 = 2.15 GB
512 KiB for each of 4,096 tokens.
- Eight of them2.15 × 8 = 17.2 GB
Each request keeps its own.
- Set it beside the weights17.2 GB > 13.5 GB
6.74 billion weights at 2 bytes each make 13.5 GB. The cache now outweighs the model.
Answer17.2 GB of cache, more than the 13.5 GB of weights.
Length and batch sit side by side in the product, so the cache only cares about the total number of tokens being served. Longer conversations and more users draw on the same memory, one for one.
KV cache size
What one token leaves behind, times the tokens in each request, times the number of requests. Every factor is a dial: double any one and the cache doubles.
cache = bytes per token × tokens × requests
The weights are paid for once, however many people use the model. The cache is paid for per token, per user. That's why serving costs climb with traffic and with context length, even though the model never changes.
Your turn
You serve 8 requests of 4,096 tokens. You switch to 4 requests of 8,192 tokens. What happens to the cache?
Part 3. Sharing keys and values across heads
Work out how much grouped-query attention shrinks the cache.
Llama 2 70B has 80 layers and 64 heads of size 128. If each head kept its own key and value, its cache would dwarf the 7B model's. Instead, its 64 query heads are split into 8 groups of 8, and each group shares one key head and one value head.
Example 3.1
Llama 2 70B, with and without sharing
Llama 2 70B keeps its cache in 16 bits. How big is one token's cache if every query head has its own key and value? And with the 8 shared pairs it really has?
- A pair per query head2 × 80 × 64 × 128 × 2
2,621,440 bytes, or 2,560 KiB a token. One 4,096-token request would hold 10.7 GB.
- A pair per group of 82 × 80 × 8 × 128 × 2
327,680 bytes, or 320 KiB a token, and 1.34 GB for the same request.
- Compare64 ÷ 8 = 8
Only the heads factor changed, from 64 to 8.
Answer320 KiB a token instead of 2,560 KiB: 8 times smaller.
Each query head still asks its own question, so the model keeps most of the benefit of many heads. Only the keys and values it searches are shared. The paper that named the idea found that grouped models come close to full multi-head quality while decoding nearly as fast as with a single shared pair.
Grouped-query attention (GQA)
Query heads split into groups, with each group sharing one key head and one value head. With a single shared pair it's called multi-query attention (MQA), and one pair per head is multi-head attention (MHA). In the cache formula, count the key-and-value heads.
bytes per token = 2 × KV heads × head size × layers × bytes
Read the config with care. num_attention_heads counts the query heads, and num_key_value_heads counts what goes into the cache. The head size still comes from the query heads: hidden_size ÷ num_attention_heads. Mixing them up puts you out by the group size, 8 for this model.
Your turn
Llama 2 7B and Llama 3 8B both have 32 layers and head size 128. Llama 2 7B has 32 KV heads, Llama 3 8B has 8. How do their per-token caches compare?
Part 4. How many requests fit
Count how many requests fit on one GPU, with and without paging.
A GPU's memory holds the weights first, and whatever is left holds caches. Take Llama 3 8B on an 80 GB H100. Requests may run to 8,192 tokens, but most end far sooner: say they average 1,500. How you hand out the cache memory decides how many fit.
Example 4.1
Reserve the maximum, or hand out blocks
Llama 3 8B's weights take 16.06 GB, and its cache is 128 KiB (131,072 bytes) a token. How many requests fit if each reserves room for 8,192 tokens? And if each gets 16-token blocks only as it grows?
- Find the free memory80 − 16.06 = 63.9 GB
A real server also keeps some back for working memory.
- Reserve 8,192 tokens each131,072 × 8,192 = 1.07 GB
Every request holds room for its longest possible length.
- Count the requests63.9 ÷ 1.07 = 59.7
Round down to whole requests: 59, each using 1,500 of its 8,192 places on average.
- Give out blocks of 161,500 ÷ 16 → 94 blocks
93.75 rounds up to 94 blocks, 1,504 tokens: 0.197 GB.
- Count again63.9 ÷ 0.197 = 324
About 5.5 times as many requests, on average.
Answer59 requests reserving the maximum, about 324 with 16-token blocks.
Reserving the maximum wastes most of the memory, because few requests reach their limit and nobody knows in advance how long a reply will be. Measuring older serving systems, the vLLM team found only 20% to 38% of cache memory held actual tokens. With blocks, the only waste is the unfilled end of each request's last block.
Paged KV cache
Cache memory cut into fixed-size blocks and handed to a request one block at a time as it grows. A small table per request lists which blocks hold its tokens, so they don't need to sit side by side. vLLM introduced it as PagedAttention (Kwon and colleagues, 2023), with 16 tokens a block as its default.
requests that fit ≈ free memory ÷ cache per request
Blocks can still run out when traffic spikes. Then the server has to pause some requests, free their blocks and redo their work later, which "Continuous batching" takes up later in this module.
Your turn
Llama 2 7B has 13.5 GB of weights, and a full 4,096-token request needs about 2.15 GB of cache. If each request reserves its full length, about how many fit on an 80 GB GPU?
Where you'll meet it
A common sizing question: can one GPU serve 50 users with 8,000-token contexts on this model? Work it out loud. Weights first: parameters × bytes per weight. Then the cache: 2 × KV heads × head size × layers × bytes, times tokens, times users. Say which config fields you used, and leave some memory spare for working space.
When vLLM starts, it loads the weights, runs a test batch to measure its working memory, and turns what's left (within the share you allow with --gpu-memory-utilization) into cache blocks. Its log then reports the cache size in tokens and how many maximum-length requests that holds at once. Divide by your typical request length for a first guess at capacity.
def kv_bytes(layers, kv_heads, head_size, tokens=1, requests=1, bytes_per=2):
per_token = 2 * layers * kv_heads * head_size * bytes_per # a key and a value
return per_token * tokens * requests
configs = { # layers, KV heads, head size
"Llama 2 7B": (32, 32, 128),
"Llama 2 13B": (40, 40, 128),
"Llama 2 70B": (80, 8, 128),
"Llama 3 8B": (32, 8, 128),
}
for name, shape in configs.items():
tok = kv_bytes(*shape)
full = kv_bytes(*shape, tokens=4096, requests=8)
print(f"{name}: {tok // 1024} KiB a token, {full / 1e9:.1f} GB for 8 x 4,096")
free = 80e9 - 8.03e9 * 2 # Llama 3 8B on an 80 GB GPU
reserved = kv_bytes(32, 8, 128, tokens=8192)
paged = kv_bytes(32, 8, 128, tokens=-(-1500 // 16) * 16)
print(int(free // reserved), int(free // paged))
# Llama 2 7B: 512 KiB a token, 17.2 GB for 8 x 4,096
# Llama 2 13B: 800 KiB a token, 26.8 GB for 8 x 4,096
# Llama 2 70B: 320 KiB a token, 10.7 GB for 8 x 4,096
# Llama 3 8B: 128 KiB a token, 4.3 GB for 8 x 4,096
# 59 324Put in any model's config numbers to size its cache before you rent the GPU.
Check yourself
Four questions from across the lesson. Try each one before you look back.
1
Written as one product, a token's cache is 2 × KV heads × head size × layers × bytes per number. What is the first 2 for?
2
Llama 3 8B keeps 128 KiB (131,072 bytes) per token. How much cache do 32 requests of 2,048 tokens hold?
3
A model has 32 layers and head size 128, and all the query heads in a layer share a single key head and value head. How big is one token's 16-bit cache?
4
A server switches from reserving 8,192 tokens per request to 16-token blocks. Requests average 1,500 tokens. Roughly how many requests can it hold at once, compared with before?
What to remember
1Bytes for one token
One token's cache is 2 × KV heads × head size × layers × bytes: 524,288 bytes (512 KiB) for Llama 2 7B.
2Times length, times batch
Multiply by tokens and requests: 8 requests of 4,096 tokens hold 17.2 GB in Llama 2 7B, more than its 13.5 GB of weights.
3Sharing keys and values across heads
Grouped-query attention shares key and value heads: Llama 2 70B's 8 for 64 query heads make its cache 8 times smaller.
4How many requests fit
Requests that fit ≈ free memory ÷ cache per request: 59 Llama 3 8B requests reserving 8,192 tokens each, about 324 with 16-token blocks.
Recap and going further
The lesson on one page
- 2 × KV heads × head size × layers × bytes
- Llama 2 7B: 512 KiB. Llama 3 8B: 128 KiB
- Times the tokens in each request, times the requests
- Length and batch trade one for one
- Query heads share key and value heads
- Llama 2 70B: 8 for 64, so 8 times smaller
- Free memory = GPU minus weights
- Divide by the cache each request needs
- Hand out 16-token blocks as requests grow
- 59 requests on one GPU become about 324
Try this
Check vLLM's memory budget by hand
Start vLLM with a small model, for example vllm serve Qwen/Qwen2.5-0.5B-Instruct, and find two lines in its startup log: how much memory model loading took, and the "GPU KV cache size" in tokens. Open the model's config.json and work out its cache per token from num_hidden_layers, num_key_value_heads and the head size (hidden_size ÷ num_attention_heads). Then restart with --gpu-memory-utilization 0.5 and read the log again.
You need: A machine with an NVIDIA GPU (a rented cloud GPU or a Colab GPU), Python, and pip install vllm.
Look for
- Whether your bytes per token × the logged token count comes close to the memory left after the weights.
- The "Maximum concurrency" the log prints, and whether your own division of cache tokens by the context length agrees.
- How much the cache shrinks when you lower the memory share.
The words, in plain language
- Head size
- The length of each head's query, key and value vectors. It's the model's width divided by its heads: 4,096 ÷ 32 = 128 for Llama 2 7B. Config files call the width
hidden_sizeand the headsnum_attention_heads, and some listhead_dimdirectly. - KV cache size
- What one token leaves behind, times the tokens in each request, times the number of requests. Every factor is a dial: double any one and the cache doubles.
- Grouped-query attention (GQA)
- Query heads split into groups, with each group sharing one key head and one value head. With a single shared pair it's called multi-query attention (MQA), and one pair per head is multi-head attention (MHA). In the cache formula, count the key-and-value heads.
- Paged KV cache
- Cache memory cut into fixed-size blocks and handed to a request one block at a time as it grows. A small table per request lists which blocks hold its tokens, so they don't need to sit side by side. vLLM introduced it as PagedAttention (Kwon and colleagues, 2023), with 16 tokens a block as its default.
Read next
- Efficient Memory Management for Large Language Model Serving with PagedAttentionKwon and colleagues, 2023The vLLM paper: blocks, block tables, and how much memory older systems wasted.
- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head CheckpointsAinslie and colleagues, 2023Grouped-query attention, and how close it stays to full multi-head quality.
Inference and Serving overviewSpotted a mistake? Tell us.