Source-linked AI summary
Zepp: Accelerating Distributed MoE Serving under Relaxed Balance Constraints
Chang Chen, Andrew Yang, Tiancheng Chen, Jiangfei Duan, Xinwei Qiang, Zhongkai Yu, Xiang Fang, Yufei Ding
TL;DR
Skewed expert workloads imbalance computation and communication, while efforts to balance one dimension can introduce other overheads. Zepp optimizes bottleneck inter-node communication subject to GPU and NIC balance constraints, coordinating placement, routing, and execution. It achieves a 1.86× geometric mean speedup over the fastest competing baseline.
Problem
Skewed expert workloads imbalance computation and communication, while balancing one dimension can introduce additional imbalance or overhead.
Source: §1 · p. 2Method
Zepp optimizes bottleneck inter-node communication under GPU and NIC balance constraints, coordinating expert placement, routing, communication restructuring, computation scheduling, and intra-node swaps for changing demand.
Source: §1 · p. 2, §4.3 · p. 8Results
1.86× geometric mean speedup over the fastest competing baseline; evaluation compares 7 baselines on 4–16 nodes, with sustained weak-scaling gains up to 32 nodes.
Source: §8 · p. 13Takeaways & Limitations
Treating GPU and NIC balance as constraints addresses communication overhead introduced by pursuing balance independently while adapting to changing demand.
Source: §8 · p. 13
Abstract
from arXiv · showhide
As Mixture-of-Experts (MoE) models continue to scale, serving them increasingly relies on expert parallelism (EP) across a growing number of devices. Yet skewed expert workloads create imbalance across computation, communication, and memory, making load balancing a central optimization objective in distributed MoE serving. We observe that balance is not free: operations introduced to balance one dimension can themselves be expensive or imbalanced. This motivates us to rethink balance as a constraint rather than an optimization objective. We present Zepp, which directly optimizes the bottleneck communication in distributed MoE serving subject to simplified balance constraints on physical resources, i.e., GPUs and NICs. Zepp progressively optimizes inter-node communication across placement, routing, and execution. It first places expert replicas to reduce token communication under GPU constraints, then reshapes communication flows through split and merge primitives under NIC constraints, and finally partitions and schedules ex- pert computation to overlap the resulting communication. To adapt to dynamic workloads, Zepp jointly coordinates computation, token communication, and expert-weight movement at each iteration. Together, these designs allow Zepp to pursue the most efficient execution rather than a single-dimension balanced one. We implement Zepp and evaluate it against 7 state-of-the-art MoE serving systems, achieving up to 6.68$\times$ MoE layer speedup and a geometric mean speedup of 1.86$\times$ over the fastest competing baseline.
1 Introduction
Skewed, time-varying expert workloads make efficient MoE serving difficult because they imbalance computation and communication across GPUs and network links. Zepp treats GPU and NIC balance as constraints while directly reducing bottleneck inter-node communication.
1 Introduction
- Zepp reframes balance as a physical-resource constraint because balancing one dimension can introduce overhead or imbalance elsewhere.The system directly minimizes bottleneck inter-node communication while constraining GPU and NIC balance.Source: §1 · p. 2
- Zepp jointly coordinates expert replication, placement, and locality-preferred routing to reduce inter-node traffic without adding communication operations.Historical demand guides replica placement across nodes, while routing favors local replicas and consolidates remaining remote requests.Source: §1 · p. 2
- Zepp uses merge to remove redundant transfers and split to expose schedulable flows and redistribute traffic, then schedules computation to overlap with communication.Merge also supports shared token transfers during dispatch and node-local pre-reduction during combine.Source: §1 · p. 2
- For changing demand, Zepp performs intra-node expert swaps that preserve node-level expert coverage and coordinates weight movement with token communication and computation.This avoids inter-node weight transfers while hiding movement costs whenever possible.Source: §1 · p. 2
- Up to 6.68× speedup over existing baselines and up to 2.78× (1.86× geometric mean) over the fastest competing baseline in end-to-end MoE layer execution.The evaluation spans multiple MoE architectures, topologies, and input token datasets.Source: §1 · p. 2
2 Background
MoE layers route each token to a dynamically selected subset of experts distributed across GPUs. The resulting traffic uses distinct intra-node and inter-node communication resources.
2 Background
- A learned router selects top-k experts per token; expert parallelism distributes experts across GPUs, requiring token dispatch and output combine.During dispatch, tokens move to GPUs hosting selected experts, and outputs return to their source for aggregation.Source: §2 · p. 3
- Intra-node communication uses high-bandwidth, low-latency fabrics such as NVLink, while inter-node transfers rely on NICs and slower scale-out networks.Peer-to-peer copies and NVSHMEM transfers use distinct data-movement resources, enabling opportunities to overlap communication with computation.Source: §2 · p. 3
- Uneven expert assignments create both uneven GPU computation and communication, and these imbalances shift across tokens and layers.In the illustrated layer, Expert 4 receives four tokens including remote ones, while Expert 3 is not activated.Source: §2 · p. 3
3 Motivation
MoE serving combines dynamic expert-load imbalance with communication planning over irregular traffic. The section explains why pursuing balance can move bottlenecks and why MoE communication requires both splitting and merging.
3 Motivation
- MoE serving must handle load imbalance from dynamic expert activation and plan communication for irregular traffic.Source: §3 · p. 3
3.1 Imbalance is Shifted, Not Eliminated
- Communication overlap can expose straggling expert computation, while replication or remapping can make expert-weight transfers the next communication bottleneck.The recurring pattern motivates optimizing the end-to-end critical path while treating resource balance as a constraint.Source: §3.1 · pp. 3–4
3.2 Split, or Merge?
- MoE routing makes splitting useful for alleviating skewed NIC traffic, while merging removes repeated shared-token transfers and reduces latency from small messages.As EP scales, individual transfers shrink and per-message latency becomes increasingly significant.Source: §3.2 · p. 4
4 System Design
Zepp progressively reduces inter-node communication through expert placement and routing, NIC-aware communication restructuring, and dependency-aware scheduling, while enforcing GPU and NIC balance as constraints. It adapts to dynamic demand through intra-node expert swaps that preserve node-level expert coverage.
4 System Design
- Zepp optimizes bottleneck inter-node communication while treating GPU compute and NIC traffic balance as resource constraints.It avoids transfers through placement and routing, restructures remaining flows with split and merge, and schedules computation around them.Source: §4 · p. 4
4.1 Expert Placement and Token Routing
- Replica placement spreads experts across nodes, while locality-preferred routing reduces remote token transfers subject to relaxed replica-load constraints.Routing preserves the model’s logical expert selections and weights while selecting physical execution locations.Source: §4.1 · p. 5, Figure 6 · p. 5
4.2 Overlapping Scheduling
- NIC-aware splitting exposes independently schedulable flows and redistributes traffic across local GPUs, moving only traffic needed to satisfy NIC constraints.The allocation minimizes added intra-node movement and keeps traffic on its original NIC whenever possible.Source: §4.2 · p. 6
- Merge shares dispatch transfers for tokens needed by multiple experts on one destination node and pre-reduces combine outputs before returning them.Compatible small flows can also be coalesced, while every expert contribution remains in the combined output.Source: §4.2 · p. 7
- Zepp overlaps communication and computation by executing dispatch groups as their inputs arrive and releasing completed combine groups for return transfer.Dispatch makes computation follow communication; combine schedules computation to make communication-ready groups early.Source: Figure 7 · p. 6, §4.2 · pp. 7–8
4.3 Three Stream Scheduling
- When dynamic demand violates GPU load constraints, intra-node swaps adjust replica locations while preserving node-level expert coverage and avoiding inter-node weight transfers.Weight movement is coordinated with token communication and computation so independent work can proceed while weights move.Source: §4.3 · pp. 8–9, Figure 8 · p. 8, Figure 9 · p. 9
5 Implementation
Zepp is implemented with C++ and CUDA, and uses grouped GEMM kernels, NVSHMEM transfers, and CUDA streams to coordinate communication and computation.
5 Implementation
- Zepp is implemented in C++ and CUDA with Python bindings and uses CUTLASS grouped GEMM kernels for expert computation.NVSHMEM provides inter-node RDMA, while CUDA streams and NVSHMEM synchronization coordinate transfers with expert computation.Source: §5 · p. 9
- Routing uses CUDA kernels and greedy waterfilling to assign tokens within per-replica load bounds, then refines assignments to consolidate remote requests.Allocation considers source-GPU replicas first, then other GPUs in the node, then remote nodes.Source: §5 · p. 9
- A greedy intra-node search swaps replicas between heavily and lightly loaded GPUs only when the exchange reduces their maximum estimated load.Routing, swap planning, and associated metadata exchange and construction account for approximately 5% of iteration latency.Source: §5 · p. 9
- NVSHMEM symmetric staging buffers hold intermediate dispatch and combine payloads and are allocated during initialization for reuse across iterations.Dispatch forwarding and combine aggregation use these buffers.Source: §5 · p. 9
6 Evaluation
Zepp is evaluated on MoE layer latency, weak scaling, end-to-end SGLang performance, component ablations, and resource timelines. Across these settings, results show gains over competing systems and illustrate how placement, routing, overlap, and expert swaps contribute under different workloads.
6.1 Experimental Setup
- The evaluation uses K2 and Qwen expert traces, BF16 workloads, and a four-GPU-per-node testbed, measuring maximum per-GPU MoE-layer latency including planning, communication, and computation.The A100 testbed uses 40 GB GPUs, with 80 GB variants for the end-to-end evaluation; H100 clusters are also used for weak scaling.Source: §6.1 · pp. 9–10, Table 2 · p. 10
6.2 MoE Performance
- At sixteen nodes, Zepp achieves 6.68× and 5.77× speedups over A2AV+GEMM for Qwen and K2 at 1 MB, and a 1.86× geometric mean speedup over the fastest competing baseline across configurations.At 16 MB, speedups over A2AV+GEMM are 4.66× for Qwen and 3.28× for K2; Figure 9 reports overall MoE-layer results.Source: §6.2 · p. 10
6.3 Weak Scalability
- Weak-scaling speedups reach 9.49× on A100 and 15.10× on H100 at 1 MB across 32 nodes, while 64 MB speedups remain 2.35–2.93× on A100 and 2.46–4.13× on H100.The evaluation holds input token volume per GPU fixed as the cluster grows from 2 to 32 nodes; Figure 10 plots K2 latency on both clusters.Source: §6.3 · pp. 10–11, Figure 10 · p. 10
6.4 End-to-End Performance
- From 1 to 16 MB per GPU, Zepp’s SGLang prefill throughput gains rise from 1.32× to 1.92×, and decode speedups rise from 1.04× to 1.57×.The comparison uses Qwen3-30B, with matched KV-cache sizes; placement uses a prompt subset and evaluation uses unseen inputs from the same dataset.Source: §6.4 · p. 11, Figure 11 · p. 11
6.5 Ablation
- With Prof. law in placement history, overlapping scheduling reaches 1.12× speedup over COMET, placement/routing raises it to 1.25×, and overlapped swaps raise it to 1.28×.Under rotating demand, the corresponding speedups are 1.05×, 1.07×, and 1.11×; swaps reduce latency by 3.0%, versus 1.9% under Prof. law-only demand.Source: §6.5 · p. 12
6.6 Case Study
- Under Prof. law demand absent from placement history, intra-node swaps overlap NVLink transfers with GEMMs and keep expert-weight traffic off inter-node links, while residual imbalance remains visible in late completion.For LiveCodeBench, placement and routing better match current demand; dispatch, expert movement, combine, and GEMMs overlap, with late-arriving contributions forming the remaining tail.Source: §6.6 · p. 12
7 Related Work
Related work spans communication optimization, computation balancing and placement, communication–computation overlap, and collective communication libraries. Zepp combines placement and locality-preferred routing with split and merge transformations and dependency-aware scheduling.
Communication optimizations
- Communication systems rebalance traffic, redirect flows, or reduce overhead, while Zepp introduces split and merge transformations to restructure token transfers and synchronization dependencies across links.FAST [22], FuseLink, and SwiftEP are cited as examples of prior communication approaches.Source: §7 · p. 12
Computation load balancing
- The described placement replicates popular experts and packs them onto GPUs according to estimated loads, while our placement spreads hot-expert replicas across nodes and couples them with locality-preferred token routing under per-GPU load constraints.Source: §7 · p. 13
Computation–communication overlap
- Prior overlap systems expose computation before communication completes; Zepp schedules expert computation around dependencies produced by split and merge.Examples include NVSHMEM, COMET [44], FlashMoE [1], and StreamEP.Source: §7 · p. 13
Collective communication libraries
- Collective libraries provide communication interfaces and execution mechanisms, while related scheduling work optimizes schedules, resource allocation, or collective work distribution.The passage names NCCL, MSCCL++ [18], DeepEP, SyCCL, HeteCCL, ForestColl, ResCCL, PReCCL, and Theseus.Source: §7 · p. 13
8 Conclusion
Zepp directly optimizes bottleneck inter-node communication while treating GPU and NIC balance as resource constraints. Evaluations report a 1.86× geometric mean speedup over the fastest competing baseline and sustained gains in weak scaling up to 32 nodes.
Conclusion
- Zepp combines expert replication, placement, routing, communication reduction, and communication–computation overlap under GPU and NIC balance constraints while adapting to changing demand.The evaluation compares Zepp against seven baselines on 4–16 nodes.Source: §8 · p. 13
- Evaluation shows a 1.86× geometric mean speedup over the fastest competing baseline, with weak-scaling gains sustained up to 32 nodes.Source: §8 · p. 13