[Feature Preview] Introduce one-sided API with ncclPut

[marksantesson]

This is an Early Access release. It contains new features and APIs that may be changed before the official General Access release. In particular, the host one-sided API is not frozen. It also may contain performance degradation that is not representative of the General Access release.

Host One-Sided RMA API

• Introduces one-sided ncclPut, ncclSignal and ncclWaitSignal APIs. Please refer to src/nccl.h.in for API definitions.
• The one-sided API support both intra-node (NVL) and inter-node (network) operations
• The implementation of the one-sided APIs does not use SM (zero-SM)
• Put/WaitSignal over network can achieve ~25% latency reduction compared to send/recv.
• The API signature, functionalities and final performance are subject to change for the official release.

Support Requirements

• The one-sided API requires NCCL window registration of the user buffers
• To enable one-sided API over the network, please set
  • NCCL_GIN_TYPE=2
  • NCCL_NET=IB
• Current one-sided APIs do not work with CUDA graph capture

Example

The example below demonstrate a ping-pong communication pattern using the ncclPut and ncclWaitSignal between two ranks.

// Configure NCCL with one-sided RMA support
ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
config.numRmaCtx = 1;  // Enable RMA with 1 context
config.blocking = 1;
NCCLCHECK(ncclCommInitRankConfig(&comm, nRanks, id, myRank, &config));

int ctx = 0;  // Use context 0 for RMA operations

// Allocate symmetric memory for RMA operations
void *sendbuff, *recvbuff;
NCCLCHECK(ncclMemAlloc((void**)&sendbuff, size));
NCCLCHECK(ncclMemAlloc((void**)&recvbuff, size));

// Register both send and receive buffers as symmetric windows for RMA operations
ncclWindow_t sendWindow, recvWindow;
NCCLCHECK(ncclCommWindowRegister(comm, sendbuff, args.end_size, &sendWindow, NCCL_WIN_COLL_SYMMETRIC));
NCCLCHECK(ncclCommWindowRegister(comm, recvbuff, args.end_size, &recvWindow, NCCL_WIN_COLL_SYMMETRIC));

// Ensure all ranks have completed window registration before proceeding
MPICHECK(MPI_Barrier(MPI_COMM_WORLD));

if (myRank == 0) {  // Rank 0: wait then put
    // Wait for signal from peer
    int nsignals = 1;
    int peer = 1;
    NCCLCHECK(ncclWaitSignal(ctx, &peer, &nsignals, 1, NCCL_SIGNAL_DISTINCT, comm, stream));

    // Put data with signal to peer's receive buffer
    NCCLCHECK(ncclPut(ctx, sendbuff, nelems, ncclInt, peer, 0, recvWindow,
                        NCCL_SIGNAL_DISTINCT, comm, stream));

} else {   // Rank 1: put then wait
    // Put data with signal to peer's receive buffer
    int peer = 0;
    NCCLCHECK(ncclPut(ctx, sendbuff, nelems, ncclInt, peer, 0, recvWindow,
                        NCCL_SIGNAL_DISTINCT, comm, stream));

    // Wait for signal from peer
    int nsignals = 1;
    NCCLCHECK(ncclWaitSignal(ctx, &peer, &nsignals, 1, NCCL_SIGNAL_DISTINCT, comm, stream));

}
NCCLCHECK(ncclCommWindowDeregister(comm, sendWindow));
NCCLCHECK(ncclCommWindowDeregister(comm, recvWindow));
NCCLCHECK(ncclMemFree(sendbuff));
NCCLCHECK(ncclMemFree(recvbuff));
CUDACHECK(cudaStreamDestroy(stream));
NCCLCHECK(ncclCommFinalize(comm));
NCCLCHECK(ncclCommDestroy(comm));