MoE routing in direct topologies is limited by unpredictable traffic. An analysis of MoX shows how static routing reduces congestion and approaches switch performance.
In classic ML clusters, the network is designed for regular collective operations. This works as long as the traffic is predictable. MoE routing breaks this assumption. Each token selects top-K experts, resulting in a sparse AllToAll-V with significant variability. The load varies from batch to batch and is distributed unevenly. In direct-connect topologies, this quickly leads to degradation: “expensive” intermediate hops (bandwidth tax) and overloaded links emerge. Importantly, execution time is determined by the most loaded link, not the average, so even moderate skew becomes a systemic issue.
MoX offers a pragmatic shift: instead of adapting the topology to traffic, it pre-optimizes MoE routing for the class of load. Instead of dynamic network reconfiguration or predicting the traffic matrix, a static policy is used. This relies on two ideas. The first is token-aware multicast: one token is spread across a tree rather than sent independently to each expert. This reduces transmission duplication. The second is balancing through precomputed per-link weights that minimize the expected load on the most congested channel. This is a trade-off: the system ignores specific runtime skew but gains stability and simplicity.
The implementation of MoE routing in MoX is built around a multicast tree for each token. The source gradually expands the set of nodes that have already received embeddings and selects paths to new destinations. Priority is given to routes where intermediate nodes are themselves recipients. This turns relay hops into “useful” deliveries. If such a path does not exist, a regular relay is used with the inevitable bandwidth tax. For balancing between equivalent paths, a weighted round robin is applied, where the probability of selection is determined by static link weights. These weights are computed offline through iterative optimization, approximating the multicast tree packing problem. Meanwhile, the runtime state remains compact: one value per link, without storing all possible routes.
A separate point is the combine phase. MoX uses a reverse traversal of the same tree and performs partial aggregation at intermediate nodes. This reduces reverse traffic since not all expert results go independently to the source. This approach offloads some work to the GPU but reduces pressure on the network.
Results show that this MoE routing yields a systemic effect. In simulations, MoX accelerates the entire MoE block by up to 1.8× compared to min-hop routing. Performance is close to that of an ideal packet switch: the deviation is in the single digits across several configurations. Importantly, this effect is achieved without dynamic network adaptation. Moreover, the static strategy outperforms a topology optimized for specific traffic if it uses naive routing. This indicates that the choice of routing policy may be more important than the physical topology itself.
In practice, a reduction in hotspot effects is also observed. In the Boardfly-level topology, the bottleneck load decreases by up to 47%. This is a direct consequence of traffic leveling across links rather than reducing the number of hops. Interestingly, the main improvement comes from the multicast structure, while the weights merely refine the balancing.
The main conclusion: MoE routing can be made predictable without traffic control. A static, load-oblivious strategy works if it considers the semantics of MoE (multicast and reduction) and optimizes for the worst link. This aligns well with the trend of simplifying networks: less dynamics in infrastructure, more intelligence in routing algorithms.
Information source
arXiv is the largest open preprint repository (since 1991, under the auspices of Cornell), where researchers quickly post working versions of papers; the materials are publicly accessible but do not undergo full peer review, so results should be considered preliminary and, where possible, checked against updated versions or peer‑reviewed journals. arxiv.org