PaperlayerView exact revision on arXiv

Source-linked AI summary

MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism

Ruidong Zhu, Ziheng Jiang, Chao Jin, Peng Wu, Cesar A. Stuardo, Dongyang Wang, Xinlei Zhang, Huaping Zhou, Haoran Wei, Yang Cheng, Jianzhe Xiao, Xinyi Zhang, Lingjun Liu, Haibin Lin, Li-Wen Chang, Jianxi Ye, Xiao Yu, Xuanzhe Liu, Xin Jin, Xin Liu

arXiv:2504.02263v4cs.DCcs.LG

TL;DR

MoE sparsity can leave FFNs underutilized at practical batch sizes, limiting the serving-cost benefits of reduced computational complexity. MegaScale-Infer separates attention and FFNs and combines tailored parallelism, pipelining, and M2N communication. It reports up to 1.9× higher per-GPU decoding throughput than state-of-the-art serving systems.

  • Problem

    In MoE decoding, sparsity can leave FFNs underutilized because practical batch sizes assign fewer tokens to each expert.

    Source: §1 · p. 2

  • Method

    MegaScale-Infer disaggregates attention and FFNs, using tailored parallelism, ping-pong micro-batch pipelining, and an M2N communication library.

    Source: §1 · pp. 2–3

  • Results

    MegaScale-Infer achieves up to 1.9× higher per-GPU decoding throughput than state-of-the-art LLM serving systems.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Evaluation across diverse models and hardware reports up to 1.9× throughput improvement over state-of-the-art systems.

    Source: §10 · p. 17

    The system assumes the model uses grouped-query attention (GQA).

    Source: §4 · p. 6

Abstract

from arXiv · show

Mixture-of-Experts (MoE) showcases tremendous potential to scale large language models (LLMs) with enhanced performance and reduced computational complexity. However, its sparsely activated architecture shifts feed-forward networks (FFNs) from being compute-intensive to memory-intensive during inference, leading to substantially lower GPU utilization and increased operational costs. We present MegaScale-Infer, an efficient and cost-effective system for serving large-scale MoE models. MegaScale-Infer disaggregates attention and FFN modules within each model layer, enabling independent scaling, tailored parallelism strategies, and heterogeneous deployment for both modules. To fully exploit disaggregation in the presence of MoE's sparsity, MegaScale-Infer introduces ping-pong pipeline parallelism, which partitions a request batch into micro-batches and shuttles them between attention and FFNs for inference. Combined with distinct model parallelism for each module, MegaScale-Infer effectively hides communication overhead and maximizes GPU utilization. To adapt to disaggregated attention and FFN modules and minimize data transmission overhead (e.g., token dispatch), MegaScale-Infer provides a high-performance M2N communication library that eliminates unnecessary GPU-to-CPU data copies, group initialization overhead, and GPU synchronization. Experimental results indicate that MegaScale-Infer achieves up to 1.90x higher per-GPU throughput than state-of-the-art solutions.

1 Introduction

MoE sparsity can leave FFNs underutilized at practical batch sizes, so reduced computational complexity does not necessarily lower serving costs. MegaScale-Infer separates attention and FFNs and combines tailored parallelism, pipelining, and M2N communication; experiments report higher throughput and lower serving costs.

1 Introduction

  • MoE sparsity assigns fewer tokens to each expert within a batch, lowering FFN GPU utilization despite the batch-size limit imposed by memory and latency constraints.The passage reports that each expert may receive less than a quarter, or even an order of magnitude fewer, tokens than in a dense model.Source: §1 · p. 2
  • MegaScale-Infer disaggregates attention and expert modules onto separate GPUs, enabling independent scaling, tailored parallelism, and heterogeneous hardware deployment.Attention uses data parallelism while FFNs use expert parallelism; consolidating requests from attention replicas increases the batch size per expert.Source: §1 · p. 2
  • Ping-pong pipeline parallelism splits request batches into micro-batches to keep attention and FFNs busy and hide communication overhead.Its effectiveness depends on conditions such as similar attention and FFN computation times.Source: §1 · pp. 2–3
  • Arbitrary attention and FFN parallelism turns token-routing All2All communication into M2N communication, for which MegaScale-Infer develops a library focused on lower overhead and greater stability.M and N denote the numbers of senders and receivers, respectively.Source: §1 · p. 3
  • Experiments report up to 1.9× higher per-GPU decoding throughput than state-of-the-art serving systems and 1.7× higher throughput per unit cost on a heterogeneous cluster.The system was evaluated with MoE models ranging from 132 to 317 billion parameters.Source: §1 · p. 3
  • 1.90× and 1.86× per-cost decoding throughput against state-of-the-art systems on homogeneous and heterogeneous clusters, respectively.The evaluation also reports up to 1.9× per-GPU decoding throughput, 4.2× higher M2N throughput and 68.2% lower M2N latency than NCCL, and a 1.5–2.0× serving-cost reduction.Source: §1 · p. 3

