Trace a Token Through TP + CP + SP + EP + DP + PP in MoE Inference

Trace a Token Through TP + CP + SP + EP + DP + PP in MoE Inference

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.

Six ways to cut the work, one color per parallelism

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.

One stage's 16 GPUs grouped as an attention mesh and as an expert mesh

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

Hop 1: GPU 11 inside DP replica 2

  1. DP: the load balancer sends the request to replica 2.
  2. CP: zigzag cuts the prompt into 4 chunks and gives CP rank r chunks r and 3 − r.
  3. SP: TP rank 0 takes chunk r, TP rank 1 takes chunk 3 − 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.

Zigzag CP split: chunk work 1, 3, 5, 7 balances to 8 per rank; chunks land on GPUs 8 to 11

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

Hop 2: all-gather with TP partner GPU 10

TP + SP on GPUs 10 and 11: all-gather to T/2 tokens, split heads, reduce-scatter to T/4

  • RMSNorm: per-token, so GPU 11 runs it alone.
  • All-gather: GPUs 10 and 11 exchange their T/4 tokens; attention needs all T/2.
  • Column-parallel QKV: each GPU holds the QKV weight columns for half the heads and projects all T/2 tokens. 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/2 KV 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

Hop 3: K/V swap with CP partners GPUs 8 and 9, then reduce-scatter with GPU 10

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.

Causal block matrix: query-key blocks GPUs 9 and 11 compute before and after the K/V swap

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/2 heads per GPU, not n_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 own T/4 tokens, 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

Hop 4: dispatch to experts on GPU 2 (via relay GPU 3) and GPU 13

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.

  1. Route: GPU 11’s router picks experts 5 and 27 (illustrative), with weights $p_5$ and $p_{27}$.
  2. 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.
  3. 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:
\[y = p_5\,E_5(x) + p_{27}\,E_{27}(x)\]

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:

One layer of prefill: every step, its collective, and the tensor shape on each GPU

  • Two layouts alternate: outside attention, the SP layout gives each GPU T/(CP·TP) whole tokens; inside, each TP rank sees T/CP tokens and 1/TP of 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

Hop 5: PP send from GPU 11 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:

Pipeline timeline: stage 1 idles until the PP send arrives, then both stages stay busy

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

Hop 6: LM head on GPUs 24 and 25; our token's 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

Occupancy map: which GPUs held the token at each step; 23 of 32 GPUs; every IB crossing is EP except the PP send

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.

Decode: GPU 11 writes the new token's K/V, copies its query to GPU 9, both attend over their KV shards, and GPU 11 merges the partial results

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.