Summary of RoutePack: Expert Placement and Attention-Aware Data Packing for MoE Reinforcement Learning

Summary (Overview)

• RoutePack is a hierarchical planner that jointly coordinates layer-wise expert placement and attention-aware data packing for Mixture-of-Experts (MoE) reinforcement learning (RL) training, addressing two coupled load-balancing problems: dense attention work in data-parallel (DP) microbatches and sparse expert work in expert-parallel (EP) ranks.

• The system leverages routing replay from the rollout phase—which reveals each sample's sequence length and layer-wise expert demand before training—to convert data layout from a routing-blind preprocessing step into a load-balancing control.

• The formulation uses a lexicographic objective that prioritizes (1) the slowest EDP shard over the optimizer-step window, (2) total projected work, and (3) the worst row-local tail, while preserving capacity-efficient row counts.

• Across Ling-3.0-Tiny (7.9B params) and Ling-3.0-Flash (124B params), RoutePack improves trainer-measured token throughput by 8.85% and 14.89% respectively over baseline, with expert rerouting contributing 3.80%/10.50% and routing-aware packing adding 4.86%/3.98%.

• The implementation uses layer-wise Longest Processing Time (LPT) placement, state-consistent expert materialization, diverse fixed-row seeding, and parallel population annealing, with a sufficient runtime condition derived for CPU packing to not extend the training-admission critical path.

Introduction and Theoretical Foundation

The Two Coupled Load-Balancing Problems

MoE RL training exposes two distinct sources of load imbalance:

1. Dense attention imbalance: Variable-length RL responses mean that equal sample counts (or even equal packed-token counts) do not imply equal attention work, especially for attention operators with sequence-pair (quadratic) components.

2. Sparse expert imbalance: Token routing determines expert work per EP rank, and the busiest physical rank at each MoE layer paces dispatch and grouped GEMM execution.

The Key Insight: Routing Replay

The rollout and training phases process the same generated token sequence, so routing replay exposes each sample's token length and layer-wise expert demand before the optimizer step is scheduled. For an optimizer-step batch of $N$ samples, the planner receives:

$$\{(t_i, A_i)\}_{i=1}^N, \qquad A_i = [a_{i,l,e}]_{l \in \mathcal{L}_{\mathrm{MoE}}, e \in \mathcal{E}_l}$$

where $t_i$ is the sequence length and $a_{i,l,e}$ is the replayed routing count for logical expert $e$ at MoE layer $l$.

Why Coordination is Necessary

The paper establishes two critical invariants:

Packing cannot remove optimizer-window expert skew: For a fixed placement $\pi_l$, the aggregate demand on physical rank $p$ is:

$$L_{l,p}^{\mathrm{agg}}(\pi_l) = \sum_i \sum_{e: \mathrm{owner}(\pi_l(e)) = p} a_{i,l,e}$$

and for any capacity-feasible packing $x$:

$$\sum_r \sum_g \max_p W_{r,g,l,p}(x; \pi_l) \geq \max_p L_{l,p}^{\mathrm{agg}}(\pi_l)$$

Expert placement does not guarantee row-local balance: A placement can make step-wide totals nearly equal while individual rows remain imbalanced due to correlated routing vectors being packed together.

Attention Work Proxy

For attention operator $A$ and DP cell $B_{r,d}$, the two-term execution proxy is:

$$F_A(B_{r,d}) = \alpha_A \sum_{i \in B_{r,d}} t_i + \beta_A \sum_{i \in B_{r,d}} t_i^2, \qquad \beta_A \geq 0$$

The linear term captures projections and recurrent/linear-attention work; the quadratic term captures token-pair interactions in full attention. KDA and Gated DeltaNet use $\beta_A = 0$, while MLA and Gated Attention use $\beta_A > 0$.

Methodology

Problem Formulation

Decision variables: Layer-local permutation $\pi_l$ maps logical experts to physical slots; binary variable $x_{i,r,d}$ indicates sample $i$ is assigned to cell at row $r$ and DP slot $d$.

Feasibility constraints (rectangular fixed-row):

$$\sum_{r=1}^{R}\sum_{d=1}^{D} x_{i,r,d} = 1, \quad \forall i$$ $$\sum_i t_i x_{i,r,d} \leq C$$

The minimum feasible row count is $R_{\min} = \min\{R: \mathcal{X}_R \neq \varnothing\}$.

Hierarchical Planning

Stage 1: Layer-wise expert placement using LPT scheduling. For candidate placement, aggregate physical-rank load is:

$$G_{l,p}(\pi_l) = \sum_{e: \mathrm{owner}(\pi_l(e)) = p} L_{l,e}$$

where $L_{l,e} = \sum_i a_{i,l,e}$. LPT greedily assigns experts in descending load order to the currently least-loaded physical rank.

Stage 2: Fixed-row data packing with the lexicographic objective:

