DISTRIBUTED INFERENCE

Tensor parallelism (Megatron-LM) splits the weight matrices themselves across GPUs. The arrangement is chosen so that communication happens as rarely as possible.

The MLP block. Two matmuls in sequence: Y = GeLU(X @ A), then Z = Y @ B.

Split A by columns across GPUs. Each GPU computes a slice of Y. Crucially, GeLU is elementwise, so each GPU can apply it to its own slice without needing anyone else's — no communication.

Then split B by rows, matching the column split of A. Each GPU computes a partial Z using its slice of Y and its rows of B. The partials sum to the correct answer, so one all-reduce finishes the block.

GPU 0:  X @ A[:, :h/2]  -> Y0 -> GeLU -> Y0 @ B[:h/2, :] -> Z_partial_0
GPU 1:  X @ A[:, h/2:]  -> Y1 -> GeLU -> Y1 @ B[h/2:, :] -> Z_partial_1
                                                 all-reduce -> Z

One all-reduce for two matmuls. Splitting the other way round would need communication between them.

The attention block splits naturally by head: give each GPU a subset of heads, let it compute those heads' attention independently, then all-reduce after the output projection. Also one all-reduce.

So: two all-reduces per transformer layer. For an 80-layer model that is 160 all-reduces per forward pass, on the critical path.

Communication volume. A ring all-reduce moves 2(N−1)/N times the data per GPU. For an activation of shape [batch, seq, d_model] in fp16:

bytes per all-reduce = 2 * batch * seq * d_model * 2 * (N-1)/N

At decode (seq = 1), batch 32, d_model = 8192, TP=4:

2 * 32 * 1 * 8192 * 2 * 0.75 = 786,432 bytes = 0.79 MB per all-reduce
x 160 all-reduces = 126 MB per token

Over NVLink at 900 GB/s that is 0.14 ms — against a decode step of a few milliseconds, so roughly 5% overhead. Over PCIe at 64 GB/s it is 1.97 ms, which can double your step time. This is the single most important practical fact in this module: TP belongs inside a node, on NVLink.

The GQA ceiling. TP splits attention by head, and a KV head cannot be split across GPUs without replicating its cache. With 8 KV heads you can run TP-8 cleanly, one KV head per GPU. At TP-16 you must duplicate KV heads, so per-GPU cache stops shrinking as you add GPUs — you pay for hardware that adds bandwidth and compute but no cache capacity. The KV head count silently caps useful TP degree.

Pipeline parallelism assigns different layers to different GPUs. With 80 layers on 4 GPUs, each holds 20 consecutive layers. Activations flow forward through the stages.

Communication is tiny. Between stages you send one activation tensor: batch × seq × d_model × 2 bytes, once per stage boundary. At decode with batch 32 and d_model = 8192 that is 0.52 MB per boundary, three boundaries — about 1.6 MB per token, versus 126 MB for TP-4. Roughly 80× less traffic.

That is why PP tolerates slow interconnects. Across nodes, over Ethernet or InfiniBand, PP works and TP does not.

The bubble is the cost. A single request must traverse all four stages sequentially. While it is in stage 2, stages 1, 3 and 4 have nothing to do. For one request in flight, utilization is 1/N.

The fix is to keep multiple microbatches in flight so every stage always has work:

time ->
stage 0:  [m1][m2][m3][m4][m5]...
stage 1:      [m1][m2][m3][m4]...
stage 2:          [m1][m2][m3]...
stage 3:              [m1][m2]...
          ^^^^^^^^^^^^ the fill bubble

Bubble fraction is (N−1)/(M + N−1) for M microbatches and N stages. With 4 stages and 32 microbatches: 3/35 = 8.6%. Acceptable. With 4 microbatches: 3/7 = 43%. Not.

In LLM serving the microbatches are your concurrent requests, so pipeline efficiency depends on having enough concurrency — which brings you back to KV cache capacity. Under light load a pipelined deployment is badly underutilized.

PP does not reduce single-request latency. A request still passes through every layer once; you have only distributed where they live. TP, by contrast, genuinely speeds up each layer because you have N× the bandwidth working on it. This is the key difference: TP improves latency, PP improves capacity.

The standard combination follows directly:

TP within a node   (NVLink, 900 GB/s, high-volume all-reduces)
PP across nodes    (InfiniBand/Ethernet, low-volume activation passing)

