Gahow Wang
859c0cc0b6
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>
83 lines
2.2 KiB
Rust
83 lines
2.2 KiB
Rust
//! Minimal NCCL FFI bindings (hand-written, like the CUDA bindings).
|
|
//! Only the collectives we need for tensor parallelism.
|
|
|
|
use std::ffi::c_void;
|
|
use std::os::raw::c_char;
|
|
use xserv_cuda::ffi::CudaStream;
|
|
|
|
/// Opaque NCCL communicator handle (`ncclComm_t`).
|
|
pub type NcclComm = *mut c_void;
|
|
|
|
/// `ncclUniqueId` is a 128-byte opaque blob shared from rank 0 to all ranks.
|
|
#[repr(C)]
|
|
#[derive(Clone, Copy)]
|
|
pub struct NcclUniqueId {
|
|
pub internal: [c_char; 128],
|
|
}
|
|
|
|
impl Default for NcclUniqueId {
|
|
fn default() -> Self {
|
|
Self { internal: [0; 128] }
|
|
}
|
|
}
|
|
|
|
// ncclDataType_t (subset)
|
|
pub const NCCL_FLOAT32: i32 = 7;
|
|
pub const NCCL_BF16: i32 = 9;
|
|
|
|
// ncclRedOp_t
|
|
pub const NCCL_SUM: i32 = 0;
|
|
|
|
// ncclResult_t
|
|
pub const NCCL_SUCCESS: i32 = 0;
|
|
|
|
unsafe extern "C" {
|
|
pub fn ncclGetUniqueId(uid: *mut NcclUniqueId) -> i32;
|
|
// ncclUniqueId is passed BY VALUE (a 128-byte struct) per the NCCL ABI.
|
|
pub fn ncclCommInitRank(comm: *mut NcclComm, nranks: i32, commid: NcclUniqueId, rank: i32) -> i32;
|
|
pub fn ncclCommDestroy(comm: NcclComm) -> i32;
|
|
pub fn ncclAllReduce(
|
|
sendbuff: *const c_void,
|
|
recvbuff: *mut c_void,
|
|
count: usize,
|
|
datatype: i32,
|
|
op: i32,
|
|
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;
|
|
}
|
|
|
|
pub fn err_string(result: i32) -> String {
|
|
unsafe {
|
|
let p = ncclGetErrorString(result);
|
|
if p.is_null() {
|
|
return format!("nccl error {result}");
|
|
}
|
|
std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned()
|
|
}
|
|
}
|
|
|
|
pub fn check(result: i32, what: &str) {
|
|
assert_eq!(result, NCCL_SUCCESS, "{what} failed: {}", err_string(result));
|
|
}
|