Inference Engineering

Serving a large model is a different engineering problem from training one. Training is throughput-bound and offline; serving is latency-bound, online, and has to hold a service-level objective while the request mix changes underneath it. This book is about that second problem.

Where the time goes

A single forward pass through a transformer splits into two phases with completely different performance characteristics. Prefill processes the whole prompt at once and is compute-bound. Decode emits one token at a time and is memory-bandwidth-bound, because every step reads the entire KV cache and the full weight matrices to produce a single token.

The practical consequence is that arithmetic intensity — FLOPs performed per byte moved — is the number worth reasoning about, not FLOPs alone.

The KV cache

Every token attends to all previous tokens, so the keys and values for the prefix are cached rather than recomputed. The cache is what makes decode tractable and also what makes it expensive.

with $L$ layers, $H_{kv}$ key/value heads, head dimension $d_{head}$, sequence length $s$, batch size $b$. The leading 2 counts keys and values separately.

Worth computing before choosing a batch size, because the cache — not the weights — is usually what runs the GPU out of memory first:

def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, batch, dtype_bytes=2):
    """Bytes held by the KV cache for one in-flight batch."""
    return 2 * layers * kv_heads * head_dim * seq_len * batch * dtype_bytes
 
 
# Llama-3-8B: 32 layers, 8 KV heads (GQA), head_dim 128, fp16
gib = kv_cache_bytes(32, 8, 128, seq_len=8192, batch=32) / 1024**3
print(f"{gib:.1f} GiB")  # -> 32.0 GiB

Batching

Static batching pads every request in a batch to the longest sequence and holds the batch until all of its requests finish, so a single long generation stalls every short one behind it. Continuous batching instead admits and retires requests at token granularity.

What to measure

Throughput and latency trade against each other, and a single average hides the trade. Report time to first token and inter-token latency separately, both at a tail percentile rather than a mean, alongside the concurrency the numbers were taken at.