2 Background and Motivation

LLM decoding has memory-intensive attention and batched FFNs, but MoE sparsity divides tokens among experts and reduces FFN GPU utilization. Disaggregating attention and FFNs offers independent scaling and heterogeneous deployment, while requiring a pipeline and specialized communication.

2.1 LLM Inference Characteristics

  • During decoding, attention accesses request-specific KV caches, whereas FFNs can reuse weights across requests, making batching more efficient for FFNs.Attention is memory-intensive during decoding; the paper notes that batching does not change its memory-access intensity.Source: §2.1 · p. 3

2.2 LLM Serving at Scale

  • MoE layers route tokens to top-k experts and combine their outputs, enabling model-size scaling without linearly increasing computation.Expert parallelism distributes experts across devices but requires token all-to-all communication and can face load imbalance.Source: §2.2 · p. 4

2.3 Problems in Large-scale MoE Serving

  • MoE sparsity sends only a fraction of a request batch to each expert, reducing FFN batch size and GPU utilization.Batch size is also constrained by latency requirements and KV-cache memory, limiting a straightforward increase in batch size.Source: §2.3 · pp. 4–5

2.4 Opportunities and Challenges

  • Disaggregating attention and FFNs enables independent scaling and cost-effective hardware choices, but creates idle periods and requires M2N/N2M communication.A ping-pong pipeline coordinated with each module’s model parallelism is proposed to address idling; existing peer-to-peer primitives can degrade performance for the M2N pattern.Source: §2.4 · p. 5

3 MegaScale-Infer Overview

MegaScale-Infer focuses on decoding and places attention and FFN modules on separate nodes. It combines tailored parallelism, a ping-pong pipeline, and customized M2N communication for the disaggregated setup.

3 MegaScale-Infer Overview

  • MegaScale-Infer decouples attention and FFN modules onto separate nodes to support independent scaling and heterogeneous deployment during decoding.The system decouples prefill and decoding into separate clusters, while this paper focuses on decoding.Source: §3 · p. 5
  • Attention parameters are replicated across attention nodes, while tensor parallelism is used within attention and expert nodes.Expert nodes form an expert-parallel group, with each expert node typically storing one expert’s parameters.Source: §3 · pp. 5–6
  • A ping-pong pipeline feeds request micro-batches through attention and expert nodes to keep them busy during communication or while awaiting results.The system selects the deployment plan using a performance model designed for disaggregated expert parallelism.Source: §3 · p. 6
  • A customized M2N communication library transfers intermediate outputs while removing unnecessary GPU-to-CPU copies, group initialization overhead, and GPU synchronization.The library also includes traffic-oriented optimizations for this communication scenario.Source: §3 · p. 6

4 Disaggregated Expert Parallelism

MegaScale-Infer chooses a deployment plan for the model, workload, hardware, and latency requirements to maximize throughput per unit cost. Its design uses ping-pong micro-batches, a performance-model-based search, and hardware choices tailored to attention and expert workloads.

4 Disaggregated Expert Parallelism

  • The deployment plan specifies module parallelism, pipeline micro-batches, maximum batch size, and hardware to maximize throughput per unit cost.The plan is selected for a given model, workload characteristics, hardware availability, and latency requirements; the paper assumes grouped-query attention.Source: §4 · p. 6

4.1 Ping-Pong Pipeline Parallelism

  • Splitting a request batch into micro-batches lets attention and expert nodes overlap forward computation with communication in a ping-pong pipeline.At least 3 micro-batches are required for fast communication, and at least 4 for relatively slower communication.Source: §4.1 · pp. 6–7

