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::{Context, Result, ensure};
11use std::ffi::c_void;
12use std::ptr;
13use std::sync::atomic::Ordering;
14use std::time::{Duration, Instant};
15
16use super::{ALL_REDUCE_DTYPE_BYTES, COLLECTIVE_TIMEOUT_SECS, NcclBackend};
17use crate::CommBackend;
18use crate::collective_diagnostics::Dtype;
19use crate::nccl::{self, NcclDataType, NcclRedOp};
20
21impl CommBackend for NcclBackend {
22    fn all_reduce(&self, ptr: u64, bytes: usize) -> Result<()> {
23        self.begin_submission("all_reduce", Dtype::Bf16, bytes, self.legacy_stream, None)?;
24        if self.world_size == 2 && self.add_kernel.load(Ordering::Relaxed) != 0 {
25            return self.all_reduce_2rank(ptr, bytes, self.legacy_stream);
26        }
27        // Fallback: ncclAllReduce reduces IN-PLACE on `ptr` and never touches
28        // `recv_buffer`, so it is not subject to the receive-buffer bound.
29        let count = bytes / ALL_REDUCE_DTYPE_BYTES;
30        let comm = *self.comm.lock();
31        let result = unsafe {
32            nccl::ncclAllReduce(
33                ptr as *const _,
34                ptr as *mut _,
35                count,
36                NcclDataType::Bfloat16,
37                NcclRedOp::Sum,
38                comm,
39                self.legacy_stream,
40            )
41        };
42        nccl::check_nccl(result, "ncclAllReduce")?;
43        self.check_async_error(comm);
44        Ok(())
45    }
46
47    fn all_reduce_async(&self, ptr: u64, bytes: usize, compute_stream: u64) -> Result<()> {
48        self.begin_submission(
49            "all_reduce_async",
50            Dtype::Bf16,
51            bytes,
52            self.comm_stream,
53            None,
54        )?;
55        if self.world_size == 2 && self.add_kernel.load(Ordering::Relaxed) != 0 {
56            // Use event-based async with 2-rank send/recv path.
57            nccl::record_event(self.compute_done_event, compute_stream)?;
58            nccl::stream_wait_event(self.comm_stream, self.compute_done_event)?;
59
60            self.all_reduce_2rank(ptr, bytes, self.comm_stream)?;
61
62            nccl::record_event(self.comm_done_event, self.comm_stream)?;
63            nccl::stream_wait_event(compute_stream, self.comm_done_event)?;
64            return Ok(());
65        }
66
67        // Fallback: in-place on `ptr`, does not use `recv_buffer`.
68        let count = bytes / ALL_REDUCE_DTYPE_BYTES;
69        let comm = *self.comm.lock();
70
71        // 1. Record "MoE compute done" on compute stream
72        nccl::record_event(self.compute_done_event, compute_stream)?;
73
74        // 2. Comm stream waits for compute to finish
75        nccl::stream_wait_event(self.comm_stream, self.compute_done_event)?;
76
77        // 3. Launch all_reduce on dedicated comm stream
78        let result = unsafe {
79            nccl::ncclAllReduce(
80                ptr as *const _,
81                ptr as *mut _,
82                count,
83                NcclDataType::Bfloat16,
84                NcclRedOp::Sum,
85                comm,
86                self.comm_stream,
87            )
88        };
89        nccl::check_nccl(result, "ncclAllReduce (async)")?;
90
91        // 4. Record "all_reduce done" on comm stream
92        nccl::record_event(self.comm_done_event, self.comm_stream)?;
93
94        // 5. Compute stream waits for all_reduce before residual_add
95        nccl::stream_wait_event(compute_stream, self.comm_done_event)?;
96
97        // Check for async errors (non-blocking, just flags unhealthy).
98        self.check_async_error(comm);
99
100        Ok(())
101    }
102
103    fn register_buffer(&self, ptr: u64, bytes: usize) -> Result<u64> {
104        let mut handle: *mut c_void = ptr::null_mut();
105        let comm = *self.comm.lock();
106        let result =
107            unsafe { nccl::ncclCommRegister(comm, ptr as *mut c_void, bytes, &mut handle) };
108        nccl::check_nccl(result, "ncclCommRegister")?;
109        self.registered_handles.lock().push(handle);
110        Ok(handle as u64)
111    }
112
113    fn deregister_buffer(&self, handle: u64) -> Result<()> {
114        let comm = *self.comm.lock();
115        let result = unsafe { nccl::ncclCommDeregister(comm, handle as *mut c_void) };
116        nccl::check_nccl(result, "ncclCommDeregister")
117    }
118
119    fn symmetric_alloc(&self, bytes: usize) -> Result<u64> {
120        // Safety: matched by the caller's symmetric_free; both go through
121        // ncclMemAlloc/ncclMemFree on the live communicator's allocator.
122        let ptr = unsafe { nccl::nccl_mem_alloc(bytes) }?;
123        Ok(ptr as u64)
124    }
125
126    fn symmetric_free(&self, ptr: u64) -> Result<()> {
127        // Safety: caller guarantees `ptr` was returned by `symmetric_alloc`.
128        unsafe { nccl::nccl_mem_free(ptr as *mut c_void) }
129    }
130
131    fn set_add_kernel(&self, handle: u64) {
132        self.add_kernel.store(handle, Ordering::Relaxed);
133        tracing::info!(
134            "NCCL backend: bf16_add_inplace kernel set \
135             (2-rank send/recv enabled)"
136        );
137    }
138
139    fn all_gather(&self, send_ptr: u64, recv_ptr: u64, bytes: usize) -> Result<()> {
140        self.begin_submission("all_gather", Dtype::U8, bytes, self.legacy_stream, None)?;
141        let comm = *self.comm.lock();
142        let result = unsafe {
143            nccl::ncclAllGather(
144                send_ptr as *const c_void,
145                recv_ptr as *mut c_void,
146                bytes,
147                NcclDataType::Uint8,
148                comm,
149                self.legacy_stream,
150            )
151        };
152        nccl::check_nccl(result, "ncclAllGather")?;
153        self.check_async_error(comm);
154        Ok(())
155    }
156
157    fn reduce_scatter(&self, send_ptr: u64, recv_ptr: u64, bytes: usize) -> Result<()> {
158        self.begin_submission("reduce_scatter", Dtype::U8, bytes, self.legacy_stream, None)?;
159        let comm = *self.comm.lock();
160        let result = unsafe {
161            nccl::ncclReduceScatter(
162                send_ptr as *const c_void,
163                recv_ptr as *mut c_void,
164                bytes,
165                NcclDataType::Uint8,
166                NcclRedOp::Sum,
167                comm,
168                self.legacy_stream,
169            )
170        };
171        nccl::check_nccl(result, "ncclReduceScatter")?;
172        self.check_async_error(comm);
173        Ok(())
174    }
175
176    fn broadcast(&self, ptr: u64, bytes: usize, root: usize) -> Result<()> {
177        self.broadcast_with_wait(ptr, bytes, root, false)
178    }
179
180    fn recv_command_u32(&self, ptr: u64, root: usize) -> Result<()> {
181        ensure!(
182            self.rank != root,
183            "idle command receive requires non-root rank"
184        );
185        self.broadcast_with_wait(ptr, 4, root, true)
186    }
187
188    fn barrier(&self) -> Result<()> {
189        self.begin_submission("barrier", Dtype::F32, 0, self.legacy_stream, None)?;
190        let comm = *self.comm.lock();
191        let result = unsafe {
192            nccl::ncclAllReduce(
193                ptr::null(),
194                ptr::null_mut(),
195                0,
196                NcclDataType::Float32,
197                NcclRedOp::Sum,
198                comm,
199                self.legacy_stream,
200            )
201        };
202        nccl::check_nccl(result, "barrier (ncclAllReduce count=0)")?;
203        self.check_async_error(comm);
204        Ok(())
205    }
206
207    fn send_to(&self, ptr: u64, bytes: usize, dest_rank: usize, stream: u64) -> Result<()> {
208        self.begin_submission("send", Dtype::U8, bytes, stream, Some(dest_rank))?;
209        let comm = *self.comm.lock();
210        let result = unsafe {
211            nccl::ncclSend(
212                ptr as *const c_void,
213                bytes,
214                NcclDataType::Uint8,
215                dest_rank as i32,
216                comm,
217                stream,
218            )
219        };
220        nccl::check_nccl(result, "ncclSend (send_to)")
221    }
222
223    fn recv_from(&self, ptr: u64, bytes: usize, src_rank: usize, stream: u64) -> Result<()> {
224        self.begin_submission("recv", Dtype::U8, bytes, stream, Some(src_rank))?;
225        let comm = *self.comm.lock();
226        let result = unsafe {
227            nccl::ncclRecv(
228                ptr as *mut c_void,
229                bytes,
230                NcclDataType::Uint8,
231                src_rank as i32,
232                comm,
233                stream,
234            )
235        };
236        nccl::check_nccl(result, "ncclRecv (recv_from)")
237    }
238
239    fn group_start(&self) -> Result<()> {
240        let result = unsafe { nccl::ncclGroupStart() };
241        nccl::check_nccl(result, "ncclGroupStart")
242    }
243
244    fn group_end(&self) -> Result<()> {
245        let result = unsafe { nccl::ncclGroupEnd() };
246        nccl::check_nccl(result, "ncclGroupEnd")
247    }
248
249    fn is_healthy(&self) -> bool {
250        if self.unhealthy.load(Ordering::Acquire) {
251            return false;
252        }
253        // Actively probe the communicator for async errors.
254        let comm = *self.comm.lock();
255        self.check_async_error(comm)
256    }
257
258    fn attempt_reconnect(&self) -> Result<()> {
259        self.reconnect_inner()
260    }
261
262    fn rank(&self) -> usize {
263        self.rank
264    }
265
266    fn world_size(&self) -> usize {
267        self.world_size
268    }
269}
270
271impl NcclBackend {
272    fn broadcast_with_wait(
273        &self,
274        ptr: u64,
275        bytes: usize,
276        root: usize,
277        idle_command: bool,
278    ) -> Result<()> {
279        self.begin_submission(
280            "broadcast",
281            Dtype::U8,
282            bytes,
283            self.legacy_stream,
284            Some(root),
285        )?;
286        let start = Instant::now();
287        let comm = *self.comm.lock();
288
289        // Broadcast raw bytes as Uint8
290        let result = unsafe {
291            nccl::ncclBroadcast(
292                ptr as *const _,
293                ptr as *mut _,
294                bytes,
295                NcclDataType::Uint8,
296                root as i32,
297                comm,
298                self.legacy_stream,
299            )
300        };
301        nccl::check_nccl(result, "ncclBroadcast")?;
302
303        // Poll instead of measuring time only AFTER an unbounded synchronize.
304        let ready = || {
305            ensure!(self.check_async_error(comm), "NCCL asynchronous failure");
306            nccl::stream_ready(self.legacy_stream)
307        };
308        let pause = || std::thread::sleep(Duration::from_millis(1));
309        let completion = if idle_command {
310            crate::collective_wait::poll_idle_command(ready, pause)
311        } else {
312            crate::collective_wait::poll_completion(
313                Duration::from_secs(COLLECTIVE_TIMEOUT_SECS),
314                || start.elapsed(),
315                ready,
316                pause,
317            )
318        };
319        crate::collective_wait::poison_on_error(completion, &self.unhealthy).with_context(
320            || {
321                format!(
322                    "NCCL broadcast rank={} world_size={} root={root} bytes={bytes}; \
323                     communicator poisoned; stop all ranks before retrying",
324                    self.rank, self.world_size,
325                )
326            },
327        )?;
328
329        Ok(())
330    }
331
332    fn begin_submission(
333        &self,
334        op: &str,
335        dtype: Dtype,
336        bytes: usize,
337        stream: u64,
338        peer: Option<usize>,
339    ) -> Result<()> {
340        crate::collective_wait::ensure_healthy(&self.unhealthy, self.rank, self.world_size, op)?;
341        self.diagnostics.submit(op, dtype, bytes, stream, peer)
342    }
343}
344
345#[cfg(test)]
346mod tests {
347    use super::*;
348
349    #[test]
350    #[allow(clippy::assertions_on_constants)]
351    fn test_collective_timeout_constant() {
352        // Sanity: timeout should be a reasonable value (not 0, not too large).
353        assert!(COLLECTIVE_TIMEOUT_SECS >= 10);
354        assert!(COLLECTIVE_TIMEOUT_SECS <= 300);
355    }
356
357    // The receive-buffer capacity invariant — which replaced a
358    // `test_recv_buffer_size_sufficient` that asserted a fixed 64 MiB constant
359    // against a 4096-token prefill chunk while the shipped default was 8192 —
360    // is tested next to the code that enforces it, in `recv_buffer.rs`.
361}
spark_comm/nccl_backend/comm_impl.rs

spark_comm/nccl_backend/
comm_impl.rs
