AI summaryⓘ
The authors focus on improving how Mixture-of-Experts (MoE) models work when they run across multiple GPUs, which can slow down because experts on different GPUs need to communicate. They found that experts often work together in ways that depend a lot on the specific task, so grouping them by average patterns hides these task-specific connections. Their solution, called Task-Aware Coactivation Grouping (TACG), groups experts based on task-related patterns and assigns them to GPUs accordingly. They also add a method called Generic Expert Shared Replication (GESR) to replicate important experts on multiple GPUs to help balance the workload. Their experiments show this approach reduces communication costs by about 31% while keeping workload fairness high, even when the tasks change a lot.
Mixture-of-Experts (MoE)conditional computationGPU communicationroutingload imbalancetask-aware groupingco-activationexpert replicationinference distribution shiftJain fairness index
Authors
Zhiyao Xu, Aoxue Liu, Zhanjie Ding, Dan Zhao, Yong Jiang, Qing Li
Abstract
Sparsely activated Mixture-of-Experts (MoE) models scale capacity via conditional computation, but distributed inference suffers from cross-GPU expert communication and routing-induced load imbalance. Existing placement methods reduce this cost by co-locating frequently co-activated experts; however, they derive a single deployment plan from globally aggregated routing traces, thereby averaging away the heterogeneous, task-specific co-activation patterns that actually drive communication in multi-task serving. We observe that expert co-activation is strongly task-conditioned: pairs tightly coupled in one task family are often uncorrelated in another, so effective deployment should group experts by task-aware co-activation rather than by a task-agnostic average. Based on this insight, we propose \emph{Task-Aware Coactivation Grouping} (TACG), a deployment-time framework that uses family-specific dispatch and co-activation traces to derive per-expert task-family preferences, reweights the co-activation graph so that intra-family locality dominates grouping, and assigns each expert to a primary GPU under exact capacity constraints. To keep the static placement robust under online workload skew, we further introduce \emph{Generic Expert Shared Replication} (GESR), a lightweight companion that identifies generic experts with consistently central co-activation profiles, replicates them across a small set of secondary GPUs, and applies locality- and load-aware selection at serving time. Experiments on three representative open-source MoE models demonstrate that our framework reduces the average communication cost by 31.39\% over the baseline, while preserving an average Jain fairness index of 0.9975. This advantage persists even under severe distribution shifts in the inference data, consistently outperforming strong baselines.