Chapter 13 · Section 13.4 · 75 min

Expert parallelism and systems

DeepSeek-V3's routed experts hold 654 billion parameters, about 1.3 TB in bfloat16. No GPU holds that, so the experts must live on many GPUs, and that changes the problem completely. On one device an MoE layer is a gather, some matrix products, and a scatter (§13.2). Across 64 devices, the gather and the scatter become network traffic. Every token has to travel to the GPUs that hold its experts and its results have to travel back, twice in the forward pass and twice more in the backward pass, in every MoE layer.

This section is about that traffic and the machinery around it. We lay out expert parallelism and check a dispatch by hand, derive how many bytes each GPU sends, find the ratio of compute to communication that decides whether the network or the arithmetic sets the speed, place the experts on a real cluster, see how an unbalanced router turns into idle GPUs and how redundant experts fix it, build the permutations and grouped matrix products that make the local work fast, and then turn to serving, where the same model needs a very different layout.

Expert parallelism

Expert parallelism (EP), made standard by GShard (Lepikhin et al. 2020), splits an MoE layer's EE experts over a group of nepn_{\text{ep}} GPUs, E/nepE/n_{\text{ep}} experts each. Everything else in the layer (attention, norms, the router) is data parallel. Each GPU runs attention on its own TlocT_{\text{loc}} tokens with a full copy of those weights. Only the experts are sharded. (The placement note at the end of §10.6 shows where the expert-parallel group sits among the other axes of the device mesh, and this section works the placement out in detail below.) One MoE layer's forward pass then runs these steps on every GPU at once.

  1. Route. Compute the router and the top-kk choices for the local tokens.
  2. Permute. Sort the TlockT_{\text{loc}}k assignments by destination GPU (expert ii lives on GPU ⌊i nep/E⌋\lfloor i\,n_{\text{ep}}/E\rfloor) and count how many go to each.
  3. Exchange counts. A tiny all-to-all of integers, so each GPU knows how many rows it will receive from each peer and can allocate buffers.
  4. Dispatch. An all-to-all (§10.2) sends each block of rows to its destination.
  5. Compute. Each GPU sorts what it received by local expert and runs its experts, one grouped matrix product.
  6. Combine. A second all-to-all, the exact reverse of the dispatch, returns every output row to the GPU its token came from.
  7. Unpermute and weight. Undo the sort and scatter-add each row into its token with its gate.

The backward pass runs the same two all-to-alls in reverse order, carrying gradients. With capacity limits, step 2 also drops or pads, so every buffer has a fixed size and step 3 can be skipped. Dropless systems (MegaBlocks, DeepSeek-V3) keep step 3 and handle variable sizes.

A dispatch you can check by hand

Take two GPUs and E=4E = 4 experts, with experts 0 and 1 on GPU 0 and experts 2 and 3 on GPU 1, top-2 routing, and three tokens per GPU. GPU 0's tokens choose the expert pairs {0,2}\{0, 2\}, {1,3}\{1, 3\}, and {0,1}\{0, 1\}. Four of its six assignments (to experts 0, 1, 0, 1) stay home and two (to experts 2 and 3) leave. GPU 1's tokens choose {2,3}\{2, 3\}, {0,2}\{0, 2\}, and {1,3}\{1, 3\}, so it keeps four assignments and sends two (to experts 0 and 1). Write entry (s,q)(s, q) of a matrix as the number of rows GPU ss sends to GPU qq,

(4224).\begin{pmatrix} 4 & 2 \\ 2 & 4 \end{pmatrix}.

Each row sums to kTloc=6kT_{\text{loc}} = 6, and always will, because every GPU ships each of its tokens kk times. Each column sums to the rows one GPU's experts must process, here 6 and 6, and the columns are what the router controls. Only the off-diagonal entries cross the network, and the combine all-to-all moves the same matrix transposed, since every row returns to where it came from. Now suppose every token had chosen {0,1}\{0, 1\}. The matrix becomes (6060)\begin{pmatrix} 6 & 0 \\ 6 & 0 \end{pmatrix}, GPU 0 processes all 12 rows while GPU 1 sits idle, and the layer takes twice as long as with balanced routing. Hold on to that picture. It returns, with numbers, when we come to load imbalance below.

