Inference — TP for latency, replicas for throughput

Same hardware, opposite goals. Training cares about the gradient AllReduce; inference cares about per-token decode time. The math behind "TP=8 or 8× replicas?" comes down to whether the workload is memory-bound or compute-bound. Beyond the TP-vs-replicas split, the other levers on single-replica decode speed — the KV cache, quantization, and speculative decoding — all fall out of the same memory-bound math, and we cover them here too.

What changes at inference time

Training is dominated by big batches doing big matmuls — comfortably compute-bound (above the roofline ridge, lesson 01). Inference splits into two regimes:

• Prefill. One forward pass on the full prompt. The matmul operates on (B, T_prefill, d); arithmetic intensity is high; compute-bound, exactly like training.
• Decode. One forward pass per generated token. The input is (B, 1, d). The KV cache from previous tokens grows but the new query is just one token. Arithmetic intensity is low: per parameter read, only one multiply-add per sample in the batch. Memory-bound: you spend all your time reading weights from HBM.

This is the core asymmetry. Training and prefill scale on compute; decode scales on HBM bandwidth. Most of inference's tokens-per-second is set by decode. So the question becomes: how fast can we read weights from HBM, per output token, per request?

What the KV cache is — and why it sets the memory wall

Decode keeps mentioning a "KV cache." Here is what it is, from first principles. Attention at step t needs the keys and values of every past token — the new query attends over all of 1..t. If you recomputed all those K and V every step, decoding a sequence of length T would cost O(T^2) work in attention alone. Instead, each layer caches the K and V it produced for every token seen so far. The new token computes its own K, V once, appends them, and does O(t) attention against the stored cache. That is what makes per-step decode O(1) in prompt length for everything except the attention dot-product itself — and it is why the weights, read once per step, dominate the cost.

The cache is not free. Its size grows linearly with sequence length and batch:

The leading 2 is the K and the V. Plug in Llama-3-70B (n_layers=80, n_kv_heads=8 with GQA, d_head=128, bf16 = 2 B):

84 GB of KV cache will not fit alongside the weights on a single 80 GB GPU. This is the punchline: on a serving GPU it is the KV cache, not the weights, that caps batch size and context length. The weights are a fixed cost; the KV cache is the variable cost that scales with how much concurrency and context you want.

The KV cache shards across TP ranks by KV head: with TP=8 and 8 KV heads, each rank holds one head's K/V, so per-rank KV drops 8×. This is the same constraint that caps honest TP at the KV-head count (the GQA wrinkle discussed later in this lesson and in lesson 06) — push TP past 8 here and you must replicate KV state, undoing the capacity win. How the cache is physically stored without fragmentation (paged attention) is in the vLLM track; how it is moved between a prefill box and a decode box is lesson 15.

The decode latency calculation

Per token, you read every weight once. For a 70B bf16 model, that's 140 GB. On one H100 at 3.35 TB/s HBM bandwidth, the absolute minimum per-token time is:

That's 24 tokens/sec, single GPU, single request, peak bandwidth utilization. In practice you achieve ~60–80% of that.

Now use TP=8. The weight read is spread across 8 GPUs' HBM bandwidths: each one reads only 1/8 of the weights, so the read time drops by 8×. The cost is a per-layer AllReduce, which on NVLink is tiny relative to the saving. New per-token time:

Roughly an 8× latency improvement. But the per-token throughput per GPU is unchanged: TP=8 uses 8× the GPUs to serve one request 8× faster. From the cluster's point of view, it's the same total work.

Animated · one decode token, bytes flowing from HBM

Pick a TP degree and watch where the bytes physically come from. At TP=1, a single GPU streams the entire 140 GB of weights through its HBM bus, layer by layer. At TP=8, each of the 8 GPUs reads only its 1/8 shard — in parallel, eight times faster. Each blue particle is a chunk of weights moving from HBM into a tensor-core compute unit.

Throughput vs latency — the choice

Now consider a serving fleet with 8 GPUs and two options:

(Both shapes batch many concurrent requests in production — we'll bring that in below. At batch=1, the table compares the single-stream extreme.)

The two options give the same total tokens-per-second. What differs is how they're allocated: TP=8 lets each request run faster; replication lets more requests run simultaneously. Pick by what matters. Latency-sensitive workloads (chat, autocomplete, coding agents) want TP high. Bulk-throughput workloads (batched embedding, document processing) want replication high.

The catch — batching changes the math

The single-stream comparison above assumes batch = 1. As you batch more requests onto the same forward pass, the per-token cost goes up only slightly (you do more arithmetic but read the same weights), so the cost per request drops. With batch B = 32 sharing the same forward:

Decode batching, plus continuous batching, is what makes shared GPU serving economical. The TP-vs-replicas question becomes: with batching, is one TP=8 replica running batch 32 cheaper than 8 TP=1 replicas each running batch 32? In wall-clock per request: the TP=8 replica is still ~8× faster (memory-bound dynamics survive batching for a while). In total cluster throughput: roughly equal again — the GPU work is the same, just spread.

Continuous batching is the mechanism that makes cost_per_request ≈ t_token / B hold in practice. Requests arrive and finish at different times, so a static batch — one fixed set of sequences run to completion together — would stall on the slowest member and leave the batch dimension half-empty for most of the run. Instead the server admits and evicts sequences every decode step: a finished sequence frees its slot immediately and a waiting request is spliced in, so B stays full and the single expensive weight read is amortized across as many live requests as the KV budget allows. That continuously-full batch is what keeps memory-bound decode cheap per token. The implementation — the scheduler, the paged KV slots it admits into — is the vLLM track; the core idea is just "never run a half-empty batch."

How long does the memory-bound regime last as B grows? Decode does ≈ 2·B·P FLOPs while reading ≈ 2·P weight-bytes, so its arithmetic intensity is ≈ B FLOP/byte. The H100 bf16 roofline ridge sits at ≈ 295 FLOP/byte (989 TFLOP/s ÷ 3.35 TB/s), so decode crosses from memory-bound to compute-bound at roughly B ≈ 200–300; below that batch the weight read dominates and TP's HBM split is the lever, above it the matmul dominates and only more FLOPs (or fp8) help.

So the right framing: at any batch size, TP and replication trade latency for parallelism, not for total throughput. The exception is when batching itself becomes a constraint:

• HBM budget for KV cache. Bigger batches need more KV-cache slots. With TP, each rank only holds 1/N of the KV cache (sharded across heads, as in lesson 09 of the vLLM series). So TP increases effective KV capacity per replica — useful for long-context workloads. The GQA wrinkle from lesson 06 applies: TP can't exceed the KV-head count, so a model with 8 KV heads (Llama-3-70B) caps at TP=8 for honest KV sharding; further TP would replicate KV state and undo the capacity win.
• Time-to-first-token (TTFT). Prefill is compute-bound and benefits from TP linearly (in compute, not memory). Workloads with very long prompts (RAG, code repos) get TP gains on both prefill and decode.

PP is rare at inference time, here's why

PP across the layers of a model adds a per-stage bubble at inference: a single request must walk all N stages sequentially, paying N - 1 stage-time of bubble before its first token comes out. For chat use cases (TTFT < 500 ms), this is too high a price. PP shows up at inference only in two cases:

• The model is so big a TP-maxed-out node still can't fit it. You add PP across nodes as a last resort.
• Throughput-only workloads with no TTFT constraint, where you can fill the pipeline with many concurrent requests.

The decision table

TP at serving differs from TP at training

Two subtle differences:

1. No backward. No gradient AllReduce. The per-step comm cost drops by 2× compared to training, since training also pays the symmetric AllReduce in backward.
2. The per-step compute is smaller (batch is smaller, sequence dim is smaller during decode). So the comm/compute ratio at decode is worse than at training — TP's AllReduces become a larger fraction of the time. Modern serving stacks (vLLM, TRT-LLM) work hard to fuse the AllReduce into the matmul kernel where possible to hide it.

This is why some serving stacks use TP=4 even when TP=8 would be available — the marginal latency improvement saturates because the AllReduce cost is no longer dwarfed by compute.

Quantization — buying decode speed by reading fewer bytes

TP and replicas move work between GPUs but never change the total bytes read. Quantization changes the bytes themselves. Recall the decode floor from this lesson: per-token time ≈ bytes_read / HBM_BW. Decode is HBM-bandwidth-bound, so halving the bytes per weight roughly halves decode latency. That makes quantization a bandwidth optimization first and a memory-capacity optimization second — the capacity win (more room for weights and KV cache) is real but secondary to the per-token speedup.

Weight-only quantization (GPTQ, AWQ)

Store the weights in INT4 or INT8 with per-group scales, and dequantize inside the matmul kernel just before the multiply — activations stay at higher precision (bf16/fp16). The matmul still happens in high precision; you have only changed how the weights are transported from HBM. Carry the number through: a 70B model at bf16 is 140 GB → 42 ms/tok floor (this lesson's own calc). At INT4 the weights are ~35 GB, so the floor drops to 35 / 3350 ≈ 21 ms/tok — roughly half, exactly because you read half the bytes.

fp8 weights + activations

On Hopper, fp8 quantizes both weights and activations, so the matmul itself runs on fp8 tensor cores. fp8 comes in two layouts — E4M3 (4 exponent / 3 mantissa bits: more precision, narrower range, used for weights and activations) and E5M2 (5 exponent / 2 mantissa bits: more range, less precision, used for gradients) — so inference forward passes pick E4M3. This wins on bandwidth and compute, which matters once batching pushes decode toward the compute-bound regime. (The fp8 formats are fully treated in lesson 17.)

KV-cache quantization

Quantize the KV cache itself to fp8 or int4. This shrinks the Section-1 memory wall directly: an fp8 KV cache halves the ~328 kB/token, so the same 80 GB GPU holds 2× the batch or 2× the context. It is the cheapest way to buy back the concurrency the KV wall takes away.

The cost of all of this is accuracy loss. It is managed, not eliminated: per-channel or per-group scales instead of one global scale, calibration on real data (AWQ picks scales that protect the most salient weight channels), and keeping the sensitive layers — notably the LM head and sometimes the first/last blocks — in higher precision. Done well, INT4/fp8 inference is within a fraction of a point of the bf16 model on most evals.

Speculative decoding — verifying many tokens for the price of one

Here is the trick that falls straight out of "decode is memory-bound." A forward pass that processes k candidate tokens reads the weights once — the same HBM traffic as processing a single token, because the weights, not the tokens, dominate the read. So if you can guess the next k tokens cheaply, the big model can verify all k in one forward pass and accept the longest correct prefix. You amortize one expensive weight read over many tokens.

Mechanism

A cheap drafter proposes k tokens — options include a small draft model (e.g. a 1B drafting for a 70B target), Medusa heads bolted onto the target, EAGLE (drafting in feature space), or even n-gram / prompt-lookup with no model at all. The target model then runs one forward over the k+1 positions, scoring all of them in parallel. You accept each drafted token that matches the target's own sampling and reject at the first mismatch, keeping the verified prefix plus one extra "free" token from the target. Expected accepted tokens per step rises with the acceptance rate α; in practice this lands around a 2–3× end-to-end decode speedup.

It is exact, not an approximation

The crucial property: with the right accept/reject rule (rejection sampling against the target's distribution), the output distribution is identical to plain autoregressive sampling from the target. Speculative decoding is a pure latency win — you get the same tokens you would have gotten, sooner. It is not a quality trade like quantization.

Animated · Amdahl ladder · TP scaling, sublinear because of comm

Pure HBM-bandwidth math says decode latency falls linearly with TP. In practice the curve flattens because (a) every TP step costs an AllReduce on NVLink, (b) at higher batch sizes some of the work becomes compute-bound and stops benefiting from more HBM lanes. The widget below traces latency as you sweep TP from 1 to 16; move the batch-size slider to watch the curve bend.

2D · choose your (TP, replicas) point on the frontier

Below is the same throughput-vs-latency frontier as the original widget, but interactive: pick a TP degree and a number of replicas with the sliders. The chosen point is highlighted; nearby alternatives are dim. Configs whose per-rank KV-cache exceeds HBM are greyed out as infeasible. The Pareto front is the chain in green.

Interactive · which shape gives you the QPS×latency point you want?

Set the cluster size, the model, and the per-request constraints. The widget lays out feasible (TP, replicas) shapes and overlays them on a "throughput × latency" plot. The Pareto front is the family of shapes that aren't dominated. Anything not on the front is a waste.