7.1 The Anatomy of LLM Inference: Prefill, Decode & The KV Cache¶

Training a large language model is a one-time (or occasional) capital expense. Serving it is a recurring operational cost that runs every second of every day, for every user, forever. Most of the money a deployed LLM ever consumes is spent here, in the inference loop. So if you want to understand why inference systems look the way they do — why vLLM and SGLang and TensorRT-LLM exist, why “continuous batching” is a phrase people say with reverence, why everyone obsesses over a thing called the KV cache — you have to start at the bottom, with the raw mechanics of how a transformer turns a prompt into tokens.

This chapter is the foundation for all of Part VII. We will dissect a single forward-generation request into its two physically distinct phases — prefill and decode — and show that they have opposite performance characteristics. We will derive the KV cache from first principles, compute its size in gigabytes for real models, and prove why decode is bandwidth-bound rather than compute-bound. From there the latency vocabulary (TTFT, TPOT, ITL), the tokens-per-second ceiling, and Little’s law for throughput will all fall out as consequences of two or three simple physical facts. By the end you should be able to look at a GPU spec sheet and a model config and estimate, on a napkin, how fast it can serve.

We assume you already understand the transformer forward pass; if attention itself is fuzzy, revisit The Attention Mechanism From Scratch. We lean heavily on the Roofline Model and on GPU Architecture & The Memory Hierarchy, so keep those mental models handy.

Two Phases, One Model¶

A decoder-only LLM generates text autoregressively: it produces one token at a time, and each new token is conditioned on every token that came before it. Given a prompt of \(S\) tokens, generating \(T\) new tokens, the model effectively runs a forward pass for positions \(S+1, S+2, \dots, S+T\). The crucial observation is that these positions are not all generated the same way.

The very first forward pass — the one that ingests the entire prompt — is special. The model sees all \(S\) prompt tokens at once and produces the first output token. This is prefill (also called the “prompt” or “context encoding” phase). Every subsequent step ingests exactly one token (the one just generated) and produces the next. This is decode (the “generation” or “extension” phase).

Why split it this way? Because of the causal attention mask. Token at position \(i\) may attend only to positions \(\le i\). During prefill, we have all \(S\) tokens in hand, so we can compute their attention in a single batched, masked matrix multiply — the whole prompt is processed in parallel, the same way a training forward pass works. During decode, we are generating position \(S+t\) and do not yet have positions beyond it, so we are forced into a strictly sequential loop: one token, then the next, then the next.

These two phases run the same weights through the same layers. What differs is the shape of the activation tensors flowing through them, and that single difference — a sequence dimension of \(S\) versus a sequence dimension of \(1\) — flips the bottleneck from compute to memory. Everything in this chapter is a downstream consequence of that flip.

A naive generation loop (and why it is quadratically wasteful)¶

Let us write the simplest possible autoregressive loop, with no cache at all, to expose the redundancy that the KV cache later eliminates.

import torch

@torch.no_grad()
def generate_naive(model, input_ids, max_new_tokens):
    """
    The simplest correct autoregressive decoder: re-run the FULL forward
    pass on the entire growing sequence at every step. Correct, but O(T^2)
    in attention work because we recompute K and V for all old tokens
    every single step. This is what the KV cache fixes.
    """
    for _ in range(max_new_tokens):
        # logits over the WHOLE sequence; we only need the last position
        logits = model(input_ids).logits        # [B, seq_len, vocab]
        next_token = logits[:, -1, :].argmax(-1, keepdim=True)  # greedy
        input_ids = torch.cat([input_ids, next_token], dim=1)   # grow seq
    return input_ids

This is correct but catastrophically slow. At decode step \(t\), the sequence has length \(S + t\), and we recompute the keys and values for all of those positions even though only the newest token is new. The cumulative attention work scales like \(\sum_{t} (S+t)^2\) — quadratic in total length and full of redundant recomputation. Almost everything from the previous step is identical and could have been saved. That observation is the entire motivation for the KV cache.

The KV Cache: Mechanism and Math¶

Recall a single attention head. For query, key, and value projections, each token’s hidden state \(x_i \in \mathbb{R}^{d}\) is mapped to:

The attention output for query position \(i\) is a softmax-weighted sum over the keys and values of all positions \(j \le i\):

Here is the key structural fact: the keys \(k_j\) and values \(v_j\) for past positions do not depend on the current query position \(i\). Once we have computed \(k_3\) and \(v_3\) for token 3, those vectors are fixed forever — they are reused, unchanged, by every future query. The query \(q_i\) changes each step, but the \(K\) and \(V\) of the past are constant.

So we cache them. The KV cache is exactly the running collection of all \(k_j\) and \(v_j\) computed so far, for every layer and every head. At each decode step we:

1. compute \(q, k, v\) for the single new token only,
2. append its \(k\) and \(v\) to the cache,
3. compute attention of the one new query against the entire cached \(K\) and \(V\).

This turns the per-step attention from “recompute everything” into “compute one new KV and read the rest.” The recomputation is gone; the price is the memory to hold the cache. We have made the classic systems trade: spend memory to save compute.

A from-scratch KV-cached attention step¶

import torch
import torch.nn.functional as F

class CachedSelfAttention(torch.nn.Module):
    """
    Single-head causal self-attention with an explicit KV cache.
    Strips away batching/multi-head bookkeeping to expose the mechanism.
    """
    def __init__(self, d_model, d_head):
        super().__init__()
        self.Wq = torch.nn.Linear(d_model, d_head, bias=False)
        self.Wk = torch.nn.Linear(d_model, d_head, bias=False)
        self.Wv = torch.nn.Linear(d_model, d_head, bias=False)
        self.scale = d_head ** -0.5
        self.k_cache = None  # [seq_so_far, d_head]
        self.v_cache = None

    def reset(self):
        self.k_cache = None
        self.v_cache = None

    def forward(self, x):
        # x: [n_new_tokens, d_model].
        #   - prefill: n_new_tokens == S (the whole prompt)
        #   - decode:  n_new_tokens == 1 (just the last token)
        q = self.Wq(x)                      # [n_new, d_head]
        k = self.Wk(x)                      # [n_new, d_head]
        v = self.Wv(x)                      # [n_new, d_head]

        # Append the new keys/values to the cache (the whole trick).
        if self.k_cache is None:
            self.k_cache, self.v_cache = k, v
        else:
            self.k_cache = torch.cat([self.k_cache, k], dim=0)
            self.v_cache = torch.cat([self.v_cache, v], dim=0)

        # Attend the new query(ies) against ALL cached keys/values.
        scores = (q @ self.k_cache.T) * self.scale  # [n_new, seq_so_far]

        # Causal mask: needed during prefill (n_new = S) so that token i
        # cannot see j > i. During decode (n_new = 1) the single new query
        # is allowed to see everything in the cache, so no mask is required.
        if q.shape[0] > 1:
            seq = self.k_cache.shape[0]
            past = seq - q.shape[0]
            i = torch.arange(q.shape[0]).unsqueeze(1) + past
            j = torch.arange(seq).unsqueeze(0)
            scores = scores.masked_fill(j > i, float("-inf"))

        attn = F.softmax(scores, dim=-1)
        return attn @ self.v_cache          # [n_new, d_head]

Notice the asymmetry baked into the shapes. In prefill, x has \(S\) rows and the attention score matrix is \(S \times S\) — a big matmul. In decode, x has a single row and the score matrix is \(1 \times (\text{seq\_so\_far})\) — a skinny matrix-times-matrix that is really a sequence of dot products. That shape difference is the whole story of why the two phases behave so differently, which we quantify next.

KV cache size: the formula that governs everything¶

The cache must store, for every transformer layer, a key and a value vector for every token seen so far. Let:

• \(L\) = number of layers,
• \(H_{kv}\) = number of key/value heads (with Multi-Query or Grouped-Query Attention this is fewer than the number of query heads — see Multi-Head Attention, MQA, GQA & MLA),
• \(d_h\) = dimension per head,
• \(S\) = sequence length (tokens cached),
• \(B\) = batch size (number of concurrent sequences),
• \(P\) = bytes per element (2 for fp16/bf16, 1 for fp8/int8).

The factor of 2 below is for storing both K and V. The total KV-cache size in bytes is:

It is worth dwelling on what this formula says and does not say. The KV cache grows linearly with sequence length and linearly with batch size. It does not depend on the model’s MLP width, vocabulary, or the number of query heads — only the KV heads matter, which is precisely why GQA/MQA were invented to shrink it. And note \(H_{kv} \cdot d_h\) is the total KV hidden dimension, often written \(d_{kv}\).

def kv_cache_bytes(batch, seq_len, n_layers, n_kv_heads, head_dim, dtype_bytes=2):
    """KV-cache size in bytes. Factor 2 = keys AND values."""
    return 2 * batch * seq_len * n_layers * n_kv_heads * head_dim * dtype_bytes

# Llama-3-style 8B: 32 layers, 8 KV heads (GQA), head_dim 128, bf16.
gib = kv_cache_bytes(batch=1, seq_len=8192, n_layers=32,
                     n_kv_heads=8, head_dim=128, dtype_bytes=2) / 2**30
print(f"{gib:.2f} GiB per 8k-token sequence")   # ~1.0 GiB

The collapse from “weights are the big thing” to “the KV cache is the big thing” is the central surprise of inference engineering. For long contexts and high concurrency, the cache can dwarf the model itself. Managing it well — packing it tightly, paging it, reusing shared prefixes — is the subject of PagedAttention & KV-Cache Memory Management and Prefix Caching & KV-Cache Reuse. This chapter just establishes why it matters so much.

Why Decode Is Memory-Bound and Prefill Is Compute-Bound¶

This is the most important concept in the whole of inference serving, so we will derive it carefully using the arithmetic intensity lens from the Roofline Model. Arithmetic intensity is the ratio of floating-point operations performed to bytes of memory moved:

A GPU has two relevant peak rates: compute throughput (FLOP/s) and memory bandwidth (bytes/s). The ridge point is their ratio, \(I^\* = \text{FLOP/s} \div \text{bytes/s}\). If a kernel’s intensity \(I > I^\*\) it is compute-bound (limited by the math units); if \(I < I^\*\) it is memory-bound (limited by how fast you can feed data from HBM). For a modern data-center GPU the ridge point is on the order of a few hundred FLOP/byte. An A100 delivers roughly 312 TFLOP/s of bf16 tensor-core compute against roughly 2.0 TB/s of HBM bandwidth, giving a ridge near \(\sim\)150 FLOP/byte. An H100 (on the order of 990 TFLOP/s dense bf16 against ~3.3 TB/s of HBM3) sits nearer \(\sim\)300, and Blackwell-generation parts raise both axes together — HBM3e bandwidth on the order of 8 TB/s alongside proportionally more tensor-core throughput — so the ridge stays in the same few-hundred range. Hold that “~100–300 FLOP/byte” number in your head: every hardware generation has moved the numerator and the denominator together, which is exactly why the conclusion we are about to draw has not changed in five years.

The decisive quantity: tokens per forward pass¶

Consider one matrix multiply that dominates the transformer: a token’s hidden vector times a weight matrix \(W \in \mathbb{R}^{d \times d}\). If we process \(N\) tokens at once, we compute \(X W\) where \(X\) is \(N \times d\). The cost is:

• FLOPs: \(2 N d^2\) (multiply-add over an \(N\times d\) by \(d \times d\) matmul).
• Bytes moved: we must read \(W\) once (\(d^2\) elements \(\times P\) bytes) plus the activations (\(N d \cdot P\)). For the large weight matrices that dominate an LLM, the weight read dominates when \(N\) is small, so bytes \(\approx d^2 P\).

Arithmetic intensity is therefore approximately:

The intensity scales with \(N\), the number of tokens processed together. That is the entire argument. Read it twice.

• Prefill processes the whole prompt at once: \(N = S\), often hundreds or thousands of tokens. Even modest \(S\) pushes \(I = 2N/P\) well above the ridge point, so prefill is compute-bound — it saturates the tensor cores. Doubling your FLOPs roughly doubles prefill throughput; the weights are read once and amortized over many tokens.

• Decode processes exactly one token per sequence: \(N = 1\) (for a single request). Then \(I \approx 2/P = 1\) FLOP/byte for fp16 — two orders of magnitude below the ridge point. Decode is therefore deeply memory-bound. The GPU spends its time streaming the entire weight matrix (and the entire KV cache) from HBM to perform a tiny amount of arithmetic, and the tensor cores sit mostly idle.

