Transformer Inference: KV Cache and Prefill vs Decode

Lesson 6, Part 1 covered latency, throughput, and batching as general serving concerns. Part 2 closes this lesson by looking at the one workload where those concerns take on a very specific, very consequential shape: transformer text generation. Batching helps throughput generally — but transformer generation has its own, more specific structure, and that structure is what this chapter derives from the ground up.

1. Why Autoregressive Generation Cannot Be Parallelized Across Time

1.1 The factorization and the causal mask

An autoregressive language model defines the probability of a sequence of tokens as a product of conditionals: the probability of token n given everything that came before it. Concretely, generating a response means repeatedly sampling

token_n ~ P(token_n | token_1, token_2, ..., token_{n-1})

and then feeding token_n back in as part of the context for producing token_{n+1}. This is not an implementation detail that a clever scheduler could work around — it is the definition of the task. The input to the computation that produces token n+1 is, in a very literal sense, not fully known until token n has been sampled. You cannot precompute token 47 of a response before token 46 exists, because token 46 might be the word that determines what token 47 should be.

Inside the network, this constraint is enforced structurally by the causal mask in self-attention. At every layer, the attention computation for position i is allowed to look at keys and values from positions 1 through i only — never from any position after i. This is what makes the model well-defined for generation in the first place: position i's output can only depend on information that would actually be available at the moment token i is being produced. It also has a second, quieter consequence that the rest of this chapter depends on entirely: because position i's key and value vectors are computed only from tokens 1..i, appending token i+1 to the sequence changes nothing about the key and value vectors already computed for positions 1..i. They are causally invariant to whatever comes later. Section 2.1 turns that invariance into the correctness argument for caching.

Contrast this with prefill — processing a prompt that already exists in full before generation starts. When the whole prompt is known upfront, there is no ordering constraint stopping the model from computing every position's representation at every layer simultaneously: token 3's key and value only depend on tokens 1 through 3, all of which are already sitting in memory, so nothing blocks computing them in parallel with token 500's key and value at the same instant. This is the fundamental asymmetry this whole chapter is built around: prefill has no temporal dependency to serialize against, decode has one built directly into the definition of the task.

1.2 What "sequential" costs you if you ignore it

The naive way to run autoregressive generation — and, historically, close to how it was first done before serving systems got sophisticated — is to treat each new token exactly like a fresh forward pass: take the full sequence of tokens produced so far, run it through the entire network end to end, and read off the model's prediction for the next token from the last position's output. This is correct. It is also enormously wasteful, and quantifying exactly how wasteful is the subject of the next section.

2. The Wasted Work of Recomputing From Scratch

2.1 What licenses caching: causal invariance

