How LLM Inference Works: From a Request to a Streamed Response
Training happens once; inference happens billions of times. This piece follows a single request through a production serving stack — tokenization, the prefill and decode phases, the KV cache, sampling, continuous batching, quantization, speculative decoding, multi-GPU serving and the metrics that matter — with a small inference engine you can run that demonstrates each optimization and checks that it doesn't change the output.
The Release Candidate
Inference starts where post-training ends: a model that has been pretrained, fine-tuned, preference-optimized and evaluated, with its weights now frozen. What ships to the serving system is a small set of artifacts: the weights (in BF16, FP16, or a quantized INT8/INT4 format), the model config (vocabulary size, number of layers, heads), the tokenizer files, the generation config (default sampling settings, maximum tokens), and the safety policies — system prompt, guardrails and filters that wrap the model at serving time.
Serving Architecture
A production deployment is more than a model on a GPU:
From the request's point of view, seven things happen: the prompt arrives through an API; it's tokenized into IDs; the prefill phase processes all prompt tokens at once; the KV cache is created and stored; the decode phase generates tokens one at a time; each token is streamed back as soon as it exists; and generation stops at an end-of-sequence token, the maximum-token limit or a stop sequence.
The Two Phases: Prefill and Decode
Generating a response has two computationally different phases, and understanding the split explains almost every optimization in this article.
Prefill — the whole prompt at once
Decode — one token at a time
Each step processes only the newest token, reusing cached keys and values for everything before it, then appends one token. Repeats until a stop condition.| Aspect | Prefill | Decode |
|---|---|---|
| Tokens per pass | All prompt tokens (10s to 100,000s) | One per sequence |
| Parallelism | Highly parallel | Sequential |
| Bottleneck | Compute (FLOPs) | Memory bandwidth: weights and KV cache |
| Determines | Time to first token | Time between tokens |
| Our run (CPU) | 512 tokens in 147 ms: 3,480 tokens/s | 11.1 ms per token: 90 tokens/s |
Why decode is memory-bound
At every decode step the GPU must read all of the model's weights from memory to produce one token per sequence. Take a 6.74B-parameter model in BF16 — 13.5 GB of weights — on an H100, which reads memory at 3.35 TB/s and computes at up to 989 dense BF16 TFLOPS:
The KV Cache
In attention, each token's query is compared against the keys of all earlier tokens and blends their values (see the transformer explainer). Those keys and values don't change once computed — the causal mask means a token's K and V never depend on what comes after. So instead of recomputing them for the whole sequence at every step, the server stores them: the KV cache. Each decode step computes Q, K and V for the new token only, appends its K and V to the cache, and attends over everything cached.
The cache's price is memory. Per token it holds a key and a value vector for every layer:
Two consequences shape modern models and servers. First, the cache often limits how many requests fit on a GPU, so architectures shrink it: multi-query attention shares one key/value head across all query heads, and grouped-query attention (GQA) shares each K/V head among a group. Llama 2 70B uses 64 query heads but only 8 K/V heads, an 8× smaller cache — try the two 70B presets above. Second, sequences grow unpredictably, and reserving memory for the maximum length wastes most of it. vLLM's PagedAttention stores the cache in fixed-size blocks that don't need to be contiguous, like an operating system's virtual memory, cutting waste to near zero and fitting more sequences per batch — 2–4× the throughput of earlier systems at the same latency.
Sampling: Choosing the Next Token
The model outputs a probability for every vocabulary token; a sampling strategy picks one. Greedy decoding takes the most likely token. Temperature sharpens (below 1) or flattens (above 1) the distribution. Top-k keeps the k most likely tokens. Top-p (nucleus) keeps the smallest set whose probabilities add up to at least p. A repetition penalty down-weights tokens already generated (introduced with Salesforce's CTRL model). Stop conditions — the end-of-sequence token, a maximum length, or a stop string — end generation.
Batching and Continuous Batching
Because a decode step reads all the weights anyway, processing many sequences in the same step is nearly free until compute becomes the limit: the weights are read once and used for every sequence in the batch. Our CPU run shows the effect, even though a CPU hits its compute limit far sooner than a GPU:
Static batching groups requests and runs them together until all finish, so short requests wait for the longest one and new arrivals wait for the whole batch. Continuous batching — introduced as iteration-level scheduling in the Orca system — makes scheduling decisions at every decode step: finished sequences leave, new requests join immediately. On GPT-3 175B, Orca reported a 36.9× throughput improvement over NVIDIA FasterTransformer at the same latency. Continuous batching is now standard in serving systems such as vLLM and TensorRT-LLM.
Memory and Compute Optimizations
Memory
Quantization stores weights in fewer bits. LLM.int8() halved inference memory for a 175B model without performance loss; GPTQ quantizes weights to 3–4 bits in about four GPU hours for a 175B model, fitting OPT-175B on a single 80 GB GPU; AWQ protects the ~1% of salient weights, chosen by activation statistics. Paged KV cache (above), KV-cache quantization, and sharding the model across GPUs round out the toolkit.Compute
FlashAttention computes exact attention in tiles that stay in fast on-chip memory, cutting slow memory traffic. Fused kernels combine several operations into one GPU launch; CUDA graphs remove per-launch CPU overhead; and optimized libraries (vLLM, TensorRT-LLM, DeepSpeed) package all of it.Our engine's int8 test shows why weight-only quantization is attractive: storing each weight row as 8-bit integers plus one scale made the linear layers 4.0× smaller (102.8 MB → 25.8 MB), changed logits by at most 0.012, and produced identical greedy output. On a GPU, smaller weights also mean fewer bytes read per decode step — directly attacking the memory bottleneck.
Speculative decoding
Decode wastes compute because each step produces one token. Speculative decoding puts that idle compute to work: a small draft model cheaply proposes several tokens, and the large target model checks all of them in a single forward pass — the same parallel pass as prefill. Every proposed token that matches what the target would have produced is kept, plus the target's own next token, so each target pass yields one or more tokens. With the right acceptance rule the output distribution is exactly the target's; the original paper reported 2–3× speedups with identical outputs.
Our run used a 148K-parameter draft for a 5.1M-parameter target (about 35× smaller), proposing 4 tokens at a time with greedy decoding. It produced exactly the target's own output, 80 tokens in 21 target passes instead of 80, accepting 59 of the draft's tokens. How much faster that is in wall-clock time depends on the draft's cost and on hardware, which is why real systems measure it rather than assume it.
Multi-GPU Serving
When a model doesn't fit on one GPU, or one GPU is too slow, it's split — the same techniques as training at scale (covered in depth in Scaling Laws, GPUs and Distributed Training), with different trade-offs:
Metrics That Matter
| Metric | Meaning | Driven by |
|---|---|---|
| Time to first token (TTFT) | Delay from request to the first streamed token | Prefill, queueing |
| Inter-token latency | Time between streamed tokens | Decode step time, batch size |
| Throughput | Tokens per second, per GPU or overall | Batching, memory, kernels |
| GPU utilization | Share of compute or memory bandwidth in use | Batching, workload mix |
| Memory usage | Weights + KV cache + activations | Quantization, context length, batch |
| Cost per request | Compute time × hardware price | All of the above |
These pull against each other. A bigger batch raises throughput and lowers cost per token, but slows every step for everyone in it; our batch-64 step took 11× longer than batch 1. Production systems tune for a latency target — say, a maximum inter-token latency — and then maximize throughput within it. Finally, tokens are streamed to the user as soon as each is generated, so the user starts reading after the time to first token instead of waiting for the whole response.
Run the Engine
A miniature inference engine in one script, about two minutes on a laptop CPU: a GPT with a KV cache (checked against full recomputation), prefill versus decode timing, batched decoding, int8 weight quantization (checked against the original output), and greedy speculative decoding with a trained draft and target model (checked against the target's own output). Needs PyTorch 2.x.
"""A miniature inference engine: KV cache, prefill vs decode, batching, int8 quantization, speculative decoding."""
import copy, itertools, random, time
import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(0); random.seed(0); torch.set_num_threads(4)
class Attention(nn.Module):
def __init__(self, d, h):
super().__init__()
self.h, self.qkv, self.proj = h, nn.Linear(d, 3 * d), nn.Linear(d, d)
def forward(self, x, cache=None):
B, T, d = x.shape
q, k, v = (t.view(B, T, self.h, d // self.h).transpose(1, 2) for t in self.qkv(x).split(d, dim=-1))
if cache is not None: # decode: append new K, V to the cache
k, v = torch.cat([cache[0], k], 2), torch.cat([cache[1], v], 2)
causal = T > 1 # prefill needs the mask; one new token sees everything
out = F.scaled_dot_product_attention(q, k, v, is_causal=causal and cache is None)
return self.proj(out.transpose(1, 2).reshape(B, T, d)), (k, v)
class Block(nn.Module):
def __init__(self, d, h):
super().__init__()
self.ln1, self.ln2, self.attn = nn.LayerNorm(d), nn.LayerNorm(d), Attention(d, h)
self.up, self.down = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)
def forward(self, x, cache=None):
a, new_cache = self.attn(self.ln1(x), cache)
x = x + a
return x + self.down(F.gelu(self.up(self.ln2(x)))), new_cache
class GPT(nn.Module):
def __init__(self, V, d, h, L, T_max=1024):
super().__init__()
self.tok, self.pos = nn.Embedding(V, d), nn.Embedding(T_max, d)
self.blocks = nn.ModuleList(Block(d, h) for _ in range(L))
self.ln_f, self.head = nn.LayerNorm(d), nn.Linear(d, V, bias=False)
def forward(self, idx, caches=None, start=0):
x = self.tok(idx) + self.pos(torch.arange(start, start + idx.shape[1]))
new = []
for i, blk in enumerate(self.blocks):
x, c = blk(x, None if caches is None else caches[i])
new.append(c)
return self.head(self.ln_f(x)), new
@torch.no_grad()
def generate(model, prompt, n, use_cache=True):
"""Greedy decoding. With the cache: one prefill pass, then one token per step."""
ids = prompt.clone()
logits, caches = model(ids) # PREFILL: whole prompt in parallel
for _ in range(n):
nxt = logits[:, -1].argmax(-1, keepdim=True)
ids = torch.cat([ids, nxt], 1)
if use_cache:
logits, caches = model(nxt, caches, start=ids.shape[1] - 1) # DECODE: only the new token
else:
logits, _ = model(ids) # no cache: recompute everything
return ids
def timed(fn, reps=3):
fn(); t = time.perf_counter()
for _ in range(reps): fn()
return (time.perf_counter() - t) / reps
# ---------------- 1. KV cache: same tokens, far less work ----------------
V, d, h, L = 1024, 512, 8, 8
model = GPT(V, d, h, L).eval()
print(f"model: {sum(p.numel() for p in model.parameters()):,} parameters, d={d}, {L} layers")
prompt = torch.randint(0, V, (1, 128))
same = torch.equal(generate(model, prompt, 64, True), generate(model, prompt, 64, False))
print(f"KV cache output identical to full recompute: {same}")
for n in (32, 128):
t_cache = timed(lambda: generate(model, prompt, n, True), 2)
t_full = timed(lambda: generate(model, prompt, n, False), 1)
print(f" {n:>3} new tokens: with cache {1000 * t_cache / n:6.1f} ms/token without {1000 * t_full / n:6.1f} ms/token")
kv_bytes = 2 * L * d * 4 # K and V, every layer, fp32
print(f"KV cache: {kv_bytes:,} bytes per token here (2 x layers x d x 4 bytes)")
# ---------------- 2. Prefill vs decode ----------------
with torch.no_grad():
long_prompt = torch.randint(0, V, (1, 512))
t_pre = timed(lambda: model(long_prompt))
_, caches = model(long_prompt)
one = torch.randint(0, V, (1, 1))
t_dec = timed(lambda: model(one, caches, start=512), 10)
print(f"prefill: 512 tokens in {1000 * t_pre:.0f} ms = {512 / t_pre:,.0f} tokens/s (this is time to first token)")
print(f"decode: 1 token in {1000 * t_dec:.1f} ms = {1 / t_dec:,.0f} tokens/s per sequence")
# ---------------- 3. Batching: decode many sequences per step ----------------
with torch.no_grad():
for B in (1, 8, 32, 64):
_, cb = model(torch.randint(0, V, (B, 128)))
step = timed(lambda: model(torch.randint(0, V, (B, 1)), cb, start=128), 5)
print(f" batch {B:>2}: {1000 * step:6.1f} ms per decode step -> {B / step:7,.0f} tokens/s total")
# ---------------- 4. Int8 weight-only quantization ----------------
def quantize_int8(w):
scale = w.abs().amax(dim=1, keepdim=True) / 127 # one scale per output row
return torch.round(w / scale).clamp(-127, 127).to(torch.int8), scale
q_model = copy.deepcopy(model)
fp32_bytes = int8_bytes = 0
for m in q_model.modules():
if isinstance(m, nn.Linear):
q, s = quantize_int8(m.weight.data)
fp32_bytes += m.weight.numel() * 4
int8_bytes += q.numel() + s.numel() * 4
m.weight.data = q.float() * s # store int8 + scales; dequantize to compute
with torch.no_grad():
a, _ = model(prompt); b, _ = q_model(prompt)
agree = (generate(model, prompt, 64) == generate(q_model, prompt, 64)).float().mean().item()
print(f"int8 linear weights: {fp32_bytes / 1e6:.1f} MB -> {int8_bytes / 1e6:.1f} MB "
f"({fp32_bytes / int8_bytes:.1f}x smaller); max logit change {(a - b).abs().max():.4f}; "
f"greedy tokens identical: {agree:.0%}")
# ---------------- 5. Speculative decoding (greedy) with a small draft model ----------------
facts = ["The cat sat on the mat and watched the rain.", "The dog chased a red ball across the park.",
"Water boils at one hundred degrees Celsius.", "The sun rises in the east every morning.",
"Bees collect nectar and turn it into honey.", "The river flows into the quiet lake."]
stream = torch.tensor([b for p in itertools.permutations(facts, 3) for b in (" ".join(p) + " ").encode()])
def train(m, steps):
opt = torch.optim.AdamW(m.parameters(), lr=2e-3)
for _ in range(steps):
i = torch.randint(0, len(stream) - 129, (16,))
x = torch.stack([stream[j:j + 128] for j in i]); y = torch.stack([stream[j + 1:j + 129] for j in i])
loss = F.cross_entropy(m(x)[0].reshape(-1, 256), y.reshape(-1))
opt.zero_grad(); loss.backward(); opt.step()
return m.eval()
target = train(GPT(256, 256, 8, 6), 300) # "large" model
draft = train(GPT(256, 64, 4, 1), 300) # ~35x smaller draft
print(f"target {sum(p.numel() for p in target.parameters()):,} params, draft {sum(p.numel() for p in draft.parameters()):,} params")
@torch.no_grad()
def speculative(prompt, n, k=4):
"""Draft proposes k tokens; the target checks all of them in ONE forward pass and keeps the agreeing prefix."""
ids, target_calls, accepted = prompt.clone(), 0, 0
while ids.shape[1] - prompt.shape[1] < n:
draft_ids = ids.clone()
for _ in range(k):
draft_ids = torch.cat([draft_ids, draft(draft_ids)[0][:, -1].argmax(-1, keepdim=True)], 1)
logits = target(draft_ids)[0]; target_calls += 1
pred = logits[:, ids.shape[1] - 1:].argmax(-1) # target's own choice at each proposed position
proposed = draft_ids[:, ids.shape[1]:]
match = int((pred[0, :k] == proposed[0]).long().cumprod(0).sum())
accepted += match
ids = torch.cat([ids, proposed[:, :match], pred[:, match:match + 1]], 1) # + the target's correction
return ids[:, :prompt.shape[1] + n], target_calls, accepted
p = torch.tensor([list("The cat sat on the mat".encode())])
spec, calls, acc = speculative(p, 80)
ref = generate(target, p, 80, use_cache=False)
print(f"speculative output identical to target's own greedy output: {torch.equal(spec, ref)}")
print(f" 80 tokens in {calls} target passes instead of 80 ({acc} draft tokens accepted); text: {bytes(spec[0].tolist()).decode()!r}")Output (PyTorch 2.14 on a 4-thread CPU; timings vary by machine and run, but the correctness checks and ratios hold):
model: 26,792,960 parameters, d=512, 8 layers KV cache output identical to full recompute: True 32 new tokens: with cache 7.3 ms/token without 34.2 ms/token 128 new tokens: with cache 6.6 ms/token without 46.5 ms/token KV cache: 32,768 bytes per token here (2 x layers x d x 4 bytes) prefill: 512 tokens in 147 ms = 3,480 tokens/s (this is time to first token) decode: 1 token in 11.1 ms = 90 tokens/s per sequence batch 1: 6.7 ms per decode step -> 150 tokens/s total batch 8: 20.6 ms per decode step -> 388 tokens/s total batch 32: 47.9 ms per decode step -> 667 tokens/s total batch 64: 73.9 ms per decode step -> 866 tokens/s total int8 linear weights: 102.8 MB -> 25.8 MB (4.0x smaller); max logit change 0.0118; greedy tokens identical: 100% target 5,132,288 params, draft 148,416 params speculative output identical to target's own greedy output: True 80 tokens in 21 target passes instead of 80 (59 draft tokens accepted); text: 'The cat sat on the mat and watched the rain. The dog chased a red ball across the park. The cat sat on'