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    fn cuStreamQuery(hStream: u64) -> i32;
259}
260
261pub fn create_stream() -> anyhow::Result<u64> {
262    let mut stream: u64 = 0;
263    let status = unsafe { cuStreamCreate(&mut stream, 1) }; // CU_STREAM_NON_BLOCKING
264    if status != 0 {
265        anyhow::bail!("cuStreamCreate failed: status {status}");
266    }
267    Ok(stream)
268}
269
270pub fn create_event() -> anyhow::Result<u64> {
271    let mut event: u64 = 0;
272    let status = unsafe { cuEventCreate(&mut event, 0x02) }; // CU_EVENT_DISABLE_TIMING
273    if status != 0 {
274        anyhow::bail!("cuEventCreate failed: status {status}");
275    }
276    Ok(event)
277}
278
279pub fn record_event(event: u64, stream: u64) -> anyhow::Result<()> {
280    let status = unsafe { cuEventRecord(event, stream) };
281    if status != 0 {
282        anyhow::bail!("cuEventRecord failed: status {status}");
283    }
284    Ok(())
285}
286
287pub fn stream_wait_event(stream: u64, event: u64) -> anyhow::Result<()> {
288    let status = unsafe { cuStreamWaitEvent(stream, event, 0) };
289    if status != 0 {
290        anyhow::bail!("cuStreamWaitEvent failed: status {status}");
291    }
292    Ok(())
293}
294
295pub fn destroy_event(event: u64) {
296    if event != 0 {
297        unsafe { cuEventDestroy_v2(event) };
298    }
299}
300
301pub fn destroy_stream(stream: u64) {
302    if stream != 0 {
303        unsafe { cuStreamDestroy_v2(stream) };
304    }
305}
306
307pub fn sync_stream(stream: u64) -> anyhow::Result<()> {
308    let status = unsafe { cuStreamSynchronize(stream) };
309    if status != 0 {
310        anyhow::bail!("cuStreamSynchronize failed: status {status}");
311    }
312    Ok(())
313}
314
315/// Allocate GPU memory backed by a symmetric memory window across the
316/// communicator. NCCL 2.28+ only — older NCCL returns `InvalidArgument`.
317///
318/// Buffers from `ncclMemAlloc` enable copy-engine collectives over NVLink
319/// and the device-side communication API. On Atlas's 2-rank Spark over
320/// RoCE, the copy-engine offload is unavailable (RoCE != NVLink), but the
321/// symmetric windows are required to compose with device-API fused kernels
322/// (TokenWeave-style AR+RMSNorm).
323///
324/// # Safety
325/// The returned pointer must be freed via [`nccl_mem_free`]. Passing the
326/// pointer to non-NCCL allocators (e.g. `cudaFree`) is undefined behavior.
327pub unsafe fn nccl_mem_alloc(size: usize) -> anyhow::Result<*mut c_void> {
328    let mut ptr: *mut c_void = std::ptr::null_mut();
329    let result = unsafe { ncclMemAlloc(&mut ptr, size) };
330    check_nccl(result, "ncclMemAlloc")?;
331    Ok(ptr)
332}
333
334/// Free a buffer previously returned by `nccl_mem_alloc`.
335///
336/// # Safety
337/// `ptr` must have been returned by [`nccl_mem_alloc`] and not yet freed.
338pub unsafe fn nccl_mem_free(ptr: *mut c_void) -> anyhow::Result<()> {
339    let result = unsafe { ncclMemFree(ptr) };
340    check_nccl(result, "ncclMemFree")
341}
342
343/// Convert NCCL result to anyhow::Result.
344pub fn check_nccl(result: NcclResult, context: &str) -> anyhow::Result<()> {
345    if result == NcclResult::Success {
346        Ok(())
347    } else {
348        let msg = unsafe {
349            let ptr = ncclGetErrorString(result);
350            if ptr.is_null() {
351                format!("{result:?}")
352            } else {
353                std::ffi::CStr::from_ptr(ptr).to_string_lossy().into_owned()
354            }
355        };
356        anyhow::bail!("NCCL error in {context}: {msg} ({result:?})")
357    }
358}
359
360/// Query completion without waiting. CUDA_ERROR_NOT_READY is 600.
361/// The caller owns the stream and its device context, as for sync_stream.
362pub fn stream_ready(stream: u64) -> anyhow::Result<bool> {
363    match unsafe { cuStreamQuery(stream) } {
364        0 => Ok(true),
365        600 => Ok(false),
366        status => anyhow::bail!("cuStreamQuery failed: status {status}"),
367    }
368}
spark_comm/nccl.rs

spark_comm/
nccl.rs
