fix(moe): equalize and align per-rank token counts for HybridEP dispatch

[HuiyingLi]

Problem

HybridEP's fused all-to-all exchanges fixed-extent buffers, and its metadata allgather asserts bytes_per_rank % 16 == 0 on a 4-byte-per-token array. That imposes two hard constraints the dispatcher did not enforce:

1. Uniformity — every rank in the EP group must dispatch the same token count; unequal extents abort or deadlock the collective.
2. Alignment — the count must be a multiple of 4; the kernel SIGABRTs even with equal counts otherwise (on main, 5&5 and 6&6 abort at hybrid_ep/extension/allgather.cu:97; 4/8/12/16 pass).

Static SFT batches satisfy both by construction (the collater emits one fixed length), but dynamic batches — variable-length or packed sequences, e.g. RL rollouts — violate them almost every step, so such workloads could not train on the hybridep backend at all. Callers could only work around this by pre-padding batches themselves, duplicating a kernel constraint outside the dispatcher.

Fix

_HybridEPManager.dispatch() all-reduces the EP-group max token count, rounds it up to the 4-token alignment, and pads this rank's hidden states, routing map, and probabilities to that extent. Padded rows route to no expert (all-False routing map, zero probs), and combine() slices them back off. deepep and uccl_ep are untouched — the added all-reduce runs only on the hybridep backend, the only one with this constraint.

Cost: one 8-byte all-reduce over the EP group per MoE layer; zero extra copies when counts are already equal and aligned.

Testing

• Unit (tests/unit_tests/moe/test_token_dispatcher.py): pad-and-slice round trip to the aligned group max, equal-but-unaligned padding, and the aligned no-pad fast path (kernel and collectives mocked).
• Functional (tests/functional_tests/moe/run_hybridep_unequal_tokens.py, new, 2×H100 through the real hybrid-ep kernels): with identity experts each token's combined output depends only on its own routing, so an unequal/unaligned run must reproduce the equal-count run exactly. Forward and backward outputs are bitwise-identical for token-count combos 8&3, 6&6, 5&3, and 13&7:

torchrun --standalone --nproc_per_node=2 tests/functional_tests/moe/run_hybridep_unequal_tokens.py
[rank 0] OK: unequal-count forward AND backward bitwise-match the equal-count run
[rank 1] OK: unequal-count forward AND backward bitwise-match the equal-count run

Nightly NeMo-CI dynamic-batch A/B

Validated the PR base (a12b237) against this head (b02855c) with the nightly container, existing finetune recipes, dispatcher=hybridep, dynamic/no-fixed-padding inputs, and 10 training steps.

Recipe	Before (base)	After (PR head)
Nemotron Omni CORD-v2	Failed on the first dynamic batch (seq_len=2047): cudaErrorIllegalAddress, then HYBRID-EP ALLGATHER TIMEOUT ... expecting 8 got 7 and SIGABRT.	Passed: 10/10 finite steps plus validation; dynamic token lengths ranged from 614 to 3543.
Qwen3.5-35B-A3B MoE / MedPix-VQA	Reproduced, then canceled after the first backward failed activation-checkpoint recomputation because dispatcher tensor extents changed (for example, saved [456, ...] vs recomputed [464, ...]).	Passed: 10/10 finite steps with varying per-step token counts.
DeepSeek-V4-Flash / HellaSwag (PP4 × EP32, 128 H100s)	Failed: no step-0 loss; downstream PP RECV timed out after 600 seconds when the upstream MoE path stopped making progress.	Passed: 10/10 finite steps for dynamic sequence lengths 71, 47, 97, 60, 76, 46, 44, 72, 41, 47.

The VLM cases used collate_fn.max_length=null; the DSV4 case used pad_to_max_length=false and pad_seq_len_divisible=null while preserving the recipe's native GBS/LBS (256/8) and PP/EP topology.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HuiyingLi]

/ok to test b02855c