Overlap AutoEP split planning with token packing (#8359)
## Summary
- add an opt-in `expert_parallel.async_split_plan` path for standard
AutoEP
- keep count AllToAll and split-size computation on the caller stream;
only pinned metadata D2H uses the cached copy stream, while token
sorting and packing stay on the caller stream
- wait only immediately before payload dispatch, preserve the existing
synchronous default, and reject unsupported AutoEP+AutoTP folding
- start asynchronous planning only for the standard communication
backend and drain newly owned pending work after failed forwards
## Validation
- [x] Local code review completed
- [x] 113 targeted CPU cases pass, including inference-mode/no-grad
warmup followed by training
- [x] Current-head two-H100 inference-to-training parity: 4/4 cases
pass, with no failures, errors, or skips
- [x] `pre-commit run --files` passed for all changed files
- [x] Current-head GPU coverage includes activation checkpointing on/off
and default/non-default caller streams, each after inference-mode warmup
- [x] parity covers loss, output, input gradient, router gradients,
expert gradients, and optimizer parameter deltas
- [x] clean E2E ABBA+BAAB completed on two upstream-stack and two
PR1+PR2 combined-stack allocations
- [x] sparse split-plan attribution completed separately from the clean
E2E runs
### Historical clean E2E (original broad-overlap design)
These timings predate the D2H-only revision and do not establish a
speedup for the current implementation.
These runs used the pre-review PR head implementation, Qwen3-30B-A3B, 48
layers, EP16, sequence length 1024, activation checkpointing, and the
same checkpoint/input/seed. `time-step-phases` and the component
profiler were completely disabled.
Each allocation discarded one fixed warm arm, then ran ABBA followed by
BAAB: four runs per arm, with 30 measured steps per run.
| Stack | Allocation | Sync median | Async median | Delta |
| --- | --- | ---: | ---: | ---: |
| upstream baseline | 1 | 975.59 ms | 974.10 ms | -1.49 ms (-0.2%) |
| upstream baseline | 2 | 989.63 ms | 978.11 ms | -11.52 ms (-1.2%) |
| PR1+PR2 combined | 1 | 990.45 ms | 976.60 ms | -13.85 ms (-1.4%) |
| PR1+PR2 combined | 2 | 983.81 ms | 972.47 ms | -11.34 ms (-1.2%) |
All four allocations favored asynchronous split planning. The median of
the two allocation medians improves by 6.51 ms (0.7%) on the upstream
stack and 12.60 ms (1.3%) on the PR1+PR2 combined stack. With only two
allocations per stack, these are performance estimates rather than a
claim of a fixed universal speedup.
### Sparse attribution
A separate combined-stack run sampled only the first MoE layer.
Activation checkpointing invoked that layer twice per step (forward and
replay), so warmup and aggregation account for two records per step.
| Metric | Median per invocation | Median summed per step |
| --- | ---: | ---: |
| count A2A | 0.690 ms | 1.501 ms |
| count A2A through D2H completion | 0.728 ms | 1.579 ms |
| split reductions/stack plus D2H | 0.108 ms | 0.175 ms |
| sort/pack | 0.250 ms | 0.397 ms |
| residual host wait before dispatch | 2.989 ms | 6.080 ms |
The original D2H range includes split reductions and stacking, rather
than isolating the metadata copy. Residual host wait also includes
outstanding GPU dependencies and rank-arrival effects. These inclusive
intervals overlap and do not isolate pure NCCL or peer-arrival cost;
this change does not remove the full historical split-sync envelope.
An alternative implementation using `async_op=True` plus `Work.wait()`
regressed a same-allocation L48 comparison by 2.1%, so that variant was
removed.
## Inference-mode review follow-up
Commit `d6874193cab3ad3a8a1c09d95f515dd7874885c8` allocates the
persistent split buffer, including pinning, inside
`torch.inference_mode(False)`. A first inference-mode validation or
warmup forward can therefore be followed by a training forward/backward
using the same cached buffer. The mode override is limited to its
one-time allocation.
The regression test reproduces the reported in-place inference-tensor
error before the fix, while its `no_grad` control passes. With the fix,
all 113 selected CPU cases and changed-file pre-commit pass. The four
GPU cases retain output, input/router/expert gradient and FP32
optimizer-update comparisons against the synchronous path after
inference-mode warmup.
GPU validation of this exact commit passed all four cases on 2 × NVIDIA
H100 80GB HBM3, using PyTorch 2.10.0.2+cu130, CUDA 13.0 and NCCL 2.28.9.
Each case uses one MoE layer, EP2/FP16, an inference-mode warmup and
three training updates, with checkpointing on/off crossed with
default/non-default caller streams. There were zero failures, errors or
skips. This is a targeted correctness regression, not a performance
measurement.
---------
Signed-off-by: yh0903 <helloyu0903@gmail.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