Section 1.1 established that, because of the causal mask, the key and value vectors for position i at every layer depend only on tokens 1..i and are therefore unaffected by anything appended after position i. This is the entire justification for caching: if a quantity will provably not change between one decode step and the next, computing it again is pure waste, not a hedge against staleness. The keys and values computed during earlier steps of a decode session are exactly the keys and values the model would recompute if it were re-run on the full sequence today — bit-for-bit, modulo floating-point non-determinism that production kernels are engineered to avoid. Query vectors do not enjoy this property (a token's query is only ever used once, to attend forward from that position, and is discarded), which is exactly why the cache stores Keys and Values and not queries — the "KV" in KV cache is not an arbitrary naming choice, it names precisely the two tensors that survive unchanged across steps.

2.2 Per-step cost: quadratic attention work without a cache, linear with one

Consider a single decode step where the context currently holds n tokens and the model is about to produce token n+1, at one attention layer with hidden size d_model, H heads, and head dimension d_h (so H × d_h = d_model).

Without a cache (naive re-run of the full sequence through the layer): the layer computes queries, keys, and values for all n positions, then the full n × n causal attention score matrix, then the full weighted sum over values for all n positions. Restricting attention just to the score-matrix and value-aggregation work (ignoring projections for a moment, which Section 2.4 covers separately):

Scores (QK^T over n×n pairs, per head):  2 × n² × d_h FLOPs
Weighted sum (softmax(scores)·V):        2 × n² × d_h FLOPs
Summed over H heads:                     4 × n² × d_h × H  =  4 × n² × d_model FLOPs

The n² term is not a loose analogy — it is the literal size of the causally-masked score matrix the naive re-run computes at that single step, even though only the last row of that matrix (token n's attention over tokens 1..n) is actually needed to produce the next token; every other row is thrown away the instant the forward pass finishes, having contributed nothing.

With a cache, the layer has one new query (for token n+1) to attend over the n already-cached keys and values. No n × n matrix is ever built — only a single 1 × n row of it:

Score (one query against n cached keys):  2 × n × d_h FLOPs
Weighted sum over n cached values:        2 × n × d_h FLOPs
Summed over H heads:                      4 × n × d_model FLOPs

The ratio of naive work to cached work at this single step is n² / n = n — it grows without bound as the conversation lengthens. Plugging in d_model = 4096 (roughly GPT-3/Llama-7B scale) and two representative context depths, for a single layer's attention:

n = 100:
  naive  = 4 × 100²  × 4096 = 4 × 10,000  × 4096 ≈ 163.8 MFLOP
  cached = 4 × 100   × 4096 ≈ 1.64 MFLOP
  ratio  ≈ 100×
 
n = 2,000:
  naive  = 4 × 2000² × 4096 = 4 × 4,000,000 × 4096 ≈ 65.54 GFLOP
  cached = 4 × 2000  × 4096 ≈ 32.77 MFLOP
  ratio  ≈ 2,000×

That is a 2,000-fold difference in attention work at a single decode step, for one layer, at a context length well within what a normal chat response reaches. Multiply by the dozens of layers in a real model and the naive approach is not merely slower — it is computationally unusable at production context lengths, which is exactly why no serving system built after roughly 2020 does it this way.

2.3 Cumulative cost across a generation: a worked walkthrough

Section 2.2 compared the two approaches at a single step. It is also worth totaling the work across an entire generated response, because that is closer to what a serving system actually pays for. Consider generating N tokens from an empty running context (ignore the prompt for a moment; Section 4 puts prefill back into the picture), one attention layer, d_model = 4096. Summing the naive per-step cost 4n²d_model from n = 1 to N uses the identity Σn² ≈ N³/3 for large N, while the cached per-step cost 4n·d_model sums to Σn ≈ N²/2:

Naive cumulative attention work  ≈ (4/3) × N³ × d_model     — grows as O(N³)
Cached cumulative attention work ≈  2      × N² × d_model     — grows as O(N²)

Caching does not just remove redundant work at each step — it removes an entire power of N from the total cost of generating a response, on top of the n versus 1 gap already visible in Section 2.2's per-step comparison. For a modest response of N = 512 tokens at d_model = 4096: naive cumulative work ≈ (4/3) × 512³ × 4096 ≈ 5.86 × 10^14 FLOP ≈ 586 TFLOP, versus cached cumulative work ≈ 2 × 512² × 4096 ≈ 2.15 × 10^9 FLOP ≈ 2.15 GFLOP — roughly a 270,000-fold difference, for attention alone, in one layer, for a single moderate-length reply.

2.4 It isn't only attention: the projection blowup

The naive re-run also redoes the linear projections — Q = XW_Q, K = XW_K, V = XW_V — for every token in the context at every step, not just the attention arithmetic. Each token's QKV projection at one layer costs 6 × d_model² FLOPs (three d_model × d_model matrix-vector products, 2d² FLOPs each). Without caching, step n reprojects all n tokens; with caching, step n projects only the one new token (the old tokens' K and V were already computed and stored — Section 2.1's invariance argument again). Summing across a generation of length N, exactly the same way as Section 2.3:

Naive cumulative projection cost  ≈ 6 × d_model² × N²/2   =  3 × d_model² × N²    — O(N²)
Cached cumulative projection cost ≈ 6 × d_model² × N                             — O(N)

With d_model = 4096 and N = 2,048 output tokens, one layer:

naive  = 3 × 4096² × 2048²  ≈ 3 × 1.678×10⁷ × 4.194×10⁶ ≈ 2.112 × 10¹⁴ FLOP ≈ 211 TFLOP
cached = 6 × 4096² × 2048   ≈ 6 × 1.678×10⁷ × 2048       ≈ 2.062 × 10¹¹ FLOP ≈ 206 GFLOP
ratio  ≈ 1,024×    (matches the general result: ratio ≈ (N+1)/2)

Scaled to a 32-layer model (roughly 7B-parameter class): naive projection cost alone would be about 6.76 PFLOP to generate a 2,048-token response, versus about 6.6 TFLOP with caching — and this is before counting the feed-forward network layers, which typically account for around two-thirds of a transformer's total FLOPs and would be recomputed just as wastefully under the naive scheme. The KV cache is often introduced as a minor bookkeeping optimization; the arithmetic above is why it is instead described, accurately, as the single most important systems-level optimization in LLM serving. Without it, autoregressive generation at real context lengths is not merely inefficient — it is asymptotically the wrong algorithm.

3. The KV Cache, Precisely

3.1 What is actually stored

The cache holds exactly the quantities Section 2.1 identified as causally invariant: the key and value tensors, per layer and per attention head, for every token position generated or ingested so far. Nothing else needs to persist across steps — not queries (used once, discarded), not attention scores (recomputed fresh each step against the current query), not intermediate activations from the feed-forward block (fully determined by that position's own hidden state, not reusable across positions the way K/V are).

Layer 0:   Head 0  K: [d_h × L]   V: [d_h × L]   ...   Head H-1  K: [d_h × L]   V: [d_h × L]
Layer 1:   Head 0  K: [d_h × L]   V: [d_h × L]   ...   Head H-1  K: [d_h × L]   V: [d_h × L]
   ⋮
Layer N-1: Head 0  K: [d_h × L]   V: [d_h × L]   ...   Head H-1  K: [d_h × L]   V: [d_h × L]
 
L = number of tokens processed so far (grows by one every decode step)

Every decode step appends exactly one new column to every one of these K and V tensors — one new key vector and one new value vector, per head, per layer:

Step 1:  K = [k1]                    V = [v1]
Step 2:  K = [k1, k2]                V = [v1, v2]
Step 3:  K = [k1, k2, k3]            V = [v1, v2, v3]
  ⋮
Step L:  K = [k1, k2, ..., kL]       V = [v1, v2, ..., vL]
                            ▲
              one new column appended per layer, per head, every step —
              every existing column is untouched (Section 2.1's invariance)

3.2 The memory footprint, derived as a formula

The total number of scalar elements resident in the cache, for a batch of B sequences each L tokens long, across every layer and head, counting keys and values separately, is:

elements = 2 × num_layers × num_heads × head_dim × L × B

The leading 2 accounts for keys and values as two separate tensors of identical shape. Multiplying by the number of bytes each element occupies (bytes_per_element — 4 for fp32, 2 for fp16/bf16, 1 for int8) gives the footprint in bytes:

KV cache bytes = 2 × num_layers × num_heads × head_dim × L × B × bytes_per_element

Every term in this formula is something the serving system controls or must accept: num_layers, num_heads, and head_dim are fixed by the model's architecture; bytes_per_element is a precision choice; B is a scheduling decision (how many requests to batch together); and L is the one term that grows on its own, silently, for the entire duration of every request the server is handling — every additional generated token, on every active sequence, adds another slice of this size to total memory in use.

3.3 Worked example: a 70B-class model at realistic context lengths

Take a concrete, well-documented architecture: Llama 3.1 70B, 80 layers, 64 query heads, head_dim = 128 (so d_model = 64 × 128 = 8192), served in bf16 (bytes_per_element = 2). Llama 3.1 70B uses grouped-query attention (GQA) with only 8 key/value heads shared across groups of query heads — Section 3.4 explains why that number, not 64, is the one that belongs in the KV-cache formula. Per token, per sequence:

KV bytes/token = 2 × 80 × 8 × 128 × 2  =  327,680 bytes  ≈  0.3125 MB/token

The model's own weights, at 70×10^9 parameters in bf16, occupy about 140 GB; at 4-bit quantization, roughly 35–42 GB depending on the exact scheme. The table below plugs the per-token figure into a batch of realistic context lengths and compares the resulting cache size against both weight footprints:

Context length (tokens)	Batch	KV cache size	vs. bf16 weights (~140 GB)	vs. int4 weights (~35–42 GB)
4,096	1	≈ 1.34 GB	≈ 1%	≈ 3–4%
32,000	1	≈ 10.5 GB	≈ 7.5%	≈ 25–30%
128,000	1	≈ 41.9 GB	≈ 30%	≈ 100–120%
4,096	32 (batched)	≈ 42.9 GB	≈ 31%	≈ 100–125%

The two bottom rows are the surprising fact worth sitting with: a single 128K-token conversation, or equivalently thirty-two concurrent 4K-token conversations, carries a KV cache that is comparable to or larger than the entire quantized model's own weights, and roughly a third of the full-precision model's weights, on top of whatever those weights already occupy. This is not a pathological worst case — 128K context windows and dozens of concurrently served users are both entirely ordinary in production LLM deployment in 2026, meaning KV cache is routinely one of the largest single consumers of accelerator memory in a serving fleet, not an afterthought line item. A back-of-envelope estimate that budgets memory for "the model" and forgets the cache will be wrong by a wide margin the moment real traffic with real context lengths shows up.

3.4 Grouped-query attention as a KV-cache lever

Standard multi-head attention (MHA) uses as many key/value heads as query heads. Had Llama 3.1 70B used plain MHA (64 KV heads instead of 8), the per-token cost would be eight times larger — 2,621,440 bytes/token — and the 128K-context row of the table above would read roughly 336 GB, more than double the model's own bf16 weight footprint, for a single conversation. GQA's entire purpose is exactly this line item: multiple query heads share one key/value head, cutting num_heads in the formula from num_query_heads down to num_kv_heads without changing how many distinct query projections the model computes. It is usually framed as a modeling choice with a small quality cost; from the KV-cache-footprint formula in Section 3.2, it is just as accurately described as a direct, linear lever on the single largest per-request memory cost in serving. Section 4.2 shows it is also a lever on decode's arithmetic intensity, not only its memory footprint — the same architectural choice pays off on both axes of the roofline at once, the way quantization did in the roofline chapter.

4. Prefill vs Decode, Through the Roofline Lens

Lesson 1, Part 2 derived that batch-1 decode is a matrix-vector multiply (GEMV) with arithmetic intensity AI = 2/S, where S is bytes per element — a result independent of model width or depth, sitting roughly two orders of magnitude below the ridge point of realistic edge and datacenter accelerators, and therefore deeply memory-bound. That chapter also showed batching turns the GEMV back into a GEMM, with AI_batched = 2B/S scaling linearly in batch size B. This section places prefill and decode on that same roofline, using exactly that machinery, and extends it to the KV cache specifically.

4.1 Prefill is prompt-length batching in disguise

Prefill processes an entire prompt of P tokens through the network in one pass. At each linear layer, this is a GEMM: a [P, d_in] activation block multiplied against a [d_in, d_out] weight matrix, not a [1, d_in] vector. Following the same derivation Lesson 1 used for the tiled matmul and for batched decode, and assuming d_in ≈ d_out ≈ d:

FLOPs        = 2 × P × d²
Ideal bytes  ≈ (2 × P × d  +  d²) × S        (activations in/out, plus the weight matrix, once each)
 
AI = 2Pd² / ((2Pd + d²)S)

For prompts of realistic length (P in the hundreds to low thousands) against realistic hidden sizes (d in the thousands), the weight term d² in the denominator, and the Pd² term in the numerator, dominate the smaller Pd cross-terms, and the expression reduces to the same shape as Lesson 1's batched-decode result:

AI_prefill ≈ 2P / S

This is not a coincidence — it is the same formula with the prompt length P standing in for the batch size B. Structurally, prefill is a very large decode batch: many token-vectors multiplied against the same resident weight matrix, exactly the condition that amortizes each streamed weight byte over many FLOPs. The difference is that decode has to assemble that batch across separate, independently-arriving requests (Section 5.2 covers exactly how), while prefill gets its batch for free, already assembled, because the whole prompt arrived as one unit. A 1,000-token prompt in fp16 gives AI_prefill ≈ 2,000/2 = 1,000 FLOP/byte — comfortably to the right of the roofline chapter's ridge points of roughly 150 for both the edge NPU and the datacenter GPU examples, i.e. solidly compute-bound. This is why prefill is universally described as GEMM-heavy and highly parallel: it is, quite literally, the large, well-tiled matmul regime from Lesson 1, Section 2.2, arrived at through a different route.

(At very long prompts, the O(P²) cost of the full causal attention score matrix — the same term identified in Section 2.2 — can itself grow large enough to matter, which is precisely the problem FlashAttention-style fused kernels are built to keep compute-bound rather than letting it degrade into a memory-bound intermediate-tensor problem. That is a kernel-implementation concern more than a prefill-versus-decode one, and is left for a later lesson on kernel-level optimization.)

4.2 Decode revisited: the cache adds a second memory-bound term

Batch-1 decode's AI = 2/S result, restated in Section 4's opening, covers only the linear projection GEMVs — the weight-streaming half of a decode step. But every decode step also runs attention against the KV cache, and that operation has its own arithmetic intensity, worth deriving with the same rigor. From Section 2.2's per-step attention cost, but now written generally with H_q query heads and H_kv key/value heads (GQA's group size is G = H_q / H_kv, and G = 1 recovers plain MHA):

FLOPs (attending 1 new query against L cached positions) = 4 × L × d_h × H_q
Bytes (reading the K and V cache for this layer)          = 2 × L × d_h × H_kv × S
 
AI_cache = (4 × L × d_h × H_q) / (2 × L × d_h × H_kv × S) = 2 × (H_q / H_kv) / S = 2G / S

Two things are immediately notable, and both echo Lesson 1's results almost exactly. First, L — the context length — cancels completely, the same way d_model cancelled out of the weight-streaming GEMV's AI = 2/S: cache-attention arithmetic intensity does not depend on how long the conversation has gotten, only on precision and GQA group size. Second, for plain MHA (G = 1) this collapses to exactly 2/S, the identical number as the weight-streaming result — decode is memory-bound on the cache for precisely the same structural reason it is memory-bound on the weights. GQA raises this specific number: Llama 3.1 70B's G = 8 in bf16 gives AI_cache = 16/2 = 8 FLOP/byte, an eight-fold improvement over plain MHA's 1 FLOP/byte — still far below any realistic ridge point, but a real, quantifiable, first-principles reason GQA earns its place in modern architectures beyond the memory-footprint savings already covered in Section 3.4.

The fact that AI_cache does not depend on L is easy to over-read, so it is worth stating precisely what does and doesn't change as context grows: the regime stays exactly as memory-bound at token 100,000 as it was at token 100 — the ridge-point comparison never flips. What grows with L is the absolute number of bytes that must move per step (2 × L × d_h × H_kv × S, linear in L), which the Roofline model's min(P_peak, B_memory × AI) bound does not capture, because that bound describes achievable throughput, not wall-clock time for a fixed amount of work. A longer context does not make decode "more memory-bound" in the roofline sense; it makes each already-memory-bound step take strictly longer, because there is strictly more cache to stream through the same bandwidth ceiling. This is the mechanism, stated with the same rigor as Lesson 1's derivations, behind the widely-observed fact that per-token decode latency creeps upward as a conversation lengthens even when the model and hardware are unchanged — and it is the same mechanism, from the opposite direction, behind why Section 3.3's memory-footprint numbers matter at all: a cache that has grown large enough to rival the model's own weights is also, by this derivation, a cache large enough to noticeably slow every subsequent decode step.

4.3 Two phases, two rooflines, plotted together

Same model, same weights, same hardware, two completely different points on the roofline chart, entered by the same request seconds apart. This is the structural fact the rest of this lesson's serving-system material is built on: a system tuned for one phase is close to worst-case for the other, because "more parallel work per byte moved" (what prefill wants) and "move fewer bytes per unit of already-scarce work" (what decode wants) are opposite optimization directions.

5. Why the Split Drives Real Serving-System Design

5.1 Disaggregated prefill/decode serving

Running prefill and decode on the same accelerator, interleaved request by request, means a compute-bound burst (a new prompt arriving) and a memory-bound trickle (existing conversations decoding one token each) are competing for the same MAC array and the same memory bus at the same time. A large incoming prefill can stall the in-flight decode steps of every other active conversation on that device — the prefill wants sustained GEMM throughput, and while it runs, decode's single-token GEMVs queue up behind it, directly inflating the per-token latency of requests that were already running.

Disaggregated serving — used in systems such as DistServe and Splitwise — addresses this by physically separating the two phases onto different accelerator pools: one set of GPUs runs prefill exclusively, handing off the freshly-built KV cache to a separate pool of GPUs that runs decode exclusively for as long as that conversation continues generating. Each pool can then be provisioned and scheduled for what it actually needs — the prefill pool tuned for peak compute throughput on large batched GEMMs, the decode pool tuned for memory bandwidth and packing as many concurrent low-arithmetic-intensity streams as memory allows — rather than a single fleet awkwardly averaged between two workloads with opposite resource profiles. The cost is the KV cache transfer between pools when a request moves from prefill to decode, which is itself an engineering problem (fast interconnects, cache formats designed for cheap serialization) that these systems address directly; the roofline analysis in Section 4 is exactly why that transfer cost is worth paying.

5.2 Continuous batching: keeping the memory-bound phase fed

Section 4.1 showed decode's arithmetic intensity scales with batch size exactly like prefill's scales with prompt length — AI_batched = 2B/S. The direct implication is that the fix for decode's memory-bound regime is the same fix Lesson 1 derived generically: batch more sequences' decode steps together so each streamed weight byte, and each streamed cache byte, gets reused across more concurrent tokens. But naive ("static") batching groups a fixed set of requests together and waits for every one of them to finish generating before admitting new requests — and different conversations finish at wildly different lengths, so a batch's effective size decays toward one as its shorter members complete, silently sliding AI_batched back down toward the unbatched, deeply memory-bound regime for whatever compute-cycles are spent waiting on the slowest remaining sequence.

Continuous batching (also called in-flight or iteration-level batching, introduced by the Orca serving system) fixes this at the scheduler level: instead of batching at the request level, it batches at the level of individual decode iterations, evicting a sequence from the batch the instant it finishes and admitting a new waiting request into the freed slot on the very next step. The batch composition changes every iteration, but its size — and therefore AI_batched — stays close to whatever the memory and compute budget can sustain, continuously, rather than sawtoothing down every time a request completes. This is not a scheduling nicety layered on top of the roofline result; it is the direct operational answer to the question the roofline model poses: decode's compute units are structurally starved unless something keeps the batch, and therefore the arithmetic intensity, high at every single step, not just on average across a request's lifetime.

5.3 PagedAttention: making the cache itself cheap to grow

Continuous batching wants to keep as many concurrent decode sequences resident as memory allows, and Section 3 established that each sequence's KV cache grows by a fixed increment every step and can, by itself, rival the size of the model's weights. Naive implementations allocate each sequence's cache as one contiguous memory block sized for the maximum sequence length up front, which wastes enormous amounts of memory on sequences that finish early, and cannot cheaply share cache memory between requests that share a common prefix (a system prompt, for instance).

PagedAttention, introduced with the vLLM serving system, borrows the fix directly from operating-system virtual memory: store each sequence's KV cache as a set of fixed-size blocks that need not be physically contiguous, with a per-sequence block table mapping logical cache positions to physical block locations — exactly the page-table abstraction an OS uses for process memory. This lets the cache grow one block at a time as a sequence generates, achieves near-zero internal fragmentation, and lets sequences that share a prefix literally share the same physical blocks for that prefix. The paper reports 2 to 4 times the throughput of prior systems at the same latency, with the improvement growing larger for longer sequences and larger models — precisely the regime Section 3.3's memory-footprint table flagged as the one where KV cache pressure is worst. It is the piece of engineering that makes continuous batching's appetite for many large, ever-growing caches actually affordable in accelerator memory, rather than a scheduling idea that memory fragmentation would otherwise undercut in practice.