gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

style: format Rust workspace

Browse Source
This commit is contained in:
Gahow Wang
2026-06-18 18:11:58 +08:00
parent 013465fc06
commit 531cd3fe08
57 changed files with 4045 additions and 1204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
crates/xserv-distributed/src/lib.rs
View File

@@ -9,8 +9,8 @@ pub mod ffi;
use std::ffi::c_void;
use ffi::{NcclComm, NcclUniqueId};
use xserv_cuda::device;
use xserv_cuda::GpuBuffer;
use xserv_cuda::device;
pub use ffi::NcclUniqueId as UniqueId;
@@ -55,7 +55,12 @@ impl TpContext {
"ncclCommInitRank",
);
ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
Self { rank, world, device, comm }
Self {
rank,
world,
device,
comm,
}
}
/// In-place AllReduce(sum) over `count` BF16 elements in `buf`.
@@ -127,7 +132,12 @@ impl PpContext {
"ncclCommInitRank",
);
ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
Self { stage, world, device, comm }
Self {
stage,
world,
device,
comm,
}
}
/// Send `count` BF16 elements at `ptr` to `peer`, on the null stream so it is
@@ -138,7 +148,16 @@ impl PpContext {
/// `ptr` must point to at least `count` BF16 elements of valid device memory.
pub fn send_bf16_ptr(&self, ptr: *const c_void, count: usize, peer: usize) {
ffi::check(
unsafe { ffi::ncclSend(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, launch_stream()) },
unsafe {
ffi::ncclSend(
ptr,
count,
ffi::NCCL_BF16,
peer as i32,
self.comm,
launch_stream(),
)
},
"ncclSend",
);
}
@@ -149,7 +168,16 @@ impl PpContext {
/// `ptr` must point to at least `count` BF16 elements of valid device memory.
pub fn recv_bf16_ptr(&self, ptr: *mut c_void, count: usize, peer: usize) {
ffi::check(
unsafe { ffi::ncclRecv(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, launch_stream()) },
unsafe {
ffi::ncclRecv(
ptr,
count,
ffi::NCCL_BF16,
peer as i32,
self.comm,
launch_stream(),
)
},
"ncclRecv",
);
}