The cell simulates the whole forward pass with four "GPUs" as entries of a Python list, checks the result against a single-device MoE, and prints the same send matrix for a random router.

Python · cell 1

The row sums are all 20, while the column sums wander, even for this random router. The widget below lets you skew the router and watch the columns, and with them the busiest GPU, pull away from the mean.

Where do tokens go in expert parallelism, and what happens when the router plays favorites?
Left, how many token copies each GPU sends to each GPU in the dispatch all-to-all. Right, how many tokens each expert receives against its capacity. Overflow beyond capacity is dropped.
Loading interactive figure…
Interview

Expert-parallel dispatch and combine

This is the expert-parallel coding question. Implement ep_moe(X, top, gate, experts, n_ep, all_to_all). GPU g holds tokens X[g] of shape (n_g, d), their chosen experts top[g] of shape (n_g, k), and gates gate[g] of the same shape. experts(e, rows) applies expert e to a 2-D array of rows. There are E = n_ep * experts.per_gpu experts, and expert e lives on GPU e // experts.per_gpu. Use exactly two calls to all_to_all(send), which takes send[src][dst] (any Python object) and returns recv[dst][src], one to dispatch and one to combine. On each destination GPU, run each local expert once on all the rows it received for that expert, from every source together (grouped, not row by row or source by source). Return (Y, counts), where Y[g] of shape (n_g, d) holds ∑jgatet,j FFN⁡topt,j(xt)\sum_{j}\text{gate}_{t,j}\,\operatorname{FFN}_{\text{top}_{t,j}}(\mathbf{x}_t) for each token and counts[s][q] is the number of rows GPU s sent to GPU q, the matrix of the worked example.

Hidden tests, hints, and AI help, from $1.99 a month.

How many bytes move

Count the dispatch for one GPU. It holds TlocT_{\text{loc}} tokens, each with kk assignments, and each assignment carries the token's dd-dimensional vector at ss bytes per number. If routing is balanced and an assignment's expert is equally likely to be on any of the nepn_{\text{ep}} GPUs, a fraction (nep−1)/nep(n_{\text{ep}}-1)/n_{\text{ep}} of the assignments leave the GPU. So the bytes each GPU sends in one dispatch, and by symmetry receives, are

Bdispatch=Tloc  k  d  sdisp  nep−1nep,B_{\text{dispatch}} = T_{\text{loc}}\; k\; d\; s_{\text{disp}}\;\frac{n_{\text{ep}}-1}{n_{\text{ep}}},

and the combine sends the same number of rows back, Bcombine=Tloc k d scomb (nep−1)/nepB_{\text{combine}} = T_{\text{loc}}\,k\,d\,s_{\text{comb}}\,(n_{\text{ep}}-1)/n_{\text{ep}}. In plain words, every token ships kk copies of itself out and gets kk results back, and almost all of them cross the network once the EP group is large.

A training step does this twice per MoE layer forward (dispatch and combine) and twice backward (the gradient of the combine is dispatched, and the gradient of the dispatch is combined), so about 2(Bdispatch+Bcombine)2(B_{\text{dispatch}} + B_{\text{combine}}) per layer per step, assuming the backward uses the same precisions. Note what does not appear. The number of experts EE and the expert width ded_e do not change the bytes, only the compute.

A worked example at DeepSeek-V3's shape

Take d=7168d = 7168 and k=8k = 8, a micro-batch of Tloc=4096T_{\text{loc}} = 4096 tokens per GPU, and EP over nep=64n_{\text{ep}} = 64 GPUs, as DeepSeek-V3 used in training. DeepSeek-V3 dispatches in FP8 (s=1s = 1) and combines in BF16 (s=2s = 2). Then

Bdispatch=4096⋅8⋅7168⋅1⋅6364=231 MB,Bcombine=462 MB.B_{\text{dispatch}} = 4096\cdot8\cdot7168\cdot1\cdot\tfrac{63}{64} = 231\text{ MB}, \qquad B_{\text{combine}} = 462\text{ MB}.

