Petit MegaMoE: Fast Multi-GPU MoE on AMD GPUs

The Challenge of Multi-GPU MoE Inference

Many frontier large language models (LLMs) use Mixture-of-Experts (MoE) layers in their feed-forward networks (FFNs). An MoE layer divides its weights among experts and activates only a small subset for each token. This sparse routing gives the model a large parameter capacity without the compute cost of a similarly sized dense FFN.

These models are too large to serve on a single GPU. Serving engines distribute them by partitioning either each expert's weights, using tensor parallelism (TP), or the experts themselves, using expert parallelism (EP). With EP, communication can become a larger bottleneck than computation: tokens must travel to the ranks that own their selected experts, and the results must return to the source ranks. Routing, synchronization, and kernel launches add further overhead. This is particularly costly during decoding, when small token batches leave little work over which to amortize those costs. In our profiles, MoE layers account for more than 30% of inference latency on the critical path.

We developed Petit MegaMoE as part of Petit to accelerate MoE inference on AMD Instinct MI350X and MI355X GPUs. In our uniform-routing microbenchmarks, Petit MegaMoE matches or exceeds vLLM/AITER at every measured token count and outperforms FlyDSL's MegaMoE implementation from 256 tokens per rank onward. In end-to-end serving, it reduces time per output token (TPOT) at every tested concurrency level from 16 through 128.

Building Petit MegaMoE required co-designing the communication algorithm and GPU kernels around AMD's CDNA 4 architecture. This article focuses on the multi-GPU data path: peer-to-peer communication, fine-grained synchronization, persistent execution, and adaptations for skewed serving workloads.

Expert Parallelism for MoE

Mathematically, an MoE FFN selects a set of top-$k$ experts $S\left(x\right)$ for each token $x$. It computes the output $y_i\left(x\right)$ of each selected expert, then combines the outputs using normalized routing weights $w_i$:

Here, $\mathrm{Act}$ is an activation function such as SiLU, while $W_i^{\text{gate}}$, $W_i^{\text{up}}$, and $W_i^{\text{down}}$ are the weights of expert $i$.

With expert parallelism, each rank owns a subset of the experts. The MoE layer dispatches tokens to the appropriate ranks, runs the selected FFNs, and returns the results to the source ranks for reduction. Figure 1 compares a conventional all-to-all implementation with Petit MegaMoE. The latency table uses A2A for all-to-all, CTA for cooperative thread array, and X-GPU for cross-GPU. The conventional path uses collectives to move token representations and expert outputs among ranks, introducing cross-GPU synchronization and additional kernel launches. The figure summarizes the global dataflow. Repeated token labels do not represent a literal broadcast of every token to every rank. The latency annotations are representative measurements from our development server with eight AMD Instinct MI350X GPUs. These measured overheads are large enough to materially affect decode latency when the token dimension $M$ is small.

Petit MegaMoE reduces this overhead in three ways: it transfers data directly over the peer-to-peer xGMI fabric, replaces layer-wide barriers with fine-grained readiness signals, and fuses communication and computation into persistent kernels. The following sections describe how we tailor these techniques to AMD CDNA 4 GPUs.

Developing MegaMoE for AMD CDNA 4

Linearizing the Symmetric Heap

GPU communication runtimes commonly use a symmetric heap to map memory from peer GPUs into each rank's virtual address space. A conventional layout exposes a separate base pointer for each rank. That indirection is a poor match for AMD's high-performance buffer instructions, which address a contiguous region of up to 4 GiB through a compact resource descriptor.

Petit MegaMoE instead arranges its peer-accessible buffers in a linear virtual address range. A kernel can derive a remote address from the destination rank without a pointer-table lookup. This layout enables efficient buffer loads and stores, hardware bounds checking, and asynchronous transfers into LDS where the data path permits them. It also provides direct control over cache policies—for example, bypassing L2 when streaming weights that are unlikely to be reused.

Replacing Global Barriers with Fine-Grained Synchronization

