Skip to content

Add CTRAN dispatch and arg validation to ncclAlltoAll - #3561

Open
YulunW wants to merge 1 commit into
meta-pytorch:mainfrom
YulunW:export-D115510404
Open

Add CTRAN dispatch and arg validation to ncclAlltoAll#3561
YulunW wants to merge 1 commit into
meta-pytorch:mainfrom
YulunW:export-D115510404

Conversation

@YulunW

@YulunW YulunW commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Summary:
NCCLX ships two all-to-all entry points with identical signatures but different
implementations, sitting ~230 lines apart in the same file:

  • ncclAlltoAll (lowercase t, collectives.cc:117) — the upstream NVIDIA API,
    added in NCCL 2.28. Dispatches straight to ncclEnqueueCheck(ncclFuncAlltoAll).
    No CTRAN, no argument validation.
  • ncclAllToAll (capital T, collectives.cc:346) — Meta-created before upstream
    had an all-to-all. Has the CTRAN switch plus argument checks, and falls back to a
    hand-rolled grouped baseSend/baseRecv loop.

This is the first diff in a stack that de-dups the two down to ncclAlltoAll. It
brings ncclAlltoAll up to parity so it is a drop-in replacement, without touching
ncclAllToAll yet — both entry points remain fully functional after this diff, so
nothing is broken at any point in the stack.

Grafted onto ncclAlltoAll, mirroring how ncclAllGather (collectives.cc:93)
already layers CTRAN on top of the upstream baseline:

  • count == 0 early-out
  • SetCudaDevRAII
  • CudaPtrCheck on sendbuff and recvbuff
  • in-place (sendbuff == recvbuff) rejection with ncclInvalidArgument
  • the NCCL_ALLTOALL_ALGO / ctranAllToAllSupport / ctranAllToAll dispatch

The non-CTRAN fallback stays on the upstream ncclEnqueueCheck(ncclFuncAlltoAll)
path rather than adopting the Meta send/recv loop. These are closer than they look:
taskAppend (enqueue.cc:3044-3051) already decomposes ncclFuncAlltoAll into the
same 2 * nRanks P2P send/recv tasks the Meta loop builds by hand — it never becomes
a ncclTaskColl and never reaches a device kernel. Staying on the upstream path also
propagates collAPI = ncclFuncAlltoAll to the net plugin (already whitelisted at
net.cc:218), emits one AlltoAll NVTX range instead of 2N Send/Recv ranges,
and keeps divergence from upstream small for future rebases.

Applied identically to v2_29 and v2_30; the two files are byte-identical here.

Differential Revision: D115510404

Summary:
NCCLX ships two all-to-all entry points with identical signatures but different
implementations, sitting ~230 lines apart in the same file:

- `ncclAlltoAll` (lowercase `t`, `collectives.cc:117`) — the upstream NVIDIA API,
  added in NCCL 2.28. Dispatches straight to `ncclEnqueueCheck(ncclFuncAlltoAll)`.
  No CTRAN, no argument validation.
- `ncclAllToAll` (capital `T`, `collectives.cc:346`) — Meta-created before upstream
  had an all-to-all. Has the CTRAN switch plus argument checks, and falls back to a
  hand-rolled grouped `baseSend`/`baseRecv` loop.

This is the first diff in a stack that de-dups the two down to `ncclAlltoAll`. It
brings `ncclAlltoAll` up to parity so it is a drop-in replacement, without touching
`ncclAllToAll` yet — both entry points remain fully functional after this diff, so
nothing is broken at any point in the stack.

Grafted onto `ncclAlltoAll`, mirroring how `ncclAllGather` (`collectives.cc:93`)
already layers CTRAN on top of the upstream baseline:
- `count == 0` early-out
- `SetCudaDevRAII`
- `CudaPtrCheck` on `sendbuff` and `recvbuff`
- in-place (`sendbuff == recvbuff`) rejection with `ncclInvalidArgument`
- the `NCCL_ALLTOALL_ALGO` / `ctranAllToAllSupport` / `ctranAllToAll` dispatch

The non-CTRAN fallback stays on the upstream `ncclEnqueueCheck(ncclFuncAlltoAll)`
path rather than adopting the Meta send/recv loop. These are closer than they look:
`taskAppend` (`enqueue.cc:3044-3051`) already decomposes `ncclFuncAlltoAll` into the
same `2 * nRanks` P2P send/recv tasks the Meta loop builds by hand — it never becomes
a `ncclTaskColl` and never reaches a device kernel. Staying on the upstream path also
propagates `collAPI = ncclFuncAlltoAll` to the net plugin (already whitelisted at
`net.cc:218`), emits one `AlltoAll` NVTX range instead of `2N` `Send`/`Recv` ranges,
and keeps divergence from upstream small for future rebases.

Applied identically to `v2_29` and `v2_30`; the two files are byte-identical here.

Differential Revision: D115510404
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 11, 2026
@meta-codesync

meta-codesync Bot commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

@YulunW has exported this pull request. If you are a Meta employee, you can view the originating Diff in D115510404.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant