Prefill and Decode: Measuring the Two Phases

Prefill and Decode: Measuring the Two Phases

Lab: Run the companion CPU exercise at /labs/track1/03-cache-and-gqa/ (prefill/decode trace, cache growth, timing).

Core idea

Prefill reads the entire prompt in one parallel forward pass. Decode generates one token at a time, each step dependent on the previous. This split is the single source of almost every latency and memory tradeoff in LLM serving.

---

One example: P=4, two output tokens

Consider a causal decoder with a 4-token prompt and two greedy output tokens. We will trace tensor shapes, cache growth, and ownership at each step.

Step	What happens	Prompt tokens	Cache length	K/V shape	First-token ownership
0	Prompt tokens t1–t4 are available	t1, t2, t3, t4	—	—	—
1	Prefill: all 4 tokens processed together in one forward pass. Final prompt-position logits predict output token g1. K/V states for all 4 tokens are stored in the cache.	t1, t2, t3, t4	P = 4	[B,H,4,D]	t4 is the last prompt token; its logits produce g1
2	Decode step 1: consume token g1. With the KV cache, the new token attends to cached K/V for t1–t4 plus its own K/V. Cache grows from 4 to 5. Logits predict output token g2.	t1, t2, t3, t4, g1	P+1 = 5	[B,H,5,D] (usable)	g1 is now the most recent token; its logits produce g2
3	Decode step 2: consume token g2. Attend to cached K/V for t1–t4 and g1, plus g2’s own K/V. Cache grows from 5 to 6.	t1, t2, t3, t4, g1, g2	P+2 = 6	[B,H,6,D] (usable)	g2 produces the final output of this trace

Key point: Each decode step processes exactly one new token. The cache extends by one column per step. There is no “extra model pass to produce output token 1”—the prefill pass already produces the logits for g1; decode step 1 consumes g1 and produces logits for g2.

Prefill and decode phase boundary

P=4 prompt tokens, two decode steps. Cache grows from 4 → 5 → 6.

Read this diagram as text

Prefill processes four prompt tokens together, creating K/V states and a cache of P columns. Decode step 1 consumes the first output token, attends to cached K/V plus its own, and extends the cache to P+1. Decode step 2 extends the cache to P+2. Tensor shapes: prefill K/V [B,H,P,D]; new K/V per step [B,H,1,D]; usable K/V after g1 [B,H,P+1,D]; usable K/V after g2 [B,H,P+2,D]. First-token ownership: t4 → g1 → g2.

Shape annotations for the example

• Prefill K/V: [B, H, P, D] — one key/value pair per token, per KV head, per head dimension, for all prompt tokens together. With B=1, H=8, P=4, D=64 this is [1, 8, 4, 64].
• First decode input: [B, 1] — a single token ID, the output from the previous step (g1 for step 2, or the sampled g1 after prefill for step 3).
• New K/V per decode step: [B, H, 1, D] — the key and value for the single new token, computed anew and written into the cache.
• Usable K/V after g1: [B, H, P+1, D] — the cache now contains the original P prompt columns plus 1 generated column. Any future attention pass can read all P+1 columns.
• Usable K/V after g2: [B, H, P+2, D] — the cache contains P prompt columns plus 2 generated columns.

These shapes assume a conventional decoder-only architecture where every attention layer stores K and V for every token in the context, across all KV heads. Grouped-query attention (GQA) or multi-query attention (MQA) would change the H dimension but the overall structure is the same.

---

Model-side duration vs. client-observed latency

Metric	What it measures	What it includes
Model-side prefill duration	Time for one forward pass over P prompt tokens.	GPU compute for Q/K/V projections and attention over the full prompt.
Model-side per-step decode duration	Time for one forward pass over one new token + cached K/V.	GPU time to read model weights, load the growing cache, compute a single attention pass, and produce logits.
Time to first token (TTFT)	Client-observed time from request submission to first generated token.	Model-side prefill duration plus queueing time, tokenization, tokenizer‑model transfer, and sampling. TTFT is not identical to prefill time.
Inter-token latency	Client-observed time between successive output tokens.	Model-side per-step decode duration plus any scheduler overhead, but typically dominated by the decode step itself.

In short: TTFT ≠ prefill time, and inter-token latency ≠ per-step decode time, because client-observed metrics fold in system-level work that the model does not control.

---

Measurement recipe

To measure the two phases separately and reliably, fix all variables that could confound the results. Here is a reproducible template; do not use it to report fabricated numbers — use it as a checklist for your own environment.

