AI Engineering Illustrated

Posted on Oct 12 Originally published at ragmehta.substack.com

How LLM Inference Serving Works: KV Cache, PagedAttention and Continuous Batching

#ai #architecture #llm #machinelearning

Why serving a model is mostly a memory-management problem, and the operating-systems idea that fixed it

Training a large language model gets the headlines, but inference is where most of the compute bill goes in production. Every chat message, every agent step and every autocomplete is a request for a GPU to generate tokens. The difference between a naive serving stack and a good one isn't 10%. The vLLM team reported up to 24× higher throughput than serving with Hugging Face Transformers when they launched in June 2023.

Where does a gap that large come from? Not from faster matrix multiplication. It comes from three things most application developers never see:

How the KV cache is stored in GPU memory
How requests are batched together
How work is scheduled between the two phases of generation

In this issue, we'll look at:

Why generation has two very different phases, prefill and decode
What the KV cache is, and how large it really gets (with exact numbers for Llama 3.1 8B)
How PagedAttention borrowed virtual memory from operating systems to stop wasting most of that memory
Why continuous batching replaced static batching
How prefix caching and chunked prefill squeeze out more performance, and what that means for how you write prompts

Two Phases: Prefill and Decode

An LLM generates text one token at a time. Each new token depends on all the previous ones. Under the hood, a single request goes through two phases with very different performance profiles.

Prefill processes the entire prompt in one forward pass. All prompt tokens go through the model in parallel, which means large matrix multiplications that keep the GPU's compute units busy. Prefill is compute-bound. Its output is the first generated token, so prefill time largely sets time-to-first-token (TTFT).

Decode generates every token after that, one forward pass per token. Each step processes only one new token per request, but to compute attention it has to read the keys and values of every previous token. So each step loads the model's weights plus the request's entire cached history from GPU memory, does relatively little math with it, and produces one token. Decode is memory-bound: the bottleneck is how fast bytes move from GPU memory (HBM), not how fast the GPU computes. vLLM's own tuning docs describe the split exactly this way: "compute-bound (prefill) and memory-bound (decode)".

This asymmetry drives almost every design decision that follows. Since decode is limited by memory bandwidth, the main way to get more out of a GPU is to decode many requests at once. Loading the weights once and using them for 64 requests is far better than loading them 64 times. Batch size is the lever, and what limits batch size is memory, specifically the KV cache.

What the KV Cache Is, and How Big It Gets

In a transformer's attention layers, each token is projected into a query, a key and a value. When generating token n, the model compares token n's query against the keys of tokens 1…n−1 and takes a weighted mix of their values. Those keys and values never change once computed, so recomputing them at every step would be wasteful. Instead, inference engines store them in the KV cache. Prefill fills the cache for the prompt, and each decode step appends one more token's K and V.

How much memory is that? It's a simple product:

KV bytes per token = 2 (K and V) × layers × KV heads × head dimension × bytes per value

Plugging in the published config.json for Llama 3.1 8B: 32 layers, 8 key/value heads (the model uses grouped-query attention, so its 32 query heads share 8 KV heads), head dimension 128, and bf16 values at 2 bytes each:

2 × 32 × 8 × 128 × 2 = 131,072 bytes = 128 KiB per token

That sounds small until you multiply by context length:

1K tokens → 128 MiB
8K tokens → 1 GiB
32K tokens → 4 GiB
128K tokens (the model's maximum context) → 16 GiB

The model's weights are about 8 billion parameters × 2 bytes ≈ 16 GB. A single conversation at full context needs about as much memory as the entire model. The weights are shared by every request. The KV cache is per request, and it grows with every token generated.

(Grouped-query attention is itself a KV-cache optimization. With 32 KV heads instead of 8, every number above would be four times larger.)

The vLLM launch post gives a sense of scale for an older model: the KV cache can take up to 1.7 GB for a single sequence in LLaMA-13B. It also describes the cache as dynamic: "its size depends on the sequence length, which is highly variable and unpredictable." That unpredictability is the real problem.

The Problem: Fragmentation

When you start generating, you don't know whether the answer will be 20 tokens or 2,000. Early serving systems handled this the obvious way. They reserved a contiguous chunk of GPU memory for each request, big enough for the maximum possible length.

The top half of the diagram shows what goes wrong:

Over-reservation. A request that ends after 200 tokens still held space for 2,048. That reserved but unused memory can't go to anyone else.
Internal fragmentation. Even within a reservation, the unused tail sits idle the whole time the request is running.
External fragmentation. Contiguous allocations of different sizes leave gaps between them that are too small to use.

The vLLM team measured this: existing systems wasted 60%–80% of KV cache memory due to fragmentation and over-reservation. Since memory caps batch size, and batch size caps decode throughput, most of the GPU's potential was being thrown away.

PagedAttention: Virtual Memory for the KV Cache

Operating systems solved this exact problem decades ago. Programs want large contiguous address spaces, but physical memory is fragmented. The OS splits memory into fixed-size pages, gives each process a contiguous virtual address space, and keeps a page table mapping virtual pages to wherever they physically sit.

PagedAttention, introduced in the vLLM paper (SOSP 2023, arXiv:2309.06180), applies the same idea to the KV cache. In the authors' words: "one can think of blocks as pages, tokens as bytes, and sequences as processes."

The bottom half of the diagram shows how it works:

KV memory is split into fixed-size blocks, each holding the keys and values for a fixed number of tokens (configurable in vLLM with --block-size).
Each request sees a sequence of logical blocks that are contiguous from its own point of view.
A block table maps each logical block to a physical block anywhere in GPU memory. Physical blocks don't need to be adjacent.
Physical blocks are allocated on demand, only when a request's current block fills up.
The attention kernel is rewritten to follow the block table, gathering each request's K and V from wherever its blocks live.

The payoff: memory waste only happens in a request's last, partially filled block. The vLLM post reports this results in "a mere waste of under 4%". Nearly all KV memory now holds real data, so far more requests fit in a batch. The paper's headline result is that vLLM "improves the throughput of popular LLMs by 2-4× with the same level of latency" compared with FasterTransformer and Orca, with bigger gains for longer sequences, larger models and more complex decoding.

Sharing comes for free

Once memory is indirected through a block table, two requests can point at the same physical block. This matters for decoding strategies that create several outputs from one prompt, like parallel sampling ("give me 4 candidate answers") or beam search. The prompt's KV blocks are stored once and shared, with reference counts. A block is copied only when one branch needs to write to it (copy-on-write, another OS technique). The vLLM team reports this cuts memory usage of these methods "by up to 55%", which can mean "up to 2.2× improvement in throughput."

Continuous Batching

Freeing memory only helps if the scheduler actually fills it with work. The classic approach, static or request-level batching, doesn't.

With static batching, the server collects a batch of requests and runs them together until all of them are done. But generation lengths vary wildly. The Orca paper (OSDI '22) describes the problem precisely: with request-level scheduling, "requests that have finished earlier than other requests in a batch cannot return to the client, while newly arrived requests have to wait until the current batch completely finishes."

Orca's fix was iteration-level scheduling, now usually called continuous batching. The scheduler re-forms the batch at every model iteration, meaning every decode step:

A request that just produced its final token leaves the batch immediately and its result is returned.
A newly arrived request joins at the very next iteration.
The batch stays full, so the GPU stays busy.

Continuous batching and PagedAttention reinforce each other. Continuous batching constantly admits and evicts requests of unpredictable length, which is exactly the workload that fragments contiguous memory. Paged blocks make admitting a new request cheap: you just need a few free blocks, not a large contiguous region.

When memory runs out anyway

Even with paging, a burst of long requests can fill the KV cache. vLLM then preempts some requests to free their blocks. In vLLM V1 the default preemption mode is recompute: the preempted request's KV is dropped and rebuilt later, rather than swapped to CPU memory, "as recomputation has lower overhead in the V1 architecture." If your logs show preemption warnings, the docs suggest giving vLLM more memory (gpu_memory_utilization, which defaults to 0.92), reducing max_num_seqs or max_num_batched_tokens, or sharding the model across more GPUs.

Prefix Caching and Chunked Prefill

Two more techniques complete the picture. The first one directly affects how you should write prompts.

Automatic prefix caching

Real traffic is highly repetitive at the start. Every request to an agent typically begins with the same system prompt, the same tool definitions, maybe the same retrieved documents. Running prefill on those identical tokens over and over is pure waste.

With prefix caching, vLLM keeps the KV blocks of previous requests and reuses them. The design docs explain how a match is found: vLLM hashes "each kv-cache block by the tokens in the block and the tokens in the prefix before the block." Because the hash includes everything before the block, a block matches only if the entire prefix up to it is identical. A lookup table maps hashes to cached physical blocks, so a new request whose first 3,000 tokens match an earlier one skips straight to prefilling only the new suffix.

Cached blocks no longer in use aren't freed immediately. They stay available for reuse and are evicted in LRU order when memory is needed. For multi-tenant deployments, vLLM supports an optional per-request cache_salt that is mixed into the first block's hash, so different tenants can't reuse, or probe, each other's cached prefixes.

The docs note that prefix caching is "widely used by many public endpoints (e.g., OpenAI, Anthropic, etc.)". This is the mechanism behind the discounted "cached input" pricing those APIs offer.

The practical lesson for prompt design: put stable content first and variable content last. System prompt, then tool definitions, then long reference documents, then the user's actual question. Changing a single token early in the prompt (a timestamp in the system prompt, say) breaks the hash chain from that point on and forces a full prefill.

Chunked prefill

Remember that prefill is compute-heavy and decode is memory-bound. If a 10,000-token prompt arrives while 60 users are streaming responses, a naive scheduler runs that huge prefill as one step, and every streaming user sees a stall.

Chunked prefill splits large prefills into pieces and mixes them into steps alongside decodes. vLLM's docs describe the V1 policy:

Each step has a token budget, max_num_batched_tokens.
The scheduler prioritizes decode requests: it "batches all pending decode requests before scheduling any prefill operations."
Any remaining budget goes to pending prefills. A prefill that doesn't fit "automatically" gets chunked.

This improves inter-token latency, because streams don't stall, and GPU utilization, because compute-bound and memory-bound work share each step. The budget is a tuning knob. Smaller values (the docs give 2048 as an example) favor smooth streaming. Larger values favor time-to-first-token, and the docs recommend going above 8192 for throughput on large GPUs. In vLLM V1, chunked prefill "is enabled by default whenever possible."

Putting It Together: The Life of a Request in vLLM

Here's what happens to one request in a modern serving engine:

The request arrives and is tokenized.
The scheduler checks the prefix cache. Matching leading blocks are reused, and only the remainder needs prefill.
At the next iteration, the scheduler admits the request if enough free KV blocks exist, preempting another request if necessary.
Its prefill runs, chunked if it's long, alongside other requests' decode steps, within the token budget.
Each decode step appends one token's K and V into the request's current block. A new physical block is allocated only when that block fills.
When the request emits its stop token, it leaves the batch immediately. Its blocks' reference counts drop, and they become cached-but-evictable.

Each step targets the same scarce resource: GPU memory for KV cache, and keeping the GPU busy with as many requests as that memory allows.

What This Means If You Don't Run GPUs

Even if you only call hosted APIs, these mechanics show up on your bill and in your latency charts:

Long prompts cost twice: compute at prefill (TTFT), and memory for every decode step that follows. Trim context you don't need.
Prompt structure affects cost. Stable prefixes hit the cache. Dynamic content at the top of the prompt defeats it.
Output length affects throughput more than you'd think. Every generated token is a full decode step holding KV memory.
TTFT and tokens-per-second are different metrics driven by different phases. A provider can be strong on one and weak on the other.

Key Takeaways

Inference has two phases: prefill (parallel, compute-bound, sets TTFT) and decode (one token per step, memory-bound).
The KV cache is the bottleneck. For Llama 3.1 8B it's 128 KiB per token, so 16 GiB at a full 128K context, about the size of the model itself.
Contiguous allocation wasted 60–80% of KV memory. PagedAttention uses fixed-size blocks and a block table, like OS virtual memory, cutting waste to under 4% and delivering 2–4× throughput in the paper.
Continuous batching (Orca's iteration-level scheduling) re-forms the batch every step, so finished requests leave and new ones join immediately.
Prefix caching reuses KV blocks for identical prefixes via chained block hashes, so put stable prompt content first.
Chunked prefill schedules decodes first and splits long prompts, so one big prompt doesn't stall everyone else.

DEV Community