PaperlayerView exact revision on arXiv

Source-linked AI summary

MegaFlux: Skew-Resilient MoE Megakernels via Pipelined Expert Replication

Jianzhu Yao, Siva Kumar Sastry Hari, Vignesh Balaji, Sana Damani, Insu Jang, Pramod Viswanath, Christos Kozyrakis

arXiv:2610.00671v1cs.DC

TL;DR

Fixed expert placement leaves MoE megakernels exposed to routing-skew stragglers, while runtime replication adds weight-transfer and gradient-reduction work. MegaFlux plans replicas and token blocks on-device and pipelines those operations with expert computation. Across eight B200 GPUs, it improves forward and backward performance over fixed placement.

  • Problem

    Fixed expert ownership in MoE megakernels leaves routing skew able to concentrate work on a few GPUs, while runtime replication adds weight transfers and training-time replica-gradient reductions.

    Source: §1 · pp. 1–2

  • Method

    MegaFlux uses an on-device planner to select replica locations and tile-aligned token-block assignments, then pipelines replica materialization and gradient reduction with expert computation.

    Source: §1 · p. 2

  • Results

    Across eight 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×.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The results show that persistent MoE execution can adapt to routing skew without being constrained by fixed expert placement.

    Source: §6 · p. 10

    The evaluation covers a single NVLink domain; topology-aware cross-node replication remains future work, and small skewed backward workloads can regress to 0.91×.

    Source: §4.6 · p. 9

Abstract

from arXiv · show

Mixture-of-experts (MoE) megakernels fuse expert-parallel communication with expert computation. However, under fixed expert placement, routing skew creates GPU stragglers: overloaded GPUs determine layer latency while others sit idle. Replicating hot experts can shift work to underloaded GPUs, but dynamic replicas introduce additional work: replicas must receive expert weights to execute and, during training, their partial weight gradients must be reduced at the expert owners. We present MegaFlux, which makes expert replication a runtime decision and pipelines the communication induced by replication within persistent MoE execution. An on-device planner jointly selects replica locations and assigns tile-aligned token blocks under a per-GPU replica budget, leaving router outputs unchanged. The forward and backward megakernels realize pipelined expert replication: replicas begin computation as their required weights arrive, while backward overlaps replica-gradient reduction with ongoing expert computation. MegaFlux extends TensorRT-LLM's CuTeDSL MegaMoE forward kernel and introduces a new backward MoE megakernel. Across 147 configurations per direction on eight NVIDIA B200 GPUs, MegaFlux achieves geometric-mean speedups of $1.45\times$ for forward and $1.28\times$ for backward over the same megakernels with fixed placement, peaking at $2.14\times$ and $2.64\times$. In ablations, pipelining hides $56$--$76$% of replica-weight transfer cost in forward and $91$--$100$% of combined weight-transfer and replica-gradient-reduction cost in backward, yielding up to $13.2$% and $26.7$% additional layer-latency reductions over the same replication plans with these operations executed separately. Integrated into vLLM for DeepSeek-V4-Pro prefill, MegaFlux delivers $1.13$--$1.26\times$ median end-to-end speedups over fixed placement.

1 INTRODUCTION

MoE megakernels pipeline communication and expert computation, but fixed expert ownership leaves routing skew to create GPU stragglers. MegaFlux plans runtime expert replication and pipelines its added communication; across eight B200 GPUs, it improves forward and backward performance over fixed placement.

Motivation

  • Dynamic replication shifts hot-expert work to underloaded GPUs, but adds weight transfers and, during training, owner-side reduction of replica gradients.Fixed expert ownership means a GPU can execute an expert only when it holds that expert’s weights.Source: §1 · pp. 1–2

Our solution

  • The on-device planner selects replicas and tile-aligned token blocks within per-GPU replica budgets while preserving token counts, expert selections, and routing weights.The plan avoids quality trade-offs from altering or dropping routed work.Source: §1 · p. 2
  • MegaFlux overlaps replica materialization and gradient reduction with tile-level expert computation instead of running them as separate stages.Forward starts input projection when its weights and an input block are ready; backward prioritizes replicated-expert gradient production so completed tiles can be reduced while other computation continues.Source: §1 · p. 2