Now the compute those bytes feed. Each GPU's experts process, on average, TlockT_{\text{loc}}k assignments, and each costs 6 d de6\,d\,d_e forward FLOPs for a SwiGLU expert with de=2048d_e = 2048. That is 4096⋅8⋅6⋅7168⋅2048=2.9×10124096\cdot8\cdot6\cdot7168\cdot2048 = 2.9\times10^{12} FLOPs, or 2.9 ms at a bfloat16 peak of 989 TFLOP/s. Moving 693 MB over a 50 GB/s InfiniBand link, the cross-node bandwidth DeepSeek reports per H800 GPU, takes 13.9 ms. Over NVLink at the 160 GB/s they report, 4.3 ms. Either way, the communication takes longer than the arithmetic it serves.

That one example explains most of the engineering in DeepSeek-V3's MoE layer, and we can make it general. Divide the FLOPs by the bytes, and TlocT_{\text{loc}}, kk, dd, and nepn_{\text{ep}} (for large groups) all cancel,

expert FLOPsbytes moved=Tloc k⋅6 d deTloc k d (sdisp+scomb)=6 desdisp+scomb.\frac{\text{expert FLOPs}}{\text{bytes moved}} = \frac{T_{\text{loc}}\,k\cdot6\,d\,d_e}{T_{\text{loc}}\,k\,d\,(s_{\text{disp}} + s_{\text{comb}})} = \frac{6\,d_e}{s_{\text{disp}} + s_{\text{comb}}} .

In plain words, the only thing that buys compute per byte is a wider expert. A GPU that does 989 TFLOP/s over a 50 GB/s link needs about 19,800 FLOPs per byte to keep the link from being the bottleneck. DeepSeek-V3's fine-grained experts give 6⋅2048/3=4,0966\cdot2048/3 = 4{,}096. Mixtral's wide experts (de=14,336d_e = 14{,}336, both directions in BF16) give 6⋅14,336/4=21,5046\cdot14{,}336/4 = 21{,}504. Fine-grained experts, which §13.5 shows are better for quality, are exactly the ones that make communication hard.

The same ratio, written as time, is what a profiler shows. With a GPU doing FF FLOP/s and a link of B\mathcal{B} bytes per second (the symbols of §10.2),

tcommtcomp=Tloc k d (sdisp+scomb)/BTloc k⋅6 d de/F=(sdisp+scomb) F6 de B,\frac{t_{\text{comm}}}{t_{\text{comp}}} = \frac{T_{\text{loc}}\,k\,d\,(s_{\text{disp}} + s_{\text{comb}})/\mathcal{B}}{T_{\text{loc}}\,k\cdot6\,d\,d_e/F} = \frac{(s_{\text{disp}} + s_{\text{comb}})\,F}{6\,d_e\,\mathcal{B}},

ignoring the factor (nep−1)/nep(n_{\text{ep}}-1)/n_{\text{ep}}, which is 63/6463/64 here. For DeepSeek-V3 over InfiniBand at peak FLOP/s this is 3⋅989×1012/(6⋅2048⋅50×109)≈4.83\cdot989\times10^{12}/(6\cdot2048\cdot50\times10^{9}) \approx 4.8, the 13.9 ms against 2.9 ms above. In plain words, a token travels the same distance whatever the expert does with it, so the narrower the expert, the more of the layer's time is spent on the wire.

A training step improves the ratio but does not change the story. The backward pass costs twice the forward FLOPs (12 d de12\,d\,d_e per assignment) and adds the same two all-to-alls, so over a whole step the ratio is 18 de/(2(sdisp+scomb))=9 de/(sdisp+scomb)18\,d_e/\big(2(s_{\text{disp}} + s_{\text{comb}})\big) = 9\,d_e/(s_{\text{disp}} + s_{\text{comb}}), which is 6,1446{,}144 for DeepSeek-V3, still about a third of the 19,800 the link needs.

Python · cell 2

What DeepSeek-V3 did about it