Fixed setup

Variable	Value (example)	Why it matters
Model checkpoint	e.g. GPT-2 small (124M) or Llama-2-7B	Different architectures have different FLOP counts and cache patterns.
Hardware	e.g. NVIDIA A100 40 GB, or Apple M2 Ultra	GPU compute capability and memory bandwidth differ dramatically.
Software stack	PyTorch 2.3, Transformers 4.41, CUDA 12.4	Kernel versions, kernel fusions, and CUDA ops affect both prefill and decode.
Precision	FP16 (half) or BF16	Mixed precision changes FLOP rate and memory traffic.
Batch size	B = 1 (single request) or a stated value	Larger batches change GPU utilization and can shift bottlenecks.
Prompt length P	e.g. P = 4 tokens	Determines the size of the prefill forward pass and the initial cache.
Output length	e.g. 2 tokens (minimum for two-phase trace)	Determines how many decode steps run.
Cache setting	use_cache=True (default) or use_cache=False	With use_cache=False the cache is not maintained; every step recomputes K/V.
Warmup runs	5–10 runs before timing	Allows GPU drivers and caches to stabilize; excludes first-run overhead.
GPU synchronization	torch.cuda.synchronize() (GPU) or performance.now() (CPU)	GPU kernels are asynchronous; synchronization is required for wall-clock accuracy.
Repeated runs	e.g. 20 runs, report median or mean	Reduces noise from system interference; report both central tendency and spread.

Separate phase timing

import statistics

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Placeholders you must set for your own run
MODEL_ID = "gpt2"          # fixed checkpoint under test
PRECISION = torch.float16  # FP16 / BF16 / FP32
WARMUP = 5                 # timing-excluded warmup passes
OUTPUT_LENGTH = 8          # decode steps to time

# 1. Load model and tokenizer (fixed checkpoint, fixed precision)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=PRECISION)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# 2. Build prompt and tokenize
prompt = "Once upon a time"
prompt_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)  # [1, P]
B, P = prompt_ids.shape

def sync():
    if torch.cuda.is_available():
        torch.cuda.synchronize()

# 3. Warmup: passes that are not timed
with torch.no_grad():
    for _ in range(WARMUP):
        model(input_ids=prompt_ids, use_cache=True)

# 4. Time prefill (all P prompt tokens in one pass)
sync()
prefill_start = torch.cuda.Event(enable_timing=True)
prefill_end = torch.cuda.Event(enable_timing=True)
prefill_start.record()
with torch.no_grad():
    out = model(input_ids=prompt_ids, use_cache=True)
prefill_end.record()
sync()
prefill_ms = prefill_start.elapsed_time(prefill_end)

# The final prompt-position logits predict output token 1 (no extra pass)
prefill_logits = out.logits[:, -1, :]        # [B, V]
past = out.past_key_values                   # per-layer (key, value)

# 5. Time each decode step (one new token + growing cache)
next_token_id = prefill_logits.argmax(-1, keepdim=True)   # [B, 1]
decode_ms_list = []

for _ in range(OUTPUT_LENGTH):
    sync()
    d_start = torch.cuda.Event(enable_timing=True)
    d_end = torch.cuda.Event(enable_timing=True)
    d_start.record()
    with torch.no_grad():
        out = model(input_ids=next_token_id, past_key_values=past, use_cache=True)
    d_end.record()
    sync()
    decode_ms_list.append(d_start.elapsed_time(d_end))
    past = out.past_key_values                                   # cache grows by 1
    next_token_id = out.logits[:, -1, :].argmax(-1, keepdim=True)  # [B, 1]

# 6. Report
cache_len = past[0][0].shape[-2]             # layer-0 key: [B, H_kv, seq, D]
print(f"Prefill (P={P}): {prefill_ms:.1f} ms (single run)")
print(f"Per-step decode (mean): {statistics.mean(decode_ms_list):.1f} ms "
      f"+/- {statistics.stdev(decode_ms_list):.1f} ms")
print(f"Cache length after {OUTPUT_LENGTH} steps: {cache_len} "
      f"(expected {P + OUTPUT_LENGTH})")

This template is illustrative: run it yourself against the checkpoint and hardware you name. No timing numbers in this lesson were measured on this site.

What to report:

• Model name / checkpoint version.
• Hardware (GPU name, CPU, RAM).
• Software (PyTorch, Transformers, CUDA versions).
• Precision (FP16, BF16, FP32).
• Batch size B.
• Prompt length P and output length (number of decode steps).
• Warmup count.
• Prefill median/mean time in ms.
• Per-step decode median/mean time in ms.
• Final cache length after all decode steps.
• Whether the measurement was on GPU or CPU, and synchronization method used.

No fabricated benchmark results are included. The table above is intentionally blank; replace the example values with measurements from your own environment.

---

Why prompt length, output length, and batch size affect the two phases differently

Parameter	Effect on prefill	Effect on decode
Longer prompt (larger P)	More tokens are processed in the single parallel pass. The matrix multiplications for Q/K/V projections and attention scale roughly with O(P²) (full attention) or O(P) (with chunked/prefix caching). Prefill time increases, and the initial cache size grows linearly with P.	The per-step decode cost is independent of P in terms of arithmetic (still one token), but the cache read cost grows linearly with P. Each decode step must read P cached key/value pairs, so per-step time slowly increases as the prompt gets longer.
Longer output (more decode steps)	No effect; prefill runs once regardless of how many tokens follow.	Total generation time multiplies the per-step decode time by the number of output tokens. A 500-token answer requires 500 sequential decode passes, making output length the dominant factor in total generation time.
Larger batch size (more concurrent requests)	Parallelism across requests increases. GPU can schedule more matrix work, but memory contention also increases. At small batch sizes, prefill may be limited by weight‑loading time; at larger batches it typically becomes more compute‑bound.	Decode per-step time per request typically increases with batch size because cached K/V for all concurrent requests compete for memory bandwidth on GPU. At very high batch sizes decode can become compute‑bound, but the usual pattern is memory‑bound contention.

Qualifications: These are general tendencies, not universal laws. The actual bottleneck depends on model size, sequence lengths, hardware, precision, and the serving system’s batching strategy. Avoid claims such as “prefill is always compute-bound and decode is always memory-bound”—the reality is a shifting balance that can be observed only by measuring.

---

Measurement illustration (blank results table)

Metric	Value
Model
Hardware
Software stack
Precision
Batch size B
Prompt length P
Output length (decode steps)
Warmup runs
Prefill time (ms)
Per-step decode time (ms, median)
Final cache length
Sync method
Notes

---

Key takeaways

• Prefill is a one-time parallel pass over the prompt that produces the first-token logits and fills the KV cache. Its duration contributes to time to first token (TTFT), but TTFT also includes queueing, tokenization, transfer, and sampling.
• Decode is a sequential loop: one forward pass per output token, each step reading the growing cache and appending one new K/V column. Per-step decode duration sets the inter-token latency, and total output length multiplies this cost.
• Tensor shapes follow a clear pattern: prefill K/V is [B, H, P, D]; each new K/V added by decode is [B, H, 1, D]; the usable cache after S steps is [B, H, P+S, D].
• Prompt length affects both phases: it increases prefill time (more parallel work) and per-step decode time (larger cache read), but the effects are different in nature.
• Output length affects only decode: total generation time scales linearly with the number of output tokens.
• Batch size affects both phases but in opposite directions: larger batches can speed up prefill through parallelism but typically slow down decode through memory‑bandwidth contention.
• A reproducible measurement recipe fixes model, hardware, software, precision, lengths, cache setting, warmup, synchronization, and repeat runs, and measures prefill and decode separately.

---

Sources and scope

• Hugging Face cache explanation documents per-layer key/value shapes and cache updates.
• vLLM benchmarking CLI shows how to measure TTFT and inter-token latency with fixed seeds and hardware.
• Pope et al., “Efficiently Scaling Transformer Inference” (arXiv:2211.05102) analyzes memory bandwidth, weight loading, KV‑cache loading, and MFU for prefill and decode.
• Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention” (arXiv:2309.06180) describes block allocation, shared prefixes, and high-request-rate behaviour.
• vLLM automatic prefix caching describes reuse of shared prompt‑prefix KV blocks across requests.

The code in this lesson is a measurement template; real serving systems may batch requests, overlap prefill and decode across calls, or use speculative decoding to accelerate generation. Those topics are covered in planned lessons orders 13–14 (KV cache and context windows) and the inference optimisation track.

Related lessons

• LLM Internals & Architecture · Inference foundationsHow LLM inference works9 min
• LLM Internals & Architecture · Inference foundationsUnderstanding the KV Cache12 min
• LLM Internals & Architecture · Generation, KV cache, and contextPretrained Weights and Text Generation12 min