From 46aeba02e72658acb8a46935a6c64e5649e8f198 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Wed, 22 Apr 2026 10:57:35 -0700 Subject: [PATCH] First pass of sum_scatter --- mlx/distributed/jaccl/jaccl.cpp | 11 +- mlx/distributed/jaccl/lib/jaccl/group.h | 3 + mlx/distributed/jaccl/lib/jaccl/mesh.cpp | 23 +++ mlx/distributed/jaccl/lib/jaccl/mesh.h | 10 ++ mlx/distributed/jaccl/lib/jaccl/mesh_impl.h | 159 ++++++++++++++++++++ mlx/distributed/jaccl/lib/jaccl/ring.cpp | 34 +++++ mlx/distributed/jaccl/lib/jaccl/ring.h | 10 ++ mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 140 +++++++++++++++++ 8 files changed, 389 insertions(+), 1 deletion(-) diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp index 01cc415a..c059069a 100644 --- a/mlx/distributed/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/jaccl.cpp @@ -145,7 +145,16 @@ class JACCLGroup : public GroupImpl { } void sum_scatter(const array& input, array& output, Stream stream) override { - throw std::runtime_error("[jaccl] sum_scatter not supported."); + auto in_ptr = input.data(); + auto out_ptr = output.data(); + size_t n_bytes = input.nbytes(); + int dtype = dtype_to_jaccl_dtype(output.dtype()); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.set_output_array(output); + encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() { + group_->sum_scatter(in_ptr, out_ptr, n_bytes, dtype); + }); } std::shared_ptr split(int color, int key = -1) override { diff --git a/mlx/distributed/jaccl/lib/jaccl/group.h b/mlx/distributed/jaccl/lib/jaccl/group.h index 495458c5..99dc8def 100644 --- a/mlx/distributed/jaccl/lib/jaccl/group.h +++ b/mlx/distributed/jaccl/lib/jaccl/group.h @@ -28,6 +28,9 @@ class Group { virtual void all_gather(const void* input, void* output, size_t n_bytes) = 0; + virtual void + sum_scatter(const void* input, void* output, size_t n_bytes, int dtype) = 0; + virtual void send(const void* input, size_t n_bytes, int dst) = 0; virtual void recv(void* output, size_t n_bytes, int src) = 0; }; diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp index 6ea1da58..bf122cee 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp @@ -176,6 +176,17 @@ void MeshGroup::all_gather(const void* input, void* output, size_t n_bytes) { static_cast(input), static_cast(output), n_bytes); } +void MeshGroup::sum_scatter( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + reduce_scatter(input, output, n_bytes, SumOp{}); + }); +} + void MeshGroup::send(const void* input, size_t n_bytes, int dst) { mesh_.send(static_cast(input), n_bytes, dst); } @@ -184,6 +195,18 @@ void MeshGroup::recv(void* output, size_t n_bytes, int src) { mesh_.recv(static_cast(output), n_bytes, src); } +template +void MeshGroup::reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op) { + auto in_ptr = static_cast(input); + auto out_ptr = static_cast(output); + int64_t count = n_bytes / sizeof(T); + mesh_.reduce_scatter(in_ptr, out_ptr, count, reduce_op); +} + template void MeshGroup::all_reduce( const void* input, diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.h b/mlx/distributed/jaccl/lib/jaccl/mesh.h index 3184f1fb..23d03220 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh.h @@ -44,10 +44,20 @@ class MeshGroup : public Group { void all_gather(const void* input, void* output, size_t n_bytes) override; + void sum_scatter(const void* input, void* output, size_t n_bytes, int dtype) + override; + void send(const void* input, size_t n_bytes, int dst) override; void recv(void* output, size_t n_bytes, int src) override; private: + template + void reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op); + template void all_reduce( const void* input, diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index 0201f7ff..f59c57ca 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -29,6 +29,165 @@ class MeshImpl { MeshImpl() : rank_(0), size_(1) {} + template + void reduce_scatter(const T* in, T* out, int64_t size, ReduceOp reduce_op) { + // Reduce-scatter for mesh topology. + // + // Each rank sends its entire input to all other ranks. + // Each rank reduces only its assigned chunk from all received inputs. + + auto [sz, buffer_size] = buffer_size_from_message(size * sizeof(T)); + int64_t N = buffer_size / sizeof(T); + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * MESH_MAX_PEERS * 2; + int64_t total = static_cast(size); + int num_peers = size_ - 1; + + // Calculate chunk for this rank + int64_t chunk_size = (total + size_ - 1) / size_; + int64_t my_chunk_start = rank_ * chunk_size; + int64_t my_chunk_size = std::min(chunk_size, total - my_chunk_start); + + // Initialize output with our own chunk + if (my_chunk_size > 0) { + std::copy_n(in + my_chunk_start, my_chunk_size, out); + } + + // A helper for convenient access to the staging buffer. + auto local_staging = [&](int buff) -> T* { + return reinterpret_cast(staging_mem_.get() + buff * MAX_BUFFER_SIZE); + }; + + // Counters to maintain the state of transfers + int in_flight = 0; + int64_t read_offset = 0; + int completed_send_count[PIPELINE] = {0}; + int recv_end[MESH_MAX_PEERS] = {0}; + int reduce_chunk = 0; + int reduce_rank = 0; + + // Total number of chunks + int64_t total_chunks = (total + N - 1) / N; + + // Prefill the pipeline + int buff = 0; + while (read_offset < total && buff < PIPELINE) { + post_recv_all(sz, buff); + + // Copy the local data to send buffer and staging buffer + int64_t elems = std::min(N, total - read_offset); + std::copy( + in + read_offset, in + read_offset + elems, local_staging(buff)); + std::copy( + in + read_offset, + in + read_offset + elems, + send_buffer(sz, buff).begin()); + recv_end[rank_]++; + post_send_all(sz, buff); + + buff++; + in_flight += 2 * num_peers; + read_offset += N; + } + + // Main loop + while (reduce_chunk < total_chunks) { + // Poll the hardware for completions. + ibv_wc wc[WC_NUM]; + int n = poll(connections_, WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + if (work_type == SEND_WR && read_offset < total) { + completed_send_count[buff]++; + if (completed_send_count[buff] == num_peers) { + int64_t elems = std::min(N, total - read_offset); + std::copy( + in + read_offset, + in + read_offset + elems, + local_staging(buff)); + std::copy( + in + read_offset, + in + read_offset + elems, + send_buffer(sz, buff).begin()); + recv_end[rank_]++; + post_send_all(sz, buff); + + completed_send_count[buff] = 0; + in_flight += num_peers; + read_offset += N; + } + } + + else if (work_type == RECV_WR) { + recv_end[rank]++; + } + } + + // Process the received chunks in order, reducing only our chunk + while (reduce_chunk < total_chunks) { + int64_t w = static_cast(reduce_chunk) * N; + if (w >= read_offset) { + break; + } + if (recv_end[reduce_rank] <= reduce_chunk) { + break; + } + int b = reduce_chunk % PIPELINE; + int64_t elems = std::min(N, total - w); + + // Check if this chunk overlaps with our output chunk + int64_t overlap_start = std::max(w, my_chunk_start); + int64_t overlap_end = + std::min(w + elems, my_chunk_start + my_chunk_size); + + if (overlap_start < overlap_end) { + int64_t out_offset = overlap_start - my_chunk_start; + int64_t in_offset = overlap_start - w; + int64_t overlap_size = overlap_end - overlap_start; + + // Data is read from the staging area for our own rank + if (reduce_rank == rank_) { + reduce_op( + local_staging(b) + in_offset, out + out_offset, overlap_size); + } + // Data is read from the recv buffers for other ranks + else { + reduce_op( + recv_buffer(sz, b, reduce_rank).begin() + in_offset, + out + out_offset, + overlap_size); + } + } + + // Check if we need to post another receive + int64_t next_chunk = static_cast(reduce_chunk) + PIPELINE; + if (next_chunk < total_chunks) { + recv_from(sz, reduce_rank, b); + in_flight++; + } + + // Move to next rank's data for this chunk + reduce_rank++; + if (reduce_rank >= size_) { + reduce_rank = 0; + reduce_chunk++; + } + } + } + + // Drain remaining in-flight completions + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = poll(connections_, WC_NUM, wc); + in_flight -= n; + } + } + template void all_reduce(const T* in, T* out, int64_t size, ReduceOp reduce_op) { // Fully connected all reduce with deterministic reduction order. diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.cpp b/mlx/distributed/jaccl/lib/jaccl/ring.cpp index 96ddb0e0..cd285a6d 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.cpp +++ b/mlx/distributed/jaccl/lib/jaccl/ring.cpp @@ -166,6 +166,17 @@ void RingGroup::all_gather(const void* input, void* output, size_t n_bytes) { n_conns_); } +void RingGroup::sum_scatter( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + reduce_scatter(input, output, n_bytes, SumOp{}); + }); +} + void RingGroup::send(const void* input, size_t n_bytes, int dst) { int right = (rank_ + 1) % size_; int left = (rank_ + size_ - 1) % size_; @@ -190,6 +201,29 @@ void RingGroup::recv(void* output, size_t n_bytes, int src) { ring_.recv(static_cast(output), n_bytes, src, n_conns_); } +template +void RingGroup::reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op) { + auto in_ptr = static_cast(input); + auto out_ptr = static_cast(output); + int64_t count = n_bytes / sizeof(T); + if (count < size_ * 2 * n_conns_) { + ring_.reduce_scatter<1, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op); + return; + } + + if (n_bytes <= 65536) { + ring_.reduce_scatter<2, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op); + return; + } + + ring_.reduce_scatter<2, T, ReduceOp>( + in_ptr, out_ptr, count, n_conns_, reduce_op); +} + template void RingGroup::all_reduce( const void* input, diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.h b/mlx/distributed/jaccl/lib/jaccl/ring.h index f0bc1ba9..55dcf578 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring.h @@ -45,10 +45,20 @@ class RingGroup : public Group { void all_gather(const void* input, void* output, size_t n_bytes) override; + void sum_scatter(const void* input, void* output, size_t n_bytes, int dtype) + override; + void send(const void* input, size_t n_bytes, int dst) override; void recv(void* output, size_t n_bytes, int src) override; private: + template + void reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op); + template void all_reduce( const void* input, diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index 3ee91b37..230cf5d3 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -45,6 +45,146 @@ class RingImpl { RingImpl() : rank_(0), size_(1), n_conns_(0) {} + template + void reduce_scatter( + const T* in_ptr, + T* out_ptr, + int64_t size, + int n_wires, + ReduceOp reduce_op) { + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS * 2 * MAX_DIR; + int64_t chunk_size = (size + size_ - 1) / size_; + int64_t size_per_wire = + (chunk_size + (MAX_DIR * n_wires) - 1) / (MAX_DIR * n_wires); + auto [sz, N] = buffer_size_from_message(size_per_wire * sizeof(T)); + N /= sizeof(T); + int64_t n_steps = (size_per_wire + N - 1) / N; + + // Counters to maintain the state of transfers + int in_flight = 0; + int64_t chunk_multiple_size = size_ * chunk_size; + int64_t send_offset[MAX_DIR]; + int64_t recv_offset[MAX_DIR]; + int64_t send_limits[MAX_DIR]; + int64_t recv_limits[MAX_DIR]; + int send_count[MAX_DIR * RING_MAX_CONNS] = {0}; + int recv_count[MAX_DIR * RING_MAX_CONNS] = {0}; + send_offset[0] = rank_ * chunk_size; + recv_offset[0] = ((rank_ + size_ - 1) % size_) * chunk_size; + if constexpr (MAX_DIR == 2) { + send_offset[1] = rank_ * chunk_size; + recv_offset[1] = ((rank_ + 1) % size_) * chunk_size; + send_limits[0] = std::min( + n_wires * size_per_wire, std::max(0, size - send_offset[0])); + send_limits[1] = + std::min(chunk_size, std::max(0, size - send_offset[1])); + recv_limits[0] = std::min( + n_wires * size_per_wire, std::max(0, size - recv_offset[0])); + recv_limits[1] = + std::min(chunk_size, std::max(0, size - recv_offset[1])); + } else { + send_limits[0] = + std::min(chunk_size, std::max(0, size - send_offset[0])); + recv_limits[0] = + std::min(chunk_size, std::max(0, size - recv_offset[0])); + } + + for (int k = 0; k < size_ - 1; k++) { + // Prefill the pipeline + int buff = 0; + while (buff < n_steps && buff < PIPELINE) { + post_recv_all(sz, buff, n_wires); + for (int lr = 0; lr < MAX_DIR; lr++) { + for (int lw = 0; lw < n_wires; lw++) { + int64_t offset = lw * N + + send_count[lr * RING_MAX_CONNS + lw] * n_wires * N + + lr * n_wires * size_per_wire; + std::copy( + in_ptr + send_offset[lr] + offset, + in_ptr + send_offset[lr] + + std::max(offset, std::min(offset + N, send_limits[lr])), + send_buffer(sz, buff, lr, lw).begin()); + send_count[lr * RING_MAX_CONNS + lw]++; + } + } + post_send_all(sz, buff, n_wires); + + buff++; + in_flight += 2 * MAX_DIR * n_wires; + } + + // Main loop + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = poll(left_, right_, WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int wire = wc[i].wr_id & 0xff; + int lr = wire / RING_MAX_CONNS; + int lw = wire % RING_MAX_CONNS; + + in_flight--; + + if (work_type == SEND_WR && send_count[wire] < n_steps) { + int64_t offset = lw * N + send_count[wire] * n_wires * N + + lr * n_wires * size_per_wire; + std::copy( + in_ptr + send_offset[lr] + offset, + in_ptr + send_offset[lr] + + std::max(offset, std::min(offset + N, send_limits[lr])), + send_buffer(sz, buff, lr, lw).begin()); + send_to(sz, buff, lr, lw); + in_flight++; + send_count[wire]++; + } + + else if (work_type == RECV_WR) { + int64_t offset = lw * N + recv_count[wire] * n_wires * N + + lr * n_wires * size_per_wire; + reduce_op( + recv_buffer(sz, buff, lr, lw).begin(), + out_ptr + recv_offset[lr] + offset, + std::max(0, std::min(N, recv_limits[lr] - offset))); + recv_count[wire]++; + if (recv_count[wire] + (PIPELINE - 1) < n_steps) { + recv_from(sz, buff, lr, lw); + in_flight++; + } + } + } + } + + send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) % + chunk_multiple_size; + recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) % + chunk_multiple_size; + if constexpr (MAX_DIR == 2) { + send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size; + recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size; + send_limits[0] = std::min( + n_wires * size_per_wire, + std::max(0, size - send_offset[0])); + send_limits[1] = + std::min(chunk_size, std::max(0, size - send_offset[1])); + recv_limits[0] = std::min( + n_wires * size_per_wire, + std::max(0, size - recv_offset[0])); + recv_limits[1] = + std::min(chunk_size, std::max(0, size - recv_offset[1])); + } else { + send_limits[0] = + std::min(chunk_size, std::max(0, size - send_offset[0])); + recv_limits[0] = + std::min(chunk_size, std::max(0, size - recv_offset[0])); + } + for (int i = 0; i < MAX_DIR * RING_MAX_CONNS; i++) { + send_count[i] = recv_count[i] = 0; + } + } + } + template void all_reduce( const T* in_ptr,