MoE 预训练中的专家耦合：通过关联放置与令牌重排减少 All-to-All 通信开销

# Expert Coupling in MoE Pretraining: Reducing All-to-All Overhead with Correlated Placement and Token Shuffling
Source: [https://arxiv.org/html/2610.09372](https://arxiv.org/html/2610.09372)
###### Abstract

Mixture\-of\-Experts \(MoE\) layers replace the feed\-forward block of a Transformer withEEexpert networks, and each token is routed tokkof these experts\. Under expert parallelism \(EP\) the experts are distributed across GPUs, and every MoE layer runs all\-to\-all collectives in the forward and backward passes to dispatch tokens to their experts and then combine the results\. On a cluster with 8 AMD Instinct MI300X GPUs per node, these collectives can take 45 % of the training step at EP32 with top\-2 routing and 60 % with top\-6 routing\. We find that early in pretraining routers have already learned to assign tokens to experts in correlated patterns, both within a layer and across layers\. At top\-2, 0\.8 % of the expert pairs in a layer are selected together by 42 % of tokens, and the experts a token selects at one layer predict the experts it selects at the next layer\. We use these correlations to keep more token–expert assignments on the token’s own GPU, which reduces communication across GPUs and across nodes\. Correlated expert placement puts experts that are often selected together on the same GPU\. Combined with a dispatcher that sends each token to each GPU once, it removes up to 58 % of dispatched rows\. Token shuffling applies when sequence parallelism shards tokens across the EP group\. It moves each token to the GPU predicted to hold its next\-layer experts during the reduce\-scatter that follows attention\. On one node this raises the share of token–expert assignments served on the token’s GPU from 12\.5 % to 59 %\. In Megatron\-LM, across EP degrees from 8 to 64 with top\-2 and top\-6 routing, the two methods reduce all\-to\-all time by 1\.16–2\.63×\\timesand end\-to\-end step time by up to 1\.41×\\times\. Neither method changes the models’ underlying routing decisions or expert parameters\.

00footnotetext:\*Corresponding author:radha@zyphra\.com## IIntroduction

Mixture\-of\-Experts \(MoE\) layers have become a central building block of large language models\([Jiang et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib11);[Liu et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib1);[Yang et al\., 2025](https://arxiv.org/html/2610.09372#bib.bib13);[Kimi Team, 2026](https://arxiv.org/html/2610.09372#bib.bib14)\)\. In an MoE layer, the single feed\-forward network \(FFN\) of a dense Transformer block is replaced by a set ofEEparallel FFNs, called*experts*, together with a lightweight router\([Shazeer et al\., 2017](https://arxiv.org/html/2610.09372#bib.bib8);[Lepikhin et al\., 2021](https://arxiv.org/html/2610.09372#bib.bib9);[Fedus et al\., 2022](https://arxiv.org/html/2610.09372#bib.bib10)\)\. After the attention sub\-layer, the router scores each token’s hidden state against all experts and selects its top\-kkexperts\. The token is processed only by those experts, and their outputs are combined using the router weights\. Model capacity therefore grows withEE, while the FLOPs per token grow only withkk\. Because each expert is updated only by the tokens routed to it, experts become specialized over the course of pretraining\([Dai et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib12);[Xue et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib15)\)\.

This capacity comes at a communication cost\. Under expert parallelism \(EP\)\([Lepikhin et al\., 2021](https://arxiv.org/html/2610.09372#bib.bib9)\), experts are sharded across GPUs, and every MoE layer requires two all\-to\-all \(A2A\) collectives\. The*dispatch*sends each token from the GPU that owns it to the GPUs holding its selected experts, and the*combine*returns the expert outputs to the owning GPU\. Both collectives run in the forward pass and again in the backward pass\. The communication volume grows with the EP degree and with top\-kk, and once the EP group spans multiple nodes, the slower inter\-node links dominate step time\. In our measurements on a cluster with 8 AMD Instinct MI300X GPUs per node \(Section[II\-A](https://arxiv.org/html/2610.09372#S2.SS1)\), all\-to\-all communication accounts for 13\.5 % and 44\.8 % of step time at EP8 and EP32 under top\-2 routing, and for 24\.0 % and 60\.4 % under top\-6 routing\.

Conventional MoE implementations place experts contiguously across the EP group, with experteeon GPU⌊e/\(E/EP\)⌋\\lfloor e/\(E/\\mathrm\{EP\}\)\\rfloor, irrespective of which experts are used together\. Depending on the routing, the all\-to\-all can also become unbalanced, and the few GPUs that receive a disproportionate share of tokens become stragglers that delay the whole collective\. Mechanisms that address this include expert replication\([Liu et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib1)\), load\-driven expert relocation\([Zhai et al\., 2023](https://arxiv.org/html/2610.09372#bib.bib19)\), and dynamically planned redundant experts with weight prefetching\([Chen et al\., 2026](https://arxiv.org/html/2610.09372#bib.bib2);[Kimi Team, 2026](https://arxiv.org/html/2610.09372#bib.bib14)\)\. SmartMoE also changes which GPU holds each expert, but it places experts to balance load, not to co\-locate experts that are selected together\. These techniques lower the*peak*per\-link traffic by distributing the traffic across GPUs\. The*total*traffic stays about the same, because the same fraction of token–expert assignments crosses a GPU boundary\. Replication also costs memory for the expert copies\. DeepSeek\-V3 additionally reduces inter\-node traffic with node\-limited routing, which restricts the experts the router can select\.

A second line of work exploits expert*affinity*for MoE inference\. GRACE\-MoE\([Han et al\., 2025](https://arxiv.org/html/2610.09372#bib.bib20)\)groups experts that are often selected together within a layer onto the same GPU, first across nodes and then across the GPUs of each node, and sends a token to each GPU only once\. Its groups differ in size, and it replicates heavily loaded experts to rebalance the GPUs, which costs memory\. We apply within\-layer grouping to pretraining with equal\-size groups and no replicas\. ExFlow\([Yao et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib3)\)uses inter\-layer expert affinity to co\-locate experts across layers\. This approach does not transfer directly to pretraining\. During autoregressive decoding, a token can be placed on any convenient GPU\. During pretraining, a token stays on the GPU that holds its part of the sequence for the attention sub\-layer, and moving it requires additional communication\. Exploiting cross\-layer affinity during training therefore requires a mechanism that ExFlow does not provide\.

Our starting point is that once experts specialize during pretraining, routing follows stable patterns\. We study expert routing across pretraining stages and load\-balancing methods, and find two such patterns\. The first is*within\-layer correlation*\. Some pairs of experts are selected together much more often than independent routing would predict\. After approximately 2\.1 B training tokens with top\-2 routing, 42 % of tokens select one of only 0\.8 % of the possible expert pairs, and this pattern is already present after about one billion tokens\. The second is*cross\-layer correlation*\. The experts a token selects at layerℓ\\ellpredict the experts it selects at layerℓ\+1\\ell\+1\. Both correlations appear with all three load\-balancing methods we test, but with differing strengths\.

We exploit these correlations to raise*token–expert locality*, the fraction of \(token, selected expert\) pairs whose expert is on the token’s GPU\. These pairs never enter the network, so higher locality reduces the total A2A volume\. Our methods require no architectural changes or expert replicas, and they do not change routing decisions or the training loss\. Correlated placement and deduplication apply to any layout that uses expert parallelism\. Token shuffling applies only to layouts that use tensor parallelism with sequence parallelism and in which the TP group and the EP group are the same set of GPUs \(TP=EP\\mathrm\{TP\}=\\mathrm\{EP\}\), such as TP8 EP8 and TP16 EP16 in our evaluation\. It does not apply to layouts without tensor parallelism \(TP1\)\. Correlated expert placement and deduplication also apply to inference, where affinity\-based grouping has already been evaluated\([Han et al\., 2025](https://arxiv.org/html/2610.09372#bib.bib20)\); we do not evaluate inference\.

We make the following contributions:

- •Correlated Expert Placement and a Deduplicating Dispatcher \(Section[V\-A](https://arxiv.org/html/2610.09372#S5.SS1)\)\.We place correlated experts on the same GPU, so that the experts a token selects are more likely to share a GPU\. The deduplicating dispatcher sends each token to a destination GPU at most once, regardless of how many of its selected experts are there\. Deduplication is an existing technique\([Zhao et al\., 2025](https://arxiv.org/html/2610.09372#bib.bib16)\)\. The all\-to\-all dispatcher in Megatron\-LM\([Shoeybi et al\., 2019](https://arxiv.org/html/2610.09372#bib.bib7)\)sends one copy of a token per selected expert, so we implement deduplication there as the basis for placement\. Together, placement and deduplication reduce both intra\-node and inter\-node communication volume\.
- •Token Shuffling \(Section[V\-B](https://arxiv.org/html/2610.09372#S5.SS2)\)\.For sequence\-parallel layouts in which the TP group equals the EP group, we \(i\) predict a token’s experts at layerℓ\+1\\ell\+1from the experts the token selected at the previous layers, and \(ii\) move the token ahead of time to the GPU that is predicted to hold those experts\. The move is folded into the reduce\-scatter that sequence parallelism\([Korthikanti et al\., 2023](https://arxiv.org/html/2610.09372#bib.bib4)\)already performs after attention\. This brings to pretraining the locality benefit that prior affinity\-based work obtained for inference\.
- •Characterization of Expert Routing \(Sections[II\-B](https://arxiv.org/html/2610.09372#S2.SS2)and[VI](https://arxiv.org/html/2610.09372#S6), and the appendix\)\.We quantify within\-layer co\-selection and cross\-layer predictability across pretraining stages, top\-kksettings and load\-balancing methods, and relate these statistics to communication volume\.
- •Evaluation in Megatron\-LM \(Section[VI](https://arxiv.org/html/2610.09372#S6)\)\.We report communication volume, A2A time and end\-to\-end step time\. With expert parallelism alone, correlated placement and deduplication reduce A2A time by1\.16×1\.16\\times–1\.38×1\.38\\timesunder top\-2 routing and by1\.35×1\.35\\times–1\.95×1\.95\\timesunder top\-6 routing\. With token shuffling added, A2A time is reduced by1\.40×1\.40\\times–1\.74×1\.74\\timesunder top\-2 and by1\.46×1\.46\\times–2\.63×2\.63\\timesunder top\-6\.

## IIMotivation

### II\-AAll\-to\-All Cost in MoE Pretraining

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/fig_step_anatomy_stacked.png)Fig\. 1:Share of training step time by component at EP8 and EP32 with top\-2 and top\-6 routing, using the Megatron\-LM all\-to\-all dispatcher and contiguous placement\. Dashed boxes mark the all\-to\-all share\.We begin by measuring how much of an MoE training step is spent in the dispatch and combine collectives\. Each MoE layer runs both collectives in the forward pass and again in the backward pass\. The cost of these collectives depends on two configuration choices\. A larger expert\-parallel degree spreads experts over more GPUs, so the fraction of a token’s experts on its own GPU falls as1/EP1/\\mathrm\{EP\}\. Beyond a single node, however, the exchange also moves from the intra\-node fabric to the inter\-node network\. A larger top\-kkraises the number of activation rows each token contributes to both collectives, so top\-6 moves three times the volume of top\-2\.

We measure a 12\-layer MoE model \(dmodel=4096d\_\{\\mathrm\{model\}\}\{=\}4096\), with 128 experts of FFN width 1024 per layer \(see Section[VI](https://arxiv.org/html/2610.09372#S6)and Table[II](https://arxiv.org/html/2610.09372#A1.T2)for the detailed configuration\)\. Figure[1](https://arxiv.org/html/2610.09372#S2.F1)shows that the two effects compound\. We measure all\-to\-all time as the wall\-clock time the step spends inside the dispatch and combine calls, which includes time spent waiting for other GPUs\. Widening the EP group from a single node \(EP8\) to four nodes \(EP32\) raises the all\-to\-all share of step time from 13\.5 % to 44\.8 % under top\-2 and from 24\.0 % to 60\.4 % under top\-6\. At fixed EP8, moving from top\-2 to top\-6 raises it from 13\.5 % to 24\.0 %\. At EP32 with top\-6, the two collectives take more of the step than expert computation and all non\-MoE computation combined\.

### II\-BExpert Correlation Within and Across Layers

Crucially, the measurements above use contiguous placement, which ignores the correlations between the experts a token selects\. Once experts have specialized, these selections follow two regular patterns, which we exploit in this paper\. We measure both on the top\-2 and top\-6 models from Section[II\-A](https://arxiv.org/html/2610.09372#S2.SS1)\. All statistics in this section are computed offline\. We record every token’s expert choices at every MoE layer on a few late training steps \(131k tokens per step\), fit each placement and predictor on some of these steps, and apply them to the recorded routes of the other steps without running the dispatcher\.

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/p2_top2_within_layer_heatmap.png)\(a\)Within layer 8
![Refer to caption](https://arxiv.org/html/2610.09372v1/images/p2_top2_across_layer_heatmap.png)\(b\)Across layers8→98\\to 9

Fig\. 2:Expert correlation in the top\-2 model, computed from the saved routing of 524k tokens at training steps 1983–1995 \(every fourth step\)\. \(a\)P⁡\(e2∈top\-2∣e1∈top\-2\)P\(e\_\{2\}\\in\\text\{top\-2\}\\mid e\_\{1\}\\in\\text\{top\-2\}\)within layer 8; \(b\)P⁡\(e2​at​9∣e1​at​8\)P\(e\_\{2\}\\text\{ at \}9\\mid e\_\{1\}\\text\{ at \}8\)across layers\. Blue marks strongly correlated pairs \(conditional probability above 0\.3\)\.TABLE I:Offline placement and prediction statistics on held\-out iterations 1983 and 1987 \(fit on 1991–1999\); all values in % \(rounded\) except*GPUs*\.*Same GPU*: tokens whose top\-kkexperts share a GPU\.*GPUs*: distinct destinations per token\.*Deduplication*: rows removed as duplicates\.*Local experts*: share of a token’s next\-layer experts on its own GPU when it is moved to the GPU predicted from its current experts \(T/EPT/\\mathrm\{EP\}tokens per GPU\) and when it stays on its home GPU \(1/EP1/\\mathrm\{EP\}\)\.Within layerAcross layers:local expertsEPPlacementSameGPUGPUsDedupli\-cationTokenmovedTokenstaystop\-2 routingEP8Contiguous121\.8863313Correlated531\.47265813EP16Contiguous61\.953286Correlated451\.5523516EP32Contiguous21\.981253Correlated361\.6418423top\-6 routingEP8Contiguous04\.49252513Correlated182\.54586013EP16Contiguous05\.2413186Correlated83\.0849496EP32Contiguous05\.676143Correlated03\.8935333##### Within a Layer, Experts Are Selected in Stable Pairs

Figure[2\(a\)](https://arxiv.org/html/2610.09372#S2.F2.sf1)shows, for layer 8, the probability that a token selecting experte1e\_\{1\}also selectse2e\_\{2\}\. The blue cells are pairs that are often chosen together: 14 pairs are chosen with at least 50 % confidence and carry 16 % of all tokens, and 64 of the 8,128 possible pairs carry 42 %, whereas independent routing with the same loads would give 1\.6 %\. For 73 of the 128 experts, one partner appears in at least 30 % of its tokens\.

A token is sent once per GPU its experts occupy, so co\-locating a frequent pair halves its traffic under a deduplicating dispatcher\. Contiguous placement does this only by chance; correlated placement does it on purpose\. Table[I](https://arxiv.org/html/2610.09372#S2.T1)\(within layer\) counts rows, with each row corresponding to one copy of a token’s hidden state sent to one GPU\. At EP8, correlated placement raises the share of top\-2 tokens whose experts share a GPU from 12 % to 53 %, and raises the fraction of rows removed by deduplication from 6 % to 26 % under top\-2 and 25 % to 58 % under top\-6\. At EP32, a GPU holds only four experts, so a top\-6 token’s experts can never all share one, yet rows removed still rise from 5\.5 % to 35\.1 %\.

##### Across Layers, the Current Experts Predict the Next Ones

Figure[2\(b\)](https://arxiv.org/html/2610.09372#S2.F2.sf2)plots, for each experte1e\_\{1\}at layer 8, the probability that a token routed toe1e\_\{1\}selects experte2e\_\{2\}at layer 9\. Most layer\-8 experts send a large share of their tokens to a single layer\-9 expert: 112 of the 128 send at least 30 % to one specific expert, and the median share is 48 %\.

This correlation can be used to move*tokens*\. Before the router of layeri\+1i\{\+\}1runs, the experts chosen at layeriialready indicate which GPU will hold most next experts, and under sequence parallelism \(Section[III\-B](https://arxiv.org/html/2610.09372#S3.SS2)\) the token can be moved there during attention\. Every GPU must own exactlyT/EPT/\\mathrm\{EP\}tokens, so not every token gets its predicted GPU\. Table[I](https://arxiv.org/html/2610.09372#S2.T1)\(across layers\) reports, under this constraint, the share of a token’s next\-layer experts that are local\. A token that stays on its home GPU has1/EP1/\\mathrm\{EP\}local\. A token moved to its predicted GPU has 58 %, 51 % and 42 % at EP8, EP16 and EP32 for top\-2 and 60 %, 49 % and 33 % for top\-6 under correlated placement\. Placement matters here too: with contiguous placement the moved token has only 33 % of its experts local at EP8, because its next experts span more GPUs\.

## IIIBackground

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/base_dedup_expert_place.png)Fig\. 3:Reducing A2A with within\-layer expert correlation\. Each tokenTiT\_\{i\}and its selected experts have the same color\.### III\-AExpert Parallelism

When the experts of a layer no longer fit in the memory of one GPU, expert parallelism \(EP\) partitions them across a group of GPUs, typically alongside data parallelism\. WithE=128E\{=\}128experts andEP=8\\mathrm\{EP\}\{=\}8, for example, each GPU hostsE/EP=16E/\\mathrm\{EP\}\{=\}16experts\. A token routed to an expert on another GPU is sent there by the dispatch all\-to\-all and returned by the combine all\-to\-all, in both the forward and the backward pass\. The first panel of Figure[3](https://arxiv.org/html/2610.09372#S3.F3)shows this pattern\. Every token–expert assignment that crosses a GPU boundary adds a row to both collectives\.

### III\-BHybrid Parallelism: Tensor–Expert Parallelism

For large models and long\-context training, expert parallelism alone is not enough to fit the model in GPU memory, and additional model parallelism is needed\. Tensor parallelism \(TP\)\([Shoeybi et al\., 2019](https://arxiv.org/html/2610.09372#bib.bib7)\)shards the computation of each layer across GPUs and is widely used in both pretraining and inference\.

In a common hybrid TP–EP configuration, attention and the MoE layer use different token layouts\. Attention is partitioned across the TP ranks, and the MoE layer operates on sequence shards distributed across the same ranks\([Korthikanti et al\., 2023](https://arxiv.org/html/2610.09372#bib.bib4);[Liu et al\., 2025](https://arxiv.org/html/2610.09372#bib.bib6);[Singh et al\., 2023](https://arxiv.org/html/2610.09372#bib.bib5)\)\. Two collectives move between these layouts\. After attention, a reduce\-scatter sums the partial TP outputs and partitions the sequence, which leaves each rank with a contiguous subset of tokens\. These tokens then pass through the dispatch all\-to\-all, the expert computation, and the combine all\-to\-all\. After the MoE layer, an all\-gather reconstructs the full sequence on every TP rank for the next Transformer block\.

Token shuffling replaces this reduce\-scatter and all\-gather with shuffled versions \(Figure[5](https://arxiv.org/html/2610.09372#S5.F5)\), and it requires the TP group and the EP group to be the same GPUs \(TP=EP\\mathrm\{TP\}=\\mathrm\{EP\}\)\. The shuffled reduce\-scatter sends each token to the rank predicted to hold most of its experts in the coming MoE layer, based on the experts it selected in the previous layers\. The shuffle\-aware all\-gather restores the canonical token order before the next attention layer\. A larger fraction of each token’s experts is then on its own GPU, and the dispatch and combine all\-to\-alls send fewer bytes across the network\. Section[V\-B](https://arxiv.org/html/2610.09372#S5.SS2)describes the mechanism in detail\. Experts are not tensor\-parallel in this work\. Some latency\-sensitive inference deployments shard experts with tensor parallelism, which we do not consider\.

## IVRelated Work

Two systems place experts by within\-layer affinity on models that are already trained\. Occult\([Luo et al\., 2025](https://arxiv.org/html/2610.09372#bib.bib21)\)measures co\-selection in a pretrained MoE model on a profiling dataset, places experts that are often selected together within a layer on the same GPU, and sends a token to each GPU only once\. It evaluates inference and fine\-tuning, and obtains further savings by restricting each token’s experts to a few GPUs, which changes routing decisions\. GRACE\-MoE\([Han et al\., 2025](https://arxiv.org/html/2610.09372#bib.bib20)\)targets inference\. It groups co\-selected experts first across nodes and then across the GPUs of each node, lets the groups differ in size, and replicates heavily loaded experts to rebalance the GPUs, which costs memory\. Correlated expert placement \(Section[V\-A](https://arxiv.org/html/2610.09372#S5.SS1)\) uses the same within\-layer signal and, like GRACE\-MoE, partitions across nodes before GPUs\. Our work differs in three ways\. We study pretraining from scratch, where the correlations do not exist at the start, and measure how early they form and how few tokens are needed to estimate them \(Section[VI\-A](https://arxiv.org/html/2610.09372#S6.SS1)\)\. We leave routing decisions unchanged and add no replicas\. We also move tokens to their experts with token shuffling \(Section[V\-B](https://arxiv.org/html/2610.09372#S5.SS2)\), which neither system does\.

## VImplementation

### V\-AWithin\-Layer Correlation: Correlated Expert Placement

As discussed in Section[II\-B](https://arxiv.org/html/2610.09372#S2.SS2)and shown in Figure[2\(a\)](https://arxiv.org/html/2610.09372#S2.F2.sf1), some experts within a layer are frequently selected together by the same tokens\. We use this structure to reduce all\-to\-all \(A2A\) communication with two techniques\.*Correlated expert placement*puts frequently co\-selected experts on the same GPU, and a*deduplicating dispatcher*sends each token to each EP rank at most once\.

Figure[3](https://arxiv.org/html/2610.09372#S3.F3)shows the A2A communication for three configurations: the baseline, with deduplication, and with both deduplication and correlated expert placement\.1shows contiguous placement with the conventional dispatcher, where each\(token,expert\)\(\\text\{token\},\\text\{expert\}\)assignment produces a separate A2A row\. When a token selects several experts on the same destination rank, its hidden state is sent there several times\.2adds deduplication\. Assignments that share the same\(token,destination rank\)\(\\text\{token\},\\text\{destination rank\}\)are merged, so the hidden state is sent to that rank once\.3adds correlated expert placement, which puts frequently co\-selected experts on the same rank\. This increases token–expert locality and gives deduplication more rows to merge\. Neither technique changes the router’s expert assignments or the expert computation\.

##### Deduplicating Dispatcher

Conventional top\-kkdispatch sends one copy of the token’s hidden state per\(token,expert\)\(\\text\{token\},\\text\{expert\}\)assignment, even when several selected experts reside on the same EP rank\. We instead send one row per unique\(token,destination rank\)\(\\text\{token\},\\text\{destination rank\}\)pair\. Each row carries the IDs and routing probabilities of the token’s experts on that rank\. The receiving rank uses this metadata to give each of its selected experts a copy of the token and to group the copies by expert for the expert computation, so no separate collective is needed for this metadata\. The combine path works the same way\. The expert rank weights the outputs of the token’s experts by their routing probabilities, sums them, and returns one partial output per\(token,rank\)\(\\text\{token\},\\text\{rank\}\)pair\. Dispatch and combine therefore send one row per token for each rank that holds at least one of its experts\. Deduplication is an existing technique\([Zhao et al\., 2025](https://arxiv.org/html/2610.09372#bib.bib16)\), and we implement it in Megatron\-LM as the basis for our correlated placement algorithms\.

##### Correlated Placement

For each MoE layer, we build a correlation graphWW, whereW⁡\[e,e′\]W\[e,e^\{\\prime\}\]counts the tokens that select both expertseeande′e^\{\\prime\}\. We look for a placement that puts the same number of experts on every GPU and maximizes the total weightW⁡\[e,e′\]W\[e,e^\{\\prime\}\]of the expert pairs that share a GPU\. Algorithm[1](https://arxiv.org/html/2610.09372#alg1)in the appendix performs a Kernighan–Lin\-style local search\([Kernighan and Lin, 1970](https://arxiv.org/html/2610.09372#bib.bib18)\)\. It starts from contiguous placement and swaps pairs of experts between GPUs whenever a swap increases this total\. Every GPU holds exactlyE/EPE/\\mathrm\{EP\}experts, so the placement does not change the memory used for expert parameters\.

For EP groups that span multiple nodes, we partition in two steps\. We first partition the experts across nodes, which keeps co\-selected experts on the same node and reduces inter\-node traffic\. We then partition each node’s experts across its GPUs, which keeps co\-selected experts on the same GPU where possible\. The resulting placement is a per\-layer permutation of global expert IDs and does not change router decisions or expert parameters\.

### V\-BCross\-Layer Correlation: Token Shuffling

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/shuffled_tp_ep_baseline.png)Fig\. 4:TP–EP baseline with correlated placement\. Each token stays on the rank of its canonical sequence slice, so its rows still cross the network even when its experts share one GPU\.![Refer to caption](https://arxiv.org/html/2610.09372v1/images/shuffled_tp_ep_e2e.png)Fig\. 5:Token shuffling end to end\. After the attention of layerℓ\\ell1, the shuffle planner predicts each token’s owner rank from its layer\-\(ℓ−1\)\(\\ell\{\-\}1\)routes, and the shuffled reduce\-scatter2delivers each row to that rank\. The dispatch3, the expert computation4and the combine5then find most of the token’s experts local\. The shuffle\-aware all\-gather6restores canonical order before the next attention block7\.Expert placement decides where experts live and does not change where tokens live\. With TP–EP parallelism, the MoE layer runs on the sequence shard that the reduce\-scatter after attention assigns to each rank\. Each rank thus owns a fixed, contiguous slice of the sequence, and a token’s rank is set by its position alone\. Even when the experts a token needs have been placed together on one GPU, the token is actually on that GPU with probability only1/EP1/\\mathrm\{EP\}\. Thus, with probability1−1/EP1\-1/\\mathrm\{EP\}the token’s row must still cross the wire in the dispatch and return in the combine\. Figure[4](https://arxiv.org/html/2610.09372#S5.F4)shows this for four ranks under correlated placement\. The conventional reduce\-scatter places T0 on GPU0, T1 on GPU1 and so on, although T0’s experts sit together on GPU1 and T1’s on GPU2, so their rows still leave the owner’s GPU\.

The cross\-layer correlation of Section[II\-B](https://arxiv.org/html/2610.09372#S2.SS2)\(Figure[2\(b\)](https://arxiv.org/html/2610.09372#S2.F2.sf2)\) lets us move the token as well\. The experts a token used at layerℓ\\ellpredict, well above chance, which GPU will host its experts at layerℓ\+1\\ell\{\+\}1, and this prediction is available a whole attention block before the layer\-\(ℓ\+1\)\(\\ell\{\+\}1\)router runs\. Token shuffling \(Figure[5](https://arxiv.org/html/2610.09372#S5.F5)\) uses this window\. We call the GPU that holds a token in the MoE layer its*owner*\. The reduce\-scatter that follows attention first permutes tokens so that each goes to its predicted owner and then scatters them, so the following MoE dispatch finds that most of the token’s experts are now local\. The rest of this section describes where ownership can change without extra communication, how the owner is predicted, and how the planner is hidden behind attention\.

##### Changing Ownership Inside Existing Collectives

The token\-shuffling construction requires the TP group and the EP group to be the same set of ranks \(TP=EP\\mathrm\{TP\}=\\mathrm\{EP\}\)\. In TP–EP, a reduce\-scatter follows the attention output projection and returns the sequence to its sharded layout\. An all\-gather precedes the next block’s QKV projection and reassembles the full sequence on every rank \(Figure[5](https://arxiv.org/html/2610.09372#S5.F5)\)\. Both collectives already move every token between ranks\. Which rank receives which token in the reduce\-scatter is set by indexing, so changing it does not change the number of bytes sent\. We therefore fuse a permutation into each collective\. The fused reduce\-scatter sends every token to its planned owner instead of the rank that holds it in the default contiguous sequence layout\. It also moves the token’s row of the residual stream, so the residual addition happens on the new owner\. The fused all\-gather restores the original token order on the way back, because attention depends on token order\. The MoE layer sees a local slice whose tokens mostly have their experts on the same GPU, and the attention layer is unchanged\. Token shuffling thus adds no additional collective\. Its two fused permutations cost a few milliseconds per step, compared with the hundreds of milliseconds removed from the all\-to\-all due to decreased communication \(Section[VI](https://arxiv.org/html/2610.09372#S6)\)\.

##### Predicting the Owner

The shuffle planner answers one question per token: given the experts it just visited, on which rank will its next experts most likely live? Routing is strongly correlated across layers, so the pair of experts a token used at the two preceding MoE layers is a good predictor of where it lands next\. Offline, from routing traces and the chosen placement, we count for every MoE layer how often each pair sends the token to each rank; the fewer observations a pair has, the more weight we give to the single preceding expert instead\. At runtime the previous layers’ routes are the only input, and the planner adds no collective\. The lookup gives each token a preferred owner and a confidence score\. Sequence parallelism requires every rank to own exactlyT/EPT/\\mathrm\{EP\}tokens, so tokens are assigned in decreasing confidence, each to its most preferred rank with capacity left\. We give the full algorithm in Algorithm[2](https://arxiv.org/html/2610.09372#alg2)\.

##### Hiding the Shuffle Planner

Prediction and quota assignment cost about a third of the forward all\-to\-all time at TP16 EP16 so, if executed sequentially, this would forfeit most of the gain\. However, crucially, the plan depends only on routes up to layeriiand is needed only at the layer\-\(i\+1\)\(i\{\+\}1\)reduce\-scatter, so we can run the planner on a side stream, concurrently with the expert computation and the next attention block, and join it with a single event wait \(Figure[6](https://arxiv.org/html/2610.09372#S5.F6)\)\. In practice the main stream waits about one millisecond per step for the plan \(0\.02 % of the step\), plus one fused permute\-plus\-residual kernel\.

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/fig_planner_overlap_preview.png)Fig\. 6:The planner runs on a side stream between the layer\-iiroutes and the layer\-\(i\+1\)\(i\{\+\}1\)reduce\-scatter\.

## VIExperimental Results

All experiments use the two models of Section[II\-A](https://arxiv.org/html/2610.09372#S2.SS1)\. Both are 12\-layer MoE models withdmodel=4096d\_\{\\mathrm\{model\}\}=4096and 128 experts of FFN width 1024 per layer, one with top\-2 and one with top\-6 routing\. Both use a sequence\-level auxiliary loss with coefficient10−210^\{\-2\}for load balancing\. We train each for 2\.1 billion tokens \(2000 steps\) on FineWeb\-Edu\([Penedo et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib17)\)with a sequence length of 8192 and a global batch size of 128\.

We take all measurements on a cluster of nodes with eight AMD Instinct MI300X GPUs each, fully connected within a node by xGMI \(one direct link between each pair of GPUs\) and connected across nodes by eight 100 Gb/s RoCE NICs per node, one per GPU\. All timing runs load the step\-2000 checkpoint of the respective model and run on held\-out data from the same corpus, so the routing decisions are those of a trained model\. We run each experiment twice\. The first run consists of 14 clean steps; we report the median over the last 10 steps, after discarding the first 4 as warm\-up \(checkpoint loading and communicator initialization\)\. The second run consists of 14 steps with the region profiler and all\-to\-all metrics enabled, which provide the time and traffic breakdowns, following the same procedure\. Speedups are relative to the baseline, which is the Megatron\-LM all\-to\-all dispatcher with contiguous placement\. We compare three configurations against it\. These are deduplication, deduplication with correlated expert placement, and, on layouts withTP=EP\\mathrm\{TP\}=\\mathrm\{EP\}, token shuffling on top of both\. The figures label the baseline and the three configurations B, D, P and S\. We implement all methods in Megatron\-LM\([Shoeybi et al\., 2019](https://arxiv.org/html/2610.09372#bib.bib7)\)\. The appendix repeats the top\-6 placement evaluation with two other load\-balancing methods \(Figure[11](https://arxiv.org/html/2610.09372#A1.F11)\)\. The correlations are present under all three, with different strengths\.

### VI\-ABuilding Expert Correlation Within and Across Layers

Correlated placement and the token\-shuffling predictor both use correlation tables fitted offline from recorded routing\. We measure how many tokens are needed for these estimates to stabilize and how quickly they go stale during training\. We test both cases on held\-out data\.

##### A Few Thousand Tokens Suffice

To measure how much routing data the tables need, we fit them on different numbers of tokens from iterations 1983–1995 and evaluate them on held\-out iteration 1999\. With tables fitted on 1,024 tokens, correlated placement removes 23\.6 % of dispatched rows at top\-2 EP8 and 57\.9 % at top\-6\. With 4,096 tokens, the Pearson correlation between the fitted and the held\-out co\-selection matrices is 0\.93 for top\-2 and 0\.98 for top\-6, and the placement and prediction results are within one percentage point of those obtained with all 524k tokens\. A single microbatch from one rank is therefore enough to fit both tables\. The tables capture a property of the trained router and do not depend on the particular tokens used to fit them\. This implies that the correlation structure is robustly discoverable from relatively few tokens\.

##### The Structure Forms Early and Remains Stable

To measure how early the correlations form, we fit the tables on routing from a single training step and evaluate them on routing from step 1999, the end of training\. Tables fitted at step 199 already give about half of the final benefit\. At top\-2 EP8, the placement built from them removes 14\.4 % of dispatched rows, compared with 26\.5 % for tables fitted at the end of training\. At top\-6 the corresponding numbers are 42\.7 % and 57\.7 %\. Tables fitted at step 1000 come within 2–3 percentage points of the final values, and tables fitted at step 1400 or later give nearly the same results as the final tables\. The correlation between the fitted and the final co\-selection matrices follows the same pattern, rising from 0\.4–0\.5 at step 199 to above 0\.93 from step 1400 on\. The predictor improves at the same rate\. At top\-2 EP8, its hit rate is 12\.7 % for tables fitted before the routing correlations have formed, close to the chance level of 12\.5 %\. It rises to 54 % for tables fitted at step 1000 and 57 % for tables fitted at step 1800\. For our 2000\-step runs, fitting the tables once in the second half of training is therefore enough\. For longer runs, we expect refitting about every 1000 steps to be sufficient, and that each refit needs only one microbatch of recorded routes\.

### VI\-BCorrelated Expert Placement

##### All\-to\-All Communication Volume

Correlated expert placement accounts for most of the reduction in all\-to\-all traffic\. Under contiguous placement, deduplication removes only the duplicates that land on the same rank by chance\. This removes 6 % of rows at top\-2 and 25 % at top\-6 for EP8, and 3 % and 12 % at EP16\. With correlated placement, the communication volume falls by 1\.36×\\timesat top\-2 EP8 and 2\.36×\\timesat top\-6 EP8 relative to the baseline\. The reduction grows with top\-kk, since more of a token’s experts can share a rank\. The effect holds as EP widens and each rank holds fewer experts, at 1\.29×\\times/1\.22×\\times/1\.12×\\timesfor top\-2 and 1\.95×\\times/1\.54×\\times/1\.24×\\timesfor top\-6 at EP16/32/64\. Dispatch and combine carry identical byte counts by construction, since every row that goes out comes back, so the byte reduction applies identically to the backward pass also\.

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/fig_place_a2a_breakdown.png)Fig\. 7:All\-to\-All Cost Under Expert Parallelism\. A2A time \(top\) and communication volume per rank \(bottom\) with EP8–EP64, for top\-2 and top\-6 routing\. B is the baseline, D adds deduplication, and P adds correlated placement to deduplication\.
##### All\-to\-All Communication Time

Across nodes, A2A time falls by 1\.34×\\times/1\.38×\\times/1\.34×\\timesfor top\-2 and 1\.95×\\times/1\.48×\\times/1\.35×\\timesfor top\-6 at EP16/32/64\. Within a node \(EP8\), where the collective runs over xGMI, it falls by 1\.16×\\timesat top\-2 and 1\.58×\\timesat top\-6\. The reduction appears in all four collectives, dispatch and combine in both the forward and the backward pass\.

##### End\-to\-End Step Time

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/fig_place_step_breakdown.png)Fig\. 8:Step\-time breakdown at EP8–EP64 for top\-2 and top\-6 routing; dashed boxes indicate A2A share\.Non\-MoE computation takes nearly the same time in all configurations \(blue in Figure[8](https://arxiv.org/html/2610.09372#S6.F8)\)\. The step speedup therefore depends on how much A2A time is saved and on how large a share of the step the A2A takes\. At top\-2 the A2A share is 13–52 %, and correlated placement speeds up the step by 1\.05×\\timesat EP8, 1\.13×\\timesat EP16, 1\.14×\\timesat EP32 and 1\.11×\\timesat EP64\. At top\-6 the share is 24–60 %, and the speedups are 1\.14×\\timesat EP8, 1\.41×\\timesat EP16, 1\.25×\\timesat EP32 and 1\.23×\\timesat EP64\. The speedup is larger at top\-6 and is largest when the A2A crosses node boundaries\.

### VI\-CToken Shuffling

Token shuffling applies where the TP group and the EP group coincide\. We evaluate it at TP8 EP8 on one node and TP16 EP16 across two nodes\. Figure[9](https://arxiv.org/html/2610.09372#S6.F9)shows its all\-to\-all results alongside deduplication and correlated placement, and Figure[10](https://arxiv.org/html/2610.09372#S6.F10)shows the end\-to\-end step time\.

##### Token–Expert Locality

Without token shuffling, a token’s expert slots are on its own GPU with probability1/EP1/\\mathrm\{EP\}, which is 12\.5 % at EP8 and 6\.3 % at EP16\. Token shuffling raises this to 59 % and 53 % for top\-2 and to 57 % and 46 % for top\-6, close to the offline hit rates in Table[I](https://arxiv.org/html/2610.09372#S2.T1)\. The dispatcher still counts these rows, but they stay on the GPU and never cross the network\. Relative to the baseline, wire bytes fall by 2\.4×\\timesat EP8 and 2\.1×\\timesat EP16 for top\-2, and by 3\.1×\\timesand 2\.4×\\timesfor top\-6\. With placement alone, the reductions are 1\.4×\\timesand 1\.3×\\timesfor top\-2 and 2\.4×\\timesand 1\.9×\\timesfor top\-6\.

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/fig_shuffle_a2a_breakdown.png)Fig\. 9:All\-to\-All Cost on theTP=EP\\mathrm\{TP\}=\\mathrm\{EP\}Layouts\. Same encoding as Figure[7](https://arxiv.org/html/2610.09372#S6.F7), with token shuffling as the fourth configuration \(S\)\. The shuffle places tokens on the GPUs that hold their experts, so more of the dispatch stays on the device and less traffic crosses the network\.![Refer to caption](https://arxiv.org/html/2610.09372v1/images/fig_shuffle_step_breakdown.png)Fig\. 10:Step\-time breakdown on theTP=EP\\mathrm\{TP\}=\\mathrm\{EP\}layouts \(TP8 EP8, TP16 EP16\) with token shuffling added as S\. Same encoding as Figure[8](https://arxiv.org/html/2610.09372#S6.F8)\.
##### All\-to\-All Communication Time

At TP16 EP16, token shuffling raises the A2A reduction from 1\.55×\\times\(placement alone\) to 1\.74×\\timesat top\-2 and from 1\.95×\\timesto 2\.63×\\timesat top\-6\. On a single node \(TP8 EP8\), token shuffling reduces bytes by a similar or larger factor, but reduces the time cost by a smaller factor\. The A2A reductions are 1\.40×\\timesat top\-2 and 1\.46×\\timesat top\-6\. At top\-6 on one node, the A2A reduction with token shuffling is slightly below that of placement alone \(1\.48×\\times\)\. On xGMI, the collective time is set by the busiest GPU pair and does not follow the aggregate bytes\. Concentrating a token’s rows on one owner raises the traffic on the busiest link even as it lowers the total\. The two fused permutations of token shuffling add a few milliseconds per step\. The planner runs behind attention, and the main stream waits about one millisecond per step for it \(Section[V\-B](https://arxiv.org/html/2610.09372#S5.SS2)\)\.

##### End\-to\-End Step Time

On the TP = EP layouts \(Figure[10](https://arxiv.org/html/2610.09372#S6.F10)\) the tensor\-parallel collectives inflate the non\-MoE part, so the A2A is a smaller share of the step: 9–21 % at top\-2 and 18–38 % at top\-6\. Token shuffling cuts the A2A time itself by 1\.40×\\timesand 1\.74×\\timesat top\-2 and by 1\.46×\\timesand 2\.63×\\timesat top\-6 \(TP8 EP8 and TP16 EP16\), but the step speedup is limited by how much of the step the A2A takes\. At top\-2 it stays within 1\.01–1\.06×\\times; at TP16 EP16 it grows from 1\.03×\\timeswith deduplication to 1\.05×\\timeswith correlated placement and 1\.06×\\timeswith token shuffling\. At top\-6 the gains are 1\.08×\\timesat TP8 EP8 and 1\.34×\\timesat TP16 EP16, where shuffling improves the speedup from 1\.28×\\timesto 1\.34×\\times\.

## VIIConclusion

We showed that expert routing in MoE pretraining is correlated both within a layer and across layers, and that these correlations form within about one billion tokens and then stay stable\. Correlated expert placement uses the within\-layer correlation to put co\-selected experts on the same GPU, and a deduplicating dispatcher sends each token at most once to its destination GPU\. Token shuffling uses the cross\-layer correlation to move each token to the GPU predicted to hold its next experts, subject to per\-rank capacity constraints\. Moreover, this shuffle occurs inside the reduce\-scatter that sequence parallelism already performs, leading to almost no overhead\. Neither method changes routing decisions or expert parameters\. On a cluster of nodes with eight MI300X GPUs each, the two methods reduce all\-to\-all time by 1\.16–2\.63×\\timesand end\-to\-end step time by up to 1\.41×\\times, with the largest gains at top\-6 when the EP group spans several nodes\. Our evaluation covers one 12\-layer model with top\-2 and top\-6 routing, trained for 2\.1B tokens\. We believe the expert correlation holds across model families and scales, and leave other parallelism layouts, such as context parallelism, and EP inference to future work\.

## References

- Chenet al\.\(2026\)Y\. Chen, C\. Li, Y\. Wang, and M\. WeiMoonEP: a perfectly balanced expert parallelism library via dynamic redundant experts\.GitHub\.Note:[https://github\.com/MoonshotAI/MoonEP](https://github.com/MoonshotAI/MoonEP)Cited by:[§A\-B](https://arxiv.org/html/2610.09372#A1.SS2.SSS0.Px2.p2.1),[§I](https://arxiv.org/html/2610.09372#S1.p3.1)\.
- Daiet al\.\(2024\)D\. Dai, C\. Deng, C\. Zhao, R\. Xu, H\. Gao, D\. Chen, J\. Li, W\. Zeng, X\. Yu, Y\. Wu,et al\.DeepSeekMoE: towards ultimate expert specialization in mixture\-of\-experts language models\.InProceedings of the 62nd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),pp\. 1280–1297\.Cited by:[§I](https://arxiv.org/html/2610.09372#S1.p1.1)\.
- Feduset al\.\(2022\)W\. Fedus, B\. Zoph, and N\. ShazeerSwitch transformers: scaling to trillion parameter models with simple and efficient sparsity\.Journal of Machine Learning Research23\(120\),pp\. 1–39\.Cited by:[§I](https://arxiv.org/html/2610.09372#S1.p1.1)\.
- Hanet al\.\(2025\)Y\. Han, L\. Pan, J\. Peng, Z\. Tao, W\. Zhang, and Y\. ZhangGRACE\-MoE: grouping and replication with locality\-aware routing for efficient distributed MoE inference\.arXiv preprint arXiv:2509\.25041\.Cited by:[§A\-B](https://arxiv.org/html/2610.09372#A1.SS2.SSS0.Px2.p2.1),[§I](https://arxiv.org/html/2610.09372#S1.p4.1),[§I](https://arxiv.org/html/2610.09372#S1.p6.1),[§IV](https://arxiv.org/html/2610.09372#S4.p1.1)\.
- Jianget al\.\(2024\)A\. Q\. Jiang, A\. Sablayrolles, A\. Roux, A\. Mensch, B\. Savary, C\. Bamford, D\. S\. Chaplot, D\. de las Casas, E\. B\. Hanna, F\. Bressand,et al\.Mixtral of experts\.arXiv preprint arXiv:2401\.04088\.Cited by:[§I](https://arxiv.org/html/2610.09372#S1.p1.1)\.
- Kernighan and Lin \(1970\)B\. W\. Kernighan and S\. LinAn efficient heuristic procedure for partitioning graphs\.The Bell System Technical Journal49\(2\),pp\. 291–307\.Cited by:[§V\-A](https://arxiv.org/html/2610.09372#S5.SS1.SSS0.Px2.p1.1)\.
- Kimi Team \(2026\)Kimi TeamKimi K3: open frontier intelligence\.arXiv preprint arXiv:2607\.24653\.Cited by:[§I](https://arxiv.org/html/2610.09372#S1.p1.1),[§I](https://arxiv.org/html/2610.09372#S1.p3.1)\.
- Korthikantiet al\.\(2023\)V\. A\. Korthikanti, J\. Casper, S\. Lym, L\. McAfee, M\. Andersch, M\. Shoeybi, and B\. CatanzaroReducing activation recomputation in large transformer models\.Proceedings of Machine Learning and Systems5,pp\. 341–353\.Cited by:[2nd item](https://arxiv.org/html/2610.09372#S1.I1.i2.p1.1),[§III\-B](https://arxiv.org/html/2610.09372#S3.SS2.p2.1)\.
- Lepikhinet al\.\(2021\)D\. Lepikhin, H\. Lee, Y\. Xu, D\. Chen, O\. Firat, Y\. Huang, M\. Krikun, N\. Shazeer, and Z\. ChenGShard: scaling giant models with conditional computation and automatic sharding\.InInternational Conference on Learning Representations,Cited by:[§I](https://arxiv.org/html/2610.09372#S1.p1.1),[§I](https://arxiv.org/html/2610.09372#S1.p2.1)\.
- Liuet al\.\(2024\)A\. Liu, B\. Feng, B\. Xue, B\. Wang, B\. Wu, C\. Lu, C\. Zhao, C\. Deng, C\. Zhang, C\. Ruan,et al\.DeepSeek\-V3 technical report\.arXiv preprint arXiv:2412\.19437\.Cited by:[§A\-A1](https://arxiv.org/html/2610.09372#A1.SS1.SSS1.p1.1),[§A\-B](https://arxiv.org/html/2610.09372#A1.SS2.SSS0.Px1.p1.1),[§A\-B](https://arxiv.org/html/2610.09372#A1.SS2.SSS0.Px2.p2.1),[§I](https://arxiv.org/html/2610.09372#S1.p1.1),[§I](https://arxiv.org/html/2610.09372#S1.p3.1)\.
- Liuet al\.\(2025\)D\. Liu, Z\. Yan, X\. Yao, T\. Liu, V\. Korthikanti, E\. Wu, S\. Fan, G\. Deng, H\. Bai, J\. Chang,et al\.MoE parallel folding: heterogeneous parallelism mappings for efficient large\-scale MoE model training with Megatron Core\.arXiv preprint arXiv:2504\.14960\.Cited by:[§III\-B](https://arxiv.org/html/2610.09372#S3.SS2.p2.1)\.
- Luoet al\.\(2025\)S\. Luo, P\. Li, J\. Peng, Y\. Zhao, Y\. Cao, Y\. Cheng, and T\. ChenOccult: optimizing collaborative communications across experts for accelerated parallel MoE training and inference\.InProceedings of the 42nd International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.267,pp\. 41235–41253\.Cited by:[§IV](https://arxiv.org/html/2610.09372#S4.p1.1)\.
- Penedoet al\.\(2024\)G\. Penedo, H\. Kydlíček, L\. Ben Allal, A\. Lozhkov, M\. Mitchell, C\. Raffel, L\. Von Werra, and T\. WolfThe FineWeb datasets: decanting the web for the finest text data at scale\.InAdvances in Neural Information Processing Systems, Datasets and Benchmarks Track,Vol\.37,pp\. 30811–30849\.Cited by:[§VI](https://arxiv.org/html/2610.09372#S6.p1.1)\.
- Shazeeret al\.\(2017\)N\. Shazeer, A\. Mirhoseini, K\. Maziarz, A\. Davis, Q\. Le, G\. Hinton, and J\. DeanOutrageously large neural networks: the sparsely\-gated mixture\-of\-experts layer\.InInternational Conference on Learning Representations,Cited by:[§I](https://arxiv.org/html/2610.09372#S1.p1.1)\.
- Shoeybiet al\.\(2019\)M\. Shoeybi, M\. Patwary, R\. Puri, P\. LeGresley, J\. Casper, and B\. CatanzaroMegatron\-LM: training multi\-billion parameter language models using model parallelism\.arXiv preprint arXiv:1909\.08053\.Cited by:[§A\-A2](https://arxiv.org/html/2610.09372#A1.SS1.SSS2.p1.1),[1st item](https://arxiv.org/html/2610.09372#S1.I1.i1.p1.1),[§III\-B](https://arxiv.org/html/2610.09372#S3.SS2.p1.1),[§VI](https://arxiv.org/html/2610.09372#S6.p2.1)\.
- Singhet al\.\(2023\)S\. Singh, O\. Ruwase, A\. A\. Awan, S\. Rajbhandari, Y\. He, and A\. BhateleA hybrid tensor\-expert\-data parallelism approach to optimize mixture\-of\-experts training\.InProceedings of the 37th International Conference on Supercomputing,pp\. 203–214\.Cited by:[§III\-B](https://arxiv.org/html/2610.09372#S3.SS2.p2.1)\.
- Xueet al\.\(2024\)F\. Xue, Z\. Zheng, Y\. Fu, J\. Ni, Z\. Zheng, W\. Zhou, and Y\. YouOpenMoE: an early effort on open mixture\-of\-experts language models\.InInternational Conference on Machine Learning,Cited by:[§I](https://arxiv.org/html/2610.09372#S1.p1.1)\.
- Yanget al\.\(2025\)A\. Yang, A\. Li, B\. Yang, B\. Zhang, B\. Hui, B\. Zheng, B\. Yu, C\. Gao, C\. Huang, C\. Lv,et al\.Qwen3 technical report\.arXiv preprint arXiv:2505\.09388\.Cited by:[§I](https://arxiv.org/html/2610.09372#S1.p1.1)\.
- Yaoet al\.\(2024\)J\. Yao, Q\. Anthony, A\. Shafi, H\. Subramoni, and D\. K\. PandaExploiting inter\-layer expert affinity for accelerating mixture\-of\-experts model inference\.In2024 IEEE International Parallel and Distributed Processing Symposium \(IPDPS\),pp\. 915–925\.Cited by:[§I](https://arxiv.org/html/2610.09372#S1.p4.1)\.
- Zhaiet al\.\(2023\)M\. Zhai, J\. He, Z\. Ma, Z\. Zong, R\. Zhang, and J\. ZhaiSmartMoE: efficiently training sparsely\-activated models through combining offline and online parallelization\.In2023 USENIX Annual Technical Conference \(USENIX ATC 23\),pp\. 961–975\.Cited by:[§A\-B](https://arxiv.org/html/2610.09372#A1.SS2.SSS0.Px2.p2.1),[§I](https://arxiv.org/html/2610.09372#S1.p3.1)\.
- Zhaoet al\.\(2025\)C\. Zhao, S\. Zhou, L\. Zhang, C\. Deng, Z\. Xu, Y\. Liu, K\. Yu, J\. Li, and L\. ZhaoDeepEP: an efficient expert\-parallel communication library\.GitHub\.Note:[https://github\.com/deepseek\-ai/DeepEP](https://github.com/deepseek-ai/DeepEP)Cited by:[1st item](https://arxiv.org/html/2610.09372#S1.I1.i1.p1.1),[§V\-A](https://arxiv.org/html/2610.09372#S5.SS1.SSS0.Px1.p1.1)\.

## Appendix ASupplementary Material

### A\-AExperimental Setup

#### A\-A1Model Architecture

Table[II](https://arxiv.org/html/2610.09372#A1.T2)summarizes the architecture and pre\-training configuration of the model used throughout our experiments\. Unless stated otherwise, every result uses this configuration, including its router \(pre\-softmax top\-kkwithk=2k=2ork=6k=6and a sequence\-level auxiliary loss of10−210^\{\-2\}\)\. The router ablation in Section[A\-B](https://arxiv.org/html/2610.09372#A1.SS2)keeps the architecture fixed and varies only the routing recipe: the pre\-softmax router with auxiliary loss10−210^\{\-2\}and10−110^\{\-1\}, and the auxiliary\-loss\-free router of DeepSeek\-V3\[[Liu et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib1)\]\(sigmoid scores with a per\-expert bias updated at rate10−310^\{\-3\}and a10−410^\{\-4\}sequence\-level loss\)\.

SymbolValueTransformerLayersLL12, every layer an MoE layerModel widthdmodeld\_\{\\text\{model\}\}4096Attention headsnheadn\_\{\\text\{head\}\}16Head dimensiondheadd\_\{\\text\{head\}\}32Attention widthnhead​dheadn\_\{\\text\{head\}\}\\,d\_\{\\text\{head\}\}512Normalization–pre\-LayerNorm,ϵ=10−5\\epsilon=10^\{\-5\}Positional encoding–learned absoluteVocabulary\|V\|\|V\|151,669 \(Qwen3 tokenizer; padded to 151,680\)Biases / dropout–none / 0Embeddings–input and output tiedMixture of expertsExperts per layerEE128, no shared expertExperts per tokenkk2 or 6Expert hidden sizedffd\_\{\\text\{ff\}\}1024 \(dmodel→dff→dmodeld\_\{\\text\{model\}\}\\to d\_\{\\text\{ff\}\}\\to d\_\{\\text\{model\}\}, GELU\)Routing–top\-kkofEE, droplessRouter–FP32 linear gate; softmax over allEE, then top\-kk;gate weights = selected softmax probabilitiesLoad\-balancing lossαaux\\alpha\_\{\\text\{aux\}\}sequence\-level auxiliary loss,10−210^\{\-2\}Dispatch–all\-to\-all, grouped GEMMSizeTotal parametersNN13\.6 B \(13\.0 B transformer, 0\.6 B embeddings\)Expert parametersNexpN\_\{\\text\{exp\}\}12\.9 B \(1\.07 B per layer\)Active per tokenNactN\_\{\\text\{act\}\}≈0\.31\\approx 0\.31B at top\-2,≈0\.71\\approx 0\.71B at top\-6 \(non\-embedding\)Pre\-trainingData–FineWeb\-Edu, 98/2 train/validation splitSequence lengthSS8192Global batchBB128 sequences \(B⋅S≈1\.05B\\cdot S\\approx 1\.05M tokens per step\)Steps / tokensTT2000 /≈2\.1\\approx 2\.1B tokensOptimizer–AdamW,β1=0\.9\\beta\_\{1\}=0\.9,β2=0\.95\\beta\_\{2\}=0\.95, weight decay 0\.1, gradient clip 1\.0Learning rateη\\eta10−510^\{\-5\}peak, 200 warm\-up steps, cosine decay toward10−610^\{\-6\}Initializationσinit\\sigma\_\{\\text\{init\}\}0\.008 \(normal\)Precision–BF16; FP32 router and optimizer stateParallelismTP / EP1 / 8, one 8\-GPU node, micro\-batch 2TABLE II:Model and pre\-training configuration used in all experiments\.
#### A\-A2Training Framework

We use Megatron\-LM\[[Shoeybi et al\., 2019](https://arxiv.org/html/2610.09372#bib.bib7)\]as our training framework and implement our methodology \(Section[V](https://arxiv.org/html/2610.09372#S5)\) on top of it, leaving the router, expert computation, optimizer, and training loop unchanged\. Specifically, our implementation modifies three components\. \(i\)*Token dispatcher*: We add a deduplicating variant of the all\-to\-all dispatcher that sends one row per unique \(token, destination rank\) pair instead of one per \(token, expert\) pair\. Expert IDs and gate weights are carried with the activation, and the receiver expands the rows into expert\-major order\. The combine direction applies the same deduplication in reverse\.

\(ii\)*Expert placement*: We extend the dispatcher’s contiguous expert\-to\-rank assignment to support an arbitrary per\-layer permutation\. The permutation is applied to the routing map, and checkpoint expert weights are keyed by global expert ID, allowing the same checkpoint to be loaded with a different expert placement\.

\(iii\)*Token shuffling*: For TP==EP layouts, we modify attention’s sequence\-parallel reduce\-scatter to assign each token to an owner rank based on the second\-order affinity table\. Canonical token order is restored around attention, the output layer, and the loss\. The route IDs the planner needs are included in the dispatcher’s existing metadata all\-gather, requiring no additional collective\.

Routing decisions used to build the offline expert placement and affinity tables are recorded through a capture hook in the router\. All variants train the same model and produce the same loss, differing only in how tokens are placed and moved across ranks\.

### A\-BAblation Over Routing Strategies

##### Load\-Balancing Methods

Our methods depend on the router producing correlated expert choices\. To test how the load\-balancing method affects this, we trained two more top\-6 models that differ from the main top\-6 model only in the balancing method\. The first uses a sequence\-level auxiliary loss ten times stronger than the main model’s \(coefficient10−110^\{\-1\}instead of10−210^\{\-2\}\)\. It gives the most balanced expert load of the three, with a maximum\-to\-mean load ratio of 1\.20\. The second uses the auxiliary\-loss\-free method of DeepSeek\-V3\[[Liu et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib1)\], which adds a per\-expert bias to a sigmoid router\. Its load ratio is 1\.37, and it reaches the lowest loss of the three \(4\.39, against 4\.46 for the main model and 4\.64 for the stronger auxiliary loss\)\. We evaluate the same configurations on all three models\.

##### Gains Follow Each Router’s Co\-Selection Structure

Correlated placement remains effective across all three routing strategies\. With correlated placement, deduplication removes 58 %, 51 %, and 47 % of dispatched rows at EP8 for the main model, the stronger auxiliary\-loss router, and the auxiliary\-loss\-free router, respectively, and 49 %, 40 %, and 36 % at EP16\. In comparison, contiguous placement removes only about 25 % at EP8 and 13 % at EP16 for every router \(Figure[11](https://arxiv.org/html/2610.09372#A1.F11), right\)\. These reductions translate into consistent end\-to\-end gains across routing strategies: at TP1 EP16, step speedups are 1\.41×\\times, 1\.30×\\times, and 1\.26×\\times, respectively, and remain at least 1\.20×\\timeson every EP16 layout \(Figure[11](https://arxiv.org/html/2610.09372#A1.F11), left\)\. Thus, the benefit of correlated placement does not depend on a particular routing strategy; its magnitude instead follows the co\-selection structure produced by each router\.

Load balance affects the realized communication time because an all\-to\-all must wait for the most heavily loaded rank\. At EP16, for example, the main model removes 49 % of dispatched rows compared with 40 % for the stronger auxiliary\-loss router, but its speedup advantage is smaller \(1\.41×\\timesversus 1\.30×\\times\), consistent with its higher load ratio \(1\.52 versus 1\.20\)\. Correlated placement and deduplication therefore target communication volume through expert co\-selection, rather than communication skew\. Methods that directly reduce skew, such as relocating or replicating heavily loaded experts\[[Zhai et al\., 2023](https://arxiv.org/html/2610.09372#bib.bib19),[Liu et al\., 2024](https://arxiv.org/html/2610.09372#bib.bib1),[Chen et al\., 2026](https://arxiv.org/html/2610.09372#bib.bib2),[Han et al\., 2025](https://arxiv.org/html/2610.09372#bib.bib20)\], address a separate source of all\-to\-all cost and are complementary to our approach\.

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/fig_ablation_routers_dedup_ep_combo.png)Fig\. 11:Router ablation on three top\-6 models that differ only in load balancing: the main recipe \(sequence auxiliary loss10−210^\{\-2\}\), a10×10\\timesstronger auxiliary loss \(10−110^\{\-1\}\), and the DeepSeek\-V3 auxiliary\-loss\-free router\. Left: step speedup of deduplication \+ correlated placement over that model’s baseline, per layout\. Right: the structure each router leaves to exploit, as the deduplication rate under correlated placement\.

### A\-CPer\-Step Measurements

Figure[12](https://arxiv.org/html/2610.09372#A1.F12)shows the ten measured steps behind every reported speedup\.

![Refer to caption](https://arxiv.org/html/2610.09372v1/images/iter_grid_all_configs.png)Fig\. 12:Step time \(left\) and A2A volume per rank \(right\) for the ten measured steps of each configuration, after the four warm\-up steps\. Data point labels give the ratio of the baseline median to the median of that configuration\.
### A\-DAlgorithms

We give algorithms for the two components described in Sections[V\-A](https://arxiv.org/html/2610.09372#S5.SS1)and[V\-B](https://arxiv.org/html/2610.09372#S5.SS2), using the notation of Table[III](https://arxiv.org/html/2610.09372#A1.T3)\.

TABLE III:Notation\.##### Correlated placement \(Alg\.[1](https://arxiv.org/html/2610.09372#alg1)\)

The placement is computed offline, once per model\. For each MoE layer we count co\-selectionsWWon the recorded routes and partition the experts into equal\-size groups that maximize the within\-group count\.Partitionstarts from contiguous groups and applies the best improving pairwise swap until none remains orImaxI\_\{\\max\}swaps are made; since every move is a swap, each rank keeps exactlyE/PE/Pexperts\. OnN\>1N\>1nodes the partition is hierarchical: experts are first split across nodes, which targets cross\-node traffic, and then across theP/NP/Nranks of each node\.

Algorithm 1Correlated placement1:Input:routing sets

\{𝒮ℓ​\(t\)\}\\\{\\mathcal\{S\}\_\{\\ell\}\(t\)\\\}, experts

EE, EP ranks

PP, nodes

NN, budget

ImaxI\_\{\\max\}
2:Output:balanced placement

πℓ\\pi\_\{\\ell\}per MoE layer

3:foreach MoE layer

ℓ\\elldo

4:

W⁡\[e,e′\]←∑t𝟏​\[e,e′∈𝒮ℓ​\(t\)\]W\[e,e^\{\\prime\}\]\\leftarrow\\sum\_\{t\}\\mathbf\{1\}\[e,e^\{\\prime\}\\in\\mathcal\{S\}\_\{\\ell\}\(t\)\]
5:if

N=1N=1then

6:

πℓ←Partition​\(W,P,Imax\)\\pi\_\{\\ell\}\\leftarrow\\textsc\{Partition\}\(W,P,I\_\{\\max\}\)
7:else

8:

ν←Partition​\(W,N,Imax\)\\nu\\leftarrow\\textsc\{Partition\}\(W,N,I\_\{\\max\}\)
9:fornode

n=0,…,N−1n=0,\\ldots,N\-1do

10:

ℳn←\{e:ν⁡\(e\)=n\}\\mathcal\{M\}\_\{n\}\\leftarrow\\\{e:\\nu\(e\)=n\\\}
11:

ρn←Partition​\(Wℳn,PN,Imax\)\\rho\_\{n\}\\leftarrow\\textsc\{Partition\}\(W\_\{\\mathcal\{M\}\_\{n\}\},\\tfrac\{P\}\{N\},I\_\{\\max\}\)
12:

πℓ​\(e\)←n​PN\+ρn​\(e\)​∀e∈ℳn\\pi\_\{\\ell\}\(e\)\\leftarrow n\\tfrac\{P\}\{N\}\+\\rho\_\{n\}\(e\)\\ \\ \\forall e\\in\\mathcal\{M\}\_\{n\}
13:endfor

14:endif

15:endfor

16:

17:procedurePartition\(

W,G,ImaxW,G,I\_\{\\max\}\)⊳\\trianglerightbalanced

18:Initialize

GGcontiguous groups of

\|W\|/G\|W\|/Gexperts

19:for

i=1,…,Imaxi=1,\\ldots,I\_\{\\max\}do

20:

\(a,b\)←\(a,b\)\\leftarrowswap with the largest gain in intra\-group

WW
21:ifthe gain is not positivethen break

22:Swap

aaand

bb
23:endfor

24:returnexpert\-to\-group assignment

25:endprocedure

##### Token shuffling \(Alg\.[2](https://arxiv.org/html/2610.09372#alg2)\)

The prediction tables are fit offline on the same traces as the placement; the planner runs at every MoE position during training\. It first predicts, for each token, how its position\-mmexpert slots will spread over ranks \(lines 3–7\): the first\- and second\-order tables are mixed per expert pair, with more weight on the second\-order table the more trace tokens support it\. It then assigns owners greedily, most confident token first, under a per\-rank cap ofT/PT/Pso the layout stays balanced \(lines 8–11\)\. The planner runs on a side stream during attention, and the main stream waits for it only at the attention output projection, where the reduce\-scatter writes each row to its owner\.

Algorithm 2Token shuffling \(TP=EP=P\\mathrm\{TP\}=\\mathrm\{EP\}=P\)1:Offline:for each MoE position

mm, tables

FmF\_\{m\},

SmS\_\{m\}, supports

nmn\_\{m\}\(optionally

QmQ\_\{m\}\) from routing traces under placement

π\\pi; back\-off

β\\beta, cap

wmaxw\_\{\\max\}\.

2:Input:top\-

kksets

𝒮m−1​\(t\)\\mathcal\{S\}\_\{m\-1\}\(t\),

𝒮m−2​\(t\)\\mathcal\{S\}\_\{m\-2\}\(t\)of all

TTtokens, gathered with the dispatcher’s metadata all\-gather at MoE position

m−1m\{\-\}1\.

3:// The shuffle planner runs on a side stream, overlapped with the expert computation atm−1m\{\-\}1and the attention that follows; see[V\-B](https://arxiv.org/html/2610.09372#S5.SS2.SSS0.Px3)\.

4:foreach token

ttdo⊳\\trianglerightpredict owner\-rank distribution

5:

dt←1k2∑e∈𝒮m−1​\(t\)∑e′∈𝒮m−2​\(t\)\[\(1−we′​e\)Fm\[e,⋅\]\+we′​eSm\[e′,e,⋅\]\]d\_\{t\}\\leftarrow\\frac\{1\}\{k^\{2\}\}\\sum\_\{e\\in\\mathcal\{S\}\_\{m\-1\}\(t\)\}\\sum\_\{e^\{\\prime\}\\in\\mathcal\{S\}\_\{m\-2\}\(t\)\}\\big\[\(1\-w\_\{e^\{\\prime\}e\}\)\\,F\_\{m\}\[e,\\cdot\]\+w\_\{e^\{\\prime\}e\}\\,S\_\{m\}\[e^\{\\prime\},e,\\cdot\]\\big\],

6:where

we′​e=min⁡\(nm​\[e′,e\]/\(nm​\[e′,e\]\+β\),wmax\)w\_\{e^\{\\prime\}e\}=\\min\\\!\\big\(n\_\{m\}\[e^\{\\prime\},e\]/\(n\_\{m\}\[e^\{\\prime\},e\]\+\\beta\),\\,w\_\{\\max\}\\big\)
7:Optionally mix in

Qm​\[tok​\(t\),⋅\]Q\_\{m\}\[\\mathrm\{tok\}\(t\),\\cdot\]with the same back\-off

8:

dt←dt/∑rdt​\[r\]d\_\{t\}\\leftarrow d\_\{t\}/\\sum\_\{r\}d\_\{t\}\[r\]
9:endfor

10:

cap⁡\[r\]←T/P\\mathrm\{cap\}\[r\]\\leftarrow T/Pfor all

rr
11:foreach

ttin decreasing order of

dt\(1\)−dt\(2\)d\_\{t\}^\{\(1\)\}\-d\_\{t\}^\{\(2\)\}do⊳\\trianglerightmost confident first

12:

owner\(t\)←argmaxr:cap⁡\[r\]\>0dt\[r\]\\mathrm\{owner\}\(t\)\\leftarrow\\arg\\max\_\{r:\\,\\mathrm\{cap\}\[r\]\>0\}d\_\{t\}\[r\];

cap⁡\[owner⁡\(t\)\]←cap⁡\[owner⁡\(t\)\]−1\\mathrm\{cap\}\[\\mathrm\{owner\}\(t\)\]\\leftarrow\\mathrm\{cap\}\[\\mathrm\{owner\}\(t\)\]\-1
13:endfor

14:Build the new token order, its inverse, and the residual\-row indices

15:// Main stream, at the attention output projection

16:Wait for the planner

17:Reduce\-scatter the projected rows to

owner⁡\(⋅\)\\mathrm\{owner\}\(\\cdot\); permute and add the residual in one fused kernel

18:MoE position

mmdispatches from this layout, which persists until the next plan

### A\-EAuthor Contributions

##### Radha Gulhane

Implemented the methods and carried out all experiments\. Led initial ideation and exploration\.

##### Quentin Anthony

Contributed to the initial ideation and shaping of the idea, helped write the paper, and provided substantial guidance throughout\.

##### Beren Millidge

Participated in technical discussions and provided suggestions throughout the project, and helped write the paper\.