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

> RoutePack jointly optimizes expert placement and data packing using routing replay, improving MoE reinforcement learning token throughput by up to 14.89%.

- **Source:** [arXiv](https://arxiv.org/abs/2608.12146)
- **Published:** 2026-09-05
- **Permalink:** https://picx.dev/p/QYaS5x
- **Whiteboard:** https://picx.dev/p/QYaS5x/image

## Summary

# 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

| Model | T/A params | Baseline | Reorder | RoutePack |
|-------|-----------|----------|---------|-----------|
| Tiny | 7.9B/1.3B | $42.86 \pm 3.05$ | $44.49 \pm 1.92$ | $46.65 \pm 2.12$ |
| Flash | 124B/5.1B | $68.50 \pm 5.02$ | $75.69 \pm 4.59$ | $78.70 \pm 3.82$ |

- **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**:

| Model | EP peak sum | Tail peak | EP balance eff. | Attention | Joint |
|-------|------------|-----------|-----------------|-----------|-------|
| Tiny | -3.13% | -11.04% | 0.891 → 0.920 | +0.09% | -1.53% |
| Flash | -3.24% | -11.62% | 0.844 → 0.869 | +0.77% | -1.35% |

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

---

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