EP: Partition Experts and Route Tokens

MoE selects experts, which are MLPs, for each token. As the number of experts grows, keeping all their weights on one GPU becomes difficult. Expert parallelism (EP) places different experts across multiple GPUs.

Not every part of a model needs the same parallelism. This article uses a simple configuration: replicate attention and other common layers with DP, and partition experts with EP. We first establish that the same GPUs use both strategies, then follow tokens as they fan out to selected experts and return to their source GPUs in their original order. We then examine how multiple expert results are combined and why equal expert counts can receive different token counts.

Replicate attention and partition experts

Suppose GPU 0 handles input sequence A and GPU 1 handles sequence B. A1 and A2 are two tokens of sequence A; B1 and B2 are two tokens of sequence B. Each small chip represents one token. What is actually computed and transferred is that token’s activation vector.

The two GPUs hold identical copies of attention and other common-layer weights. GPU 0 computes A and GPU 1 computes B using the same weights: the data parallelism approach introduced earlier. In this example, a sequence is not split across GPUs. Attention handles relationships among tokens within each GPU’s own sequence.

Expert placement is different. E0 and E1 reside on GPU 0; E2 and E3 reside on GPU 1. These are four MLPs with different weights, not replicas of the same weights. The router, which chooses experts for each token, has replicated weights on the two GPUs. Each GPU selects from all four experts for its own tokens.

Each GPU stores only some expert weights, but attention and router weights remain replicated. Halving the number of experts per GPU therefore does not halve total memory usage.

GPU 0 and GPU 1 stay in the same positions from top to bottom. After computing attention, they send tokens to the selected experts in the MoE block. Once results return, each GPU continues computing its own inputs. The GPUs are not divided into separate DP and EP sets: weight placement and data flow change by model part on the same GPUs.

Using DP for common layers is this article’s example configuration. EP does not prescribe one attention parallelism strategy. vLLM’s EP documentation distinguishes replicated attention from attention partitioned with TP. To focus on the token round trip, we omit TP and other partitioning strategies for the common layers here.

Send tokens to experts and return their results

Now zoom in on the MoE block. We begin with Top-1, selecting one expert per token. To focus on dispatch and return, this example sets the combination weight to one. Suppose the router selects E0 for A1, E2 for A2, E1 for B1, and E2 for B2.

Use Next in the figure to follow input, expert selection, dispatch, computation, return, and order restoration. Token colors and names remain consistent, while GPUs and experts stay in place.

The router determines destinations. On GPU 0, A1’s E0 is local, but A2’s E2 is on GPU 1. On GPU 1, B2’s E2 is local, while B1’s E1 is on GPU 0. A single central router does not gather every input to make these decisions. Each GPU computes expert-selection information from its own tokens.

Dispatch groups token vectors by the selected expert. A2’s vector goes to GPU 1, and B1’s vector goes to GPU 0. A1 and B2 are placed in their expert inputs within the same GPU. Expert weights stay in place while token vectors move, rather than moving weights toward the inputs each time. Experts receive the token input vectors used for selection, not the router scores.

E2 receives A2 and B2 from different sequences. This works because the expert MLP transforms each token independently. Their vectors can be stacked as rows for matrix operations with the same E2 weights, but B2’s values are not mixed into A2’s output. Grouping tokens under one expert batches their computation; it does not combine information across tokens as attention does.

Results return to their source GPUs and original tokens. E2’s result for A2 returns to GPU 0, and E1’s result for B1 returns to GPU 1. Locally computed results for A1 and B2 stay on their GPUs. GPU 0 arranges outputs in A1, A2 order, and GPU 1 in B1, B2 order, then passes them to the next operation. Each token has one result here, so the multi-result weighted sum needed in Top-2 is absent.

This round trip changes the arrangement from grouped by input → grouped by expert → grouped by input again. It connects to All-to-All as an exchange by destination, but destinations may receive different token counts. The association between each result and its original token must also be tracked. Megatron’s MoE documentation describes this dispatch-and-return relationship and several dispatchers. The central idea of EP is that selected experts receive inputs and their results return to the correct tokens, rather than the name of a particular communication function.

Return multiple expert results to the same token

With Top-2, each token selects two experts, so one input vector is used in two expert computations. Focus on A2 from GPU 0. Suppose its router selects local E0 and E2 on GPU 1, assigning combination weights of 0.7 and 0.3. Other tokens are omitted.

E0 and E2 receive the same A2 vector but use different weights, producing two results. In the figure, u0 is E0’s output and u2 is E2’s output. GPU 0 computes u0 locally; GPU 1 computes u2 and sends it back to GPU 0.

At the source, GPU 0 forms a weighted sum of only the two results for A2: y(A2) = 0.7 × u0 + 0.3 × u2. If both outputs are d-dimensional, the result is also d-dimensional and can be connected to A2’s original position. This is why expert-selection information and weights used for result combination must be kept distinct and available.

Selecting more experts increases the number of token–expert computations and the results to collect. However, input delivery can be shared when two experts reside on the same GPU, so network traffic is not always directly proportional to the number selected. Actual transfer volume depends on expert placement and implementation.

Equal expert counts can receive different token counts

Even when GPUs start with equal numbers of inputs, their routers need not select destinations evenly. Increase the input to four tokens per GPU: A1 through A4 on GPU 0 and B1 through B4 on GPU 1. The four equally sized experts still occupy two GPUs, two experts each, with Top-1 selection.

In the figure, E0 receives A1, E1 receives B1, and E2 receives the other six tokens. GPU 0 therefore processes two tokens and GPU 1 processes six. Evenly partitioning inputs for common layers does not guarantee an even workload in the expert block.

An expert selected frequently can concentrate computation, input-buffer demand, and communication on its GPU. Tokens originating on other GPUs also have to wait for that expert’s results. Conversely, too few tokens per expert can create many small matrix operations that use GPUs inefficiently.

Three times as many tokens does not mean exactly three times the execution time. The figure shows that dividing weights into equal counts and balancing incoming work are different tasks. Expert placement must therefore be considered together with the distribution produced by routing.

With EP, the same GPUs can use DP for common layers and exchange the tokens needed by their experts. The next article will examine how to choose and combine parallelism strategies.