distributed: NCCL P2P primitives (PpContext + send/recv)

Add ncclSend/ncclRecv FFI and a PpContext that initializes a NCCL
communicator across P pipeline stages and hands the hidden state to
neighbour stages on the null stream. Mirrors TpContext; the collective
differs (point-to-point hand-off vs in-layer AllReduce).

tests/sendrecv.rs: 2-GPU stage0->stage1 send/recv smoke test.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-05-29 18:45:42 +08:00
parent c2362df1f1
commit 859c0cc0b6
3 changed files with 143 additions and 0 deletions

View File

@@ -45,6 +45,23 @@ unsafe extern "C" {
comm: NcclComm,
stream: CudaStream,
) -> i32;
// Point-to-point primitives for pipeline parallelism (Phase 18).
pub fn ncclSend(
sendbuff: *const c_void,
count: usize,
datatype: i32,
peer: i32,
comm: NcclComm,
stream: CudaStream,
) -> i32;
pub fn ncclRecv(
recvbuff: *mut c_void,
count: usize,
datatype: i32,
peer: i32,
comm: NcclComm,
stream: CudaStream,
) -> i32;
pub fn ncclGroupStart() -> i32;
pub fn ncclGroupEnd() -> i32;
pub fn ncclGetErrorString(result: i32) -> *const c_char;