[QST] Is bounded load imbalance within MoonEP’s design scope?

[Marsssqqq]

Hi MoonEP team, thanks for open-sourcing this work.

MoonEP currently provides exact load balance, where every rank receives
exactly S * K token assignments.

As discussed in #6, the trade-off between eliminating compute imbalance and
introducing expert-weight prefetch overhead may depend on workload shape,
especially token load and expert size.

Before investing further in this direction, I would like to understand the
intended design scope of MoonEP:

1. Is exact S * K balance considered a fundamental MoonEP contract, or is it
   mainly the policy implemented by the current planner?
2. Would an optional planner policy that allows a small, bounded amount of
   residual imbalance be considered in scope, if it reduces weight-prefetch
   and backward gradient-reduction traffic?
3. Are there any public plans, experimental branches, or ongoing efforts in
   this area that external contributors should coordinate with?

Thanks!

[asp0ex]

Hi @Marsssqqq, thanks for the question — this is a trade-off we studied carefully. Short answer: exact $S \times K$ balance is a deliberate design choice.
The math. We formally analyzed the trade-off between balance exactness and expert redundancy. With $R$ ranks, $m = E/R$ local experts per rank, and $B$ redundant remote-expert slots per rank, let $\beta = B/m$. We proved that the worst-case per-rank load over adversarial router outputs is
$$x_{\max}(B) = SK \cdot \frac{(1+\beta)^2}{4\beta}$$

Reference points:

$B/m$	1/8	1/4	1/3	1/2	3/4	1
$x_{\max}/SK$	2.53	1.56	1.33	1.125	1.021	1.0

So half the redundancy budget already caps worst-case overflow at 12.5% ($B = m/2$), and $B = m$ gives exactly $SK$. Bounded imbalance is indeed well-behaved.
Why we still chose exact $S \times K$ balance.

1. Compute balance — exact balance eliminates stragglers in expert compute by construction, rather than bounding them to $1 + \varepsilon$.
2. Planner performance — under the exact $S \times K$ constraint our planner comes very close to the ILP optimum, but once the constraint is relaxed, the max load it produces deviates significantly from the ILP optimum. Better algorithms for the relaxed setting may well exist, but they appear to be considerably more complex.
   On your questions: (1) a bounded-imbalance planner is a legitimate direction — the main cost would be designing a good planning algorithm for the relaxed setting; (2) no public work in this direction yet — if you prototype something, we're happy to discuss numbers here.