metrale_comm/nccl_backend/
comm_impl.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `impl CommBackend for NcclBackend`, plus the broadcast wait and
4//! the per-submission check. The `unsafe` contract is the one in the
5//! `nccl_backend` module header.
6//!
7//! Owner: metrale-comm.
8//! Invariants: none beyond the types.
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 metrale_gpu_sys::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        // 2026-09-26: In place on `ptr`; `recv_buffer` is not used, so its
28        // capacity does not bound this path.
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            nccl::record_event(self.compute_done_event, compute_stream)?;
57            nccl::stream_wait_event(self.comm_stream, self.compute_done_event)?;
58
59            self.all_reduce_2rank(ptr, bytes, self.comm_stream)?;
60
61            nccl::record_event(self.comm_done_event, self.comm_stream)?;
62            nccl::stream_wait_event(compute_stream, self.comm_done_event)?;
63            return Ok(());
64        }
65
66        // 2026-09-26: In place on `ptr`; `recv_buffer` is not used.
67        let count = bytes / ALL_REDUCE_DTYPE_BYTES;
68        let comm = *self.comm.lock();
69
70        nccl::record_event(self.compute_done_event, compute_stream)?;
71
72        nccl::stream_wait_event(self.comm_stream, self.compute_done_event)?;
73
74        let result = unsafe {
75            nccl::ncclAllReduce(
76                ptr as *const _,
77                ptr as *mut _,
78                count,
79                NcclDataType::Bfloat16,
80                NcclRedOp::Sum,
81                comm,
82                self.comm_stream,
83            )
84        };
85        nccl::check_nccl(result, "ncclAllReduce (async)")?;
86
87        nccl::record_event(self.comm_done_event, self.comm_stream)?;
88
89        nccl::stream_wait_event(compute_stream, self.comm_done_event)?;
90
91        self.check_async_error(comm);
92
93        Ok(())
94    }
95
96    fn register_buffer(&self, ptr: u64, bytes: usize) -> Result<u64> {
97        let mut handle: *mut c_void = ptr::null_mut();
98        let comm = *self.comm.lock();
99        let result =
100            unsafe { nccl::ncclCommRegister(comm, ptr as *mut c_void, bytes, &mut handle) };
101        nccl::check_nccl(result, "ncclCommRegister")?;
102        self.registered_handles.lock().push(handle);
103        Ok(handle as u64)
104    }
105
106    fn deregister_buffer(&self, handle: u64) -> Result<()> {
107        let comm = *self.comm.lock();
108        let result = unsafe { nccl::ncclCommDeregister(comm, handle as *mut c_void) };
109        nccl::check_nccl(result, "ncclCommDeregister")
110    }
111
112    fn symmetric_alloc(&self, bytes: usize) -> Result<u64> {
113        // 2026-09-26: SAFETY: `nccl_mem_alloc` only requires that the pointer
114        // is later freed with `nccl_mem_free`, which `symmetric_free` does.
115        let ptr = unsafe { nccl::nccl_mem_alloc(bytes) }?;
116        Ok(ptr as u64)
117    }
118
119    fn symmetric_free(&self, ptr: u64) -> Result<()> {
120        // 2026-09-26: SAFETY: The trait contract: `ptr` came from
121        // `symmetric_alloc` and has not been freed.
122        unsafe { nccl::nccl_mem_free(ptr as *mut c_void) }
123    }
124
125    fn set_add_kernel(&self, handle: u64) {
126        self.add_kernel.store(handle, Ordering::Relaxed);
127        tracing::info!(
128            "NCCL backend: bf16_add_inplace kernel set \
129             (2-rank send/recv enabled)"
130        );
131    }
132
133    fn all_gather(&self, send_ptr: u64, recv_ptr: u64, bytes: usize) -> Result<()> {
134        self.begin_submission("all_gather", Dtype::U8, bytes, self.legacy_stream, None)?;
135        let comm = *self.comm.lock();
136        let result = unsafe {
137            nccl::ncclAllGather(
138                send_ptr as *const c_void,
139                recv_ptr as *mut c_void,
140                bytes,
141                NcclDataType::Uint8,
142                comm,
143                self.legacy_stream,
144            )
145        };
146        nccl::check_nccl(result, "ncclAllGather")?;
147        self.check_async_error(comm);
148        Ok(())
149    }
150
151    fn reduce_scatter(&self, send_ptr: u64, recv_ptr: u64, bytes: usize) -> Result<()> {
152        self.begin_submission("reduce_scatter", Dtype::U8, bytes, self.legacy_stream, None)?;
153        let comm = *self.comm.lock();
154        let result = unsafe {
155            nccl::ncclReduceScatter(
156                send_ptr as *const c_void,
157                recv_ptr as *mut c_void,
158                bytes,
159                NcclDataType::Uint8,
160                NcclRedOp::Sum,
161                comm,
162                self.legacy_stream,
163            )
164        };
165        nccl::check_nccl(result, "ncclReduceScatter")?;
166        self.check_async_error(comm);
167        Ok(())
168    }
169
170    fn broadcast(&self, ptr: u64, bytes: usize, root: usize) -> Result<()> {
171        self.broadcast_with_wait(ptr, bytes, root, false)
172    }
173
174    fn recv_command_u32(&self, ptr: u64, root: usize) -> Result<()> {
175        ensure!(
176            self.rank != root,
177            "idle command receive requires non-root rank"
178        );
179        self.broadcast_with_wait(ptr, 4, root, true)
180    }
181
182    fn barrier(&self) -> Result<()> {
183        self.begin_submission("barrier", Dtype::F32, 0, self.legacy_stream, None)?;
184        let comm = *self.comm.lock();
185        let result = unsafe {
186            nccl::ncclAllReduce(
187                ptr::null(),
188                ptr::null_mut(),
189                0,
190                NcclDataType::Float32,
191                NcclRedOp::Sum,
192                comm,
193                self.legacy_stream,
194            )
195        };
196        nccl::check_nccl(result, "barrier (ncclAllReduce count=0)")?;
197        self.check_async_error(comm);
198        Ok(())
199    }
200
201    fn send_to(&self, ptr: u64, bytes: usize, dest_rank: usize, stream: u64) -> Result<()> {
202        self.begin_submission("send", Dtype::U8, bytes, stream, Some(dest_rank))?;
203        let comm = *self.comm.lock();
204        let result = unsafe {
205            nccl::ncclSend(
206                ptr as *const c_void,
207                bytes,
208                NcclDataType::Uint8,
209                dest_rank as i32,
210                comm,
211                stream,
212            )
213        };
214        nccl::check_nccl(result, "ncclSend (send_to)")
215    }
216
217    fn recv_from(&self, ptr: u64, bytes: usize, src_rank: usize, stream: u64) -> Result<()> {
218        self.begin_submission("recv", Dtype::U8, bytes, stream, Some(src_rank))?;
219        let comm = *self.comm.lock();
220        let result = unsafe {
221            nccl::ncclRecv(
222                ptr as *mut c_void,
223                bytes,
224                NcclDataType::Uint8,
225                src_rank as i32,
226                comm,
227                stream,
228            )
229        };
230        nccl::check_nccl(result, "ncclRecv (recv_from)")
231    }
232
233    fn group_start(&self) -> Result<()> {
234        let result = unsafe { nccl::ncclGroupStart() };
235        nccl::check_nccl(result, "ncclGroupStart")
236    }
237
238    fn group_end(&self) -> Result<()> {
239        let result = unsafe { nccl::ncclGroupEnd() };
240        nccl::check_nccl(result, "ncclGroupEnd")
241    }
242
243    fn is_healthy(&self) -> bool {
244        if self.unhealthy.load(Ordering::Acquire) {
245            return false;
246        }
247        let comm = *self.comm.lock();
248        self.check_async_error(comm)
249    }
250
251    fn attempt_reconnect(&self) -> Result<()> {
252        self.reconnect_inner()
253    }
254
255    fn rank(&self) -> usize {
256        self.rank
257    }
258
259    fn world_size(&self) -> usize {
260        self.world_size
261    }
262}
263
264impl NcclBackend {
265    fn broadcast_with_wait(
266        &self,
267        ptr: u64,
268        bytes: usize,
269        root: usize,
270        idle_command: bool,
271    ) -> Result<()> {
272        self.begin_submission(
273            "broadcast",
274            Dtype::U8,
275            bytes,
276            self.legacy_stream,
277            Some(root),
278        )?;
279        let start = Instant::now();
280        let comm = *self.comm.lock();
281
282        let result = unsafe {
283            nccl::ncclBroadcast(
284                ptr as *const _,
285                ptr as *mut _,
286                bytes,
287                NcclDataType::Uint8,
288                root as i32,
289                comm,
290                self.legacy_stream,
291            )
292        };
293        nccl::check_nccl(result, "ncclBroadcast")?;
294
295        // 2026-09-26: Poll `cuStreamQuery` with a 1 ms pause, so the deadline
296        // applies while waiting, not after an unbounded synchronise.
297        let ready = || {
298            ensure!(self.check_async_error(comm), "NCCL asynchronous failure");
299            nccl::stream_ready(self.legacy_stream)
300        };
301        let pause = || std::thread::sleep(Duration::from_millis(1));
302        let completion = if idle_command {
303            crate::collective_wait::poll_idle_command(ready, pause)
304        } else {
305            crate::collective_wait::poll_completion(
306                Duration::from_secs(COLLECTIVE_TIMEOUT_SECS),
307                || start.elapsed(),
308                ready,
309                pause,
310            )
311        };
312        crate::collective_wait::poison_on_error(completion, &self.unhealthy).with_context(
313            || {
314                format!(
315                    "NCCL broadcast rank={} world_size={} root={root} bytes={bytes}; \
316                     communicator poisoned; stop all ranks before retrying",
317                    self.rank, self.world_size,
318                )
319            },
320        )?;
321
322        Ok(())
323    }
324
325    fn begin_submission(
326        &self,
327        op: &str,
328        dtype: Dtype,
329        bytes: usize,
330        stream: u64,
331        peer: Option<usize>,
332    ) -> Result<()> {
333        crate::collective_wait::ensure_healthy(&self.unhealthy, self.rank, self.world_size, op)?;
334        self.diagnostics.submit(op, dtype, bytes, stream, peer)
335    }
336}
337
338#[cfg(test)]
339mod tests {
340    use super::*;
341
342    #[test]
343    #[allow(clippy::assertions_on_constants)]
344    fn test_collective_timeout_constant() {
345        assert!(COLLECTIVE_TIMEOUT_SECS >= 10);
346        assert!(COLLECTIVE_TIMEOUT_SECS <= 300);
347    }
348}
metrale_comm/nccl_backend/comm_impl.rs

metrale_comm/nccl_backend/
comm_impl.rs
