Start fixing ring

mlx/backend/cpu/distributed.cpp

@@ -98,6 +98,15 @@ void Recv::eval_cpu(
 void ReduceScatter::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  throw std::runtime_error("[ReduceScatter] Not implemented yet.");
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto [in, copied] = ensure_row_contiguous(inputs[0], stream());
+  outputs[0].set_data(allocator::malloc(outputs[0].nbytes()));
+  distributed::detail::sum_scatter(group(), in, outputs[0], stream());
+  if (copied) {
+    auto& enc = cpu::get_command_encoder(stream());
+    enc.add_temporary(in);
+  }
 }
 } // namespace mlx::core::distributed

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<char>();
+    auto out_ptr = output.data<char>();
+    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<GroupImpl> split(int color, int key = -1) override {

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;
 };

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<const char*>(input), static_cast<char*>(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<T>(input, output, n_bytes, SumOp<T>{});
+  });
+}
+
 void MeshGroup::send(const void* input, size_t n_bytes, int dst) {
   mesh_.send(static_cast<const char*>(input), n_bytes, dst);
 }
@@ -184,6 +195,18 @@ void MeshGroup::recv(void* output, size_t n_bytes, int src) {
   mesh_.recv(static_cast<char*>(output), n_bytes, src);
 }
 
+template <typename T, typename ReduceOp>
+void MeshGroup::reduce_scatter(
+    const void* input,
+    void* output,
+    size_t n_bytes,
+    ReduceOp reduce_op) {
+  auto in_ptr = static_cast<const T*>(input);
+  auto out_ptr = static_cast<T*>(output);
+  int64_t count = n_bytes / sizeof(T);
+  mesh_.reduce_scatter(in_ptr, out_ptr, count, reduce_op);
+}
+
 template <typename T, typename ReduceOp>
 void MeshGroup::all_reduce(
     const void* input,

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 <typename T, typename ReduceOp>
+  void reduce_scatter(
+      const void* input,
+      void* output,
+      size_t n_bytes,
+      ReduceOp reduce_op);
+
   template <typename T, typename ReduceOp>
   void all_reduce(
       const void* input,

mlx/distributed/jaccl/lib/jaccl/mesh_impl.h

@@ -29,6 +29,165 @@ class MeshImpl {
 
   MeshImpl() : rank_(0), size_(1) {}
 
+  template <typename T, typename ReduceOp>
+  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<int64_t>(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<T*>(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<T>());
+      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<T>());
+            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<int64_t>(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<T>() + in_offset,
+                out + out_offset,
+                overlap_size);
+          }
+        }
+
+        // Check if we need to post another receive
+        int64_t next_chunk = static_cast<int64_t>(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 <typename T, typename ReduceOp>
   void all_reduce(const T* in, T* out, int64_t size, ReduceOp reduce_op) {
     // Fully connected all reduce with deterministic reduction order.

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<T>(input, output, n_bytes, SumOp<T>{});
+  });
+}
+
 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<char*>(output), n_bytes, src, n_conns_);
 }
 
