SIGN IN SIGN UP

[Distributed] Add zero-SM all_gather and all-to-all (#79686)

* [Distributed] Add zero-SM all_gather and all-to-all

Build the collectives on top of the symmetric-memory window registration: a
group created with cta_policy=2 runs them on the Copy Engines and the RMA CPU
proxy, so they consume no SM.

- add paddle.distributed.zero_sm with empty(), all_gather() and
  alltoall_single(), mirroring the signatures of paddle.distributed.stream. The
  output is allocated from the registered pool and any input that is not already
  pooled is staged into it
- take the single-call ncclAlltoAll when the splits are uniform and both buffers
  are registered: a group of ncclSend/ncclRecv always launches a point-to-point
  device kernel, while ncclAlltoAll is a symmetric collective and can reach the
  zero-SM path. Everything else keeps the Send/Recv implementation

* [Distributed] Address review on the zero-SM collectives

- gate the single-call ncclAlltoAll on the even split instead of on locally
  uniform split sizes. Explicit splits cannot be judged rank-locally: rank 0 may
  see [1, 1] while rank 1 legally sees [1, 2], which would send one rank into
  ncclAlltoAll and the others into Send/Recv and hang the group. With the even
  split a legal all-to-all forces one count group-wide, so the choice needs no
  agreement
- let the list form of all_gather take a 0-D input, which stream.all_gather
  answers with one scalar per rank, rather than indexing an empty shape
- settle the write-back in is_completed() as well as in wait(), so that a caller
  polling the task and then reading its list sees the gathered values
- drop the per-call votes from _ZeroSMPool: every decision is local now and the
  caller owes the pool a rank-symmetric call sequence, which is what the
  docstrings spell out

Verified on one node with two GPUs against NCCL 2.30.7: the even-split call
still reaches the symmetric path and the uneven one keeps using Send/Recv.

* address review opinion

* address review opinion

* address review opinion

* retrigger AI review on current head
A
adam-xiaoyao committed
4ac7a7572b1697b538340e6a38944cd060d791ee
Parent: 80001ff
Committed by GitHub <noreply@github.com> on 9/9/2026, 1:18:14 PM