Four techniques bring the problem under control, and each maps to a term above.

  • Lower precision on the wire. Dispatching in FP8 halves sdisps_{\text{disp}}. The combine stays in BF16 for accuracy.
  • Node-limited routing. Each token may use experts on at most M=4M = 4 of the 8 nodes in the EP group. The router scores each node by the sum of the token's two highest affinities there, keeps the best four nodes, and picks its 8 experts among them. A token is sent across InfiniBand at most once per destination node, and then forwarded to the right GPUs inside the node over the faster NVLink. Because NVLink is about 3.2 times faster than InfiniBand (160 against 50 GB/s), the paper notes that one InfiniBand arrival can feed about 3.2 experts in that node at no extra cost, so the same traffic could serve up to 13 routed experts per token (4 nodes times 3.2), though the model uses 8.
  • Overlap. Communication runs while other work computes. DeepSeek's DualPipe schedule interleaves the forward of one micro-batch with the backward of another, so one's all-to-all hides behind the other's matrix products, and custom kernels use only 20 of the GPU's 132 streaming multiprocessors to drive the network.
  • No tensor parallelism for experts. Each expert lives whole on one GPU, so no all-reduce is needed inside it.

How much of the problem is left after the first two? The report says that cross-node expert parallelism still gave DeepSeek-V3 a computation-to-communication ratio of about one to one, and the formula above reproduces that order of magnitude in two steps. First, the 4.8 above sent every one of a token's nearly 8 remote copies over InfiniBand. With the limit, a token crosses InfiniBand once per destination node, at most 4 nodes, and if nodes were picked evenly its own node would be among them half the time, so about 3.5 copies. That cuts the InfiniBand bytes, and the ratio, by more than half, from 4.8 to about 4.8⋅3.5/7.9≈2.14.8\cdot3.5/7.9 \approx 2.1 at peak FLOP/s. Second, the arithmetic does not run at peak. The whole run achieved about 35 percent of the BF16 peak, some 350 TFLOP/s (§8.9 derives it from the reported GPU hours). If the expert products ran at that rate, compute would take about three times as long as at peak, and the ratio would fall to about 2.1⋅0.35≈0.752.1\cdot0.35 \approx 0.75. In plain words, even with FP8 and the node limit, moving tokens takes about as long as computing on them, which is why overlap, the third technique, is not optional.

Bytes are not the whole cost. An all-to-all over nepn_{\text{ep}} GPUs is nep−1n_{\text{ep}} - 1 separate messages from each GPU, so in the latency and bandwidth model of §10.2 it pays a fixed latency per message on top of the bytes divided by the bandwidth. In training, with thousands of tokens per GPU, the bandwidth term dominates. In decoding, with a handful of tokens per GPU per step, each message is a few kilobytes and latency dominates. That is why DeepSeek's open-source DeepEP library ships two kinds of kernels, high-throughput ones for training and prefill that cross InfiniBand once and forward over NVLink, and low-latency ones for decoding that send directly over RDMA.

Is an expert-parallel MoE layer limited by arithmetic or by the network?
One MoE layer, forward pass, balanced routing. Compute is the experts' FLOPs at 989 TFLOP/s. Communication is the dispatch plus the combine all-to-all, per GPU.
Loading interactive figure…
Interview

All-to-all bytes per GPU

Implement alltoall_bytes(T_loc, k, d, n_ep, s_disp, s_comb, training=False) returning the bytes one GPU sends over the network in one MoE layer, assuming balanced routing with experts spread evenly over the nepn_{\text{ep}} GPUs of the EP group (assignments to a local expert cost nothing). Count the dispatch and the combine of the forward pass, and if training is true, the two all-to-alls of the backward pass as well, at the same precisions. Return a float.

Hidden tests, hints, and AI help, from $1.99 a month.
Research

Node-limited top-k routing

Implement DeepSeek-V3's group-limited selection, node_limited_topk(scores, k, n_groups, top_groups). The E experts in scores of shape (T, E) (already bias-adjusted affinities) are split into n_groups consecutive groups of E // n_groups experts, one group per node. For each token, score every group by the sum of its two largest expert scores, keep the top_groups best groups (ties to the smaller group index), and choose the token's k experts as the largest scores among experts in the kept groups only (ties to the smaller expert index). Return idx of shape (T, k), in descending order of score. No loop over tokens.

