# MegaFlux: Skew-Resilient MoE Megakernels via Pipelined Expert Replication

> MegaFlux dynamically replicates experts across GPUs, pipelining weight transfers and gradient reductions to achieve 1.45x forward and 1.28x backward speedups over fixed placement.

- **Source:** [arXiv](https://arxiv.org/abs/2610.00671)
- **Published:** 2026-10-03
- **Permalink:** https://picx.dev/p/e4jAzZ
- **Whiteboard:** https://picx.dev/p/e4jAzZ/image

## Summary

## Summary (Overview)

- **MegaFlux** is a system that enables **dynamic expert replication** within persistent Mixture-of-Experts (MoE) megakernels to address GPU stragglers caused by routing skew, without altering router outputs or token assignments.
- An **on-device planner** jointly selects replica locations and assigns tile-aligned token blocks under a per-GPU replica budget, using bounded greedy search for load balancing.
- The **forward megakernel** pipelines replica-weight transfers with expert computation by exploiting phase-level weight readiness (FC1 can start before FC2 weights arrive).
- The **backward megakernel** pipelines data-gradient computation, weight-gradient reduction, requantization, and token movement, prioritizing gradient production for replicated experts to enable early reduction.
- On eight NVIDIA B200 GPUs, MegaFlux achieves **geometric-mean speedups of 1.45× (forward) and 1.28× (backward)** over fixed placement, peaking at 2.14× and 2.64×, with end-to-end prefill speedups of 1.13–1.26× when integrated into vLLM.

---

## Introduction and Theoretical Foundation

### Background and Motivation

Mixture-of-Experts (MoE) has become a common scaling strategy in large language models (DeepSeek-V3, Qwen3, GLM-4.5, Kimi K2). By activating only a few experts per token, MoE increases parameter capacity without proportionally increasing per-token computation. Expert parallelism (EP) partitions experts across GPUs, introducing two system challenges:

1. **Communication** for moving routed activations
2. **Load imbalance** when input-dependent routing concentrates work on experts hosted by only a few GPUs

### Fixed-Placement Megakernels

Recent MoE megakernels fuse token dispatch, expert computation, and token return into persistent kernels, enabling fine-grained pipelining. However, **expert ownership remains fixed**: a GPU can execute an expert only if it holds that expert's weights. Megakernels make each GPU's assigned work run efficiently but do not change where that work can execute.

### Routing Skew Problem

Input-dependent routing can concentrate work on a few experts, leaving some GPUs underutilized while hot-expert GPUs determine layer latency. The paper quantifies routing skew as:

$$\kappa = E \sum_{e=1}^{E} \left(\frac{c_e}{\sum_j c_j}\right)^2$$

where $c_e$ is the token count routed to expert $e$ and $E$ is the total number of experts. Uniform routing gives $\kappa = 1$; larger values indicate more uneven loads.

In captured Qwen3 and OLMoE workloads, routing skew slows fixed-placement MegaMoE by **7.3–41.0%**, motivating the need for dynamic expert replication.

### Key Insight

Dynamic replication introduces new communication: replicas must receive expert weights to execute, and during training, their partial weight gradients must be reduced at expert owners. The challenge is realizing dynamic replication within the fine-grained communication-computation pipeline of MoE megakernels.

---

## Methodology

### 1. Planner: Replica Placement and Token-Block Assignment

#### Problem Formulation

The planner partitions each expert's routed rows into token blocks of at most $\beta$ rows, matching the megakernel's tiled computation. The optimization problem is:

$$\min_{q \in \mathbb{Z}^{E \times P}_{\geq 0}} \max_r \ell_r, \quad \ell_r = \sum_e q_{er}$$

subject to:

$$\sum_r q_{er} = b_e \quad (\forall e), \quad \sum_{e: h(e) \neq r} \mathbb{1}[q_{er} > 0] \leq s \quad (\forall r)$$

where $q_{er}$ is the number of expert $e$'s blocks assigned to GPU $r$, $b_e = \lceil c_e/\beta \rceil$ is the block count, $h(e)$ is the owner of expert $e$, and $s$ is the per-GPU replica budget.

#### Bounded Greedy Search

Four greedy searches run in parallel, each starting from the no-replica assignment and shifting blocks from overloaded owners toward $\tau = \lceil \sum_e b_e / P \rceil$ blocks per GPU. Searches combine:
- **Donor priority**: decreasing total overload or largest remaining home-expert load
- **Fan-out**: at most one new copy per expert per round, or multiple copies

For backward, a **reduction-aware variant** accounts for owner-side cost of replica-gradient reduction using the proxy:

$$C_\lambda(q) = \max_r (\ell_r + \lambda n_r)$$

where $n_r$ counts remote replicas of experts owned by GPU $r$, and $\lambda = 12$ is a fixed penalty.

### 2. Executor: Pipelined Expert Replication

#### Forward

A replica can begin useful work before all weights arrive:
- **FC1** requires only $W_1$ and an input token block
- **FC2** requires $W_2$ and the completed activation block from FC1

Separate readiness for the two weight matrices overlaps FC1 with $W_2$ transfer. Communication warps progressively join the replica-weight queue after draining their token blocks.

#### Backward

The backward megakernel overlaps both replica materialization and gradient reduction with the existing DGrad-WGrad pipeline. Key mechanisms:

1. **Prioritized gradient production**: WGrad preparation and gradient-tile production for replicated experts are prioritized across all participants
2. **Early gradient reduction**: Owner-side workers pull and accumulate partials in FP32 in a fixed order as soon as they become available
3. **Overlapped requantization**: Token-axis $D_e$ and $dZ_e$ requantization is scheduled in otherwise idle communication intervals

The expert computation follows:

$$Z_e = X_e W_{1,e}, \quad A_e = f(Z_e), \quad O_e = A_e W_{2,e}, \quad y_t = \sum_{e \in S_t} p_{t,e} O_e[t]$$

where $Z_e = [G_e, U_e]$, $f(Z_e) = \text{SiLU}(G_e) \odot U_e$, and $O_e[t]$ denotes expert $e$'s output for token $t$.

Backward gradients:

$$dA_e = D_e W^\top_{2,e}, \quad dX_e = dZ_e W^\top_{1,e}, \quad dW_{2,e} = A^\top_e D_e, \quad dW_{1,e} = X^\top_e dZ_e$$

---

## Empirical Validation / Results

### Performance Relative to Fixed Placement

| Metric | Forward | Backward |
|--------|---------|----------|
| Geometric-mean speedup | 1.45× | 1.28× |
| Peak speedup (real-median) | 2.14× | 1.43× |
| Peak speedup (real-high) | 1.99× | 1.46× |
| Peak speedup (real-extreme) | 1.93× | 1.92× |
| Peak speedup (Zipf-κ3) | 1.53× | 1.75× |
| Peak speedup (Zipf-κ6) | 1.86× | 2.64× |

- **Balanced routes** remain near parity (0.985× forward, 0.999× backward geometric mean)
- Small skewed backward workloads can regress to 0.91×
- Gains strengthen as computation amortizes planning and replica-operation overhead

### Decomposing Pipelined Expert Replication

Separate-stage replication reduces latency by up to **42.4% forward** and **51.2% backward**; pipelining provides an additional **13.2% and 26.7%** reduction.

**Hidden replica-operation cost** (real-high routes):

| Operation | Hidden Fraction |
|-----------|----------------|
| Forward weight transfer | 56–76% |
| Backward transfer + reduction | 91–100% |

### Comparison to Prior Art (BF16, E=256)

| Baseline | Forward Speedup | Backward Speedup |
|----------|----------------|------------------|
| Megatron-Core | 2.31× | 1.53× |
| UltraEP | 1.74× | 2.35× |
| Mixture-of-Kittens (unmodified) | 1.84× | 1.37× |
| MoK (comparable-work estimate) | 1.70× | 1.26× |

### Online Planner

- Planning takes **42.6–65.9 μs** (E=128) and **49.8–72.0 μs** (E=256)
- Planning-to-forward-latency ratio decreases from 9.17% to 0.14% as tokens/rank increase from 1K to 128K
- Planner is **1.31–2.35× faster** than UltraEP's planner with CUDA-graph replay

### End-to-End Evaluation (vLLM, DeepSeek-V4-Pro)

| Batch | 16K chunk | 32K chunk |
|-------|-----------|-----------|
| 8 | 1.128× | 1.133× |
| 16 | 1.227× | 1.238× |
| 32 | 1.165× | 1.260× |

---

## Theoretical and Practical Implications

### Theoretical Contributions

1. **Block-granular runtime expert replication**: Formulates replication as a constrained optimization problem with tile-aligned token blocks, preserving token counts, expert selections, and routing weights.

2. **Phase-level weight readiness**: Demonstrates that replicas need not be fully materialized before computation begins—FC1 can overlap with $W_2$ transfer.

3. **Pipelined gradient reduction**: Shows that prioritizing gradient production for replicated experts across all participants enables early reduction that overlaps with remaining computation.

### Practical Implications

- **System design**: MegaFlux demonstrates that persistent MoE execution can adapt to routing skew without being constrained by fixed expert placement, achieving significant speedups on real hardware.
- **Workload amortization**: Larger workloads amortize planning and replica-operation overhead, making replication most effective when substantial compute exists.
- **Balanced workloads**: Near-balanced assignments do not guarantee lower latency; exposed replica-operation costs may outweigh rebalancing benefits.

### Limitations

- Evaluation is limited to a single NVLink domain; topology-aware, cross-node replication is left to future work
- A runtime gate to enable replication only when predicted savings exceed costs is suggested but not implemented
- The reduction-aware penalty $\lambda = 12$ is a heuristic; a more principled, workload-adaptive cost model is left to future work

---

## Conclusion

MegaFlux resolves GPU stragglers due to work imbalance in persistent MoE execution via dynamic expert replication. An on-device planner assigns replicas and token blocks, while forward and backward megakernels pipeline replica-weight transfers and gradient reductions with expert computation. Across diverse routing skew and workload sizes on eight NVIDIA B200 GPUs, MegaFlux achieves geometric-mean speedups of 1.45× in forward and 1.28× in backward execution over fixed placement, peaking at 2.14× and 2.64×.

**Future directions** include:
- Topology-aware, cross-node replication
- Runtime gating to enable replication only when beneficial
- More principled cost models for reduction-aware planning
- Extending to additional workload types and model architectures

---

_Markdown view of https://picx.dev/p/e4jAzZ, served by PicX — AI-generated visual whiteboard summaries of research papers._