Trait CommBackend 

pub trait CommBackend: Send + Sync {
    // Required methods
    fn all_reduce(&self, ptr: u64, bytes: usize) -> Result<()>;
    fn all_gather(
        &self,
        send_ptr: u64,
        recv_ptr: u64,
        bytes: usize,
    ) -> Result<()>;
    fn reduce_scatter(
        &self,
        send_ptr: u64,
        recv_ptr: u64,
        bytes: usize,
    ) -> Result<()>;
    fn broadcast(&self, ptr: u64, bytes: usize, root: usize) -> Result<()>;
    fn barrier(&self) -> Result<()>;
    fn send_to(
        &self,
        ptr: u64,
        bytes: usize,
        dest_rank: usize,
        stream: u64,
    ) -> Result<()>;
    fn recv_from(
        &self,
        ptr: u64,
        bytes: usize,
        src_rank: usize,
        stream: u64,
    ) -> Result<()>;
    fn rank(&self) -> usize;
    fn world_size(&self) -> usize;

    // Provided methods
    fn all_reduce_async(
        &self,
        ptr: u64,
        bytes: usize,
        compute_stream: u64,
    ) -> Result<()> { ... }
    fn register_buffer(&self, _ptr: u64, _bytes: usize) -> Result<u64> { ... }
    fn deregister_buffer(&self, _handle: u64) -> Result<()> { ... }
    fn symmetric_alloc(&self, _bytes: usize) -> Result<u64> { ... }
    fn symmetric_free(&self, _ptr: u64) -> Result<()> { ... }
    fn set_add_kernel(&self, _handle: u64) { ... }
    fn group_start(&self) -> Result<()> { ... }
    fn group_end(&self) -> Result<()> { ... }
    fn is_healthy(&self) -> bool { ... }
    fn attempt_reconnect(&self) -> Result<()> { ... }
}

Communication backend trait for distributed operations.

All collective operations take raw device pointer + byte count. The pointer type is u64 (matching CUDA CUdeviceptr) to avoid coupling this crate to spark-runtime’s DevicePtr.

Required Methods

fn all_reduce(&self, ptr: u64, bytes: usize) -> Result<()>

All-reduce: sum across all ranks, result on all ranks.

fn all_gather(&self, send_ptr: u64, recv_ptr: u64, bytes: usize) -> Result<()>

All-gather: each rank contributes a chunk, all ranks get full buffer.

fn reduce_scatter( &self, send_ptr: u64, recv_ptr: u64, bytes: usize, ) -> Result<()>

Reduce-scatter: reduce + scatter (inverse of all-gather).

fn broadcast(&self, ptr: u64, bytes: usize, root: usize) -> Result<()>

Broadcast from root rank to all ranks.

fn barrier(&self) -> Result<()>

Barrier: block until all ranks reach this point.

fn send_to( &self, ptr: u64, bytes: usize, dest_rank: usize, stream: u64, ) -> Result<()>

Send tokens to a specific rank (for EP token dispatch).

Sends bytes from ptr on this rank to dest_rank. Must be paired with a matching recv_from on the destination rank. stream is the CUDA stream on which the operation is enqueued.

fn recv_from( &self, ptr: u64, bytes: usize, src_rank: usize, stream: u64, ) -> Result<()>

Receive tokens from a specific rank (for EP token combine).

Receives bytes into ptr on this rank from src_rank. Must be paired with a matching send_to on the source rank. stream is the CUDA stream on which the operation is enqueued.

fn rank(&self) -> usize

This rank’s index (0-based).

fn world_size(&self) -> usize

Total number of ranks.

Provided Methods

fn all_reduce_async( &self, ptr: u64, bytes: usize, compute_stream: u64, ) -> Result<()>

Async all-reduce using GPU-side event synchronization.

Replaces gpu.synchronize(stream) + all_reduce(ptr, bytes). Uses a dedicated comm stream + CUDA events so the CPU never blocks. compute_stream is where MoE kernels ran and where residual_add will run.

fn register_buffer(&self, _ptr: u64, _bytes: usize) -> Result<u64>

Pre-register a GPU buffer with the communication backend.

For NCCL over IB/RoCE, this caches the IB memory registration (ibv_reg_mr), avoiding per-call overhead in all_reduce. Returns an opaque handle for deregistration.

fn deregister_buffer(&self, _handle: u64) -> Result<()>

Deregister a previously registered buffer.

fn symmetric_alloc(&self, _bytes: usize) -> Result<u64>

Allocate a GPU buffer in NCCL’s symmetric-memory window (NCCL ≥ 2.28 / ncclMemAlloc). Returns the device pointer as u64.

Symmetric-memory allocations are the substrate for:

1. Copy-engine offload of NVLink collectives (frees SMs for compute).
2. Device-side communication API (kernels invoke collectives in-kernel), which TokenWeave-style fused AR+RMSNorm+Residual builds on.

On Atlas’s 2-rank Spark over RoCE, the copy-engine offload itself does not apply (RoCE is not NVLink), but the symmetric windows are still required for future device-API fusions and to reduce per-call setup. Returns an error if the linked NCCL is < 2.28; backends that don’t support symmetric memory return Err and callers must fall back.

fn symmetric_free(&self, _ptr: u64) -> Result<()>

Free a buffer previously returned by symmetric_alloc.

fn set_add_kernel(&self, _handle: u64)

Provide a kernel handle for the BF16 in-place addition kernel.

Used by the 2-rank send/recv all-reduce path. The kernel is loaded by the model layer (which has access to AtlasRegistry) and passed to the comm backend at init time.

fn group_start(&self) -> Result<()>

Begin a group of point-to-point operations (send_to/recv_from).

All send_to/recv_from calls between group_start and group_end are batched into a single NCCL launch for efficiency.

fn group_end(&self) -> Result<()>

End a group of point-to-point operations.

fn is_healthy(&self) -> bool

Check if the communicator is healthy (no async errors, no timeouts).

Returns true if the communicator is operational. Implementations may actively probe the underlying transport (e.g., ncclCommGetAsyncError).

fn attempt_reconnect(&self) -> Result<()>

Attempt to recover a degraded communicator.

For NCCL, this aborts the dead communicator and re-initializes via TCP bootstrap. Both ranks must call this concurrently. Returns Ok(()) on successful recovery, Err if recovery failed.

Implementors

impl CommBackend for NcclBackend

impl CommBackend for SingleGpuBackend