The conventional path in Figure 1 introduces cross-GPU barriers. Petit MegaMoE avoids layer-wide barriers by dividing tokens into groups of 32 and starting expert computation as soon as a group is ready. Persistent kernels remain resident on the GPU and dynamically schedule ready work, reducing both idle time and kernel-launch overhead.

The kernels use ticketed readiness counters instead of a contended global flag, which reduces synchronization traffic. For remote writes, they also use buffer-instruction cache controls rather than flushing the entire L1 and L2 hierarchy. Our microbenchmarks show that a full cache flush is slower even when caching combines some writes to remote ranks.

Dispatching Tokens Directly over xGMI

In an eight-GPU AMD CDNA 4 system, the GPUs are connected by a full-mesh xGMI fabric, with up to 64 GB/s of bandwidth in each direction on every link. Petit MegaMoE pushes token values directly from their source rank to the ranks that own the selected experts. Distributing transfers across the source GPUs makes better use of the mesh than having destination ranks pull their inputs.

The push design gives up asynchronous loads from remote GPU memory directly into LDS. Even so, it is faster because it avoids an additional wait in the receive path and uses aggregate peer-to-peer bandwidth more effectively. For the small transfers common during decoding, those benefits outweigh the loss of the remote-to-LDS load path.

Adapting to Routing Skew

Token routing is rarely uniform in model serving. Some experts receive far more tokens than others, creating imbalanced communication and compute loads across ranks. We profiled all 1,319 GSM8K problems while serving DeepSeek V4 Pro. A perfectly uniform eight-rank workload would assign 12.5% of token-expert pairs to each rank, but in some layers we observed one rank receiving as much as 44%. The imbalance is especially pronounced during decoding, where fewer tokens are available to smooth out the distribution. Figure 2 shows the largest share of token-expert assignments received by any rank in each layer during prefill and decode.

Petit MegaMoE adapts by changing the number of compute units (CUs) assigned to cross-rank transfers and by selecting matrix-multiplication tile sizes based on the available work. These choices prevent a configuration tuned for uniform routing from becoming a bottleneck when one rank is disproportionately busy.

Evaluation

We compare Petit MegaMoE with two baselines: vLLM 0.27.1 using its default AITER MoE all-to-all backend, which we call vLLM/AITER, and FlyDSL's MegaMoE implementation. The microbenchmark uses an expert-parallelism degree of eight and reports effective TFLOPS across a range of token counts per rank. We calculate effective TFLOPS by dividing the floating-point operation count for DeepSeek V4 Pro's expert shapes by measured execution time. The end-to-end benchmark serves DeepSeek V4 Pro over all 1,319 GSM8K problems. We ran both benchmarks on one server with eight AMD Instinct MI350X GPUs.

FlyDSL's MegaMoE implementation has not yet been integrated into vLLM, so the end-to-end comparison includes only Petit MegaMoE and the vLLM/AITER baseline.

Figure 3 shows that Petit MegaMoE matches vLLM/AITER at 16 and 32 tokens per rank and outperforms it at the other measured sizes. The advantage is largest at 1,024 tokens per rank, where Petit delivers nearly 1.5 times the effective throughput. Petit matches FlyDSL at 8 tokens per rank, trails it from 16 through 128, and outperforms it from 256 through 1,024.

Figure 4 shows a small low-concurrency tradeoff: at concurrency 8, vLLM/AITER has lower TPOT. From concurrency 16 through 128, however, Petit MegaMoE is faster, reaching an advantage of approximately 21% at concurrency 128. This pattern reflects the design's goal of overlapping communication and computation when enough work is available.

Conclusion

Petit MegaMoE demonstrates that AMD hardware can run demanding multi-GPU MoE workloads efficiently when software is designed around the architecture. Its gains on MI350X come from co-designing memory layout, xGMI communication, synchronization, cache behavior, persistent scheduling, and matrix-multiplication tiling for real routing patterns.

The broader lesson is simple: understand the hardware, measure the full critical path, and optimize computation and data movement together. We hope the Petit MegaMoE implementation provides a useful reference for building high-performance kernels on AMD GPUs.

Acknowledgments

We thank AMD for providing the server used to develop and evaluate Petit MegaMoE.