LLM Inference: KV Cache, Batching & Scheduling

KV cache mechanics, memory management, and paging (PagedAttention). Continuous batching (Orca), chunked prefill, prefix caching, and latency vs. throughput tradeoffs in LLM serving.

Overview

Inference is a different problem from training

Training is compute-bound: large batches, weight reuse, a backward pass. Production inference balances two competing goals, throughput (tokens per second across all users) and latency (time to first token and inter-token delay per user). The KV cache is the central data structure, and every serving optimization revolves around managing it.

KV cache — Recomputing keys and values for every prior token
PagedAttention — Fragmentation from pre-allocated KV memory
Continuous batching — Idle slots and head-of-line blocking in static batches
Chunked prefill — Long prompts stalling other requests' decode
Prefix caching — Recomputing a shared prompt for every request

The KV cache

Attention needs the keys and values of all prior positions. Recomputing them for each new token would be O(N^2) in sequence length, so the cache stores K and V for every past position and appends the new token's K and V each step.

Memory cost per layer: 2 * seq_len * num_heads * head_dim * sizeof(dtype) bytes.

Layers, KV heads (GQA), head dim — 80, 8, 128
Per layer — 2 * seq_len * 8 * 128 * 2 = 4096 * seq_len bytes
Per token, all layers — 327,680 bytes
One 4096-token request — 1.28 GB of HBM
HBM left after ~42 GB of weights on an 80 GB GPU — ~38 GB
Concurrent 4096-token requests — About 30

> [!WARNING] Running out of KV cache memory stalls new requests or forces expensive preemption of in-flight ones. Size the KV cache budget before planning a deployment.

PagedAttention and vLLM

Naive allocation reserves the maximum sequence length for each request, even one that ends after 100 tokens.

Internal fragmentation — Reserved but unused space inside a request
External fragmentation — Enough free memory in total, but not in one piece for a new request

PagedAttention, the core of vLLM, stores KV in fixed-size blocks (typically 16 tokens). Logical pages map to physical pages like OS virtual memory, pages are allocated on demand, and a finished request's pages return to the pool at once.

> [!NOTE] Near-zero fragmentation lets vLLM fit more concurrent requests in the same memory, for 2-4x higher throughput than Hugging Face TGI without paged attention.

Continuous batching (Orca)

Static batching waits for a fixed number of requests and returns them together. Early requests wait for the batch to fill, and GPU cycles are wasted when requests generate different numbers of tokens.

Continuous batching (Orca, 2022) batches each token step across all in-flight requests, which keeps the GPU busy, removes head-of-line blocking and returns short requests right away.

Chunked prefill

Work — Computes the KV cache for the prompt — Generates one token at a time
Bound by — Compute (for a long prompt) — Memory bandwidth

A long prefill can monopolize the GPU and delay other requests' decode steps, raising tail latency. Chunked prefill splits the prompt into fixed-size chunks (e.g., 512 tokens) and interleaves them with ongoing decode steps.

> [!NOTE] The tradeoff is slightly higher TTFT for the new request.

Prefix caching

Many requests share a prefix: a system prompt, few-shot examples or RAG context. Prefix caching computes that prefix's KV once, stores it and lets later requests pay only for the tokens after it.

> [!TIP] For a 2048-token system prompt shared across all requests, prefix caching removes ~50% of each request's TTFT.

Interview questions

Explain PagedAttention in detail. How does it manage KV cache memory and what are its limits?
How do you architect an LLM serving system to handle both latency-sensitive and throughput-sensitive workloads simultaneously?