This is why batching helps decode enormously but barely helps prefill. Prefill already saturates compute, so cramming more sequences in does not raise per-token throughput much (and can even hurt latency). Decode, by contrast, reads the same weights regardless of how many sequences share the step — so if you decode \(B\) sequences together, you read \(W\) once and reuse it across all \(B\) tokens, lifting \(N\) from \(1\) to \(B\) and intensity from \(\approx 2/P\) to \(\approx 2B/P\). Batching decode is essentially free FLOPs riding on memory traffic you were paying for anyway. This single insight is the entire reason continuous batching exists — see Continuous Batching & Request Scheduling.

The catch: batching amortizes weights, not the KV cache¶

There is an asymmetry hidden in that argument, and missing it is the most common error in napkin models of decode. Batching amortizes the weight read — one copy of \(W\) serves all \(B\) tokens — but it cannot amortize the KV-cache read, because every sequence owns a private cache that no other sequence in the batch can reuse. Count the attention work in one decode step, per layer, per sequence: reading \(K\) and \(V\) costs \(2 \cdot S \cdot H_{kv} \cdot d_h \cdot P\) bytes, while the two matmuls (\(q K^\top\) and \(\text{attn}\,V\)) cost about \(4 \cdot S \cdot H_q \cdot d_h\) FLOPs. The arithmetic intensity of decode attention is therefore

which is a small constant that does not grow with batch size at all — the \(S\) cancels, and \(B\) never appears. For a Llama-3-8B-style model (\(H_q = 32\), \(H_{kv} = 8\), bf16) that is \(2 \cdot 32 / (8 \cdot 2) = 4\) FLOP/byte; for full MHA (\(H_q = H_{kv}\)) it collapses to \(2/P = 1\). Three consequences follow, and they justify a large fraction of Part VII:

• Decode attention is permanently memory-bound, no matter how large the batch. Only the weight matmuls escape into the compute-bound regime as \(B\) grows.
• GQA/MLA speed up decode, not just shrink it. Cutting \(H_{kv}\) by the group size \(g\) raises \(I_{\text{attn}}\) by exactly \(g\) and cuts the bytes streamed by \(g\) — a benefit beyond the capacity argument of the last section. Reading the same key once for \(g\) query heads is only possible if the kernel is GQA-aware, which modern FlashAttention/FlashDecoding kernels are.
• At long context the KV traffic overtakes the weight traffic, and the “weights / bandwidth” ceiling below becomes an over-estimate. Exercise 5 works this out numerically: at \(B=32\), \(S=8192\) the cache read is roughly twice the weight read for an 8B model.

Batch-1 long-context decode has a further problem: a single query row against one cache gives the GPU almost no parallelism to fill its SMs. FlashDecoding fixes this by splitting the KV cache along the sequence axis, attending in parallel chunks, and combining the partial softmax statistics — the same online-softmax rescaling trick used in training, applied along a different axis. See FlashAttention 2 & 3.

A measurement you can run¶

You can feel the memory-bound nature of decode directly. The following microbenchmark shows that decode latency per token is nearly flat as batch size grows (because you are bandwidth-limited, not compute-limited), until the batch gets large enough to finally start saturating compute.

import torch, time

def bench_decode(model, batch_sizes, d_model, n_steps=50, device="cuda"):
    """
    Single-token (decode-shaped) forward passes at several batch sizes.
    Memory-bound behavior => per-step latency stays ~flat as batch grows,
    so tokens/sec scales almost linearly with batch. Compute-bound behavior
    would instead show latency rising proportionally with batch.
    """
    model.eval()
    for B in batch_sizes:
        x = torch.randn(B, 1, d_model, device=device, dtype=torch.bfloat16)
        # warmup (CUDA lazy init, autotuning, caches)
        for _ in range(5):
            with torch.no_grad():
                model(x)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(n_steps):
            with torch.no_grad():
                model(x)
        torch.cuda.synchronize()
        dt = (time.perf_counter() - t0) / n_steps
        print(f"B={B:4d}  {dt*1e3:6.2f} ms/step  "
              f"{B/dt:8.0f} tok/s  {dt*1e3/B:6.3f} ms/tok")