Hidden tests, hints, and AI help, from $1.99 a month.

Where the experts sit on the cluster

Expert parallelism needs no GPUs of its own. Take the ndpn_{\text{dp}} data-parallel GPUs that hold one pipeline stage (§10.5). For attention, norms, and every other dense weight they are ordinary replicas. For the experts they are cut into ndp/nepn_{\text{dp}}/n_{\text{ep}} expert-parallel groups of nepn_{\text{ep}} GPUs, and each group holds one full set of experts, E/nepE/n_{\text{ep}} of each MoE layer per GPU. Every expert therefore exists in

nedp=ndpnepn_{\text{edp}} = \frac{n_{\text{dp}}}{n_{\text{ep}}}

copies, one per group, and those copies form that expert's own data-parallel group. In plain words, a GPU is a full replica for the dense weights and one shard of nepn_{\text{ep}} for the expert weights, so dense gradients are averaged over all ndpn_{\text{dp}} GPUs and an expert's gradients only over the nedpn_{\text{edp}} GPUs that hold a copy of it.

Getting that group wrong is a silent bug. With 4 experts per GPU, GPU 0 of an expert-parallel group holds experts 0 to 3 and GPU 1 holds experts 4 to 7 in the same memory slots, so an all-reduce over all ndpn_{\text{dp}} GPUs would add expert 0's gradient to expert 4's. Frameworks build separate process groups for the two cases, and ZeRO-style sharding of optimizer states (§10.3) follows the same split, the dense states over ndpn_{\text{dp}} GPUs and the expert states over nedpn_{\text{edp}}.

DeepSeek-V3's layout, counted

DeepSeek-V3 trained on 2,048 H800 GPUs with 16 pipeline stages, 64-way expert parallelism spanning 8 nodes, and ZeRO-1 data parallelism, with no tensor parallelism at all. Each stage has ndp=2048/16=128n_{\text{dp}} = 2048/16 = 128 GPUs, which form two expert-parallel groups of 64, so nedp=2n_{\text{edp}} = 2 and each GPU holds 256/64=4256/64 = 4 routed experts of every MoE layer in its stage. The routed experts are 58⋅256⋅44,040,192=653.958\cdot256\cdot44{,}040{,}192 = 653.9B parameters (one expert is 3⋅7168⋅20483\cdot7168\cdot2048, §13.1), about 40.9B per stage and 0.64B per GPU. Everything else (attention, the shared experts, the three dense layers, routers, norms, embeddings) is the remaining 17.1B, about 1.07B per stage, and it is replicated on all 128 GPUs of the stage. So on average a GPU holds 1.71B parameters, 3.4 GB in BF16, of a model that is 1.3 TB in BF16. Without expert parallelism, the same 16-stage layout would put 41.9B parameters on every GPU, 84 GB of weights alone. (Stages do not split the layers exactly evenly, so these are averages.)

Python · cell 3

Expert parallelism has a ceiling. With whole experts on each GPU, nep≤En_{\text{ep}} \le E, so Mixtral's 8 experts fill at most 8 GPUs. Past that, each expert must be split too, with the column and row split of §10.4 applied inside every expert (Megatron calls it expert tensor parallelism). That adds an all-reduce of each expert's output and makes every expert's matrices thinner, so it suits models with a few wide experts. Models with hundreds of narrow experts, like DeepSeek-V3, use pure expert parallelism, because their experts are already small and there are enough of them to give every GPU a few.

Warm-up

An expert-parallel layout

Implement ep_layout(n_gpu, n_stages, n_ep, E, n_moe_layers, expert_params, other_params) for a cluster of n_gpu GPUs split into n_stages pipeline stages, with expert parallelism of degree n_ep inside each stage's data-parallel GPUs and no tensor parallelism. The model has n_moe_layers MoE layers of E routed experts with expert_params parameters each, and other_params parameters that are not routed experts. Return a dict with n_dp (GPUs per stage), n_edp (copies of each expert), experts_per_gpu (routed experts of one MoE layer per GPU), and params_per_gpu (the average parameters held by one GPU, as a float, assuming layers split evenly over stages). Raise ValueError if n_stages does not divide n_gpu, or if n_ep does not divide both n_dp and E.

