Cooling-Aware Routing in Mixture-of-Experts Models: Optimizing Expert Assignment Under Routing-Deviation Constraints

Abstract

Mixture-of-Experts (MoE) models reduce computational demand by only activating a subset of available experts for each input token, but uneven expert selection often produces imbalanced workloads when experts are distributed across physical accelerators. This study investigates whether token-level routing flexibility can be used to reduce modeled direct-to-chip liquid-cooling demand while limiting departure from the model’s original router preferences. Routing data were collected from the Switch-Base-8 using five separately seeded 1,000-document samples from the C4 validation dataset, including 978,748 non-padding token positions and 5,872,488 token-layer routing decisions. Expert assignments were mapped only to an 8-GPU configuration and connected to a reduced-order thermal and hydraulic model. Next, a constrained optimization was used to redistribute selected token assignments while minimizing peak GPU workload and routing deviation. Across the five samples, baseline maximum GPU workload was approximately 21% above the equal-share reference. At a target corresponding to 75% of the maximum achievable workload balancing, an average on just 3.16% of token-layer assignments were rerouted, with the mean routing deviation coming to 0.009857, while the maximum GPU workload decreased by 12.96%. The result was reproducible across separately seeded samples and stayed present under alternative expert-to-GPU placements. Under modeled hydraulic assumptions, the corresponding mean pump-power reduction ranged from 12.96% to 24.18%. These results indicate that token-level MoE routing flexibility may provide a useful mechanism for reducing modeled peak routing-dependent cooling demand, even though hardware validation and direct evaluation of model-quantity effects remain necessary.