Evaluation

  • 1.45×/1.28× geometric-mean forward/backward speedups over fixed placement, peaking at 2.14×/2.64×, are reported across eight B200 GPUs.Ablations report up to 13.2%/26.7% further latency reduction from pipelining over separate-stage replication; vLLM DeepSeek-V4-Pro prefill achieves 1.13–1.26× median end-to-end speedups over fixed placement.Source: §1 · p. 2

2 MEGAMOE EXECUTION AND MOTIVATION

The section describes fixed-placement MegaMoE execution and measures routing skew in real workloads. Even with communication and computation pipelined, skewed routes expose substantial latency when expert placement is fixed.

2.1 FIXED-PLACEMENT MEGAMOE EXECUTION

  • Expert parallelism distributes experts across GPUs, which store subsets of parameters and receive routed activations; persistent megakernels keep CTAs resident and expose tile-level communication–computation overlap.This contrasts with conventional separate dispatch, computation, and return stages.Source: §2.1 · pp. 2–3
  • Fixed-placement forward MegaMoE overlaps token transfers with tiled FC1/FC2 computation, but every expert tile still runs on its owner.Backward streams token blocks through DGrad and WGrad, with token return and requantization overlapping expert computation; the final top-k sum remains separate.Source: §2.1 · p. 3

2.2 ROUTING SKEW IN REAL WORKLOADS

  • Routing profiles cover Qwen3 and OLMoE inference and training workloads; κ is 1 for uniform routing, with larger values indicating more uneven expert loads.The profiles include serving and training datasets, plus OLMoE inference routes on Finance and Contrast.Source: §2.2 · pp. 3–4
  • 7.3–41.0% higher latency occurs on P95-skew routes than on balanced counterparts across OLMoE and Qwen3 workloads.The comparison uses native MoE configurations and checkpoint weights on eight B200 GPUs; the balanced route changes expert IDs while preserving token counts.Source: Figure 1 · p. 3, §2.2 · p. 4

3 MEGAFLUX

MegaFlux plans replica placement and token-block assignments from current routing, then realizes that plan by overlapping replica-related communication with expert computation. Its planner balances block-level GPU load subject to replica capacity.

3.1 PLANNER: REPLICA PLACEMENT AND TOKEN-BLOCK ASSIGNMENT

  • The planner assigns tile-aligned token blocks and expert replicas across GPUs under per-GPU replica budgets, without changing expert selections or routing weights.Partitioning each expert’s routed rows into blocks of at most β rows leaves at most one partial block per expert regardless of replication.Source: §3.1 · p. 4
  • Block counts proxy compute load, and the planner’s bounded greedy searches move blocks from overloaded owners toward eligible GPUs while minimizing maximum GPU load.Ties favor fewer replicas, followed by pruning and assignment refinement; backward planning additionally accounts for owner-side replica-gradient reduction cost.Source: §3.1 · pp. 4–5

3.2 EXECUTOR: PIPELINED EXPERT REPLICATION

  • Forward replicas can start FC1 when W1 and an input block are ready, overlapping FC1 with W2 transfer and avoiding a replica-wide readiness barrier.The executor reuses communication warps for replica-weight movement while preserving token movement.Source: §3.2 · p. 5, Figure 2 · p. 5
  • Backward overlaps replica-weight transfer and requantization with DGrad, then reduces completed replica-gradient tiles at their owners while other tiles are still produced.MegaFlux prioritizes WGrad preparation and tile production across participating GPUs so the matching partials needed for reduction become ready earlier.Source: §3.2 · pp. 5–6, Figure 2 · p. 5

4 EVALUATION

The evaluation tests dynamic replication and pipelining across routing distributions and workload sizes, then compares MegaFlux with prior systems, its planner alternative, and fixed placement end-to-end. Results show strong gains under skew, with benefits shaped by workload size and bounded by the evaluation’s single-NVLink-domain scope.

4.1 EXPERIMENTAL SETUP

  • The evaluation uses eight B200 GPUs in one NVLink domain, with 25% spare expert capacity and fixed-placement MegaMoE as the primary baseline.Trace-derived routes preserve observed expert-load distributions for controlled sweeps, but are not end-to-end replays of source models.Source: §4.1 · p. 6