Hidden tests, hints, and AI help, from $1.99 a month.

Load imbalance becomes idle time

The all-to-alls synchronize the expert-parallel group, so the expert phase of an MoE layer lasts as long as the busiest GPU's. Write ℓg\ell_g for the rows GPU gg's experts process (a column sum of the dispatch matrix) and ℓˉ\bar\ell for the mean over the group. The fraction of the group's expert compute that does useful work is

η=ℓˉmax⁡gℓg=11+MaxVio⁡GPU,MaxVio⁡GPU=max⁡gℓg−ℓˉℓˉ,\eta = \frac{\bar\ell}{\max_g \ell_g} = \frac{1}{1 + \operatorname{MaxVio}_{\text{GPU}}}, \qquad \operatorname{MaxVio}_{\text{GPU}} = \frac{\max_g\ell_g - \bar\ell}{\bar\ell},

the maximal violation of §13.3 measured per GPU instead of per expert. In plain words, a GPU that receives 50% more than its share makes the whole group run at two thirds of its speed. It also receives the most bytes, so the all-to-all slows by the same factor. This is what a capacity factor (§13.2) buys in expert-parallel training. Capping every expert at CF\text{CF} times its fair share caps every GPU at CF\text{CF} times the balanced load, so a step is never more than CF\text{CF} times slower than with perfect balance, and the price is paid in dropped tokens. Dropless systems give up that guarantee and rely on balancing instead.

Holding several experts per GPU helps, for two reasons you can derive in a line each. A GPU's load is a sum, ℓg=∑i on gci\ell_g = \sum_{i \text{ on } g} c_i over its m=E/nepm = E/n_{\text{ep}} experts, where cic_i is expert ii's load. First, the busiest GPU carries at most mmax⁡icim\max_i c_i while the mean is ℓˉ=mcˉ\bar\ell = m\bar c, so MaxVio⁡GPU≤MaxVio⁡\operatorname{MaxVio}_{\text{GPU}} \le \operatorname{MaxVio}, never worse than the busiest expert. Second, if the expert loads fluctuate independently around cˉ\bar c with standard deviation σ\sigma, a sum of mm of them has standard deviation m σ\sqrt m\,\sigma (variances of independent terms add) and mean mcˉm\bar c, so its relative spread is (σ/cˉ)/m(\sigma/\bar c)/\sqrt m. DeepSeek-V3's 4 experts per GPU halve the relative spread of one expert. Mixtral at EP8, with one expert per GPU, gets no averaging at all.

What averaging cannot fix is an expert that is hot every step. Its GPU is the straggler every time, and no amount of randomness evens it out. The fix DeepSeek-V3 uses in serving (below) is redundant experts. Give every GPU a few spare slots, put extra copies of the hottest experts in them, and split each expert's tokens evenly among its copies. Choosing the copies and placing them is a small scheduling problem, and a greedy method does well.

  1. Replicate. Hand out the spare slots one at a time, each to the expert whose load per copy, ci/qic_i/q_i with qiq_i its number of copies, is currently largest.
  2. Pack. Sort all copies by their load per copy, largest first, and put each on the least-loaded GPU that still has a free slot.

Step 2 is the longest processing time first rule from scheduling, which Graham (1969) proved finishes within 4/34/3 of the best possible makespan on identical machines when there is no slot limit. DeepSeek's open-source expert load balancer, EPLB, is built on the same replicate-then-pack pattern.

Here is a case to check by hand. Eight experts with loads (90,10,10,10,10,10,10,10)(90, 10, 10, 10, 10, 10, 10, 10) sit on 4 GPUs with 3 slots each, so there are 4 spare slots. The plain layout puts experts 2g2g and 2g+12g+1 on GPU gg, and GPU 0 carries 100100 against a mean of 160/4=40160/4 = 40, so η=0.40\eta = 0.40. The greedy plan gives all four spare slots to expert 0 (its load per copy goes 90→45→30→22.5→1890 \to 45 \to 30 \to 22.5 \to 18, always above 10), packs the five copies of 18 and the seven tens, and ends with GPU loads (46,38,38,38)(46, 38, 38, 38), so η=0.87\eta = 0.87. It is not optimal. Four copies of expert 0 (22.5 each) plus a second copy of one cold expert (5 and 5) give (42.5,42.5,37.5,37.5)(42.5, 42.5, 37.5, 37.5) and η=0.94\eta = 0.94. Greedy is a heuristic, a good one, and the widget lets you watch it work on realistic loads.

