Inference Engineering
Module 05·Serving an LLM·3 min read·2 drills

KV cache math

Decode is a memory budget problem before it is a model problem. Compute the KV cache cost per token, then turn VRAM into concurrency.

By the end

Estimate KV-cache memory per token and use it to reason about max context, batch size, and serving capacity.

Autoregressive generation has a strange asymmetry: prefill processes the prompt in parallel, then decode produces one token at a time.

The KV cache is what makes decode possible. Every generated token stores key and value vectors for every transformer layer, so the next token can attend to the context without recomputing the whole past.

That cache is also what turns serving into a memory-management problem.

The formula

For a decoder-only transformer, a useful KV-cache estimate is:

bytes_per_token =
  layers
  * 2              # K and V
  * kv_heads
  * head_dim
  * bytes_per_element

Then:

request_kv_bytes = bytes_per_token * total_tokens_in_request

Total tokens means prompt tokens plus generated tokens that remain in the cache.

Grouped-query attention changes the formula by reducing kv_heads. The query heads can be many; the KV heads are what matter for the cache.

A concrete example

Suppose:

layers = 32
kv_heads = 8
head_dim = 128
dtype = bfloat16 = 2 bytes

Then:

32 * 2 * 8 * 128 * 2 = 131,072 bytes/token

That is 128 KiB per token. A single 8K-token request wants roughly 1 GiB of KV cache. Ten such requests want roughly 10 GiB before you count weights, activations, fragmentation, allocator overhead, or runtime workspace.

This is why long-context serving gets weird fast. The weights are fixed. The cache grows with traffic.

Capacity is not just model size

A naive deployment question is:

Does the model fit on the GPU?

The serving question is:

After weights and runtime overhead, how many live tokens fit?

Live tokens are the unit of concurrency. A system with 12 GiB free for KV cache and 128 KiB per token can hold about 98,000 live tokens. That could be:

  1. 98 requests at 1K tokens.
  2. 12 requests at 8K tokens.
  3. 3 requests at 32K tokens.

Same GPU. Same model. Very different product experience.

Fragmentation and paging

Real serving engines do not just allocate one huge perfect cache per request. Request lengths vary, requests finish at different times, and beam search or parallel sampling can share prefixes. Waste matters.

PagedAttention, introduced with vLLM, treats KV cache memory in fixed-size blocks and maps logical request blocks to physical memory blocks. The operating-systems analogy is virtual memory paging: the logical sequence does not need to be physically contiguous.

The practical lesson is not "always use vLLM." The lesson is that KV cache allocation strategy can decide how many requests fit before the GPU is out of memory.

Prefill and decode pressure different things

Prefill is usually compute-heavy because the model processes many prompt tokens at once. Decode is often memory-bandwidth-heavy because each step reads model weights and existing cache to produce one new token.

The KV cache sits directly in decode's path. More live tokens mean more cache memory. More concurrent requests mean more scheduler work. Longer contexts mean attention touches more cached history unless the architecture or kernel changes the story.

That is why a serving benchmark needs to report prompt length, output length, concurrency, batch policy, and context limits. A token-per-second number without those dimensions is mostly atmosphere.

References

Checkpoint

  1. Write the KV-cache bytes-per-token formula from memory.
  2. Why does grouped-query attention reduce cache size?
  3. A GPU has 20 GiB available for KV cache, and a model needs 160 KiB per live token. Roughly how many live tokens fit?

Practice this lesson

The reading is the model. These drills are the hours — 2 problems that force the numbers onto paper before the next lesson.