Workload-adaptive reduction-cost model

Derive and validate a principled workload-adaptive cost model that converts replica-gradient-reduction and expert-computation costs into an appropriate penalty for reduction-aware expert-replica planning.

Background

For backward execution, MegaFlux uses a reduction-aware planner with a fixed heuristic penalty of λ=12 token blocks per remote replica. This penalty is intended to approximate owner-side replica-gradient-reduction cost, but it is not derived from measured execution costs and is kept constant across workloads.

The paper explicitly leaves unresolved the development of a more principled model that adapts the penalty to workload characteristics and measured reduction and computation costs. Such a model would improve planning decisions when expert loads, token counts, hardware conditions, or overlap opportunities differ.

References

The penalty is not intended as a precise conversion between reduction and compute time. A more principled, workload-adaptive cost model that derives this penalty from measured reduction and computation costs is left to future work.

— MegaFlux: Skew-Resilient MoE Megakernels via Pipelined Expert Replication  (2610.00671 - Yao et al., 30 Sep 2026) in Appendix, Section “Planner Details,” subsection “Reduction-aware Planning,” subsection “Reduction-aware penalty”