AutoEP: share one DeepEP buffer across MoE layers instead of one per layer (#8557)
## Summary
Share a DeepEP `ElasticBuffer` across compatible MoE layers within each
AutoEP-converted model, instead of constructing one buffer per layer.
This removes redundant buffer construction and unblocks the 48-layer
EP32 configuration tested here.
## Motivation
In the tested four-node EP32 setup, buffer construction failed at the
twenty-eighth buffer inside `ncclDevCommCreate`. An isolated
buffer-construction probe reproduced the failure without a model,
pointing to a resource limit associated with creating many buffers
rather than model depth itself. This is an observed limit in this
environment, not a general DeepEP limit.
With sharing enabled, the 48-layer EP32 benchmark completes
successfully.
## Design
- Sharing is scoped to one converted model. Separately converted models
retain independent buffers, even when their EP group and configuration
match.
- Layers share only when their EP group, expert count, top-k, hidden
size, token capacity, SM budget, and QP margin match.
- References are released during engine teardown. The final release
destroys the buffer and removes its registry entry.
- Dispatch outputs are separate tensors, and each autograd node retains
its own handle for backward. Sharing relies on the current serialized
DeepEP execution path; this PR does not introduce communication overlap.
## Testing Done
- 43 unit tests pass, covering sharing, configuration separation, model
isolation, and teardown.
- Four-GPU parity testing: 4 passed, 2 failed. The same two
optimizer-delta comparisons also fail on the pre-change control; the
detailed comparison is documented in the testing comment.
- The 48-layer EP32 paired benchmark completes with the shared-buffer
implementation.
- EP16 with an 8192-token sequence also completes.
## Scope
This change addresses redundant communication-buffer construction and
the observed EP32 initialization failure. It does not claim a
steady-state throughput improvement or a reduction in PyTorch peak
allocated memory: these DeepEP buffers are allocated outside the PyTorch
caching allocator.
---------
Signed-off-by: Han Yu <hanyu@linkedin.com>
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>