4.2 Deployment Plan Search

  • The deployment search balances attention and expert computation, then compares feasible micro-batch plans using simulated throughput per unit cost under the SLO.Its performance model estimates attention, expert, and communication times; the search also accounts for memory limits and uses binary search to find the maximum global batch size that meets the SLO.Source: §4.2 · pp. 7–8

4.3 Heterogeneous Deployment

  • Heterogeneous deployment assigns memory-bandwidth- and capacity-oriented GPUs to attention and compute-cost-effective GPUs to experts.The paper identifies H20 as more suitable for attention and L40S as more cost-effective for experts, and notes heterogeneous deployment can reduce energy consumption.Source: §4.3 · p. 8

5 High-Performance M2N Communication

MoE token routing turns attention-to-FFN communication into M2N transfers, where existing NCCL operations add latency and become unstable at higher percentiles. MegaScale-Infer’s M2N library removes unnecessary copies and coordination overhead, with traffic-specific tuning for stability.

5 High-Performance M2N Communication

  • NCCL’s latency exceeds the perftest baseline and its P99 latency surges particularly at 32 or more receivers.The comparison uses a single sender transmitting 128K bytes to each receiver.Source: §5 · p. 8
  • MegaScale-Infer’s M2N library targets unnecessary GPU-to-CPU copies, group-operation overhead, and GPU synchronization and memory accesses.The sender uses RDMA writes with immediate and polls completion queues; the receiver polls completions and uses GDRCopy with a flush operation.Source: §5 · pp. 9–10
  • High-priority ACK queues address sender bottlenecks from delayed acknowledgements, while congestion-control tuning targets latency degradation in unbalanced transfers.The ACK and congestion-control settings are fine-tuned empirically to isolate ACK packets and allow faster rate convergence.Source: §5 · p. 10
  • The design uses CPU-to-CPU inter-node communication rather than DeepEP’s GPU-to-GPU approach, avoiding GPU compute allocation but potentially sacrificing throughput for very small packets.The paper says a single-threaded CPU is sufficient to saturate bandwidth at its typical several-hundred-kilobyte per-connection transfer volume; smaller volumes may favor GPU parallelism.Source: §5 · p. 10

6 Implementation

The implementation combines fused kernels with a traffic-informed expert placement plan. Kernel fusion reduces launch and memory-access overhead, while the placement plan uses expert popularity to balance computation across nodes.

6 Implementation

  • Fusing tensor-parallel communication with adjacent GEMMs and combining sequential memory-intensive operations reduces kernel-launch and memory-access overhead.The example combines gating with top-k expert selection, token-count and weight computation, and token scattering.Source: §6 · p. 10
  • The expert deployment plan distributes M experts across N nodes to minimize the maximum node computation cost, using a greedy approximation based on traffic from a previous time period.The cost model accounts for each expert’s allocation fraction and active-token computation cost, with a floor for cold experts.Source: §6 · pp. 10–11

7 Evaluation

MegaScale-Infer is evaluated against vLLM and TensorRT-LLM across MoE models, GPU configurations, M2N microbenchmarks, and ablations. Results show higher decoding throughput, lower M2N latency, and gains from disaggregation and pipeline tuning, with performance depending on deployment balance.

7.1 Experimental Setup

  • The evaluation compares vLLM and TensorRT-LLM across Ampere and heterogeneous H20/L40S clusters, using production workload traces and a 150 ms time-between-tokens requirement.Models include Mixtral 8x22B, DBRX, and Scaled-MoE; decoding measurements exclude prefill.Source: §7.1 · pp. 11–12

7.2 End-to-end Experiment

  • 1.90× higher per-GPU decoding throughput than TensorRT-LLM is reported for Scaled-MoE, while end-to-end throughput improves by up to 1.18× in homogeneous deployment.The paper reports comparable time-between-tokens latency to baselines; Figure 9 also presents heterogeneous deployment performance.Source: §7.2 · pp. 12–13, Figure 9 · p. 13
  • 3.24× and 1.86× higher unit-cost decoding throughput than vLLM and TensorRT-LLM on H20, respectively, is reported for heterogeneous deployment.MegaScale-Infer assigns H20 GPUs to attention and L40S GPUs to experts; Figure 9 presents decoding throughput per unit cost.Source: §7.2 · p. 13, Figure 9 · p. 13