On real hardware you will see the per-token time barely move from \(B=1\) to \(B=16\) or \(B=32\) — you are getting those extra tokens “for free” off the same weight reads — and only at large batch do you finally cross into the compute-bound regime where latency climbs. That flat region is the memory-bound signature, and exploiting it is the single biggest lever in throughput-oriented serving.

The Tokens-Per-Second Ceiling¶

Because single-stream decode is memory-bound, we can compute a remarkably tight upper bound on generation speed using nothing but the model size and the GPU’s memory bandwidth. Every decode step must read every weight of the model from HBM at least once (each weight participates in some matmul for the new token). For a Mixture-of-Experts model, replace “every weight” with “every active weight” — the dense trunk plus only the experts routed to this step’s tokens — which is precisely why a sparse model decodes far faster than a dense model of the same total parameter count, and why nearly every 2026 frontier model is an MoE; see Serving Mixture-of-Experts. For a dense model the time for one decode step is bounded below by the time to stream all the weights:

and the per-stream decode throughput is bounded above by the reciprocal:

This “weights divided by bandwidth” estimate is one of the most useful back-of-envelope tools in the field. (At large context lengths you must also add the KV-cache bytes read per step; for short contexts the weight read dominates.)

The flip side: this ceiling is per stream. If you batch \(B\) independent sequences, you still read the weights only once per step and serve \(B\) tokens, so aggregate throughput is roughly \(B \times\) the single-stream rate, right up until you exhaust KV-cache memory or finally saturate compute. Single-stream latency and aggregate throughput are different axes, and batching trades one for the other. Hold this thought — it is the seed of Little’s law below.

The Latency Vocabulary: TTFT, TPOT, ITL¶

Streaming LLM applications (a chat UI typing tokens at you) are not characterized by a single “latency” number. The user’s experience splits along the prefill/decode boundary, and the industry has settled on a small, precise vocabulary. Learn these exact definitions — they appear in every serving SLA and every interview.

The decomposition is clean and worth memorizing:

TTFT is what the user feels as “responsiveness” — how long the cursor blinks before text appears. It is governed by prefill, so it scales with prompt length \(S\) (longer prompts take longer to encode) and by queueing delay under load. Because prefill is compute-bound, TTFT is improved by more FLOPs, by chunked prefill (slicing a long prompt so it interleaves with ongoing decodes), and by prefix caching (skipping prefill entirely for a shared prompt prefix). See Disaggregated Prefill/Decode & Chunked Prefill.

TPOT is what the user feels as “typing speed” — the steady cadence once text starts flowing. It is governed by decode, so it is memory-bound and follows the tokens/sec ceiling above. A common product target is to keep TPOT below the human reading rate (very roughly 5–10 tokens/sec is readable, so a TPOT under ~100–200 ms feels smooth). Crucially, batching raises aggregate throughput but can raise per-stream TPOT, because a larger decode batch takes slightly longer per step. This tension — throughput versus per-user latency — is the central knob of serving.

def measure_streaming_latency(stream_fn, prompt):
    """
    Wrap a streaming generation call and report TTFT, mean TPOT, and the
    per-token ITL series. `stream_fn(prompt)` must yield one token at a time.
    """
    import time
    t_start = time.perf_counter()
    t_prev = None
    ttft = None
    itls = []                      # inter-token latencies (seconds)
    n_tokens = 0
    for _token in stream_fn(prompt):
        now = time.perf_counter()
        if ttft is None:
            ttft = now - t_start   # arrival -> first token
        else:
            itls.append(now - t_prev)
        t_prev = now
        n_tokens += 1
    tpot = sum(itls) / len(itls) if itls else float("nan")
    return {
        "ttft_ms": ttft * 1e3,
        "tpot_ms": tpot * 1e3,                # mean inter-token latency
        "p99_itl_ms": sorted(itls)[int(0.99 * len(itls)) - 1] * 1e3 if itls else None,
        "output_tokens": n_tokens,
        "decode_tps": (n_tokens - 1) / sum(itls) if itls else float("nan"),
    }

In practice you rarely write that timing harness yourself, because the serving engines ship one that already reports exactly this vocabulary. vLLM’s benchmark client drives an OpenAI-compatible endpoint at a chosen arrival rate and prints TTFT, TPOT and ITL percentiles alongside aggregate throughput:

# Terminal 1: start the engine (OpenAI-compatible server on :8000).
vllm serve meta-llama/Llama-3.2-1B-Instruct --max-model-len 8192

# Terminal 2: load-test it and get the metrics defined in the table above.
# (Recent vLLM exposes this as `vllm bench serve`; older releases ship the
#  same tool as benchmarks/benchmark_serving.py. Check --help for your
#  version, since flag names move between releases.)
vllm bench serve \
  --model meta-llama/Llama-3.2-1B-Instruct \
  --dataset-name random --random-input-len 1024 --random-output-len 256 \
  --request-rate 4 --num-prompts 200

Sweeping --request-rate from well below saturation to well above it and plotting p99 TTFT against achieved throughput draws the latency–throughput curve for your deployment; the knee of that curve is the operating point you are actually choosing when you set a batch size. SGLang ships an equivalent harness (python -m sglang.bench_serving), and the two are close enough in interface that you can benchmark both engines with nearly the same command line.

When you benchmark a serving system, you report these as distributions, not point values — p50 and p99 TTFT, p50 and p99 TPOT — because tail latency under load is what breaks SLAs. A system with great median TPOT but a terrible p99 (caused by, say, a giant prompt prefill blocking the decode batch) feels janky and stutters. This is exactly the pathology chunked prefill was designed to cure.

Little’s Law: From Latency to Throughput¶

To reason about a fleet serving many users — not a single request — we borrow a result from queueing theory. Little’s law states that for any stable system in steady state, the average number of in-flight requests equals the arrival rate times the average time each request spends in the system:

where \(L\) is the average number of concurrent requests resident in the system, \(\lambda\) is the throughput (requests completed per second, which in steady state equals the arrival rate), and \(W\) is the average latency (time in system). It is astonishingly general — it assumes nothing about the arrival distribution or service discipline, only stability.

For LLM serving we apply it at the token level, which makes the bandwidth story precise. Let \(L\) be the number of sequences we can hold concurrently (capped by KV-cache memory), let \(W\) be the average time a sequence stays resident (roughly its total generation time, \(\approx T \cdot \text{TPOT}\)), and let \(\lambda\) be the sequence completion rate. Rearranged:

This ties the whole chapter together. The numerator \(L\) is set by KV-cache memory — the formula from earlier directly bounds how many sequences fit. The denominator \(W\) is set by decode speed — the bandwidth-bound TPOT ceiling. So your serving throughput is fundamentally governed by the two physical facts we derived: how much KV cache fits in HBM, and how fast HBM bandwidth lets you decode.

The practical consequence: to raise throughput you either increase \(L\) (shrink the KV cache via GQA/MLA, quantize the cache, page it more tightly, share prefixes) or decrease \(W\) (faster decode via quantized weights, speculative decoding, better kernels). Every serving optimization in Part VII is, underneath, an attack on one of these two terms. Continuous batching keeps \(L\) as full as possible at all times; PagedAttention lets you pack \(L\) tighter without fragmentation; speculative decoding shrinks \(W\) by emitting multiple tokens per forward pass.

Putting It Together: An Annotated Generation Loop¶

To consolidate everything, here is a complete, cache-using generation loop with the prefill/decode split made explicit and annotated with the performance regime of each phase. This is the conceptual skeleton inside every production engine; what vLLM and SGLang add is scheduling many of these loops together and managing the KV memory, which later chapters cover.

import torch

