fix(moe): equalize and align per-rank token counts for HybridEP dispatch (#3641)
HybridEP's fused all-to-all exchanges fixed-extent buffers, and its metadata allgather asserts 16-byte alignment on a 4-byte-per-token array. Unequal per-rank token counts deadlock the collective, and counts that are not multiples of 4 abort it outright — so dynamic batches (variable-length or packed sequences) could not train on the hybridep backend at all. _HybridEPManager.dispatch() now all-reduces the EP-group maximum token count, rounds it up to the 4-token kernel alignment, and pads this rank's hidden states, routing map, and probabilities to that extent; padded rows route to no expert, and combine() slices them back off. 2-GPU parity: forward and backward outputs of unequal/unaligned runs bitwise-match the equal-count run. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
H
Huiying committed
97451037f9d9627df0c240a73487b01620d68eb1
Parent: fc3162d
Committed by GitHub <noreply@github.com>
on 8/26/2026, 1:59:44 AM