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

• 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):

Comparison to Prior Art (BF16, E=256)

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)

---

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