7.3 Performance of M2N Communication

  • At 256 KB, M2N communication improves throughput 4.2× and reduces median and tail latency by 68.2% and 92.9% versus NCCL.Across varying sender and receiver counts, the library reduces tail latency by 54.7%–96.9% and improves throughput by 3.3×–5.8×.Source: §7.3 · p. 14, Figure 11 · p. 14, Figure 12 · p. 14

7.4 Ablation Study

  • Disaggregation with NCCL improves throughput up to 4.66× over the colocated baseline, and optimized M2N communication adds up to 1.53×.The reported gains are on Ampere GPUs; optimization brings single-micro-batch communication time below computation time, allowing overlap in the ping-pong pipeline.Source: §7.4 · p. 14, Figure 13 · p. 15
  • Increasing micro-batches from 1 to 2 improves throughput by 1.9×; increasing from 2 to 3 adds 1.10×, 1.28×, and 1.38× for Mixtral 8x22B, DBRX, and Scaled-MoE.The gains from 2 to 3 micro-batches are greater for larger models; further increases yield only marginal improvement in the high-bandwidth testbed.Source: §7.4 · p. 15, Figure 14 · p. 15
  • With DBRX, normalized decoding throughput peaks at DP degree 8, where attention and FFN computation times are roughly equal; larger degrees shift the bottleneck to experts.Latency stays constant from DP 1 to 4 while normalized throughput scales linearly; Figure 15 shows latency and throughput across DP degrees.Source: §7.4 · p. 15, Figure 15 · p. 15

8 Deployment Experience

Production observations reveal both expert-side and attention-side load imbalance. Expert loads are steadier during decoding than prefill, while attention-node sequence-length variation can create synchronization bubbles.

8 Deployment Experience

  • On a production cluster of nearly 10,000 GPUs, heterogeneous deployment reduces the cost of serving the same traffic by 1.5–2.0×, depending on workload characteristics.Source: §8 · p. 15
  • Expert load is substantially imbalanced within a batch, but remains relatively stable across decoding batches and fluctuates more during prefill.These observations motivate static or periodic expert balancing during decoding and more frequent expert-plan adjustments during prefill.Source: §8 · pp. 15–16, Figure 16 · p. 16
  • Sequence-length variation can leave attention nodes with different computation times, creating synchronization bubbles; batches are composed to match a target execution time.The system profiles key operators across sequence lengths and batch sizes to estimate execution time.Source: §8 · p. 16

9 Related Work

Prior LLM-serving work improves scheduling, cache management, batching, and load balancing, but largely targets dense models. MegaScale-Infer addresses MoE sparsity by disaggregating attention and FFN modules and identifying M2N communication overheads.

LLM serving

  • Existing serving systems improve scheduling, KV-cache management, batching, and request migration, but their approaches primarily focus on dense models and often overlook large-scale MoE sparsity.Examples include Orca, vLLM [51], Sarathi-Serve [23], LoongServe [75], LLM-nix [68], ServerlessLLM [37], and dLoRA.Source: §9 · p. 16

Resource disaggregation

  • MegaScale-Infer disaggregates attention and FFN modules during decoding, while prior systems offload attention to cheaper devices but leave MoE sparsity challenges unresolved.Prior systems include FASTDECODE [42], Lamina [31], and MoE-Lightning [28].Source: §9 · p. 16

MoE optimization

  • MoE optimization work addresses load balancing, communication, and resource offloading; this work focuses on sparsity-related inefficiencies in large-scale distributed serving.The cited optimization concerns include load balancing [53], communication [46] [53] [62], and offloading and preloading [28] [47].Source: §9 · p. 16

Collective communication for ML

  • Disaggregating attention and FFN modules requires M2N collective communication, and MegaScale-Infer targets overhead and instability in existing solutions.Prior communication work includes NCCL, SCCL [27], TACCL, TE-CCL, CoCoNET [48], and Centauri.Source: §9 · pp. 16–17

10 Conclusion

MegaScale-Infer disaggregates attention and FFN modules and uses ping-pong parallelism with an M2N communication library. Across diverse models and hardware, it achieves up to 1.9× throughput improvement over state-of-the-art systems.

10 Conclusion

  • Up to 1.9× throughput improvement over state-of-the-art systems across diverse models and hardware is reported for MegaScale-Infer.The system combines attention–FFN disaggregation, ping-pong parallelism, and a high-performance M2N communication library.Source: §10 · p. 17

Related papers

Loading 2504.02263v4…