MegaScale-Infer M2N Library

• MegaScale-Infer M2N Library is a specialized GPU–RDMA runtime designed for M-to-N token routing in MoE architectures with disaggregated attention and expert modules.
• It replaces traditional all-to-all NCCL collectives with a custom sender-receiver design, significantly reducing overhead and tail latency.
• Integration with ping-pong pipeline parallelism and optimized deployment strategies boosts inference throughput and cost efficiency in large-scale MoE models.

MegaScale-Infer M2N Library is the specialized communication layer of MegaScale-Infer, a serving system for large Mixture-of-Experts (MoE) LLMs that implements disaggregated expert parallelism by placing attention and FFN expert modules on different GPU nodes and scaling them independently. In this setting, token routing no longer matches the fixed, symmetric all-to-all regime assumed by conventional MoE runtimes. The M2N library instead provides an $M$-senders-to-$N$-receivers communication runtime for repeated token dispatch and aggregation between attention GPUs and expert GPUs, using a specialized GPU–RDMA design rather than standard cross-node NCCL collectives (Zhu et al., 3 Apr 2025).

1. Architectural setting and motivation

MegaScale-Infer targets Transformer MoE inference in which each layer combines an attention module and an MoE FFN module. During decoding, the attention module reads KV caches and is memory-intensive, whereas the MoE FFN becomes sparsely activated through top-$K$ gating. The paper formalizes the utilization shift from dense FFNs to MoE FFNs as

$util = \min\left(\frac{B}{F} b, 1\right)$

for a dense FFN, and

$$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$

