[EazyReal]

Summary

Sequence packing appends loss-masked tokens to align training batches. With router replay enabled, those tokens still pass through MoE dispatch, but their replay payload was filled with zero bytes. Every synthetic route therefore selected expert 0, concentrating padding-only forward work on one expert and potentially one expert-parallel rank.

Fix

• Read the model's routed-expert count from the same configuration already loaded by the batch packer, including composite text configs.
• Distribute each layer's complete padding route set as full expert cycles plus an evenly spaced remainder, then rotate assignments between layers. The cycle order keeps every row's top-k distinct across cycle boundaries, bounds per-expert count skew to one assignment, and spreads the unavoidable remainder optimally across any contiguous expert-parallel partition.
• Widen compact replay payloads when the model's expert range requires it, preserving all sampled route values. For example, a uint8 payload for a 512-expert model becomes uint16 before padding is appended.
• Fail closed before mutating the batch when routing padding is required but the expert count is unavailable.

Non-padding replay routes and loss-mask behavior are unchanged.

Verification

Validated at 893604c3f10d6153eb261dd7e51e98cbf38f8146, based on main 91f58bd9d958860a8414284ea7ec7619f50a23a7. The current sampling-mask padding and multimodal batch paths are preserved.

• uv run --locked pytest -q tests/unit/orchestrator/test_batch.py: 36 passed on Python 3.12.
• A 256-expert, top-k 8, one-row padding case assigns one route to each of eight contiguous EP ranks per layer.
• A 256-expert, top-k 2, seven-row padding case produces 14 distinct routes per layer; per-expert counts differ by at most one and the eight-rank histogram is [1, 1, 2, 2, 2, 2, 2, 2] up to rank ordering.
• A 256-expert, top-k 8, 33-row padding case crosses a full expert cycle while retaining distinct routes per row, per-expert skew of one, and exactly 33 assignments per EP rank.
• A 512-expert case verifies uint8 to uint16 widening, exact preservation of sampled routes, row/layer rotation, and balanced aggregate EP assignment.
• A supplemental CPU probe passed three cases combining routing replay, sampling masks, and objective-weight streams: original payloads remained intact, sampling-mask padding counts were zero, and padding contributed no objective weight.
• Ruff check, Ruff format check, and git diff --check pass for the changed files.

No multi-GPU or DeepEP end-to-end training run was performed; the regression coverage validates the production packing payload and its dispatch-facing invariants on CPU.

[mikasenghaas]

uh this is cool, we will check!

[MarioSieg]

Amazing - looks good so far :)
Can you rebase onto main and fix the conflicts?

[EazyReal]

Done!

[EazyReal]

@MarioSieg rebased onto current main and resolved the conflicts, preserving the new sampling-mask and multimodal paths.Could you take another look?

[EazyReal]

Closing as superseded by #4184 for the original zero-filled routing-replay padding hotspot. That PR repeats real routes and tokens, while this PR synthesizes balanced expert assignments; the stronger balance guarantees here could be evaluated separately if benchmarks justify a follow-up.

The original diagnosis and proposed fix remain documented here. I’ve requested explicit attribution and consideration of coauthor credit on #4184: #4184 (comment). Closing this PR does not resolve that attribution request.