Module nccl_backend 

2026-09-26: NcclBackend, the multi-rank CommBackend over NCCL. Bootstrap: rank 0 creates an NCCL unique id and sends it over TCP to every other rank, then all ranks call ncclCommInitRank.

At world_size == 2, once the add kernel is set, all-reduce is a grouped ncclSend/ncclRecv into a registered receive buffer plus a local BF16 add; otherwise it is ncclAllReduce. After each collective the backend queries ncclCommGetAsyncError; a broadcast waits for completion for at most COLLECTIVE_TIMEOUT_SECS, an idle command receive without a deadline, and a failure in either marks the communicator unhealthy. attempt_reconnect aborts the communicator and bootstraps again. METRALE_COMM_DIAGNOSTICS=1 logs every host submission.

Owner: metrale-comm. Invariants:

• Every collective and point-to-point operation first refuses an unhealthy communicator and validates its byte count and peer (begin_submission).
• A 2-rank send/recv all-reduce never receives more than recv_capacity bytes: all_reduce_2rank checks before any NCCL call.

Each unsafe block wraps one NCCL or CUDA driver call. The declarations follow nccl.h and cuda.h. Callers must pass device pointers valid for bytes on the communicator’s device and streams that outlive the enqueued work.

Structs

NcclBackend

2026-09-26: NCCL communicator plus the streams, events and receive buffer its operations use.

Constants

ALL_REDUCE_DTYPE_BYTES

2026-09-26: Bytes per element of the all-reduce: BF16. NcclBackend divides byte counts by it to get NCCL element counts, and required_model_recv_bytes sizes the buffer with it.

Functions

required_model_recv_bytes

2026-09-26: Receive-buffer bytes for tokens rows of BF16 activations or logits: tokens × max(hidden, vocab) × 2. Serve passes max_batch_tokens, hidden_size and vocab_size.

required_recv_bytes

2026-09-26: max_batch_tokens × hidden_size × dtype_bytes, in bytes.