for an MoE FFN. For the same global batch $b$, FFN utilization is therefore scaled down by $\frac{topk}{\#expert}$, making expert computation memory-bound and under-utilized (Zhu et al., 3 Apr 2025).

MegaScale-Infer addresses this by disaggregating attention and expert FFN modules across different GPU nodes. Attention nodes hold all attention parameters and the KV cache, use data parallelism, and remain memory-bound. Expert nodes each hold exactly one expert, optionally tensor-parallel across GPUs within a node, and collectively form an expert parallel group. This arrangement aggregates tokens from many attention replicas onto the expert side, turning experts into compute-bound units.

This architectural separation creates the specific problem that motivates M2N. Traditional MoE implementations route tokens through all-to-all within a fixed expert-parallel group. Disaggregation breaks that symmetry: the number of attention GPUs $M$ and expert GPUs $N$ differs, the mapping is chosen independently, and deployment may be heterogeneous in GPU type and network bandwidth. The M2N library exists to make this communication regime practical.

2. Communication semantics and runtime role

Within MegaScale-Infer, M2N is a specialized GPU–RDMA communication runtime implemented as a PyTorch extension of approximately 4.9K LOC in C/C++ and 5K LOC in Python. It sits below the MegaScale-Infer runtime and ping-pong scheduler, above low-level RDMA and GPUDirect primitives, and replaces NCCL for cross-node token movement while coexisting with NCCL for intra-node tensor-parallel collectives over NVLink or PCIe (Zhu et al., 3 Apr 2025).

The library serves a repeated two-phase routing pattern. After gating on the attention side, tokens, weights, and indices are packed into per-expert buffers and dispatched from attention GPUs to expert GPUs. After expert FFN computation, outputs are returned to the originating attention GPUs and reassembled in token order. MegaScale-Infer describes these as “M2N and N2M communication between $M$ attention GPUs and $N$0 expert GPUs, replacing the traditional All-to-All communication used in each MoE layer.”

For each MoE layer and each micro-batch, the sequence is fixed. Attention nodes run attention and gating, then invoke an M2N sender to transmit routed tensors. Expert nodes use an M2N receiver to await incoming buffers from all attention GPUs, execute FFN GEMMs on the received tokens, and then use M2N again to send outputs back. The runtime’s ping-pong scheduler interleaves these communications with computation across micro-batches.

The topology is correspondingly split. Within a node, GPUs form tensor-parallel groups of size $N$1 on attention nodes and $N$2 on expert nodes, with NCCL used for TP collectives. Across nodes, attention nodes form an attention cluster under data parallelism, expert nodes form an expert cluster under expert parallelism, and M2N performs peer-to-peer RDMA writes over InfiniBand or RoCE NICs. The paper characterizes this pattern as akin to batched, dynamic point-to-point writes tuned for MoE routing, where per-pair message sizes are typically hundreds of kilobytes per micro-batch and destination sets vary with top-$N$3 gating.

3. Sender–receiver design and core optimizations

The paper identifies two empirical deficiencies when NCCL is applied to this pattern: excessive overhead and instability at high-percentile latency. NCCL introduces GPU–CPU proxy copies, batches send/recv operations in groups of at most eight, and incurs general group setup and internal checking overhead. It also exhibits high variance due to GPU synchronization and device-memory accesses, with instability worsening as the number of receivers $N$4 grows (Zhu et al., 3 Apr 2025).

M2N removes these costs through a custom sender–receiver architecture. On the send path, GPU memory is registered as an RDMA buffer with the NIC, and a custom Core Sender uses RDMA write with immediate plus GPUDirect so that data flows directly from GPU memory to the NIC and into remote GPU memory or pinned host memory without device–host copies. A low-utilization GPU kernel on a CUDA stream enforces ordering: once preceding attention, gating, and packing kernels have completed, the kernel marks readiness and the CPU side issues RDMA writes from the GPU buffer.

On the receive path, incoming writes target pre-registered RDMA buffers. A CPU daemon polls the RDMA completion queue to detect arrival. Data may then be moved into final user tensors by a custom GPU copy kernel or GDRCopy, and this movement can be overlapped with computation. Completion tracking remains on the CPU through RDMA completion queues; the design avoids GPU polling on completion and minimizes GPU-side coordination after readiness has been established.

The library further applies traffic-oriented optimizations. High-priority ACK queues separate acknowledgements from data packets so that control traffic is not delayed behind bulk data under ping-pong bidirectional communication. Congestion control is refined for imbalanced traffic patterns so that rate limiting converges faster and latency remains stable when $N$5 and $N$6 are large. Together, these changes remove the group initialization, GPU synchronization, and proxy-copy behaviors that the paper identifies as the major sources of tail-latency inflation.

4. Interaction with ping-pong pipeline parallelism and formal models

Disaggregation is exploited through ping-pong pipeline parallelism. A global batch of size $N$7 is split into $N$8 micro-batches, and at each MoE layer tokens flow Attention $N$9 Experts $K$0 Attention. Multiple micro-batches remain in flight so that expert computation on micro-batch $K$1, attention computation on $K$2, and communication on $K$3 can overlap. The pipeline is designed to hide communication time $K$4 behind attention compute $K$5 and expert compute $K$6 (Zhu et al., 3 Apr 2025).

The paper states the overlap conditions as

$K$7

which imply a minimum micro-batch count

$K$8

Reducing $K$9 through M2N therefore permits smaller $util = \min\left(\frac{B}{F} b, 1\right)$0, typically $util = \min\left(\frac{B}{F} b, 1\right)$1–$util = \min\left(\frac{B}{F} b, 1\right)$2, which avoids excessively small expert GEMMs and preserves expert MFU.

The iteration-time model makes the same dependence explicit. For $util = \min\left(\frac{B}{F} b, 1\right)$3 MoE layers, decoding iteration latency for one micro-batch satisfies

$util = \min\left(\frac{B}{F} b, 1\right)$4

and the total iteration latency for the global batch is

$util = \min\left(\frac{B}{F} b, 1\right)$5

Communication time per micro-batch is modeled as

$util = \min\left(\frac{B}{F} b, 1\right)$6

where $util = \min\left(\frac{B}{F} b, 1\right)$7 and $util = \min\left(\frac{B}{F} b, 1\right)$8 are micro-batch sizes on attention and expert nodes, $util = \min\left(\frac{B}{F} b, 1\right)$9 is hidden size, $$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$0 is the top-$$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$1 expert count, $$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$2 and $$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$3 are tensor-parallel sizes, $$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$4 and $$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$5 are available network bandwidths per GPU, and $$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$6 is an empirical bandwidth-utilization curve obtained by profiling M2N.

The deployment planner uses these quantities together with profiled linear models

$$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$7

and chooses

$$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$8

to satisfy $$util = \min\left(\frac{topk}{\#expert} \frac{B}{F} b, 1\right)$$9. In this formulation, M2N is not merely a transport primitive; it is an explicit term in the deployment search that determines feasible overlap and throughput-per-dollar operating points.

5. Implementation details, data layout, and integration boundaries

The paper does not provide a fully documented public API, but it describes several implementation-level interfaces. As a PyTorch extension, M2N is invoked through custom op bindings from the attention-side and expert-side MoE implementations. Attention modules pass GPU tensors and routing metadata after gating; expert modules receive packed buffers, run FFN GEMMs, and send outputs back through the same runtime (Zhu et al., 3 Apr 2025).

On attention nodes, the MoE layer computes top-$b$0 experts and weights for each token, determines token counts per expert node, and rearranges tokens into per-expert packed buffers for RDMA transfer. The paper states that these steps are fused with gating so as to reduce both kernel launch overhead and memory access. This fused kernel produces the packed token buffers and routing metadata in the layout expected by M2N, minimizing extra copies.

Memory management relies on pre-registered RDMA buffers that are registered once and reused across iterations and micro-batches. This avoids repeated registration overhead and dynamic allocation. On the receive side, GDRCopy may be used when copying from RDMA buffers into final user tensors. Tensor parallelism remains orthogonal to M2N and is handled intra-node through fused NCCL-plus-GEMM kernels, while expert parallelism is expressed as one expert per node with TP inside the node.

The same design supports heterogeneous deployment. The deployment section assigns memory-rich, high-bandwidth but lower-TFLOP GPUs such as H20 to attention, and computation-rich, cost-efficient GPUs such as L40S to experts. M2N is the communication mechanism that connects these two clusters via GPUDirect RDMA without host-memory staging. A plausible implication is that M2N’s abstraction is as much about decoupled hardware specialization as about raw transport efficiency.

6. Evaluation, comparisons, and scope

The paper includes dedicated M2N micro-benchmarks. In the benchmark setup, each sender and receiver is a GPU with a 200 Gbps NIC, and message size is the number of bytes sent from one sender to one receiver. For a representative 256 KB message with fixed $b$1, M2N reduces median latency by 68.2%, reduces P99 latency by 92.9%, and improves throughput by $b$2. Across sizes, throughput improvements reach up to $b$3, median-latency reductions up to 80.8%, and P99-latency reductions up to 96.2%. With fixed 256 KB messages and growing $b$4, tail-latency reduction ranges from 54.7% to 96.9% and throughput improvements from $b$5 to $b$6 (Zhu et al., 3 Apr 2025).

End-to-end evaluation attributes the system-level gains to the combination of ping-pong pipeline parallelism, optimized deployment, and M2N. On an Ampere A800 cluster, MegaScale-Infer achieves up to $b$7 higher per-GPU decoding throughput than TensorRT-LLM for a scaled 317B MoE and up to $b$8 higher throughput than vLLM for Mixtral 8x22B. On a heterogeneous H20+L40S cluster, it achieves up to $b$9 and $\frac{topk}{\#expert}$0 higher throughput per cost compared with vLLM and TensorRT-LLM on H20. The paper also notes that there is no end-to-end ablation that isolates “use NCCL vs use M2N with the same pipeline,” so the library’s system impact is established indirectly through micro-benchmarks and the viability of the resulting deployment plans.

Its comparison point is also specific. M2N is not a general replacement for NCCL or MPI. It does not implement all-reduce and other traditional collectives; those remain necessary elsewhere in the stack. This distinction matters in light of broader multi-node inference work, where all-reduce is often the dominant bottleneck for tensor-parallel dense models and specialized hierarchical collectives such as NVRAR are proposed instead (Singhania et al., 12 Nov 2025). M2N addresses a different communication primitive: repeated $\frac{topk}{\#expert}$1 and $\frac{topk}{\#expert}$2 token routing under disaggregated attention and expert clusters.

The principal limitations follow directly from this specialization. M2N assumes RDMA-capable NICs and GPUDirect support; benefits may not materialize on commodity Ethernet without GPUDirect. It is tailored to MoE token routing rather than arbitrary multi-GPU workloads. Scaling beyond the tested $\frac{topk}{\#expert}$3 values may introduce congestion and scheduling issues not evaluated in the paper. Integration also requires adoption of its RDMA buffer management, completion-queue polling, and fused kernel layouts. These constraints define the library’s scope precisely: it is a communication runtime for disaggregated expert parallelism, not a universal distributed-systems substrate.