spark_comm/nccl_backend/
comm_impl.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `impl CommBackend for NcclBackend` — collective operations.
4//!
5//! See `crate::nccl_backend` module-level docs for the global SAFETY
6//! contract that applies to every `unsafe { nccl*(...) }` call here:
7//! comm/buffers/sizes/streams come from valid prior allocations on this
8//! rank's device, and the `extern "C"` ABI matches NCCL 2.28+.
9
10use anyhow::Result;
11use std::ffi::c_void;
12use std::ptr;
13use std::sync::atomic::Ordering;
14use std::time::Instant;
15
16use super::{ALL_REDUCE_DTYPE_BYTES, COLLECTIVE_TIMEOUT_SECS, NcclBackend};
17use crate::CommBackend;
18use crate::nccl::{self, NcclDataType, NcclRedOp};
19
20impl CommBackend for NcclBackend {
21    fn all_reduce(&self, ptr: u64, bytes: usize) -> Result<()> {
22        if self.world_size == 2 && self.add_kernel.load(Ordering::Relaxed) != 0 {
23            return self.all_reduce_2rank(ptr, bytes, self.legacy_stream);
24        }
25        // Fallback: ncclAllReduce reduces IN-PLACE on `ptr` and never touches
26        // `recv_buffer`, so it is not subject to the receive-buffer bound.
27        let count = bytes / ALL_REDUCE_DTYPE_BYTES;
28        let comm = *self.comm.lock();
29        let result = unsafe {
30            nccl::ncclAllReduce(
31                ptr as *const _,
32                ptr as *mut _,
33                count,
34                NcclDataType::Bfloat16,
35                NcclRedOp::Sum,
36                comm,
37                self.legacy_stream,
38            )
39        };
40        nccl::check_nccl(result, "ncclAllReduce")?;
41        self.check_async_error(comm);
42        Ok(())
43    }
44
45    fn all_reduce_async(&self, ptr: u64, bytes: usize, compute_stream: u64) -> Result<()> {
46        if self.world_size == 2 && self.add_kernel.load(Ordering::Relaxed) != 0 {
47            // Use event-based async with 2-rank send/recv path.
48            nccl::record_event(self.compute_done_event, compute_stream)?;
49            nccl::stream_wait_event(self.comm_stream, self.compute_done_event)?;
50
51            self.all_reduce_2rank(ptr, bytes, self.comm_stream)?;
52
53            nccl::record_event(self.comm_done_event, self.comm_stream)?;
54            nccl::stream_wait_event(compute_stream, self.comm_done_event)?;
55            return Ok(());
56        }
57
58        // Fallback: in-place on `ptr`, does not use `recv_buffer`.
59        let count = bytes / ALL_REDUCE_DTYPE_BYTES;
60        let comm = *self.comm.lock();
61
62        // 1. Record "MoE compute done" on compute stream
63        nccl::record_event(self.compute_done_event, compute_stream)?;
64
65        // 2. Comm stream waits for compute to finish
66        nccl::stream_wait_event(self.comm_stream, self.compute_done_event)?;
67
68        // 3. Launch all_reduce on dedicated comm stream
69        let result = unsafe {
70            nccl::ncclAllReduce(
71                ptr as *const _,
72                ptr as *mut _,
73                count,
74                NcclDataType::Bfloat16,
75                NcclRedOp::Sum,
76                comm,
77                self.comm_stream,
78            )
79        };
80        nccl::check_nccl(result, "ncclAllReduce (async)")?;
81
82        // 4. Record "all_reduce done" on comm stream
83        nccl::record_event(self.comm_done_event, self.comm_stream)?;
84
85        // 5. Compute stream waits for all_reduce before residual_add
86        nccl::stream_wait_event(compute_stream, self.comm_done_event)?;
87
88        // Check for async errors (non-blocking, just flags unhealthy).
89        self.check_async_error(comm);
90
91        Ok(())
92    }
93
94    fn register_buffer(&self, ptr: u64, bytes: usize) -> Result<u64> {
95        let mut handle: *mut c_void = ptr::null_mut();
96        let comm = *self.comm.lock();
97        let result =
98            unsafe { nccl::ncclCommRegister(comm, ptr as *mut c_void, bytes, &mut handle) };
99        nccl::check_nccl(result, "ncclCommRegister")?;
100        self.registered_handles.lock().push(handle);
101        Ok(handle as u64)
102    }
103
104    fn deregister_buffer(&self, handle: u64) -> Result<()> {
105        let comm = *self.comm.lock();
106        let result = unsafe { nccl::ncclCommDeregister(comm, handle as *mut c_void) };
107        nccl::check_nccl(result, "ncclCommDeregister")
108    }
109
110    fn symmetric_alloc(&self, bytes: usize) -> Result<u64> {
111        // Safety: matched by the caller's symmetric_free; both go through
112        // ncclMemAlloc/ncclMemFree on the live communicator's allocator.
113        let ptr = unsafe { nccl::nccl_mem_alloc(bytes) }?;
114        Ok(ptr as u64)
115    }
116
117    fn symmetric_free(&self, ptr: u64) -> Result<()> {
118        // Safety: caller guarantees `ptr` was returned by `symmetric_alloc`.
119        unsafe { nccl::nccl_mem_free(ptr as *mut c_void) }
120    }
121
122    fn set_add_kernel(&self, handle: u64) {
123        self.add_kernel.store(handle, Ordering::Relaxed);
124        tracing::info!(
125            "NCCL backend: bf16_add_inplace kernel set \
126             (2-rank send/recv enabled)"
127        );
128    }
129
130    fn all_gather(&self, send_ptr: u64, recv_ptr: u64, bytes: usize) -> Result<()> {
131        let comm = *self.comm.lock();
132        let result = unsafe {
133            nccl::ncclAllGather(
134                send_ptr as *const c_void,
135                recv_ptr as *mut c_void,
136                bytes,
137                NcclDataType::Uint8,
138                comm,
139                self.legacy_stream,
140            )
141        };
142        nccl::check_nccl(result, "ncclAllGather")?;
143        self.check_async_error(comm);
144        Ok(())
145    }
146
147    fn reduce_scatter(&self, send_ptr: u64, recv_ptr: u64, bytes: usize) -> Result<()> {
148        let comm = *self.comm.lock();
149        let result = unsafe {
150            nccl::ncclReduceScatter(
151                send_ptr as *const c_void,
152                recv_ptr as *mut c_void,
153                bytes,
154                NcclDataType::Uint8,
155                NcclRedOp::Sum,
156                comm,
157                self.legacy_stream,
158            )
159        };
160        nccl::check_nccl(result, "ncclReduceScatter")?;
161        self.check_async_error(comm);
162        Ok(())
163    }
164
165    fn broadcast(&self, ptr: u64, bytes: usize, root: usize) -> Result<()> {
166        let start = Instant::now();
167        let comm = *self.comm.lock();
168
169        // Broadcast raw bytes as Uint8
170        let result = unsafe {
171            nccl::ncclBroadcast(
172                ptr as *const _,
173                ptr as *mut _,
174                bytes,
175                NcclDataType::Uint8,
176                root as i32,
177                comm,
178                self.legacy_stream,
179            )
180        };
181        nccl::check_nccl(result, "ncclBroadcast")?;
182
183        // Synchronize to measure wall-clock time for timeout detection.
184        nccl::sync_stream(self.legacy_stream)?;
185
186        let elapsed = start.elapsed();
187        if elapsed.as_secs() >= COLLECTIVE_TIMEOUT_SECS {
188            tracing::error!(
189                "NCCL broadcast took {:.1}s (threshold: {}s) \
190                 — marking communicator unhealthy",
191                elapsed.as_secs_f64(),
192                COLLECTIVE_TIMEOUT_SECS,
193            );
194            self.unhealthy.store(true, Ordering::Release);
195        }
196
197        // Also check for async errors.
198        self.check_async_error(comm);
199
200        Ok(())
201    }
202
203    fn barrier(&self) -> Result<()> {
204        let comm = *self.comm.lock();
205        let result = unsafe {
206            nccl::ncclAllReduce(
207                ptr::null(),
208                ptr::null_mut(),
209                0,
210                NcclDataType::Float32,
211                NcclRedOp::Sum,
212                comm,
213                self.legacy_stream,
214            )
215        };
216        nccl::check_nccl(result, "barrier (ncclAllReduce count=0)")?;
217        self.check_async_error(comm);
218        Ok(())
219    }
220
221    fn send_to(&self, ptr: u64, bytes: usize, dest_rank: usize, stream: u64) -> Result<()> {
222        let comm = *self.comm.lock();
223        let result = unsafe {
224            nccl::ncclSend(
225                ptr as *const c_void,
226                bytes,
227                NcclDataType::Uint8,
228                dest_rank as i32,
229                comm,
230                stream,
231            )
232        };
233        nccl::check_nccl(result, "ncclSend (send_to)")
234    }
235
236    fn recv_from(&self, ptr: u64, bytes: usize, src_rank: usize, stream: u64) -> Result<()> {
237        let comm = *self.comm.lock();
238        let result = unsafe {
239            nccl::ncclRecv(
240                ptr as *mut c_void,
241                bytes,
242                NcclDataType::Uint8,
243                src_rank as i32,
244                comm,
245                stream,
246            )
247        };
248        nccl::check_nccl(result, "ncclRecv (recv_from)")
249    }
250
251    fn group_start(&self) -> Result<()> {
252        let result = unsafe { nccl::ncclGroupStart() };
253        nccl::check_nccl(result, "ncclGroupStart")
254    }
255
256    fn group_end(&self) -> Result<()> {
257        let result = unsafe { nccl::ncclGroupEnd() };
258        nccl::check_nccl(result, "ncclGroupEnd")
259    }
260
261    fn is_healthy(&self) -> bool {
262        if self.unhealthy.load(Ordering::Acquire) {
263            return false;
264        }
265        // Actively probe the communicator for async errors.
266        let comm = *self.comm.lock();
267        self.check_async_error(comm)
268    }
269
270    fn attempt_reconnect(&self) -> Result<()> {
271        self.reconnect_inner()
272    }
273
274    fn rank(&self) -> usize {
275        self.rank
276    }
277
278    fn world_size(&self) -> usize {
279        self.world_size
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286
287    #[test]
288    #[allow(clippy::assertions_on_constants)]
289    fn test_collective_timeout_constant() {
290        // Sanity: timeout should be a reasonable value (not 0, not too large).
291        assert!(COLLECTIVE_TIMEOUT_SECS >= 10);
292        assert!(COLLECTIVE_TIMEOUT_SECS <= 300);
293    }
294
295    // The receive-buffer capacity invariant — which replaced a
296    // `test_recv_buffer_size_sufficient` that asserted a fixed 64 MiB constant
297    // against a 4096-token prefill chunk while the shipped default was 8192 —
298    // is tested next to the code that enforces it, in `recv_buffer.rs`.
299}
spark_comm/nccl_backend/comm_impl.rs

spark_comm/nccl_backend/
comm_impl.rs
