Expert Parallelism › Worked Example: Tracing a Token Through Expert Parallelism

To make the concepts concrete, let us trace a single token through the full expert-parallel MoE layer step by step. Suppose we have 4 devices, each holding 4 experts (16 experts total). Our token starts on device 0 and uses top-2 routing.

Step 1: Routing decision. The gating network on device 0 processes the token's hidden representation (a vector of dimension $D$) and produces logits for all 16 experts. After softmax normalization, the top-2 selection identifies expert 2 (local, on device 0) and expert 11 (remote, on device 2) as the chosen experts. The routing weights are, say, 0.6 and 0.4 respectively.

Step 2: Dispatch phase. The token creates two dispatch entries: one for expert 2 (destination: device 0) and one for expert 11 (destination: device 2). The dispatch entry for expert 2 is a "local" routing that does not require network transfer. The dispatch entry for expert 11 must travel across the network. During the all-to-all, device 0 sends a copy of this token's representation to device 2 along with a tag identifying the target expert (expert 11, which is local expert index 3 on device 2, computed as $11 \mod 4 = 3$). At the same time, device 0 may receive token representations from other devices that need to be processed by experts 0-3.

Step 3: Expert computation. Device 0 processes all tokens routed to experts 0-3, including the token with expert 2. Device 2 processes all tokens routed to experts 8-11, including the copy of our token for expert 11. Each device processes its tokens independently and in parallel, with no inter-device communication during this phase. The expert transforms the token representation through its two-layer feed-forward network.

Step 4: Combine phase. Device 2 sends the processed result for expert 11 back to device 0. During the second all-to-all, all processed results travel back to their originating devices. Device 0 now has two processed representations: the output from expert 2 (which never left device 0) and the output from expert 11 (returned from device 2). The final output is the weighted sum: $0.6 \times \text{output}_2 + 0.4 \times \text{output}_{11}$.

Step 5: Residual connection. The combined output is added to the original token representation via the residual connection, and the result flows to the next transformer layer. From this layer's perspective, the computation looks identical to a dense feed-forward layer: input goes in, output comes out. The distributed routing machinery is completely transparent to subsequent layers.

This trace illustrates several key design choices. The routing weights (0.6 and 0.4) determine how much each expert contributes to the final representation, letting the model to blend expert knowledge. The local routing for expert 2 is free from a communication standpoint, which is why some implementations add incentives for local routing (though this can hurt model quality if taken too far). And the two-phase all-to-all structure ensures that every token gets both its expert results without the ordering information being lost.