How much does one popular expert slow an expert-parallel step, and how much do redundant copies help?
Each bar is one GPU's expert load for one step, split into the experts it holds. The dashed line is the mean, and the step lasts as long as the tallest bar.
Loading interactive figure…
Research

Redundant experts and balanced placement

Implement place_experts(load, n_gpu, slots) for expert loads load of shape (E,) (tokens per expert in a recent window), n_gpu GPUs, and slots expert slots per GPU, with n_gpu * slots >= E. Every slot must hold one copy of some expert, every expert must have at least one copy, and an expert's tokens are split evenly over its copies. Return (copies, placement), where copies of shape (E,) counts each expert's copies and placement is a list of n_gpu lists of expert ids, slots per GPU. Use the greedy method of the text (ties to the smaller expert index, then the smaller GPU index), or anything better. The tests check the layout is valid and that its busiest GPU is no busier than the greedy plan's. Raise ValueError if the slots cannot hold every expert.

Hidden tests, hints, and AI help, from $1.99 a month.

Token permutation and grouped GEMMs

Inside each GPU, the received rows have to be grouped by local expert before the expert products can run, and the outputs ungrouped afterward. Both are permutations. If order is the stable argsort of the expert ids, then X[order] puts the rows in expert order, and the inverse permutation, inv[order] = arange(n), puts them back, Y_sorted[inv]. Production kernels fuse the permutation into the all-to-all or into the GEMM's loads, because a separate gather of a (Tk,d)(Tk, d) tensor costs a full read and write of HBM, the same reason for fusion as in §11.7.

Then come the products themselves. Expert ii has nin_i rows, and the nin_i differ from expert to expert and from step to step. There are three ways to run EE matrix products of different heights.

  • A loop of GEMMs. One kernel launch per expert. Simple, but the launches add overhead and small nin_i give small, inefficient products.
  • Pad to capacity and batch. Make every block exactly cap⁡\operatorname{cap} rows, so one batched GEMM of shape (E,cap⁡,d)×(E,d,de)(E, \operatorname{cap}, d)\times(E, d, d_e) does everything. Fast and static, but padding rows cost real FLOPs and memory, and overflowing rows are dropped.
  • A grouped GEMM. One kernel that receives a list of problem sizes and computes all EE products of different heights in one launch, scheduling the output tiles of §11.5 from all of them across the GPU. No padding beyond the tile size, no dropping. CUTLASS provides grouped GEMMs, and MegaBlocks (Gale et al. 2022) reached the same goal with block-sparse kernels, reporting training up to 40% faster than the Tutel MoE library while never dropping tokens.

The cell measures what padding costs on a realistic, slightly unbalanced load.

Python · cell 4

With the busiest expert at about 1.5 times the mean, capacity padding forces a bad choice. At CF=1\text{CF} = 1 it drops 7% of assignments and still pads the idle experts' buffers, and at CF=1.25\text{CF} = 1.25 it drops half a percent while computing a quarter more rows than needed. A grouped GEMM with 128-row tiles pads under 3% and drops nothing. That is why modern training stacks are dropless.

Core

Permute and unpermute

Implement permute(expert_ids, E) for a 1-D array of the expert id of each of nn rows. Return (order, counts, offsets, inverse). order is the stable sort of the rows by expert id, counts of shape (E,) the rows per expert (zeros included), offsets of shape (E + 1,) the start of each expert's block with offsets[E] = n, and inverse the permutation that undoes order, so that X[order][inverse] equals X.

Hidden tests, hints, and AI help, from $1.99 a month.

In PyTorch

On real hardware the two all-to-alls are calls to torch.distributed.all_to_all_single, which takes a flat send buffer and per-destination split sizes. The sketch below shows the dispatch and combine for one EP group. It cannot run in a notebook without several GPUs, so it is shown as static code.