An 8-GPU node running TP-8, with PP across nodes. Sequence or context parallelism (splitting the sequence dimension, as in Ring Attention) is a third axis used for very long context, where even one sequence's activations and KV cache exceed a single GPU.

Work through TP for Llama-3-70B on H100s, 8k context, so the effects are concrete.

Weights per GPU: 141.2 / N GB. KV cache per GPU: 320 KiB per token / N — because each GPU holds only its own heads' K and V, provided N ≤ n_kv_heads. Aggregate bandwidth: 3.35 × N TB/s.

TP   weights/GPU   KV/tok/GPU   agg BW      fits in 80GB?   TPOT floor
 1     141.2 GB      320 KiB    3.35 TB/s        NO             --
 2      70.6 GB      160 KiB    6.70 TB/s        yes         21.1 ms
 4      35.3 GB       80 KiB   13.40 TB/s        yes         10.5 ms
 8      17.7 GB       40 KiB   26.80 TB/s        yes          5.3 ms
16       8.8 GB       40 KiB   53.60 TB/s        yes          2.6 ms

TPOT floor is 141.2 GB / aggregate bandwidth — the weights are read once collectively per step.

Three things to read out of this table.

TP improves latency roughly linearly. From TP-2 to TP-8, TPOT falls 4×. You are not just fitting the model, you are decoding it faster, because you added memory bandwidth.

KV cache per GPU stops shrinking at TP-8. Llama-3-70B has 8 KV heads. At TP-16 two GPUs share each KV head and must both hold its cache, so per-GPU KV stays at 40 KiB/token. You doubled the hardware and got no additional cache capacity — only bandwidth. That is the GQA ceiling biting.

Communication grows with TP degree. The all-reduce volume scales as (N−1)/N, approaching a constant, but the number of participants grows and ring all-reduce latency grows with hop count. Beyond a node, where you leave NVLink, the cost jumps discontinuously.

Now the effect that surprises people: TP degree affects TTFT and throughput differently.

TTFT improves with TP because prefill is compute-bound and you added compute. Communication during prefill is larger in absolute terms (activations are seq times bigger) but it is amortized over far more arithmetic, so the fraction of time spent communicating is lower during prefill than decode.

Throughput per GPU usually falls with TP. You are paying communication overhead on every layer and getting sub-linear scaling. Two independent TP-4 replicas will generally serve more total tokens per second than one TP-8 deployment — but each request will be slower.

That is the real decision:

minimize latency        ->  maximize TP (up to the KV head count / node boundary)
maximize throughput/$   ->  minimum TP that fits the model, then replicate

Most production deployments choose the smallest TP that fits the model with acceptable latency, then scale out with replicas.

An MoE layer replaces the single MLP with E expert MLPs and a router that sends each token to the top k of them, typically k = 2 out of 8 to 256 experts.

The inference consequence follows straight from this course's thesis. You only read the weights of the experts you activate. A model with 8 experts and top-2 routing has roughly 4× the parameters of its dense equivalent but moves about the same bytes per token.

For decode, which is memory-bound, this is close to free quality. Mixtral-8x7B has ~47B total parameters but activates ~13B per token: it decodes at roughly the speed of a 13B model while being substantially better than one. MoE is the most direct possible exploitation of the fact that bytes moved, not parameters held, is what costs you.

The problems are all about capacity and routing.

Every expert must be resident. You do not know in advance which experts a token will need, so all of them must be in memory somewhere. Mixtral-8x7B needs 94 GB at fp16 even though only 26 GB is read per token. MoE trades memory capacity for memory bandwidth, which is a good trade when capacity is what you have.

Expert parallelism distributes experts across GPUs, so each holds E/N of them. Now routing requires an all-to-all: send each token's hidden state to whichever GPUs hold its chosen experts, compute, send results back. Two all-to-alls per MoE layer, and all-to-all is the most demanding collective — every GPU talks to every other.

Load imbalance is the hard part. Routing is learned and data-dependent, so nothing guarantees even distribution. If 40% of tokens in a batch pick expert 3, the GPU holding expert 3 does 40% of the work while others idle. Since the layer cannot complete until every GPU finishes, the slowest GPU sets the pace for all of them.

Mitigations, none complete:

• Capacity factor. Cap tokens per expert; drop or reroute the overflow. Bounds the imbalance and costs quality on dropped tokens.
• Auxiliary load-balancing loss during training. Helps on average, not per batch.
• Expert replication. Duplicate hot experts across GPUs. Costs memory.
• Larger batches. Averaging over more tokens smooths the distribution — one case where batching helps for a reason unrelated to arithmetic intensity.

Combining parallelism. Production MoE deployments typically use expert parallelism for the MoE layers and tensor parallelism for attention and the dense layers, which means switching communication patterns within each layer. This is genuinely intricate, and it is why MoE serving is harder than dense serving even though the arithmetic is friendlier.

This is the conclusion the whole course has been building toward.

Prefill is compute-bound, latency-critical at the request level, and benefits from high TP degree and large FLOP capacity. Decode is memory-bound, latency-critical per token, benefits from bandwidth and cache capacity, and wants large batches. Module 5 showed that co-locating them means one interferes with the other, and chunked prefill only mitigates it.

Disaggregation puts them on separate machine pools:

  request
     |
     v
  [ PREFILL POOL ]        compute-optimized, high TP, small batches
     |                    produces the prompt's KV cache
     |  transfer KV
     v
  [ DECODE POOL ]         bandwidth-optimized, large batches, big KV budget
     |
     v
  streamed tokens

What this buys:

No interference. A 100k-token prefill cannot stall anyone's token stream, because it is not on the same GPU.

Independent tuning. Each pool gets its own parallelism strategy, batch size and even hardware. Prefill wants FLOPs; decode wants bandwidth and capacity. You can buy different GPUs for each — an H200 or a large-memory part for decode, a compute-dense part for prefill.

Independent scaling. Traffic mix shifts. A day of long-document summarization is prefill-heavy; a day of chat is decode-heavy. Separate pools scale separately.

The cost is the KV transfer. After prefill you must move the prompt's entire KV cache to the decode pool. For Llama-3-70B with a 4k prompt that is 320 KiB × 4096 = 1.25 GiB. Over InfiniBand at ~50 GB/s that is 25 ms; over NVLink where available, ~1.4 ms.

Whether disaggregation wins depends on comparing that transfer against the interference it removes. DistServe's argument, and it is convincing, is that for most realistic workloads the transfer is cheaper — especially because it can be overlapped layer by layer: start sending layer 0's KV as soon as it is computed, rather than waiting for the whole prefill to finish. Done well, most of the transfer hides behind the remaining prefill compute.

Splitwise reached similar conclusions independently, and Mooncake built a production system around a disaggregated KV store shared across the cluster — which also makes prefix caching global rather than per-node, a significant additional benefit.

When it is not worth it: short prompts (little to transfer, but also little interference), low load (no contention to eliminate), or a slow interconnect between pools. Like everything in this course, it is a trade you should compute rather than assume.

Collect the numbers, because every choice in this module reduces to them:

NVLink is 3.7× slower than HBM. PCIe is 52× slower. Ethernet is 280× slower.

Now map the strategies onto them:

The design rule that falls out:

within a node (NVLink):    TP, EP, SP -- the chatty strategies
across nodes (IB/Eth):     PP, disaggregation, replication -- the quiet ones

Violating it is the most common and most expensive mistake in multi-GPU deployment. Running TP-8 across two 4-GPU nodes connected by PCIe means every one of your 160 all-reduces per token crosses a link 52× slower than HBM. Deployments have been observed running slower on 8 GPUs than on 4 for exactly this reason.

A few practical notes:

Check your topology before choosing TP degree. nvidia-smi topo -m tells you which GPUs have NVLink and which are behind PCIe switches. A "4-GPU machine" may be two NVLink pairs connected by PCIe, in which case TP-4 is a mistake and TP-2 with PP-2 is correct.

All-reduce cost is not purely bandwidth. Ring all-reduce has 2(N−1) communication steps, so latency grows with participant count even at fixed volume. At decode, where messages are small (under a megabyte), you are often *latency*-bound on the collective rather than bandwidth-bound — which is why decode suffers proportionally more from high TP degree than prefill does.

Overlap what you can. Communication and computation can overlap if the schedule allows it. Fine-grained overlapping of all-reduce with the next matmul is standard in mature implementations and is worth a meaningful fraction of the communication cost.