@torch.no_grad()
def generate(model, tokenizer, prompt, max_new_tokens=256,
             eos_id=None, device="cuda"):
    """
    KV-cached autoregressive generation, prefill/decode split made explicit.
    Uses a HF-style model whose forward() accepts past_key_values and
    returns updated ones (use_cache=True). Greedy for clarity; swap in
    your sampler from the decoding chapter for real use.
    """
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    S = input_ids.shape[1]

    # ---- PREFILL ----------------------------------------------------------
    # One big forward pass over all S prompt tokens. COMPUTE-BOUND: the SxS
    # attention and the S-row matmuls saturate the tensor cores. This builds
    # the initial KV cache and produces the first output token. Its cost is
    # what the user experiences as TTFT.
    out = model(input_ids=input_ids, use_cache=True)
    past = out.past_key_values                      # the KV cache, all layers
    next_id = out.logits[:, -1, :].argmax(-1, keepdim=True)
    generated = [next_id]

    # ---- DECODE -----------------------------------------------------------
    # Loop, one token at a time. MEMORY-BOUND: each step feeds a single token,
    # so it streams the full weights + KV cache from HBM to do tiny math. The
    # per-step time is TPOT. We pass only the NEW token plus the cache, never
    # the whole sequence again -- that is the KV cache doing its job.
    for _ in range(max_new_tokens - 1):
        out = model(input_ids=next_id, past_key_values=past, use_cache=True)
        past = out.past_key_values                  # cache grew by one token
        next_id = out.logits[:, -1, :].argmax(-1, keepdim=True)
        generated.append(next_id)
        if eos_id is not None and next_id.item() == eos_id:
            break

    return tokenizer.decode(torch.cat(generated, dim=1)[0])

Three things to internalize from this skeleton. First, prefill runs once, decode runs hundreds of times — for a typical chat turn with a 1,000-token prompt and a 300-token answer, almost all wall-clock time is in the decode loop even though prefill touches more tokens, because decode is the slow, memory-bound part repeated 300 times. Second, the cache (past) is the entire state — its size is the formula we derived, and managing it is the hardest part of real serving. Third, the only input to each decode step is one token plus the cache — which is exactly why many sequences can share a decode step (continuous batching) and why the weight read can be amortized across them.

From toy loop to real engines: the Cache API and CUDA graphs¶

The loop above is the real thing, but the object it passes around has evolved. Modern transformers returns a Cache object rather than the legacy tuple-of-tuples, and which subclass you pick is a performance decision, not a formality:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

name = "meta-llama/Llama-3.2-1B-Instruct"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(
    name, torch_dtype=torch.bfloat16).to("cuda").eval()

past = DynamicCache()          # grows by concatenation -- exactly our toy cache
ids = tok("The capital of France is", return_tensors="pt").to("cuda")

with torch.no_grad():
    out = model(**ids, past_key_values=past, use_cache=True)        # PREFILL
    next_id = out.logits[:, -1:, :].argmax(-1)                      # [B, 1]
    for _ in range(32):                                             # DECODE
        # `past` is updated in place by the Cache object each step.
        out = model(input_ids=next_id, past_key_values=past, use_cache=True)
        next_id = out.logits[:, -1:, :].argmax(-1)

print(past.get_seq_length())   # tokens currently cached, per layer

DynamicCache is the right default for a research loop and the wrong one for a latency-critical server: it reallocates and copies on every step, and its ever-changing tensor shapes defeat graph capture. StaticCache instead preallocates a buffer of max_cache_len and writes each new key/value in place at the step’s cache_position, so every decode step is the same fixed-shape graph. That is what lets model.generate(..., cache_implementation="static") compose with torch.compile(model, mode="reduce-overhead"), which captures the step as a CUDA graph and replays it — see Kernel Fusion, torch.compile, CUDA Graphs & Compilers. vLLM and SGLang arrive at the same design from the other direction: their KV cache is a preallocated block pool (never a torch.cat) and their decode step is a captured graph, which is why a first-principles PyTorch loop typically leaves a large factor on the table against them even at batch 1.

Key Takeaways¶

Further Reading¶

• Vaswani et al., Attention Is All You Need (2017) — the original transformer and the attention computation whose K/V we cache.
• Pope et al., Efficiently Scaling Transformer Inference (Google, 2022) — the canonical analysis of inference arithmetic, memory-bound decode, and the prefill/decode distinction; the source of much of this chapter’s framing.
• Dao et al., FlashAttention and FlashAttention-2 — IO-aware attention kernels that make prefill (and long-context decode) memory-efficient; see FlashAttention I.
• Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM, 2023) — the paper that reframed KV-cache memory as a paging problem; see vLLM Internals.
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (2019) — introduces Multi-Query Attention, the first major attack on KV-cache size; see MHA, MQA, GQA & MLA.
• Williams, Waterman & Patterson, Roofline: An Insightful Visual Performance Model (2009) — the arithmetic-intensity model underpinning the compute- vs memory-bound argument; see The Roofline Model.
• Little, A Proof for the Queuing Formula \(L = \lambda W\) (1961) — the throughput law used for fleet sizing.