The runnable version simulates an EP group inside one process. Each "rank" is a slot in a list, the all-to-all is a transpose of a list of lists, and autograd flows through it as it would through the real collective, because all_to_all_single is differentiable in the same way (its backward is the reverse all-to-all).

🔥 In PyTorchfor a machine with PyTorch

Serving an MoE model

Inference flips the cost structure. A decoding step processes one new token per sequence, so the matrix products are thin and the step is limited by reading weights from memory, not by arithmetic (§6.5, §11.2). For a dense model, every step reads every weight once, and a batch of BB sequences shares that read. For an MoE model, a step reads only the experts that some token in the batch chose, and each chosen expert sees only the tokens that chose it.

Count the arithmetic intensity of one expert's product. With nn tokens routed to it, the product does 2 n d de2\,n\,d\,d_e FLOPs per weight matrix and reads d ded\,d_e weights of 2 bytes, so the intensity is about nn FLOPs per byte. An H100 needs about 989/3.35≈295989/3.35 \approx 295 FLOPs per byte (its ridge point, §11.2) to stop being memory-bound. So each expert needs roughly 300 tokens per step, and with EE experts and kk choices that means a batch of about 300E/k300E/k tokens. For DeepSeek-V3, 300⋅256/8≈9,600300\cdot256/8 \approx 9{,}600 tokens per step per model replica. No single server with a few dozen users comes close, which is why MoE serving is organized around very large batches spread over many GPUs.

DeepSeek-V3's own deployment shows the pattern, with prefill and decoding served by separate deployments, each laid out for its own shape of work. For prefill (processing prompts, many tokens per request), the minimum unit is 4 nodes of 8 GPUs. Attention runs with 4-way tensor parallelism plus sequence parallelism and 8-way data parallelism, the MoE layers with 32-way expert parallelism, and 32 redundant experts (extra copies of the most heavily loaded experts, chosen from serving statistics and periodically rearranged, the replicate-and-pack problem of the load-imbalance section above) are added so each GPU hosts its 8 experts plus one duplicate. Prefill also runs two micro-batches of similar size at once, so the attention and expert computation of one hides the dispatch and combine of the other. For decoding, the unit is 40 nodes, 320 GPUs. Attention is 4-way tensor parallel with 80-way data parallelism, and the MoE layers use 320-way expert parallelism with one expert per GPU, and 64 GPUs hosting redundant and shared experts. One expert per GPU means each GPU's weight read is shared by every token in the whole 320-GPU batch that chose that expert.

At the other extreme, a single-GPU or laptop deployment has to hold all the experts somewhere. Mixtral in bfloat16 is 93 GB, so it needs two 80 GB GPUs, or 4-bit quantization (about 23 GB), or keeping most experts in CPU memory and copying them in when chosen. Offloading works better than you might expect because consecutive tokens often reuse experts (the temporal locality Mixtral's authors measured, §13.6), but it is slow, since a PCIe link moves a few tens of GB/s against an HBM's few TB/s.

layouthow experts are splitcommunication per MoE layerstrengthsweaknesses
tensor parallelevery expert split across nn GPUsall-reduce of (T,d)(T, d)no load imbalance across GPUs, simpleevery GPU touches every expert, all-reduce grows with TT
expert parallelwhole experts per GPUtwo all-to-alls of (Tk,d)(Tk, d)each expert's weights read once per batch, scales to hundreds of GPUsload imbalance becomes GPU idle time, needs large batches
EP plus redundant expertshot experts duplicatedsame, better balancedevens out real trafficmore memory, placement must track load
offloadingexperts in CPU memoryPCIe copies of chosen expertsruns on small hardwarevery slow, latency depends on cache hits

Read more

Watch next

🎬 Ultra-scale playbook, ch.6 - "Expert Parallelism" · Little ML book clubA reading group's walk through the Ultra-Scale Playbook's chapter on expert parallelism.
🎬 Lecture 48: The Ultra Scale Playbook · GPU MODE · starts at 149:00Nouamane Tazi on expert parallelism and how it combines with the other parallelisms on a real cluster.