SIGN IN SIGN UP

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