4.2 PERFORMANCE RELATIVE TO FIXED PLACEMENT

  • 1.45× forward and 1.28× backward geometric-mean speedups over fixed placement; peaks reach 2.14× and 2.64×, respectively.Across Zipf settings, all 84 improve; balanced routes remain near parity, and small skewed backward workloads can regress to 0.91×.Source: §4.2 · p. 7

4.3 DECOMPOSING PIPELINED EXPERT REPLICATION

  • Pipelining adds up to 13.2% forward and 26.7% backward latency reduction over separate-stage replication using identical plans.It hides an estimated 56–76% of forward weight-transfer cost and 91–100% of backward weight-transfer plus gradient-reduction cost.Source: §4.3 · pp. 7–8, Figure 4 · p. 7

4.4 MEGAFLUX COMPARED TO PRIOR ART

  • MegaFlux’s geometric-mean speedups over Megatron-Core, UltraEP, and unmodified MoK are 2.31×/1.53×, 1.74×/2.35×, and 1.84×/1.37× for forward/backward.For comparable work, adjusted speedups over MoK are 1.70× forward and 1.26× backward.Source: §4.4 · p. 8

4.5 ONLINE PLANNER

  • Planner assignments yield 0.981–1.102× UltraEP-plan F+B performance on MegaFlux across 1K–32K tokens/rank, and planner execution is 1.31–2.35× faster.At 1K and 4K, the reduction-aware planner creates no replicas to avoid transfers and gradient reductions that are harder to amortize.Source: §4.5 · pp. 8–9

4.6 END-TO-END EVALUATION AND DISCUSSION

  • 1.13–1.26× median end-to-end prefill speedups over fixed placement in vLLM for DeepSeek-V4-Pro.The evaluation measures batch latency until every request produces one output token.Source: Table 1 · p. 9, §4.6 · p. 9
  • MegaFlux is most effective with substantial routing imbalance and enough expert computation to amortize planning and replica operations.Balanced workloads remain near parity, small skewed backward workloads can regress to 0.91×, and cross-node replication is left for future work.Source: §4.6 · p. 9

5 RELATED WORK

Prior MoE megakernels fuse communication and computation but optimize for fixed expert placement. MegaFlux combines persistent execution with runtime replication and pipelines the resulting transfer and gradient-consolidation work.

5 RELATED WORK

  • FLUX and COMET overlap distributed communication with computation, while later MoE systems fuse communication and expert computation into persistent or megakernel execution.MPK provides a more general in-kernel task runtime.Source: §5 · p. 9
  • Prior systems adapt expert replication, placement, or routed-work assignment to skew, while some aggregate replica gradients during training.Related planners also model costs beyond token counts, including activated experts, weight streaming, and tile-rounded computation.Source: §5 · p. 9
  • MegaFlux differs from fixed-placement megakernels by integrating runtime replication and its communication dependencies into persistent expert computation.Forward progress is gated by weight availability, while backward gradient consolidation overlaps expert computation.Source: §5 · p. 9

6 CONCLUSION

MegaFlux uses dynamic expert replication and pipelined forward and backward megakernels to address work imbalance in persistent MoE execution. Across routing skew and workload sizes on eight B200 GPUs, it outperforms fixed placement in both directions.

6 CONCLUSION

  • MegaFlux addresses GPU stragglers from work imbalance through dynamic expert replication in persistent MoE execution.Its planner assigns replicas and token blocks, while forward and backward megakernels pipeline replica-weight transfers and gradient reductions with expert computation.Source: §6 · p. 9
  • 1.45× forward and 1.28× backward geometric-mean speedups over fixed placement, peaking at 2.14× and 2.64×, respectively.Evaluated across diverse routing skew and workload sizes on eight NVIDIA B200 GPUs.Source: §6 · pp. 9–10
  • The reported results show that persistent MoE execution can adapt to routing skew without being constrained by fixed expert placement.Source: §6 · p. 10

A.1 GREEDY PLACEMENT AND REFINEMENT

MegaFlux uses bounded greedy searches and refinement to place replicas and redistribute token blocks, then applies a reduction-aware pruning pass for backward plans. Assignment preserves routed token identities and routing weights while distributing blocks across selected GPUs.

