MoX: Efficient MoE Routing on Direct-Connect Topologies

2026-07-22

AI Summary

This work addresses the challenge of efficiently mapping sparse, runtime-dependent Mixture-of-Experts (MoE) communication onto static direct-connect optical switching networks, which often leads to bandwidth underutilization and scheduling bottlenecks. To overcome this, the authors propose a load-oblivious yet token-aware static multicast routing scheme that formulates traffic balancing as a constrained multicast tree packing problem, eliminating the need for dynamic topology reconfiguration or real-time traffic information. By leveraging token-level multicast trees, statically precomputed link weights, and a tailored packing algorithm, the approach enables efficient simulation in ASTRA-sim. Evaluated on the Boardfly topology, the method reduces bottleneck link load by up to 47% and achieves up to 1.8× speedup in MoE training and inference, closely approaching the ideal performance of packet-switched networks.

Abstract

Optically switched networks suit the regular communication of dense ML models, but MoE introduces sparse, runtime-dependent traffic. We show that efficient offline-optimized routing enables efficient MoE training and inference on direct-connect topologies without the need for MoE traffic matrix or dynamic topology reconfiguration. MoX constructs token-aware multicast trees to reduce bandwidth tax, then uses static, precomputed link weights to balance traffic by solving a restricted multicast tree-packing problem. Using recorded traffic from large MoE models, token-level traces, and ASTRA-sim, we find that MoX accelerates the full MoE block -- dispatch, expert computation, and combine -- by up to 1.8x over min-hop routing. Moreover, it attains nearly ideal packet-switched network performance in random expander topologies. On a 1,024-TPU model of Google's Boardfly topology, MoX reduces the dispatch bottleneck link load by up to 47%. These results show that high-performance MoE on static direct-connect fabrics can be achieved via optimized load-oblivious routing without demand-driven reconfiguration.

Problem

• Mixture of Experts
• direct-connect topologies
• efficient routing
• sparse traffic
• static networks

Innovation

• MoE routing
• direct-connect topology
• multicast tree packing
• static load balancing
• optical switching