spark_comm/
nccl.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Raw NCCL FFI bindings (minimal surface for EP all-reduce + broadcast).
4//!
5//! Only the functions Atlas actually calls are bound here — no attempt
6//! at complete coverage. Type sizes match NCCL 2.28+ on aarch64
7//! (symmetric memory `ncclMemAlloc`/`ncclMemFree` require NCCL ≥ 2.28).
8//!
9//! ## Safety
10//!
11//! Every `unsafe` block wraps a single NCCL FFI call. Invariants:
12//! - The `NcclComm`/`NcclUniqueId` arguments come from prior `Self`
13//!   constructors that called the matching NCCL init function.
14//! - GPU buffers are valid `DevicePtr`s alive on the device that owns
15//!   the comm.
16//! - Counts × dtype-size match the buffer byte count.
17//! - The `extern "C"` declarations match the NCCL header ABI for the
18//!   linked library version.
19
20use std::ffi::c_void;
21
22/// Opaque NCCL communicator handle.
23pub type NcclComm = *mut c_void;
24
25/// NCCL unique ID for bootstrapping (128 bytes, passed by value in C ABI).
26#[repr(C)]
27#[derive(Copy, Clone)]
28pub struct NcclUniqueId {
29    pub internal: [u8; 128],
30}
31
32/// NCCL config for ncclCommInitRankConfig.
33/// Set blocking=0 for non-blocking mode (ncclGroupEnd returns immediately).
34#[repr(C)]
35pub struct NcclConfig {
36    /// Size of this struct (for versioning).
37    pub size: usize,
38    /// Magic number (0x4d43434e = "NCCM" in LE for NCCL 2.27+).
39    pub magic: u32,
40    /// Version (NCCL_VERSION_CODE).
41    pub version: u32,
42    /// 1 = blocking (default), 0 = non-blocking.
43    pub blocking: i32,
44    /// CGA cluster size (0 = default).
45    pub cga_cluster_size: i32,
46    /// Min CTAs for launch (0 = default).
47    pub min_ctas: i32,
48    /// Max CTAs for launch (0 = default).
49    pub max_ctas: i32,
50    /// Network name (null-terminated, or all zeros for default).
51    pub net_name: [u8; 8],
52    /// Split share (0 = default).
53    pub split_share: i32,
54}
55
56impl NcclConfig {
57    /// Create a default config with non-blocking mode enabled.
58    /// Must match NCCL_CONFIG_INITIALIZER from nccl.h:
59    ///   { sizeof(ncclConfig_t), 0x4e43434c, NCCL_VERSION(MAJOR,MINOR,PATCH),
60    ///     NCCL_CONFIG_UNDEF_INT, ... }
61    pub fn non_blocking() -> Self {
62        // NCCL_CONFIG_UNDEF_INT = -1 means "use default"
63        Self {
64            size: std::mem::size_of::<Self>(),
65            magic: 0x4e43434c,    // "NCCL" in LE (not "NCCM")
66            version: 22907,       // NCCL 2.29.7 (major*10000 + minor*100 + patch)
67            blocking: 0,          // NON-BLOCKING
68            cga_cluster_size: -1, // NCCL_CONFIG_UNDEF_INT
69            min_ctas: -1,
70            max_ctas: -1,
71            net_name: [0; 8],
72            split_share: -1,
73        }
74    }
75}
76
77/// NCCL result code.
78#[repr(C)]
79#[derive(Debug, Copy, Clone, PartialEq, Eq)]
80pub enum NcclResult {
81    Success = 0,
82    UnhandledCudaError = 1,
83    SystemError = 2,
84    InternalError = 3,
85    InvalidArgument = 4,
86    InvalidUsage = 5,
87    RemoteError = 6,
88    InProgress = 7,
89}
90
91/// NCCL data types (matches nccl.h enum values).
92#[repr(C)]
93#[derive(Debug, Copy, Clone)]
94#[allow(dead_code)]
95pub enum NcclDataType {
96    Int8 = 0,
97    Uint8 = 1,
98    Int32 = 2,
99    Uint32 = 3,
100    Int64 = 4,
101    Uint64 = 5,
102    Float16 = 6,
103    Float32 = 7,
104    Float64 = 8,
105    Bfloat16 = 9,
106}
107
108/// NCCL reduction operations.
109#[repr(C)]
110#[derive(Debug, Copy, Clone)]
111#[allow(dead_code)]
112pub enum NcclRedOp {
113    Sum = 0,
114    Prod = 1,
115    Max = 2,
116    Min = 3,
117    Avg = 4,
118}
119
120#[link(name = "nccl")]
121unsafe extern "C" {
122    pub fn ncclGetUniqueId(id: *mut NcclUniqueId) -> NcclResult;
123
124    pub fn ncclCommInitRank(
125        comm: *mut NcclComm,
126        nranks: i32,
127        id: NcclUniqueId,
128        rank: i32,
129    ) -> NcclResult;
130
131    /// Non-blocking variant: set config.blocking=0 so ncclGroupEnd returns immediately.
132    /// Poll ncclCommGetAsyncError to check completion and detect hangs with timeout.
133    pub fn ncclCommInitRankConfig(
134        comm: *mut NcclComm,
135        nranks: i32,
136        id: NcclUniqueId,
137        rank: i32,
138        config: *const NcclConfig,
139    ) -> NcclResult;
140
141    pub fn ncclAllReduce(
142        sendbuf: *const c_void,
143        recvbuf: *mut c_void,
144        count: usize,
145        datatype: NcclDataType,
146        op: NcclRedOp,
147        comm: NcclComm,
148        stream: u64, // cudaStream_t
149    ) -> NcclResult;
150
151    pub fn ncclBroadcast(
152        sendbuf: *const c_void,
153        recvbuf: *mut c_void,
154        count: usize,
155        datatype: NcclDataType,
156        root: i32,
157        comm: NcclComm,
158        stream: u64,
159    ) -> NcclResult;
160
161    pub fn ncclCommDestroy(comm: NcclComm) -> NcclResult;
162
163    pub fn ncclGetErrorString(result: NcclResult) -> *const std::ffi::c_char;
164
165    // Buffer registration (pre-registers with IB HCA, avoids per-call ibv_reg_mr).
166    pub fn ncclCommRegister(
167        comm: NcclComm,
168        buff: *mut c_void,
169        size: usize,
170        handle: *mut *mut c_void,
171    ) -> NcclResult;
172
173    pub fn ncclCommDeregister(comm: NcclComm, handle: *mut c_void) -> NcclResult;
174
175    // NCCL 2.28+ symmetric memory window APIs.
176    //
177    // Buffers allocated via `ncclMemAlloc` participate in symmetric memory
178    // windows across the communicator, enabling:
179    //   1. Copy-engine offload on NVLink-connected ranks (frees SMs for
180    //      compute during AllReduce/AllGather).
181    //   2. The device-side communication API (kernels can issue collectives
182    //      directly without host round-trip), which is the substrate for
183    //      fused AllReduce+RMSNorm+Residual kernels (TokenWeave-style).
184    //
185    // For Atlas's 2-rank Spark over RoCE, the copy-engine path doesn't
186    // apply (RoCE is not NVLink), but the symmetric-memory windows are
187    // still needed to compose with future device-API fusions and reduce
188    // NCCL setup overhead via pre-registered handles.
189    //
190    // Returns `InvalidArgument` if the linked NCCL is < 2.28.
191    pub fn ncclMemAlloc(ptr: *mut *mut c_void, size: usize) -> NcclResult;
192
193    pub fn ncclMemFree(ptr: *mut c_void) -> NcclResult;
194
195    // Point-to-point (for custom 2-rank all-reduce).
196    pub fn ncclSend(
197        sendbuf: *const c_void,
198        count: usize,
199        datatype: NcclDataType,
200        peer: i32,
201        comm: NcclComm,
202        stream: u64,
203    ) -> NcclResult;
204
205    pub fn ncclRecv(
206        recvbuf: *mut c_void,
207        count: usize,
208        datatype: NcclDataType,
209        peer: i32,
210        comm: NcclComm,
211        stream: u64,
212    ) -> NcclResult;
213
214    // Collective: all-gather (each rank sends `count`, recv gets `world_size * count`).
215    pub fn ncclAllGather(
216        sendbuf: *const c_void,
217        recvbuf: *mut c_void,
218        sendcount: usize,
219        datatype: NcclDataType,
220        comm: NcclComm,
221        stream: u64,
222    ) -> NcclResult;
223
224    // Collective: reduce-scatter (send has `world_size * count`, each rank recvs `count`).
225    pub fn ncclReduceScatter(
226        sendbuf: *const c_void,
227        recvbuf: *mut c_void,
228        recvcount: usize,
229        datatype: NcclDataType,
230        op: NcclRedOp,
231        comm: NcclComm,
232        stream: u64,
233    ) -> NcclResult;
234
235    // Group API (batch multiple send/recv into one launch).
236    pub fn ncclGroupStart() -> NcclResult;
237    pub fn ncclGroupEnd() -> NcclResult;
238
239    // Health check: retrieve asynchronous errors from the communicator.
240    pub fn ncclCommGetAsyncError(comm: NcclComm, async_error: *mut NcclResult) -> NcclResult;
241
242    // Abort: destroy a communicator that is in a failed state.
243    // Unlike ncclCommDestroy, this does not block and cleans up immediately.
244    pub fn ncclCommAbort(comm: NcclComm) -> NcclResult;
245}
246
247// CUDA driver API for inter-stream synchronization (async all-reduce).
248// spark-comm already links libcuda transitively via NCCL.
249#[link(name = "cuda")]
250unsafe extern "C" {
251    fn cuStreamCreate(phStream: *mut u64, flags: u32) -> i32;
252    fn cuEventCreate(phEvent: *mut u64, flags: u32) -> i32;
253    fn cuEventRecord(hEvent: u64, hStream: u64) -> i32;
254    fn cuStreamWaitEvent(hStream: u64, hEvent: u64, flags: u32) -> i32;
255    fn cuEventDestroy_v2(hEvent: u64) -> i32;
256    fn cuStreamDestroy_v2(hStream: u64) -> i32;
257    fn cuStreamSynchronize(hStream: u64) -> i32;
258}
259
260pub fn create_stream() -> anyhow::Result<u64> {
261    let mut stream: u64 = 0;
262    let status = unsafe { cuStreamCreate(&mut stream, 1) }; // CU_STREAM_NON_BLOCKING
263    if status != 0 {
264        anyhow::bail!("cuStreamCreate failed: status {status}");
265    }
266    Ok(stream)
267}
268
269pub fn create_event() -> anyhow::Result<u64> {
270    let mut event: u64 = 0;
271    let status = unsafe { cuEventCreate(&mut event, 0x02) }; // CU_EVENT_DISABLE_TIMING
272    if status != 0 {
273        anyhow::bail!("cuEventCreate failed: status {status}");
274    }
275    Ok(event)
276}
277
278pub fn record_event(event: u64, stream: u64) -> anyhow::Result<()> {
279    let status = unsafe { cuEventRecord(event, stream) };
280    if status != 0 {
281        anyhow::bail!("cuEventRecord failed: status {status}");
282    }
283    Ok(())
284}
285
286pub fn stream_wait_event(stream: u64, event: u64) -> anyhow::Result<()> {
287    let status = unsafe { cuStreamWaitEvent(stream, event, 0) };
288    if status != 0 {
289        anyhow::bail!("cuStreamWaitEvent failed: status {status}");
290    }
291    Ok(())
292}
293
294pub fn destroy_event(event: u64) {
295    if event != 0 {
296        unsafe { cuEventDestroy_v2(event) };
297    }
298}
299
300pub fn destroy_stream(stream: u64) {
301    if stream != 0 {
302        unsafe { cuStreamDestroy_v2(stream) };
303    }
304}
305
306pub fn sync_stream(stream: u64) -> anyhow::Result<()> {
307    let status = unsafe { cuStreamSynchronize(stream) };
308    if status != 0 {
309        anyhow::bail!("cuStreamSynchronize failed: status {status}");
310    }
311    Ok(())
312}
313
314/// Allocate GPU memory backed by a symmetric memory window across the
315/// communicator. NCCL 2.28+ only — older NCCL returns `InvalidArgument`.
316///
317/// Buffers from `ncclMemAlloc` enable copy-engine collectives over NVLink
318/// and the device-side communication API. On Atlas's 2-rank Spark over
319/// RoCE, the copy-engine offload is unavailable (RoCE != NVLink), but the
320/// symmetric windows are required to compose with device-API fused kernels
321/// (TokenWeave-style AR+RMSNorm).
322///
323/// # Safety
324/// The returned pointer must be freed via [`nccl_mem_free`]. Passing the
325/// pointer to non-NCCL allocators (e.g. `cudaFree`) is undefined behavior.
326pub unsafe fn nccl_mem_alloc(size: usize) -> anyhow::Result<*mut c_void> {
327    let mut ptr: *mut c_void = std::ptr::null_mut();
328    let result = unsafe { ncclMemAlloc(&mut ptr, size) };
329    check_nccl(result, "ncclMemAlloc")?;
330    Ok(ptr)
331}
332
333/// Free a buffer previously returned by `nccl_mem_alloc`.
334///
335/// # Safety
336/// `ptr` must have been returned by [`nccl_mem_alloc`] and not yet freed.
337pub unsafe fn nccl_mem_free(ptr: *mut c_void) -> anyhow::Result<()> {
338    let result = unsafe { ncclMemFree(ptr) };
339    check_nccl(result, "ncclMemFree")
340}
341
342/// Convert NCCL result to anyhow::Result.
343pub fn check_nccl(result: NcclResult, context: &str) -> anyhow::Result<()> {
344    if result == NcclResult::Success {
345        Ok(())
346    } else {
347        let msg = unsafe {
348            let ptr = ncclGetErrorString(result);
349            if ptr.is_null() {
350                format!("{result:?}")
351            } else {
352                std::ffi::CStr::from_ptr(ptr).to_string_lossy().into_owned()
353            }
354        };
355        anyhow::bail!("NCCL error in {context}: {msg} ({result:?})")
356    }
357}
spark_comm/nccl.rs

spark_comm/
nccl.rs
