Trace a Token Through TP + CP + SP + EP + DP + PP in MoE Inference
Table of Contents
- Six Cuts
- The Setup: One Model, 32 GPUs, Two Meshes
- Hop 1: Land on GPU 11
- Hop 2: All-Gather With GPU 10
- Hop 3: Swap K/V With GPU 9
- Hop 4: Dispatch to Two Experts
- One Layer at a Glance
- Hop 5: PP Send to GPU 27
- Hop 6: Leave K/V on GPUs 10, 11, 26, 27
- The Whole Trip: Map and Traffic
- The Token’s Afterlife
- Three Answers
Let’s trace one prompt token through prefill of a Mixture-of-Experts (MoE) model on 32 GPUs with six parallelisms: data (DP), context (CP), sequence (SP), tensor (TP), expert (EP), and pipeline (PP). Prefill is the forward pass over the prompt that fills the KV cache before the first output token.
Each is easy alone. The hard part is how they compose: which GPUs form which group, why the layout is ordered as it is, and why data-parallel replicas stop being independent once you add experts. Following one token answers all three.
No one deploys all six; on 4 nodes × 8 GPUs, DeepSeek prefills with attention DP32 + EP32. This mesh puts every axis in one trace.
The map above is the route; its six numbered steps are the six hops below. Prefill discards everything it computes for the token except its K/V.
Six Cuts
Each parallelism cuts one axis of the work. It pays in communication wherever a later step needs something the cut put on another GPU.
| Parallelism | Cuts | Buys | Pays with | Hops |
|---|---|---|---|---|
| DP | requests | throughput | nothing directly | 1, 4 |
| CP | one long prompt’s tokens | long-prompt latency and memory | K/V swap inside attention | 1, 3 |
| SP | tokens, outside attention | each token handled once outside attention | nothing extra: reuses TP’s two collectives | 1–4 |
| TP | attention heads (and the LM head’s vocabulary) | weight memory and latency | all-gather before attention, reduce-scatter after | 2, 3, 6 |
| EP | experts | expert memory and bigger per-expert batches | all-to-all dispatch and combine | 4 |
| PP | layers | weight memory with the least communication | point-to-point send per stage boundary | 5 |
All-gather: every GPU gets every slice. Reduce-scatter: sum across GPUs, each keeps its slice. All-to-all: every GPU sends a different slice to every other.
SP and CP both split tokens, but SP’s split, inside a TP pair, is undone by an all-gather before attention; CP’s never is (CP ranks swap K/V). SP here is Megatron-style; DeepSpeed-Ulysses and several ring-attention papers use the name for splitting one long sequence across GPUs, which is closer to CP.
The Setup: One Model, 32 GPUs, Two Meshes
| Component | Value |
|---|---|
| Layers | 8 |
| Hidden size | d = n_h·d_h |
| Attention | n_h query and n_kv KV heads, both even (GQA: several query heads share one KV head); head size d_h |
| MoE FFN | 32 routed experts per layer, top-2 (k = 2), no shared expert |
| Vocabulary | V |
| Prompt | T tokens, one prefill chunk (engines prefill long prompts in chunks) |
| Cluster | 4 nodes × 8 H100; NVLink 450 GB/s per direction inside a node, InfiniBand (IB) 50 GB/s per GPU across nodes |
MoE inference lays two meshes over the same GPUs (Megatron’s MoE Parallel Folding): attention uses TP2 × CP2 × DP4 × PP2, experts use EP16 × PP2. DP copies attention, not experts: the four replicas’ 16 GPUs share one copy of the 32 experts, 2 each, so each expert sees every replica’s tokens.
Numbering: rank = tp + 2·cp + 4·dp + 16·pp. Stage 0 is GPUs 0–15 (nodes 0–1, layers 1–4); stage 1 is GPUs 16–31 (nodes 2–3, layers 5–8).
Why this order: TP and CP partners talk every layer, so they share one node’s NVLink. DP replicas share one EP group, so they fill a stage. PP talks once per stage boundary, so it goes outermost, across IB. EP crosses nodes on purpose: spanning both of a stage’s nodes leaves 2 experts per GPU instead of 4 (a copy per node), freeing memory for KV cache.
Meet Our Token
Our token, at position t in the prompt’s third quarter, lives on GPU 11.
| Group | Rule | GPU 11’s partners |
|---|---|---|
| DP | same dp (no traffic) | replica 2: 8–11 (stage 1: 24–27) |
| CP | stride 2 (NVLink) | GPU 11 ↔ GPU 9, GPU 10 ↔ GPU 8 (stage 1: 27 ↔ 25, 26 ↔ 24) |
| SP | same pair as TP (NVLink) | GPU 10 (stage 1: 26) |
| TP | stride 1 (NVLink) | GPU 10 (stage 1: 26) |
| EP | whole stage (NVLink + IB); GPU g hosts experts 2s, 2s+1, s = g mod 16 | 0–15 (stage 1: 16–31) |
| PP | stride 16 (IB) | 11 → 27 |
A trip log counts GPUs that held any piece of the token, node crossings (IB transfers of its data), and comm ops.
Hop 1 happens once; Hops 2–4 are one layer, repeated in all 8; Hop 5 crosses to stage 1; Hop 6 ends prefill.
Hop 1: Land on GPU 11
- DP: the load balancer sends the request to replica 2.
- CP: zigzag cuts the prompt into 4 chunks and gives CP rank
rchunksrand3 − r. - SP: TP rank 0 takes chunk
r, TP rank 1 takes chunk3 − r.
This is the SP layout: each GPU owns T/4 whole tokens.
Zigzag balances the causal mask (Llama 3, §3.3.2). In half-chunk blocks, chunk i attends to i full earlier chunks plus half of itself, so chunks cost 1, 3, 5, 7 and each rank gets 8. With 2·CP chunks, rank r gets r and 2·CP − 1 − r.
The embedding lookup is local: this toy gives every stage-0 GPU the full table, though real engines split it by vocabulary across the TP pair, like the LM head (Hop 6).
Gotcha: RoPE (rotary position embeddings) and the causal mask must use the global position t, not the local index on GPU 11.
Hop 2: All-Gather With GPU 10
- RMSNorm: per-token, so GPU 11 runs it alone.
- All-gather: GPUs 10 and 11 exchange their
T/4tokens; attention needs allT/2. - Column-parallel QKV: each GPU holds the QKV weight columns for half the heads and projects all
T/2tokens. Our token’s first-half heads land on GPU 10, its second-half heads on GPU 11. - KV cache write: our token’s K/V splits the same way,
n_kv/2KV heads per GPU.
Constraint: TP must divide n_h. It must also divide n_kv, unless n_kv < TP; then TP must be a multiple of n_kv, and each KV head and its cache are replicated TP/n_kv times. Both are even here, so nothing is replicated.
Trip log: 2/32 GPUs (+10) · 0 node crossings · 1 comm op (+all-gather)
Hop 3: Swap K/V With GPU 9
GPUs 10 and 11 hold chunks 1 and 2, but our query also needs chunk 0. GPUs 8 and 9 hold chunks 0 and 3 after their all-gather (before it, chunk 0 was only on GPU 8).
Ring attention keeps queries in place and moves K/V. At CP = 2 the ring is one swap between same-head partners, GPU 11 ↔ GPU 9 and GPU 10 ↔ GPU 8, equal to Llama 3’s K/V all-gather; ring differs only at larger CP. Our K/V goes to GPUs 8 and 9 to serve chunk 3, whose queries all come after it.
GPU 11 now holds two partial outputs for our query: $O_1$ over its local K/V and $O_2$ over the partner’s chunk 0. Each is already divided by its softmax sum $\ell_i = \sum_j e^{s_j - m_i}$ and carries its max score $m_i$, so GPU 11 merges them locally:
\[\begin{aligned} O &= \frac{w_1 O_1 + w_2 O_2}{w_1 + w_2} \\ w_i &= e^{m_i - m}\,\ell_i, \quad m = \max(m_1, m_2) \end{aligned}\]- GQA shrinks the swap: K and V cover
n_kv/2heads per GPU, notn_h/2. - Reduce-scatter: the row-parallel output projection (each GPU holds its heads’ rows) leaves each GPU a
[T/2, d]share of the sum; reduce-scatter sums them and hands each GPU its ownT/4tokens, so our token is whole again on GPU 11.
Adding SP costs TP no extra bytes (Korthikanti et al.). Plain TP would end attention with an all-reduce, leaving every token on both GPUs. A ring all-reduce is a reduce-scatter then an all-gather; SP runs the per-token work in between.
Trip log: 4/32 GPUs (+8, 9) · 0 node crossings · 3 comm ops (+K/V swap, reduce-scatter)
Hop 4: Dispatch to Two Experts
After the reduce-scatter, each token lives on one GPU and is dispatched once, not twice: SP’s real job in MoE inference. vLLM’s sequence-parallel MoE is a simpler version: it keeps the all-reduce and splits tokens just before dispatch.
- Route: GPU 11’s router picks experts 5 and 27 (illustrative), with weights $p_5$ and $p_{27}$.
- Dispatch: an all-to-all over the EP group. Expert 27 is on GPU 13, one NVLink hop away. Expert 5 is on GPU 2, node 0: the token crosses IB once, to GPU 3, the relay, which has GPU 11’s local index and so the same IB rail; GPU 3 forwards it over NVLink. DeepSeek-V3 routes its all-to-all this way, as do DeepEP’s high-throughput kernels for prefill.
- Combine: GPUs 2 and 13 run their experts and return results by the reverse path, each relay first summing its node’s weighted outputs so one copy crosses IB. GPU 11 adds the pieces:
This is where DP replicas stop being independent: each has its own attention, but all four share one copy of the experts, so every MoE layer is a rendezvous. An idle replica still runs a dummy batch so the others don’t hang. The slowest replica sets the pace for all 16 GPUs, so the load balancer must even out attention work.
Trip log: 7/32 GPUs (+13, 3, 2) · 2 node crossings (+2) · 5 comm ops (+dispatch, combine)
One Layer at a Glance
Each layer runs the steps above the dashed line:
- Two layouts alternate: outside attention, the SP layout gives each GPU
T/(CP·TP)whole tokens; inside, each TP rank seesT/CPtokens and1/TPof the heads. - TP’s all-gather stops at
[T/CP, d]: CP swaps K/V instead of gathering tokens, so no attention step materializes[T, d]. - Five comm ops per layer: all-gather, K/V swap, reduce-scatter, dispatch, combine.
Hop 5: PP Send to GPU 27
- Layers 2–4: experts on GPUs 11 (itself) and 4 (node 0, via relay GPU 3: one IB crossing each way); then 15 and 8, both on node 1, so no crossing (GPU 8 already held our K/V); then 1 and 6, both on node 0, still one crossing each way.
- The PP send: after layer 4, GPU 11 sends its
[T/4, d]slice to GPU 27 on node 3. - Layers 5–8: the same on nodes 2 and 3 (TP partner 26, CP partners 24 and 25, relay GPU 19). Each layer sends the token to one expert on node 2 and one on node 3: GPUs 23 and 28, 16 and 25 (already a CP partner), 20 and 29, 21 and 31.
A pipeline pays off only with several micro-batches in flight: other requests, or successive chunks of one long prompt (chunked pipeline parallelism). With one, each stage idles while the other works. Because DP replicas meet at every MoE layer, each micro-batch (A–C below) spans all four replicas, which move in lockstep:
Trip log: after layer 4, 11/32 GPUs · 6 node crossings · 20 comm ops; after the PP send, 12/32 · 7 · 21.
Hop 6: Leave K/V on GPUs 10, 11, 26, 27
Only the last prompt position, T − 1, on GPU 25 (chunk 3), feeds the LM head, unless prompt logprobs are requested. Stage 1’s TP pairs each hold the LM head split by vocabulary, but only GPUs 24 and 25 use theirs: GPU 25 copies the row to GPU 24, each computes V/2 logits, and the halves are gathered for sampling.
For every token but T − 1, only K/V outlives prefill. Ours is 2·n_kv·d_h·8 elements: the first half of the KV heads on GPUs 10 (layers 1–4) and 26 (layers 5–8), the second half on GPUs 11 and 27.
Layer 8 computes its K/V before attention and has no next layer, so its attention and MoE serve only T − 1: our token’s layer-8 expert trip (GPUs 21 and 31, 2 of its 15 node crossings) is dead work. Engines compute it anyway, and with EAGLE or MTP (multi-token prediction) speculative decoding it isn’t wasted: the draft head reads every final hidden state.
The Whole Trip: Map and Traffic
Trip log: 23/32 GPUs · 15 node crossings (14 EP, 1 PP) · 41 comm ops.
Elements sent per GPU per layer (g = n_h/n_kv; times at 450 GB/s NVLink, 50 GB/s IB):
| Parallelism | Elements per GPU | Link | Relative time |
|---|---|---|---|
| DP | 0 | — | — |
| CP | 2·(T/2)·(n_kv/2)·d_h = Td/(2g) | NVLink | 1/g |
| SP | 0 extra: reuses TP’s collectives | — | — |
| TP | all-gather + reduce-scatter: 2·(T/4)·d = Td/2 | NVLink | 1× (baseline) |
| EP | 2·k·(T/4)·d = Td before dedup, ≈ 0.38·Td of it over IB | IB (+ NVLink, overlapped) | ≈ 7× |
| PP | (T/4)·d once per stage boundary (every 4 layers) | IB | ≈ 1.1× per layer (4.5× per send) |
With uniform routing, 1 − (16/32)·(15/31) ≈ 76% of tokens have at least one expert on the other node, each crossing once each way: 2·0.76·(T/4)·d ≈ 0.38·Td. The table also assumes BF16, peak link rates, and no compute-communication overlap. Skewed routing makes EP worse; FP8 dispatch and overlapping the all-to-all with compute make it better. Either way, EP moves the most data over the slowest link.
Swapping PP and DP (rank = tp + 2·cp + 4·pp + 8·dp) would spread each stage over all four nodes. The PP send would move onto NVLink, but three quarters of a token’s candidate experts would sit across IB instead of half, raising EP’s IB traffic from about 0.38·Td to 0.67·Td per layer. Four layers of that outweigh one (T/4)·d send.
The Token’s Afterlife
Suppose this mesh also runs decode, and the first output token, position T, lands on GPU 11. In every layer, its query reads our token’s K/V.
Prefill moved K/V to the queries; decode moves the query to the K/V (“pass-Q” in Meta’s CP inference paper). After their all-gather, GPUs 10 and 11 each hold half the query’s heads and send them to their CP partners, GPUs 8 and 9; the partial outputs come back for merging. The new token’s K/V is written before attention, on GPUs 10 and 11. Keeping every new token there would grow CP rank 1’s cache, so implementations assign new tokens to CP ranks round-robin.
vLLM’s decode context parallelism also shards the KV cache by token, but carves the shards out of the TP group; for GQA it needs TP > n_kv, which this mesh doesn’t meet. Other axes change too:
| Prefill | Decode | |
|---|---|---|
| DP | one long prompt per replica | many requests, one new token each per step, still meeting at MoE layers |
| CP | splits the prompt; K/V travels | K/V stays put; the query goes out, partial outputs come back |
| SP | splits prompt tokens | splits the batch’s decode tokens (padded to a multiple of TP) |
| TP | compute-bound matmuls; big collectives | memory-bound matmuls; tiny, latency-bound collectives |
| EP | big per-expert batches | few tokens per expert; pool tokens from many DP replicas; latency-bound, so the all-to-all sends straight to expert GPUs with no NVLink relay |
| PP | filled by other micro-batches | each generated token waits on the previous one, so every stage’s latency is on the critical path |
Because the bottlenecks differ, large deployments run prefill and decode on separate pools with different meshes (DeepSeek: prefill EP32, decode EP144). In that setup, our token’s K/V makes one last trip, with the rest of the KV cache, from the prefill pool to the decode pool.
Three Answers
- Which GPUs form which group:
rank = tp + 2·cp + 4·dp + 16·pp. TP and CP partners sit inside the token’s 4-GPU replica; its EP group spans the 16-GPU stage. - Why this order: CP, SP, and TP talk every layer, so they never leave the node. EP crosses nodes on purpose and carries 14 of 15 node crossings. PP, which carries the 15th, goes outermost because putting DP outside it would nearly double EP’s IB traffic.
- Why DP replicas stop being independent: they share one set of experts, so every MoE layer is a rendezvous.