$$\mathrm{Score}(x; \pi) = \left(\max_g U_g, \sum_g U_g, \max_{r,g} J_{r,g}\right)$$

where $U_g = \sum_r J_{r,g}$ is the projected work accumulated by EDP shard $g$, and $J_{r,g} = A_{r,g} + E_{r,g}$ combines attention and expert costs.

The attention cost per row-shard is:

$$A_{r,g} = \sum_{s \in \mathcal{A}} \max_{d \in \mathcal{D}_g} \left(\alpha_s \bar{T}_{r,d} + \beta_s \bar{Q}_{r,d}\right)$$

with normalized statistics $\bar{T}_{r,d} = \frac{1}{C}\sum_i t_i x_{i,r,d}$ and $\bar{Q}_{r,d} = \sum_i (t_i/C)^2 x_{i,r,d}$.

The expert cost is:

$$E_{r,g} = \frac{1}{C} \sum_{l \in \mathcal{L}_{\mathrm{MoE}}} \max_p W_{r,g,l,p}(\pi, x)$$

where $W_{r,g,l,p}$ is the routed-token load on physical EP rank $p$ at layer $l$.

Population Annealing Solver

The solver combines:

• Diverse fixed-row seeding: Window Shuffle, RandomizedBestFit, and EDPAwarePair operators generate structurally diverse feasible seeds
• Fixed-temperature population annealing: Parallel chains with Metropolis acceptance based on scalar energy $\mathcal{E}(x) = s_0(x) + 10^{-6}s_1(x) + 10^{-9}s_2(x)$
• Systematic resampling between temperature levels with Boltzmann weights
• Trajectory-best archive for non-regressing global incumbent

Empirical Validation / Results

Trainer-Measured Throughput

• Tiny: Expert reordering improves throughput by 3.80%; packing adds 4.86% → total 8.85% over baseline
• Flash: Reordering improves 10.50%; packing adds 3.98% → total 14.89% over baseline
• All pairwise comparisons significant after Bonferroni correction (largest adjusted p-value: 0.0110)

Load-Balance Decomposition

Global EP-rank CV: LPT reduces mean global EP-rank coefficient of variation by more than 99% on both models.

Online load metrics from length-only FFD to guided packing:

Key finding: Guided packing accepts a small increase in projected attention work (+0.09% to +0.77%) in exchange for larger reductions in EP peaks (3.13–3.24%) and the joint bottleneck (1.35–1.53%).

Theoretical and Practical Implications

Design Requirements Established

1. Capacity efficiency must precede load-balance optimization — otherwise, lower peaks from extra microbatches increase total work
2. Correct parallel scope matters: attention is per DP cell, expert work is per EDP shard (as induced by the deployed communicator)
3. Complete-step optimization: independent EDP shards accumulate work at different rates; the primary tail objective must consider the slowest shard over the window

Framework Parallel Schedule

The modeled makespan from routing readiness is:

$$T_{\text{target}} = \max\left\{T_{\mathrm{LPT}} + T_{\mathrm{pack}}, \max(T_{\mathrm{LPT}}, T_{\mathrm{aux}}) + T_{\mathrm{actor}}\right\} + T_{\mathrm{admit}}$$

CPU packing introduces no additional admission straggler if:

$$T_{\mathrm{LPT}} + T_{\mathrm{pack}} \leq \max(T_{\mathrm{LPT}}, T_{\mathrm{aux}}) + T_{\mathrm{actor}}$$

Complexity Bounds

With saturated workers ($W \geq \max\{M, P_{\mathrm{pop}}\}$), total host-side wall complexity is polynomial:

$$T_{\mathrm{CPU}}^{\mathrm{sat}} = O\left(A_{\mathrm{route}} + L(E\log E + EP) + nB\log B + B^2 GLP + qR^*(F + \log R^*) + P_{\mathrm{pop}}(n + qR^*) + N_T K(R^*GLP + R^*D)\right)$$

Conclusion

RoutePack demonstrates that routing replay can serve as both an expert-side and a data-placement signal in sparse RL systems. The key contributions are:

1. Joint control: State-consistent, layer-wise expert rerouting coordinated with whole-sample data packing
2. Joint objective: Linear-quadratic attention work combined with per-layer physical EP-rank tails within each projected EDP shard
3. Practical implementation: Layer-wise LPT, state-consistent materialization, diverse seeding, and parallel population annealing with exact lexicographic scoring

Future directions identified by the authors:

• Scaling studies across datasets, larger EDP shards, multiple nodes, and nontrivial PP/CP degrees
• End-to-end instrumentation to quantify planning overhead in real deployments
• Communication-heavy deployments requiring calibrated topology-aware cost terms
• Kernel microbenchmark suites for shape-aware latency surfaces
• Extending evaluation to coding, STEM, instruction-following, and mixed-domain workloads