LLM inference optimization: prefill, decode, KV cache, batching and speculative decoding
Generation has two phases: prefill processes your whole prompt in parallel and is limited by compute, and decode generates one token at a time and is limited by memory bandwidth, because every step reads the model weights and the KV cache. The main optimisations are continuous batching to keep the GPU busy, PagedAttention to manage KV cache memory, quantization to shrink what each step reads, prefix caching to reuse shared prompts, and speculative decoding to produce several tokens per expensive step. For most teams, the practical bottleneck isn't training a model, it's serving one reliably and cheaply enough for real users.
What happens during inference
Inference is when a model processes a request and generates an answer. It's the most expensive step in an LLM system, because every request runs through billions of parameters on a GPU.
Prefill
The model reads your whole prompt at once. All the prompt's tokens are processed in parallel, which keeps the GPU's compute units busy, so prefill is compute-bound. Prefill produces the first output token, so it decides your time to first token.
Decode
The model then generates the answer one token at a time. Each new token needs a full pass through the model, and each pass reads all the weights from GPU memory, plus the KV cache. The computation per token is small, so the GPU spends most of its time waiting for data to move. That's why decode is memory-bound.
KV cache
During attention, each token produces key and value vectors. Without a cache, generating token 500 would recompute the keys and values for the previous 499 tokens. The KV cache stores them, so each step only computes the new token's.
You can think of the KV cache as the model's working memory, like RAM for a CPU. It grows with the number of layers, the context length, the number of requests being served together and the precision.
def kv_cache_gb(layers: int, kv_heads: int, head_dim: int, tokens: int, batch: int, bytes_per_value: int = 2) -> float:
return 2 * layers * kv_heads * head_dim * tokens * batch * bytes_per_value / 1e9
print(f"{kv_cache_gb(32, 8, 128, tokens=8_192, batch=1):.1f} GB for one 8K-token request")
print(f"{kv_cache_gb(32, 8, 128, tokens=8_192, batch=32):.1f} GB for 32 of them")
For an 8 billion parameter model in this shape (32 layers, 8 key-value heads, head size 128), one 8,000-token request needs about 1 GB of KV cache, and 32 concurrent requests need about 34 GB, more than the model's own weights at 16-bit. That is why memory, not compute, limits how many users a GPU can serve. See the KV cache.
Techniques
| Technique | What it does |
|---|---|
| Continuous batching | Adds and removes requests from the running batch at every step, instead of waiting for a whole batch to finish |
| PagedAttention (vLLM) | Stores the KV cache in fixed-size blocks, like virtual memory pages, so memory isn't wasted on fragmentation |
| Prefix caching | Reuses the KV cache for prompt prefixes shared across requests, like a long system prompt |
| Quantization | Smaller weights and KV cache, so each decode step moves less data |
| Speculative decoding | A small draft model proposes several tokens, the large model checks them in one pass |
| Tensor and pipeline parallelism | Splits a large model across several GPUs |
| Chunked prefill | Splits long prompts so they don't block decode for other users |
Continuous batching
A GPU serving one request at a time is mostly idle during decode. Batching serves many requests together, so each read of the weights produces tokens for all of them. Static batching waits until every request in the batch finishes, so short requests wait for long ones. Continuous batching swaps requests in and out at every step. Anyscale reported up to 23 times higher throughput with continuous batching and memory optimisations in their benchmark.
PagedAttention
Before PagedAttention, servers reserved one contiguous block of memory per request for its maximum possible length. The vLLM paper found existing systems wasted 60 to 80% of KV cache memory this way. PagedAttention stores the cache in small blocks allocated as needed, which reduced waste to under 4% and improved throughput 2 to 4 times over the systems they compared against.
Speculative decoding
A small, fast draft model guesses the next few tokens. The large model checks all of them in one forward pass, keeps the ones it agrees with, and continues from the first disagreement. The output is the same as the large model alone would produce, and when the draft guesses well, generation is several times faster.
What to measure
- Time to first token (TTFT): mostly prefill and queueing.
- Time per output token or tokens per second per user: decode speed.
- Throughput: total tokens per second per GPU, across all users.
- GPU memory use and KV cache utilisation.
- Cost per million tokens.
Benchmark your server at 1, 8 and 32 concurrent requests, and plot throughput and p95 latency. Throughput rises with concurrency until memory runs out, and latency per user rises with it, so you choose the point that meets your latency target.
From the floor
In the inference zone of the production pipeline I teach, the prefill engine loads the whole prompt into the GPU in parallel and fills the KV cache, managed with PagedAttention, then the decode engine generates tokens one by one, and quantized weights reduce the memory each step needs, with vLLM and NVIDIA TensorRT-LLM as the examples I use. I describe this as the most expensive step in the whole pipeline, which is why every zone before it tries to avoid reaching it.
Karan Mittal, CTO at Dextar, ran a guest session for the cohort on inference engines, and his framing was that for most organisations the bottleneck isn't training a frontier model, it's serving models reliably and cheaply enough for real users, measured through latency, throughput, GPU utilisation and cost per token.