Add direct IB reduce-scatter collective (#3058)#3058
Open
goelayu wants to merge 6 commits into
Open
Conversation
meta-cla
Bot
added
the
CLA Signed
Contributor
|
@goelayu has exported this pull request. If you are a Meta employee, you can view the originating Diff in D110544920. |
goelayu
pushed a commit
to goelayu/torchcomms
that referenced
this pull request
Jul 6, 2026
Summary: Pull Request resolved: meta-pytorch#3058 Direct all-pairs IB reduce-scatter for GB300, cleaned down to the production-ready algorithm and its validation surface. This diff keeps the asym 128/384 kernel and the `asym_cpasync` variant. The cp.async variant uses `CpAsyncSmemReduce<float, SumOp, 8192, 384, 2>` in the recv role, with 131072 B dynamic shared memory, to hide reduce-side global-load latency without the register pressure and spills seen in register-prefetch variants. The shared-memory async pipeline is intentionally retained as an A/B alongside the production asym path because it is the only measured variant that is both register-frugal and faster across the useful mid-size range. The committed benchmark is intentionally one clean reproducible regression source: NCCL Ring+Simple vs asym 128/384 vs asym-cpasync, 1KB through 2GB, fixed 8-block GB300 budget. The MCCL adapter / nccl-test-suite integration is useful follow-up work, but this diff keeps the standalone benchmark so future changes can reproduce the direct-kernel numbers without also depending on MCCL adapter plumbing. The cleanup removed the experimental sweep surface from the diff: timed trace hooks, fused transport hooks, local-reduce / reduce-ceiling microbenches, unused kernel variants, and temporary launcher/enqueue instrumentation. Those artifacts were archived under `rsdirect/claude/archive/prod_cleanup_20260706/` and summarized in `rsdirect/claude/RUNNING_LOG.md` for future reference. Differential Revision: D110544920
goelayu
force-pushed
the
export-D110544920
branch
from
a2a3f29 to
5080f09
Compare
meta-codesync
Bot
changed the title
added 5 commits
July 7, 2026 10:14
Summary: Redesigns the NVLink tile-path protocol state around fixed channels. Replaces the per-peer `NvlinkTransportTileState` with `NvlChannelState`: one 128-byte-aligned state record per channel per peer, with IPC-shared per-peer channel-state arrays and staging buffers owned by the transport. Migrates NVL send/recv/forward and their callers to fixed channel geometry from `MultiPeerNvlTransportConfig::max_num_channels` and `per_channel_slot`, instead of sizing channel state from each call's active block count. Non-tile benchmark transports explicitly set `max_num_channels = 0` to stay on legacy behavior. Fixes the AllReduce-NVL timeout introduced by the migration by bounding `P2pNvlTransportDevice::pipeline_window()` to the per-channel staging ring capacity on the fixed-channel path. This prevents a single send from wrapping its channel ring and waiting on `slot_free` before the peer has entered recv. Caps fixed-channel NVL per-channel staging at 128MB and collapses `ctranPipesNvlMaxNumChannels()` to the real grid sizing inputs: `NCCL_CTRAN_MAX_NBLOCKS`, plus `NCCL_CTRAN_HIER_AG_NVL_NUM_BLOCKS` when hierarchical all-gather overlap is enabled. The hier-AG term is intentionally retained because hier-AG can launch more tile NVL groups than `NCCL_CTRAN_MAX_NBLOCKS`; dropping it would leave a `group.total_groups > max_num_channels` device-trap path. The Triton send/recv wrapper intentionally removes the `active_blocks` argument to match the fixed-channel NVL API. IBGDA still defaults to `active_blocks = 0` through that wrapper. Differential Revision: D110094643
Summary:
D1 (D110094643) migrated every production caller of the SIMPLE-protocol cooperative path (`send_group`/`recv_group`/`forward_group` on `P2pNvlTransportDevice`) over to the per-channel tile path (`send`/`recv`/`forward`). This diff deletes the now-unused cooperative path and all of its supporting infrastructure.
Deletions on `P2pNvlTransportDevice`:
- `send_group`, `recv_group`, `forward_group` method bodies
- `receiverStateBuffer` + `senderStateBuffer` fields on `LocalState` and `RemoteState`
- `useDualStateBuffer` field on `P2pNvlTransportOptions` and its (long) dual/single state-machine doc block
Deletions on the host transport (`MultiPeerNvlTransport`):
- `useDualStateBuffer` field on `MultiPeerNvlTransportConfig`
- `stateBufferHandler_` allocation + the dual/single `LocalState`/`RemoteState` selection branch in `getP2pTransportDevice` (collapses to a single literal without the chunk-state fields)
- `perPeerChunkStateBufferSize_` member
Deletions on the synchronisation primitive itself:
- `fbcode/comms/prims/core/ChunkState.cuh` file
- `chunk_state` BUCK rule + references from sibling rules
Cross-tree plumbing removed:
- `pipesUseDualStateBuffer` from `ctranPipesConfig` (`fbcode/comms/ctran/CtranComm.h`)
- `NCCL_CTRAN_PIPES_USE_DUAL_STATE_BUFFER` cvar read in `CtranPipes.cc` + the cvar declaration in `nccl_cvars.yaml`
- `pipesUseDualStateBuffer` plumb-through in `NcclxConfig.{h,cc}` + `MetaFactory.cc`
- ChunkState references in `CtranAlgo.cc` (`getPerPeerChunkStatesSize`, `remoteChunkStatesMap`/`localChunkStatesMap`, `DeviceSpan<ChunkState>` constructions for `receiverStateBuffer`)
- Mention of `pipesUseDualStateBuffer` in `comms/ncclx/meta/design_docs/per_comm_config_summary.md`
Triton extern wrapper:
- `torchcomms_transport_send_groups` + `torchcomms_transport_recv_groups` (+ their .h decls, Python `_send_groups_extern`/`_recv_groups_extern` bindings, `send_groups`/`recv_groups` jit wrappers, and the corresponding `transport_sendrecv_kernel` + `test_transport_sendrecv_ring_*` tests) — wholesale removed since wrapping cooperative semantics around tile `send`/`recv` would be a behavioural change rather than a no-op rename, and no in-tree caller remains.
Test surface:
- Cooperative-path GTests in `P2pNvlTransportTest.{cc,cu}` (testSend/testRecv/testSendRecv/testRecvSend/testWeightedPartition/testForward + all parametrized `useDualStateBuffer` rows) deleted; the tile-path GTests added in D1 cover the new API.
- Cooperative-path tests in `MultiPeerNvlTransportIntegrationTest`, `MultiPeerTransportKernelTest`, `PipesTransportApiTest` deleted/trimmed (LL128 sites retained).
- `BarrierTest.cc`, `P2pNvlTransportDeviceTest.cc`, `ExternalStagingBuffersTest.cc`, `BarrierBench.cu` swept to drop now-removed fields from `LocalState`/`RemoteState` literals.
- `comms/prims/benchmarks/P2pSyncBench.{cc,cu,cuh}` benchmark binary deleted entirely (its only purpose was benchmarking the deleted `ChunkState`).
Doc:
- `tile_sendrecv.md`: dual-state mode + ChunkState references removed; NVL channel-state and tile API content retained.
Out of scope (deferred to a follow-up diff):
- LL128 cooperative variants (`ll128_*_group`) — all callers + tests intact
- LL path (`ll_send`/`ll_recv`/`ll_forward`) — intact
- Pipes-library → Prims-library naming sweep (some artefacts of an in-progress rename were dropped from this diff so it stays focused on cooperative-path removal)
Reviewed By: saifhhasan
Differential Revision: D110109646
Summary: Add the shared IB fixed-channel vocabulary and config helpers that will be used by the IBGDA and IBRC channel migration. This introduces channel/direction QP state structs and explicit fixed-channel config fields while preserving existing transport behavior. Reviewed By: saifhhasan Differential Revision: D110239640
Summary: Normalize IB sendrecv configuration onto fixed channel geometry, allocate IBGDA main and companion QPs by channel and direction, and move IBGDA device-side lane selection and flush bookkeeping from physical block state to channel-direction state. Existing active_blocks arguments are retained as compatibility shims but no longer change staging layout; IBRC gets direction-aware method shims so shared sendrecv code continues to compile until the proxy backend is migrated in the next diff. Reviewed By: saifhhasan Differential Revision: D110242335
Summary: Migrate the IBRC CPU-proxy transport to the fixed IB channel resource model. Host-side IBRC now allocates command queues and main QPs using `max_num_channels * 2 * qpsPerConnection` per NIC, stores per-peer `IbLocalChannel` state for device-side QP cursors, and publishes fixed-channel geometry into the P2P device handle. Device-side IBRC now selects proxy queues from `(group_id, IbDirection, NIC, qp_index)`, uses Send as the default raw API direction, and lets send/recv internals route recv-side `SLOT_FREE` operations through the Recv direction. The proxy descriptor ABI stays unchanged; the GPU still selects the queue and the CPU proxy executes the queue's bound QP. Reviewed By: saifhhasan Differential Revision: D110244557
goelayu
pushed a commit
to goelayu/torchcomms
that referenced
this pull request
Jul 7, 2026
Summary: Pull Request resolved: meta-pytorch#3058 Direct all-pairs IB reduce-scatter for GB300, cleaned down to the production-ready collective implementation and validation surface. This diff adds the direct IB reduce-scatter launch path through `launch_direct_reduce_scatter_ib` and the CUDA kernel `direct_reduce_scatter_ib_kernel`. The kernel maps one CTA to one IB channel, splits the CTA into send and receive groups, and uses the non-smem `TileReduceStaged` receive/reduce path as the parent baseline. `CopyOp.cuh` preserves the existing `TileReduce` policy and adds `TileReduceStaged` as a separate policy for this direct IB path. The retained benchmark is one reproducible regression source: `DirectIbReduceScatterBenchmark.DirectIbVsNcclRingSweep`, comparing NCCL Ring+Simple, NCCL PAT, and RSDirect from 1KB through 2GB with the fixed 8-block GB300 budget and the default RSDirect tuning. The MCCL adapter / nccl-test-suite integration remains useful follow-up work, but this diff keeps the standalone benchmark so future kernel changes can reproduce the direct-collective numbers without depending on MCCL adapter plumbing. The shared-memory / `cp.async` receive pipeline is intentionally excluded from this parent diff and will be measured in a child diff so its correctness and performance delta stay isolated. Differential Revision: D110544920
goelayu
force-pushed
the
export-D110544920
branch
from
5080f09 to
76ce2f3
Compare
goelayu
pushed a commit
to goelayu/torchcomms
that referenced
this pull request
Jul 7, 2026
Summary: Pull Request resolved: meta-pytorch#3058 Direct all-pairs IB reduce-scatter for GB300, cleaned down to the production-ready collective implementation and validation surface. This diff adds the direct IB reduce-scatter launch path through `launch_direct_reduce_scatter_ib` and the CUDA kernel `direct_reduce_scatter_ib_kernel`. The kernel maps one CTA to one IB channel, splits the CTA into send and receive groups, and uses the non-smem `TileReduceStaged` receive/reduce path as the parent baseline. `CopyOp.cuh` preserves the existing `TileReduce` policy and adds `TileReduceStaged` as a separate policy for this direct IB path. The retained benchmark is one reproducible regression source: `DirectIbReduceScatterBenchmark.DirectIbVsNcclRingSweep`, comparing NCCL Ring+Simple, NCCL PAT, and RSDirect from 1KB through 2GB with the fixed 8-block GB300 budget and the default RSDirect tuning. The MCCL adapter / nccl-test-suite integration remains useful follow-up work, but this diff keeps the standalone benchmark so future kernel changes can reproduce the direct-collective numbers without depending on MCCL adapter plumbing. The shared-memory / `cp.async` receive pipeline is intentionally excluded from this parent diff and will be measured in a child diff so its correctness and performance delta stay isolated. Differential Revision: D110544920
goelayu
force-pushed
the
export-D110544920
branch
from
76ce2f3 to
b61acf4
Compare
Summary: Pull Request resolved: meta-pytorch#3058 Direct all-pairs IB reduce-scatter for GB300, cleaned down to the production-ready collective implementation and validation surface. This diff adds the direct IB reduce-scatter launch path through `launch_direct_reduce_scatter_ib` and the CUDA kernel `direct_reduce_scatter_ib_kernel`. The kernel maps one CTA to one IB channel, splits the CTA into send and receive groups, and uses the non-smem `TileReduceStaged` receive/reduce path as the parent baseline. `CopyOp.cuh` preserves the existing `TileReduce` policy and adds `TileReduceStaged` as a separate policy for this direct IB path. The retained benchmark is one reproducible regression source: `DirectIbReduceScatterBenchmark.DirectIbVsNcclRingSweep`, comparing NCCL Ring+Simple, NCCL PAT, and RSDirect from 1KB through 2GB with the fixed 8-block GB300 budget and the default RSDirect tuning. The MCCL adapter / nccl-test-suite integration remains useful follow-up work, but this diff keeps the standalone benchmark so future kernel changes can reproduce the direct-collective numbers without depending on MCCL adapter plumbing. The shared-memory / `cp.async` receive pipeline is intentionally excluded from this parent diff and will be measured in a child diff so its correctness and performance delta stay isolated. Reviewed By: saifhhasan Differential Revision: D110544920
goelayu
force-pushed
the
export-D110544920
branch
from
b61acf4 to
c0f479f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Direct all-pairs IB reduce-scatter for GB300, cleaned down to the production-ready collective implementation and validation surface.
This diff adds the direct IB reduce-scatter launch path through
launch_direct_reduce_scatter_iband the CUDA kerneldirect_reduce_scatter_ib_kernel. The kernel maps one CTA to one IB channel, splits the CTA into send and receive groups, and uses the non-smemTileReduceStagedreceive/reduce path as the parent baseline.CopyOp.cuhpreserves the existingTileReducepolicy and addsTileReduceStagedas a separate policy for this direct IB path.The retained benchmark is one reproducible regression source:
DirectIbReduceScatterBenchmark.DirectIbVsNcclRingSweep, comparing NCCL Ring+Simple, NCCL PAT, and RSDirect from 1KB through 2GB with the fixed 8-block GB300 budget and the default RSDirect tuning. The MCCL adapter / nccl-test-suite integration remains useful follow-up work, but this diff keeps the standalone benchmark so future kernel changes can reproduce the direct-collective numbers without depending on MCCL adapter plumbing.The shared-memory /
cp.asyncreceive pipeline is intentionally excluded from this parent diff and will be measured in a child diff so its correctness and performance delta stay isolated.Reviewed By: saifhhasan
Differential Revision: D110544920