perf(moe): sync-free HybridEP dispatch (capacity mode, equal token counts)

[yisongbetter]

What does this PR do ?

Removes the two per-dispatch host synchronisations on the HybridEP MoE path under dynamic routing, both as opt-in BackendConfig knobs that default to today's behaviour. (1) dispatcher_capacity_factor: HybridEP's blocking dispatch learns the permuted row count by draining the compute stream (torch.cuda.current_stream().synchronize() after metadata preprocessing) on every MoE layer, forward and activation-checkpoint recompute; its non-blocking mode takes num_permuted_tokens from the caller and only uses it to size the output buffers, with the real counts left on the device and an overflow flag in the handle. The knob runs one blocking calibration dispatch per layer, sizes every later dispatch to ceil(rows × factor) (EP-group max, 4-token aligned) and runs non-blocking; torch._assert_async on the overflow flag fails loudly instead of training on truncated tokens. The combine's backward dispatch takes the same host int, so it needs no drain either. (2) dispatcher_equal_token_counts: the #3641 pad size is an EP-group MAX all-reduce of a host-known shape plus an int() host sync; for fixed-shape batches every rank holds the same count, so the knob replaces the collective with the aligned local count (keep it off for variable-length or in-batch-packed inputs). With both on, the CPU runs ahead of the GPU through the whole MoE layer and no layer acts as an EP-group barrier.

Changelog

• nemo_automodel/components/models/common/utils.py: BackendConfig.dispatcher_capacity_factor (float | None, default None) and dispatcher_equal_token_counts (bool, default False), with docstrings.
• nemo_automodel/components/moe/experts.py: GroupedExpertsDeepEP reads both knobs from the backend, forwards them to TokenDispatcherConfig, and skips its per-microbatch count_nonzero host read in capacity mode (rows are never empty).
• nemo_automodel/components/moe/megatron/token_dispatcher.py: TokenDispatcherConfig.moe_hybridep_capacity_factor / moe_hybridep_equal_token_counts; _HybridEPManager calibration (_calibrate_hybridep_capacity), non-blocking dispatch with the capacity, _assert_no_hybridep_overflow, and the equal-count pad size.
• nemo_automodel/components/moe/megatron/fused_a2a.py: HybridEPCombine.backward passes non_blocking=True when the count is a host int; HybridEPDispatchReplayRecorder.record keeps a host-side extent (capacity mode, static-routing pin) so perf(moe): reuse the DeepEP dispatch layout on activation-checkpoint recompute #3684's activation-checkpoint replay redispatches to the same extent the forward was sized to, and finalize reduces tokens_per_expert only for blocking dispatches.
• Tests: tests/unit_tests/moe/test_token_dispatcher.py (calibration → capacity passed non-blocking, exact count kept for the calibration dispatch's combine, overflow guard, static routing / no factor unchanged; equal counts skip the all-reduce and align, default keeps it), tests/unit_tests/moe/test_backend_config.py (defaults); tests/unit_tests/moe/test_fused_a2a.py (a recorded host extent is kept without reducing tokens_per_expert; a checkpointed dispatch with a capacity extent replays with that extent).

Validation

HybridEP API check (one GB200 node, EP4, 32 experts top-8, random dynamic routing, 8192 rows × 7168): with capacity 1×, 1.5× and 2× the actual count, the dispatched rows, probs, device-side tokens_per_expert and the combine round trip are torch.equal to the blocking path on all ranks, overflow flag 0; with capacity = actual/2 the flag is set (silent truncation, no fault) — hence the device-side assert. An end-to-end check through _HybridEPManager with autograd, 12 dynamic-routing steps, gave forward outputs and input gradients bit-identical to blocking on all 4 ranks; the dispatch call's host time fell from 1.30 ms to 0.21 ms (the blocking value is bounded by whatever GPU work is queued ahead — in training that is the whole backlog).

Throughput (8 × GB200 NVL72, EP16 × PP2, Kimi-K3 16-layer twin, 2k-token rows, mock data, benchmark_static_routing: false = dynamic routing, same job, A/A band 0.08 %):

config	tok/s/GPU	step s	vs control
control (blocking dispatch, per-dispatch pad all-reduce)	1957.9	1.046	—
dispatcher_capacity_factor: 1.5 + dispatcher_equal_token_counts: true	2124.0	0.964	+8.5 %
control repeated (A/A)	1959.5	1.045	+0.08 %
reference: benchmark_static_routing: true (the existing pins)	2089.2	0.980	+6.7 %

Loss curves identical (last-step 12.4403–12.4404 across the four rungs); one calibration per layer (33152 rows × 1.5 → capacity 49728), zero overflows. Isolating the pad collective under static routing at the same scale: +7.4 % (5.8 points from the host sync, 1.6 from the barrier); a side-stream variant of the same collective on the dynamic path recovers nothing (+0.29 %) because the blocking dispatch still drains the stream — which is why both knobs ship together. The 64-node measurement of the dynamic-routing stack will be added here.

Tests: tests/unit_tests/moe/test_token_dispatcher.py, tests/unit_tests/moe/test_fused_a2a.py, tests/unit_tests/moe/test_backend_config.py, tests/unit_tests/moe/test_static_routing_m_splits.py — 102 passed in 6.5 s on one GB200 with a cold Triton cache (job 3090506, this branch on main 28a79d0, tree e48796d1300c).

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? (calibration then capacity passed as a host int with non-blocking dispatch; the calibration dispatch keeps its exact count; overflow flag trips the guard; static routing and the default ignore the factor; equal counts skip the all-reduce and align, default keeps the per-dispatch all-reduce; the replay recorder keeps a host extent and the checkpoint recompute redispatches to it)
• Did you add or update any necessary documentation? (BackendConfig docstrings for both knobs; comments at the dispatch call explaining the two syncs)

Additional Information

• Three commits: one knob each (both default off; whether the K3 GB200 yaml enables them follows the 64-node measurement) and the replay-recorder adaptation to perf(moe): reuse the DeepEP dispatch layout on activation-checkpoint recompute #3684.
• Interplay with perf(moe): reuse the DeepEP dispatch layout on activation-checkpoint recompute #3684 (merged 2026-09-17, this branch is rebased onto it): under activation checkpointing with ignore_router_for_ac the recorder replays the forward's HybridEP layout and sizes the redispatch to a host integer. In capacity mode the forward output is capacity rows, so the recorder now records that integer instead of int(tokens_per_expert.sum().item()); without this the recompute would produce the actual row count and fail the checkpoint metadata check. Recording the known integer also keeps capacity mode free of the per-layer device-to-host copy that the reduction adds for blocking dispatches. Static-routing pins (a host int since fix(moe): equalize and align per-rank token counts for HybridEP dispatch #3641) get the same treatment.
• Memory: capacity mode holds dispatch output buffers of factor × the calibrated rows per live MoE layer (≈ +0.5 GiB per layer at the mini shape, ≈ +0.9 GiB at 64 nodes with factor 1.5).
• dispatcher_equal_token_counts is the "just remove the all-gather" option discussed for fix(moe): equalize and align per-rank token counts for HybridEP dispatch #3641's padding: it is safe exactly when every rank's row count is equal, which the caller knows and HybridEP does not, hence a declaration rather than a default. Variable-length inputs must keep the all-reduce (or pin a static cap).
• Related to # (issue)

[yisongbetter]

/ok to test 52fb62e

[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.

[yisongbetter]

/ok to test c0cf685