fix: distribute routing replay padding across experts

[EazyReal]

Summary

Rollout routing replay appends synthetic expert selections for per-sample, context-parallel, and batch-alignment padding. These rows are loss-masked, but they still pass through MoE expert dispatch and therefore must not concentrate on one expert block or expert-parallel rank.

The original fill flattens [pad, num_layers, topk] before taking modulo num_experts, so a fixed layer advances by num_layers * topk and every padding call restarts at zero. Numbering pad rows after CP fixes those aliases, but assigning ((pad_row + layer) * topk + column) % num_experts still produces contiguous expert IDs. At layer 0 with contiguous expert sharding, 1-5 rows at topk=8 can activate as few as 1/1/2/2/3 of 8 EP ranks for 128 experts and 1/1/1/1/2 for 256 experts. At EP=32 the corresponding counts are 2/4/6/8/10 ranks for 128 experts and 1/2/3/4/5 for 256 experts.

Fix

Keep padding marked as -1 while the per-sample and context-parallel layouts are constructed. After CP slicing is complete, number only the pad rows that remain. For pad-row ordinal p, layer l, route column c, expert-parallel size P, top-k K, and Q = num_experts / P experts per rank, fill:

a = p * K + c + l * (P * K)
expert_rank = a % P
local_expert = (a // P) % Q
expert = expert_rank * Q + local_expert

This is a rank-first permutation of expert IDs. Every fixed-layer row reaches min(P, K) EP ranks, consecutive pad rows give an optimal prefix histogram across ranks, and every num_experts assignments visit every expert once. Layer blocks retain the same per-layer EP balance. Sequence parallelism then partitions these already distributed assignments.

The EP size comes from Megatron's initialized mpu, which already owns the TP and CP geometry used by this helper. Assignment arithmetic uses int64 and casts once to the routed-expert payload dtype, avoiding overflow in intermediate products. Non-padding replay data is unchanged.

Verification

• The production helper is exercised for every 1-5-pad case at 128/256 experts, topk=8, and EP=8/32. EP=8 now activates all 8 ranks for every row count; EP=32 activates 8/16/24/32/32 ranks. Every per-layer rank histogram differs by at most one assignment, and all 8 * pad_rows expert IDs are distinct in these cases.
• Non-padding rows are compared exactly before and after the repair. Multi-row batch-alignment padding continues the same post-CP ordinal sequence.
• Ordinary CP and allgather_cp are checked across CP ranks, then reconstructed from sequence-parallel slices to verify that padding is assigned after CP and before SP.
• A synthetic Qwen3-30B-A3B CP=1 layout uses 32 mixed sample token lengths summing to 19,969, TP=4, EP=8, and the default 128x data-padding multiplier. Per-sample padding contributes 32 terminal rows; alignment to 4 * 128 = 512 rows adds 511 more because 19,969 = 39 * 512 + 1. The resulting pre-SP layout has 20,480 rows, including 543 padding rows and 4,344 synthetic assignments per layer. Upstream's original flattened fill uses 8/128 experts and 1/8 EP ranks with 16.0x expert and 8.0x EP peak/mean. The repaired fill uses 128/128 experts and 8/8 ranks; its expert peak is 34 over a 33.9375 mean (1.001842x), and every EP rank receives exactly 543 assignments (1.0x peak/mean).
• python -m pytest -q tests/test_accelerator.py tests/test_cp_utils.py: 18 passed on CPU-only PyTorch.
• The repository's complete pre-commit hook set passes for both changed files.

The quantitative case validates routing metadata mechanics rather than end-to-end training throughput; a multi-GPU Megatron run can be added if maintainers want that integration evidence.