Exercises¶

1. (Conceptual) A colleague proposes speeding up your serving system by increasing the decode batch size \(B\) from 1 to 32, and separately by increasing the prefill batch size from 1 to 32. Using the arithmetic-intensity argument \(I \approx 2N/P\), explain why the first change gives a near-32x aggregate throughput win while the second gives almost nothing. What single physical quantity (\(N\), the tokens processed per forward pass) is different between the two phases before you batch, and why does that make batching act so differently on each?

2. (Quantitative) Consider a hypothetical model with \(L = 80\) layers, \(H_{kv} = 8\) key/value heads (GQA), head dimension \(d_h = 128\), served in bf16 (\(P = 2\) bytes). Compute: (a) the KV-cache cost per token in KiB; (b) the KV-cache size of one 4096-token context in GiB; © how many such 4096-token contexts fit in a 50 GiB KV-cache budget.

3. (Quantitative) You serve a 13B-parameter model on an A100 with 2.0 TB/s of HBM bandwidth. Using the single-stream decode ceiling \(\text{tokens/sec} \le \text{BW} / (N_{\text{params}} \cdot P)\), compute the optimistic ceiling in (a) fp16 (\(P = 2\)) and (b) int4 (\(P = 0.5\)). By what factor does quantization raise the ceiling, and why - does it add FLOPs?

4. (Quantitative) Reuse the fleet from Exercise 2: \(L \approx 40\) concurrent sequences fit in KV memory. Suppose each request generates on average \(T = 1000\) tokens, and batched decode achieves a per-stream TPOT of 25 ms. Ignoring TTFT, use Little’s law to estimate (a) the sustainable request throughput \(\lambda\) and (b) the aggregate output-token throughput. © What happens if the arrival rate exceeds \(\lambda\)?

5. (Implementation) The chapter’s tokens/sec ceiling assumes the weight read dominates, which is only true for short contexts. At long context the decode step must also stream the growing KV cache from HBM. Using the chapter’s kv_cache_bytes helper, implement decode_ceiling_tps(...) that returns the aggregate tokens/sec ceiling for a batch of \(B\) sequences, each at sequence length \(S\), accounting for both the weight read and the KV-cache read per step. Then evaluate it for an 8B model (16 GB fp16 weights) with \(L=32\), \(H_{kv}=8\), \(d_h=128\), bf16, at \(B=32\) and \(S=8192\) on a 2.0 TB/s GPU, and compare against the weight-only ceiling.

def kv_cache_bytes(batch, seq_len, n_layers, n_kv_heads, head_dim, dtype_bytes=2):
    """KV-cache size in bytes. Factor 2 = keys AND values."""
    return 2 * batch * seq_len * n_layers * n_kv_heads * head_dim * dtype_bytes

def decode_ceiling_tps(weight_bytes, bw_bytes_per_s, batch, seq_len,
                       n_layers, n_kv_heads, head_dim, dtype_bytes=2):
    """
    Aggregate tokens/sec ceiling for a decode step of `batch` sequences,
    each at length `seq_len`, accounting for BOTH the one-time weight read
    and the KV-cache read for the whole batch. Memory-bound: step time is
    set by bytes streamed from HBM divided by bandwidth.
    """
    kv_bytes = kv_cache_bytes(batch, seq_len, n_layers, n_kv_heads,
                              head_dim, dtype_bytes)
    bytes_per_step = weight_bytes + kv_bytes       # weights + KV, per step
    t_step = bytes_per_step / bw_bytes_per_s       # seconds (lower bound)
    return batch / t_step                          # tokens/sec aggregate

W_BYTES = 16e9          # 8B params * 2 bytes (fp16)
BW      = 2.0e12        # 2.0 TB/s

# Weight-only baseline (chapter's ceiling), aggregated over the batch:
baseline = 32 / (W_BYTES / BW)         # B / t_weight_only
full     = decode_ceiling_tps(W_BYTES, BW, batch=32, seq_len=8192,
                              n_layers=32, n_kv_heads=8, head_dim=128,
                              dtype_bytes=2)
print(f"weight-only ceiling : {baseline:8.0f} tok/s")
print(f"with KV-cache reads : {full:8.0f} tok/s")