Inference Parallelism — Tensor Parallel, Expert Parallel & Prefill/Decode Disaggregation
Serving scale-out is not the training mesh. This lesson is inference only: tensor parallel shards weights and KV and pays a collective on the decode path, expert parallel routes MoE tokens with an all-to-all, and prefill/decode disaggregation splits pools when TTFT and tokens per second fight.
- 1Gist
- 2Maps
- 3Q&A
- 4Sandbox
Voice readout needs Web Speech Synthesis in this browser.
Which scale-out are you buying?
Prefer
Match the collective to the bottleneck
Tensor parallel shards every dense layer and synchronizes often. Expert parallel moves only the tokens a router chose. A prefill/decode split moves KV between pools so the two phases stop stealing each other’s GPUs.
- Fast NVLink and a model that does not fit: serving tensor parallel.
- Sparse experts with a balance plan: expert parallel.
- TTFT SLO fighting decode concurrency: separate pools.
Alternative
Reuse the training mesh because it already exists
Training tensor parallel hides collectives behind huge steps and gradient sync. Serving pays a small collective on every token. Eight-way shards on PCIe for a 7B chat model get slower.
- FSDP is the wrong tool for serving latency.
- Hot experts strand the all-to-all.
- KV transfer you did not measure regresses TTFT.
Overview
Scaling serving is not scaling training. Distributed PyTorch — DDP, FSDP, Tensor & Pipeline Parallelism already covers data parallel, sharded training state, and the 3D training grid. This page is inference only.
One paragraph on training tensor parallel
Training tensor parallel shards matmuls inside a larger plan, often with FSDP and pipeline stages, to push throughput of a huge global batch and to sync gradients. Inference tensor parallel shards the same column and row linears so a request’s latency drops and the weights fit. KV heads are split across GPUs. Every decode step may all-reduce partial logits or activations. Those collectives sit on the time-per-token path. They are not hidden behind a long training step. If you need training memory tricks, leave this page and open the Distributed PyTorch study. Stop there.
Serving tensor parallel
- Column-parallel linear. Split output features across GPUs. Each rank holds a weight shard and produces a partial output.
- Row-parallel linear. Split input features. Partial outputs must be reduced.
- Attention. QKV projections are sharded. KV-cache heads are partitioned so each GPU stores a subset of heads, not a full copy of every head.
- MLP / SwiGLU. Sharded the same way, column then row, or the reverse, depending on the implementation.
- Collective. All-reduce or all-gather on activations or logits, per layer or at the end. That latency is a floor for small batches.
Helps when the model does not fit on one GPU, or when single-stream latency is the goal and NVLink or NVSwitch is fast.
Hurts when prompts are tiny, decode is already memory-bound, and the interconnect is slow. The collective tax exceeds the compute you saved.
Expert parallel
Mixture-of-experts models route each token to a few experts (top-k gates). Expert parallel places different experts on different ranks. Tokens travel with an all-to-all, run the chosen experts, and travel back.
Popular experts become hotspots. Stragglers set the step time. Capacity factors, training-time auxiliary losses, and runtime balancing exist because of that. Treating a MoE like a dense tensor-parallel-only model either melts the interconnect or leaves experts idle.
Large MoE servers often combine expert parallel with tensor parallel inside an expert. The hybrid is normal. The failure mode is still imbalance plus an all-to-all you did not size.
Prefill / decode disaggregation
One pool that both prefills long prompts and decodes many streams will interfere with itself. Prefill wants compute for time-to-first-token. Decode wants memory bandwidth and a fat batch for tokens per second. A long prefill stalls decodes. A decode-optimized batch delays the first token.
Split the pools. Prefill workers produce KV. Decode workers consume KV and stream tokens. KV has to move or be shared over the network or NVLink, and the block metadata has to match the pager from PagedAttention & Continuous Batching.
Use the split when the TTFT SLO is strict, decode concurrency is high, or you have heterogeneous hardware (large prefill GPUs, decode-sized GPUs). You pay transfer latency, a larger ops surface, and a new partial-failure domain. If you never measure transfer round-trip, TTFT can get worse than the colocated pool.
Industry writeups, including TensorRT-LLM’s disaggregated serving notes, move quickly. The mechanism is the split plus the KV path. Engine flags stay in TensorRT-LLM — Engine Build, In-Flight Batching & Quantization.
Teaching order, not a required pipeline
A service might use only one lever. The diagram reads top to bottom so the three stories stay in one column.
Flow
- 1
1 GPU0 weight shard and KV heads
- next2 GPU1 weight shard and KV heads
- 2
2 GPU1 weight shard and KV heads
- next3 NCCL all-reduce on activations
- 3
3 NCCL all-reduce on activations
- next4 Router picks top-k experts
- 4
4 Router picks top-k experts
- next5 All-to-all moves tokens
- 5
5 All-to-all moves tokens
- next6 Experts live on other ranks
- 6
6 Experts live on other ranks
- next7 Prefill pool writes KV
- 7
7 Prefill pool writes KV
- next8 Transfer KV to the decode pool
- 8
8 Transfer KV to the decode pool
- next9 Decode pool streams tokens
- 9
9 Decode pool streams tokens
Lesson map
Inference Parallelism — Tensor Parallel, Expert Parallel & Prefill/Decode
Serving scale-out is not the training mesh. This lesson is inference only: tensor parallel shards weights and KV and pays a collective on the decode path, expert parallel routes MoE tokens with an all-to-all, and prefill/decode disaggregation splits pools when TTFT and tokens per second fight.
Architecture. Architecture
Select a node to see why it exists, or an edge to see the protocol, direction, effect, and consequence.
Mermaid export
flowchart TB a["1 GPU0 weight shard and KV heads"] b["2 GPU1 weight shard and KV heads"] c["3 NCCL all-reduce on activations"] d["4 Router picks top-k experts"] a -->|1 GPU0 weight shard and KV heads| b b -->|2 GPU1 weight shard and KV heads| c c -->|3 NCCL all-reduce on activations| d
Pros and cons
| Lever | Gain | Cost |
|---|---|---|
| Serving tensor parallel | Fit and latency on a fast fabric | Collective floor. Weak on tiny batches |
| Expert parallel | Sparse capacity | All-to-all and imbalance |
| Prefill/decode split | Isolate TTFT from tokens per second | KV fabric and ops |
| Single-GPU continuous batching | Simplest | Model-size and peak-throughput caps |
| Training FSDP | Training memory | Wrong tool for serving latency. See Distributed PyTorch |
Sandbox: collective sketch and expert heat
The step model is T_compute + T_allreduce. Ring intuition: bytes moved scale like 2 · (n - 1) / n times the payload. Constants are ignored on purpose. The slow-link case is worse than the fast-link case. It can still beat one GPU when the shard’s compute drop is large. The hurt case is when compute barely drops and the link is slow.
Press Run. Snippets must be self-contained — no network, files, or native modules.
Press Run. Snippets must be self-contained — no network, files, or native modules.
A max load above 1.5 times the mean is a teaching flag, not a production threshold. The point is that one hot expert sets the iteration.
Pitfalls
Box 1: 70B dense, NVSwitch, latency SLO on one user. Box 2: MoE, expert 3 takes 40 percent of tokens. Box 3: 8k prompts and 64 concurrent decodes missing both TTFT and tokens per second on the same eight GPUs. Which collective or transfer do you introduce, and which study do you refuse to reopen for FSDP?
Interview Q&A
Where does NCCL show up in serving tensor parallel?
Answer
After sharded linears and attention, ranks synchronize partial activations or logits. That is often an all-reduce, and it sits on every decode step’s critical path. It is not a gradient bucket at the end of a training iteration.
Why disaggregate prefill and decode?
Answer
Prefill wants compute for time-to-first-token. Decode wants memory bandwidth and batching for tokens per second. Colocation makes them interfere. Separate pools scale those knobs independently, at the cost of moving KV. See also the phase split in Prefill vs Decode & KV Cache Mechanics.
How is expert parallel different from tensor parallel?
Answer
Tensor parallel shards every dense layer’s weights, and every token touches every shard. Expert parallel places different experts on different ranks and moves a token only to the experts the router picked, via all-to-all. You can combine them. You should not pretend they are the same collective.
Column-parallel versus row-parallel, briefly.
Answer
Column-parallel splits the output dimension. Each GPU writes a slice of the features and you often all-gather afterward if the next op needs the full vector. Row-parallel splits the input dimension and all-reduces partial sums. Attention and MLP blocks alternate these so the next layer sees a full activation.
Why can tensor parallel make a small model slower?
Answer
Local compute shrinks, but the collective does not. On a slow link, or on a tiny batch that was already bandwidth-bound, the all-reduce exceeds the compute you removed. Eight-way on PCIe for a 7B chat service is the usual story.
What does KV-head partitioning change?
Answer
Each rank stores a subset of KV heads, so per-GPU cache bytes drop, and attention is local to those heads until a collective reassembles activations. Your single-GPU byte formula still holds for the whole model. Per-rank memory is that total divided by the tensor-parallel degree, plus overhead, not the full cache on every GPU.
How do hot experts show up?
Answer
One rank’s expert queue is much deeper than the mean. The all-to-all waits on that rank. Tokens per second and p99 follow the straggler. Capacity factors and runtime balance exist to cap that skew. A dense tensor-parallel plan does not fix a router that always picks the same expert.
What has to be true for the KV transfer to be worth it?
Answer
Transfer time plus decode start must beat the interference you removed. Long prompts and moderate generations are the usual win, because prefill was blocking a lot of decode work. Short prompts with a slow network lose. Measure round-trip before you split the on-call rotation in two.
Is this the Distributed PyTorch lesson?
Answer
No. That study owns DDP, FSDP, and the training tensor and pipeline grid. The only overlap is the shared idea of sharding a matrix. Goals, collectives, and memory (KV versus optimizer state) differ. Open Distributed PyTorch — DDP, FSDP, Tensor & Pipeline Parallelism if the question is about training.
Where do product stack choices go?
Answer
After the mechanism. vLLM vs SGLang — Runtime Choice (comparative) is the thin matrix. Comparative Playbook — Serving Stack Choice, Metrics & Failure Modes is the metrics narrative. Multi-GPU engine flags stay in the TensorRT-LLM study.
Go Deeper
- Megatron-LM — public tensor-parallel linear patterns
- vLLM documentation and SGLang documentation — multi-GPU serving
- TensorRT-LLM disaggregated serving
- Distributed PyTorch — DDP, FSDP, Tensor & Pipeline Parallelism — training contrast only
- TensorRT-LLM — Engine Build, In-Flight Batching & Quantization
- Next: vLLM vs SGLang — Runtime Choice (comparative)