+template <typename T, typename ReduceOp>
+void RingGroup::reduce_scatter(
+    const void* input,
+    void* output,
+    size_t n_bytes,
+    ReduceOp reduce_op) {
+  auto in_ptr = static_cast<const T*>(input);
+  auto out_ptr = static_cast<T*>(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 <typename T, typename ReduceOp>
 void RingGroup::all_reduce(
     const void* input,

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 <typename T, typename ReduceOp>
+  void reduce_scatter(
+      const void* input,
+      void* output,
+      size_t n_bytes,
+      ReduceOp reduce_op);
+
   template <typename T, typename ReduceOp>
   void all_reduce(
       const void* input,

mlx/distributed/jaccl/lib/jaccl/ring_impl.h

@@ -45,6 +45,174 @@ class RingImpl {
 
   RingImpl() : rank_(0), size_(1), n_conns_(0) {}
 
+  template <int MAX_DIR, typename T, typename ReduceOp>
+  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_;
+    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 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_ + size_ - 1) % size_) * chunk_size;
+    recv_offset[0] = ((rank_ + size_ - 2) % size_) * chunk_size;
+    if constexpr (MAX_DIR == 2) {
+      send_offset[1] = ((rank_ + 1) % size_) * chunk_size;
+      recv_offset[1] = ((rank_ + 2) % size_) * chunk_size;
+      send_limits[0] = std::min(
+          n_wires * size_per_wire, std::max<int64_t>(0, size - send_offset[0]));
+      send_limits[1] =
+          std::min(chunk_size, std::max<int64_t>(0, size - send_offset[1]));
+      recv_limits[0] = std::min(
+          n_wires * size_per_wire, std::max<int64_t>(0, size - recv_offset[0]));
+      recv_limits[1] =
+          std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[1]));
+    } else {
+      send_limits[0] =
+          std::min(chunk_size, std::max<int64_t>(0, size - send_offset[0]));
+      recv_limits[0] =
+          std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[0]));
+    }
+
+    for (int k = 0; k < size_ - 1; k++) {
+      const T* read_ptr = (k == 0) ? in_ptr : out_ptr;
+      int64_t read_offset[MAX_DIR];
+      read_offset[0] = (k == 0) ? send_offset[0] : 0;
+      if constexpr (MAX_DIR == 2) {
+        read_offset[1] = (k == 0) ? send_offset[1] : 0;
+      }
+
+      // Prefill the pipeline
+      int buff = 0;
+      while (buff < n_steps && buff < PIPELINE) {
+        // Post receives
+        post_recv_all<MAX_DIR>(sz, buff, n_wires);
+        in_flight += MAX_DIR * n_wires;
+
+        // Copy to send buffers and also to the output buffer
+        for (int lr = 0; lr < MAX_DIR; lr++) {
+          for (int lw = 0; lw < n_wires; lw++) {
+            // send
+            int64_t offset = lw * N +
+                send_count[lr * RING_MAX_CONNS + lw] * n_wires * N +
+                lr * n_wires * size_per_wire;
+            std::copy(
+                read_ptr + read_offset[lr] + offset,
+                read_ptr + read_offset[lr] +
+                    std::max(offset, std::min(offset + N, send_limits[lr])),
+                send_buffer(sz, buff, lr, lw).begin<T>());
+            send_count[lr * RING_MAX_CONNS + lw]++;
+            in_flight++;
+            send_to(sz, buff, lr, lw);
+
+            // copy for in-place recv reduce
+            offset -= n_wires * N;
+            std::copy(
+                in_ptr + recv_offset[lr] + offset,
+                in_ptr + recv_offset[lr] +
+                    std::max(offset, std::min(offset + N, recv_limits[lr])),
+                out_ptr + offset);
+          }
+        }
+
+        buff++;
+      }
+
+      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) {
+            int64_t offset = lw * N + send_count[wire] * n_wires * N +
+                lr * n_wires * size_per_wire;
+
+            // More stuff to send
+            if (send_count[wire] < n_steps) {
+              std::copy(
+                  read_ptr + read_offset[lr] + offset,
+                  read_ptr + read_offset[lr] +
+                      std::max(offset, std::min(offset + N, send_limits[lr])),
+                  send_buffer(sz, buff, lr, lw).begin<T>());
+              send_to(sz, buff, lr, lw);
+              in_flight++;
+              send_count[wire]++;
+            }
+
+            // Copy the input chunk into the output
+            offset -= n_wires * N;
+            std::copy(
+                in_ptr + recv_offset[lr] + offset,
+                in_ptr + recv_offset[lr] +
+                    std::max(offset, std::min(offset + N, recv_limits[lr])),
+                out_ptr + offset);
+          }
+
+          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<T>(),
+                out_ptr + recv_offset[lr] + offset,
+                std::max<int64_t>(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] + size - chunk_size) % size;
+      recv_offset[0] = (recv_offset[0] + size - chunk_size) % size;
+      if constexpr (MAX_DIR == 2) {
+        send_offset[1] = (send_offset[1] + chunk_size) % size;
+        recv_offset[1] = (recv_offset[1] + chunk_size) % size;
+        send_limits[0] = std::min(
+            n_wires * size_per_wire,
+            std::max<int64_t>(0, size - send_offset[0]));
+        send_limits[1] =
+            std::min(chunk_size, std::max<int64_t>(0, size - send_offset[1]));
+        recv_limits[0] = std::min(
+            n_wires * size_per_wire,
+            std::max<int64_t>(0, size - recv_offset[0]));
+        recv_limits[1] =
+            std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[1]));
+      } else {
+        send_limits[0] =
+            std::min(chunk_size, std::max<int64_t>(0, size - send_offset[0]));
+        recv_limits[0] =
+            std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[0]));
+      }
+      for (int i = 0; i < MAX_DIR * RING_MAX_CONNS; i++) {
+        send_count[i] = recv_count[i] = 0;
+      }
+    }
+  }
+
   template <int MAX_DIR, typename T, typename ReduceOp>
   void all_reduce(
       const T* in_ptr,