A.1 GREEDY PLACEMENT AND REFINEMENT

  • Four greedy searches start from the no-replica assignment and combine donor priority with one-copy or multiple-copy fan-out strategies.A new replica goes to the least-loaded eligible GPU below τ with a free replica slot that does not already hold the expert; searches stop after s + 2 rounds or a round without movement.Source: §A.1 · p. 13
  • Refinement selects the candidate with the smallest maximum GPU load, breaks ties by replica count, then prunes redundant replicas and moves single blocks between GPUs sharing an expert.These refinement steps neither open new replicas nor increase maximum compute load.Source: §A.1 · p. 13

A.2 REDUCTION-AWARE PLANNING

  • Remote replicas produce full weight-gradient partials regardless of token count, so balancing computation can concentrate reduction work at expert owners.The pruning proxy counts remote replicas owned by each GPU and does not explicitly model transfer time or overlap.Source: §A.2 · p. 13
  • The reduction-aware penalty is fixed at λ = 12 token blocks of augmented owner load per remote replica, not a precise conversion to reduction time.A workload-adaptive penalty derived from measured reduction and computation costs is left to future work.Source: §A.2 · p. 13
  • Pruning removes a remote replica when its positive allocation v is below λ, returning its blocks to the owner and repeating while the selected owner has such a replica.The plan reverts to q0 unless pruning strictly improves Cλ; unlike load-only refinement, pruning may increase maximum compute load for lower owner-side reduction cost.Source: §A.2 · pp. 13–14
  • Inference omits reduction-aware pruning, while backward reuses the reduction-aware plan selected before forward without changing replica locations or routed-assignment destinations.Planning cost is charged once, to forward.Source: §A.2 · p. 14
  • For each expert, routed assignments are ordered by source GPU and local token and top-k position, partitioned into blocks, and assigned according to the planner’s per-GPU block counts.Token identities, top-k positions, and routing weights are retained for output combination; replica gradients are reduced at the expert owner.Source: §A.2 · p. 14

B.1 TRACE-DERIVED WORKLOAD CONSTRUCTION

The workload construction resizes source expert-load histograms to the evaluated expert counts and builds routes that realize the resulting counts. Evaluation aggregates repeated runs and seeded realizations into configuration latencies and speedups.

B.1 TRACE-DERIVED WORKLOAD CONSTRUCTION

  • Three distinct source events are selected near the 25th, 50th, and 75th percentile positions within each routing-skew band.Forward uses Qwen3 retrieval histograms and backward uses OLMoE training histograms, restricted to captures with 32K input tokens per rank; skew labels are relative within each source and do not imply matched κ across directions.Source: §B.1 · p. 14
  • Resized expert counts are scaled to PMk assignments, constrained by top-k feasibility, rounded to preserve the total, and permuted over expert IDs to construct routes.All methods use identical routes, routing weights, and initial ownership; the procedure preserves trace-derived load distributions subject to rebinning and feasibility, but not original token-level expert combinations.Source: §B.1 · p. 14
  • Configuration latency is the median across six realization–run measurements, and speedup is fixed-placement latency divided by MegaFlux latency.Each realization uses two fresh processes with five warmup and 20 timed iterations; forward adds measured planner latency, while backward reuses the saved plan without another planning charge.Source: §B.1 · p. 14

B.2 REDUCTION IN ASSIGNMENT IMBALANCE

The section quantifies how MegaFlux redistributes assignment imbalance, reporting median reductions in busiest-GPU excess assignments while distinguishing assignment redistribution from latency reduction.

B.2 REDUCTION IN ASSIGNMENT IMBALANCE

  • The imbalance calculation counts valid token–expert assignment rows, excluding padded blocks and routing-weighted assignments; the per-GPU mean is Mk.Rebalancing preserves all routed work.Source: §B.2 · p. 14
  • η measures the fraction of the busiest GPU’s excess assignments removed.Cases with zero fixed-placement excess are excluded.Source: §B.2 · pp. 14–15
  • 99.6% for forward and 98.2% for backward are the reported median assignment-redistribution values, not latency reductions.The medians are over 126 nonbalanced settings per direction.Source: §B.2 · p. 15
Loading 2610.00671v1…