Tentative JACCL examples

mlx/backend/cpu/distributed.cpp

@@ -98,15 +98,6 @@ void Recv::eval_cpu(
 void ReduceScatter::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  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);
-  }
+  throw std::runtime_error("[ReduceScatter] Not implemented yet.");
 }
 } // namespace mlx::core::distributed

mlx/distributed/jaccl/jaccl.cpp

@@ -145,16 +145,7 @@ class JACCLGroup : public GroupImpl {
   }
 
   void sum_scatter(const array& input, array& output, Stream stream) override {
-    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);
-    });
+    throw std::runtime_error("[jaccl] sum_scatter not supported.");
   }
 
   std::shared_ptr<GroupImpl> split(int color, int key = -1) override {

mlx/distributed/jaccl/lib/examples/CMakeLists.txt

@@ -35,6 +35,8 @@ endfunction()
 # Examples
 build_example(minimal_env.cpp)
 build_example(minimal_cfg.cpp)
+build_example(monte_carlo_pi.cpp)
+build_example(file_broadcast.cpp)
 
 # Benchmarks
 build_example(allreduce_bench.cpp)

mlx/distributed/jaccl/lib/examples/file_broadcast.cpp

@@ -0,0 +1,360 @@
+// Copyright © 2025 Apple Inc.
+//
+// File Broadcast with JACCL
+//
+// This example demonstrates distributed file transfer using JACCL's all_sum
+// operation to broadcast a file from any rank to all other machines.
+//
+// The algorithm:
+// 1. The sender rank reads the file into memory
+// 2. All other ranks allocate zero-filled buffers of the same size
+// 3. Use all_sum to broadcast: sender has data, others have zeros
+// 4. After all_sum, all ranks have the file data
+// 5. All ranks write the file to disk
+//
+// For large files, the transfer is chunked to manage memory efficiently.
+//
+// Usage:
+//   Set environment variables (see README.md), then run:
+//
+//     ./jaccl_file_broadcast -f <file> [-s <sender_rank>] [-o <output_dir>]
+//
+//   Or with mlx.launch:
+//
+//     mlx.launch --hostfile hosts.json ./jaccl_file_broadcast -f myfile.bin
+//
+// Example output (4 ranks, sender rank 2):
+//   Rank 0 of 4: Received 10485760 bytes from rank 2 (982.5 MB/s)
+//   Rank 1 of 4: Received 10485760 bytes from rank 2 (985.2 MB/s)
+//   Rank 2 of 4: Sent 10485760 bytes (980.1 MB/s)
+//   Rank 3 of 4: Received 10485760 bytes from rank 2 (978.9 MB/s)
+
+#include <jaccl/jaccl.h>
+#include <jaccl/types.h>
+
+#include <sys/stat.h>
+#include <atomic>
+#include <chrono>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <string>
+#include <thread>
+#include <vector>
+
+static void usage(const char* prog) {
+  std::cerr
+      << "Usage: " << prog << " [options]\n"
+      << "  -f <file>       File to broadcast (required)\n"
+      << "  -s <rank>       Sender rank (default: 0)\n"
+      << "  -o <dir>        Output directory (default: current dir)\n"
+      << "  -c <bytes>      Chunk size in bytes (default: 67108864 = 64MB)\n"
+      << "  -v              Verbose output\n"
+      << "  -h              Show this help\n";
+}
+
+static bool file_exists(const std::string& path) {
+  struct stat buffer;
+  return (stat(path.c_str(), &buffer) == 0);
+}
+
+static std::int64_t file_size(const std::string& path) {
+  struct stat buffer;
+  if (stat(path.c_str(), &buffer) != 0) {
+    return -1;
+  }
+  return static_cast<std::int64_t>(buffer.st_size);
+}
+
+static bool create_directory(const std::string& path) {
+  if (path.empty() || path == ".") {
+    return true;
+  }
+  return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST;
+}
+
+static std::string basename(const std::string& path) {
+  size_t pos = path.find_last_of("/\\");
+  return (pos == std::string::npos) ? path : path.substr(pos + 1);
+}
+
+struct BroadcastStats {
+  std::int64_t total_bytes;
+  std::int64_t chunks_sent;
+  std::int64_t chunks_received;
+  double total_time_ms;
+  int sender_rank;
+};
+
+int main(int argc, char** argv) {
+  std::string input_file;
+  std::string output_dir = ".";
+  int sender_rank = 0;
+  std::int64_t chunk_size = 67108864;
+  bool verbose = false;
+
+  for (int i = 1; i < argc; i++) {
+    std::string arg = argv[i];
+    if (arg == "-h" || arg == "--help") {
+      usage(argv[0]);
+      return 0;
+    } else if (arg == "-f" && i + 1 < argc) {
+      input_file = argv[++i];
+    } else if (arg == "-s" && i + 1 < argc) {
+      sender_rank = std::atoi(argv[++i]);
+    } else if (arg == "-o" && i + 1 < argc) {
+      output_dir = argv[++i];
+    } else if (arg == "-c" && i + 1 < argc) {
+      chunk_size = std::atoll(argv[++i]);
+    } else if (arg == "-v" || arg == "--verbose") {
+      verbose = true;
+    } else {
+      std::cerr << "Unknown option: " << arg << "\n";
+      usage(argv[0]);
+      return 1;
+    }
+  }
+
+  if (input_file.empty()) {
+    std::cerr << "Error: Input file is required (-f <file>)\n";
+    usage(argv[0]);
+    return 1;
+  }
+
+  auto group = jaccl::init();
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  int rank = group->rank();
+  int nranks = group->size();
+
+  if (sender_rank < 0 || sender_rank >= nranks) {
+    std::cerr << "Error: Sender rank " << sender_rank << " is out of range [0, "
+              << nranks << ")\n";
+    return 1;
+  }
+
+  std::int64_t total_file_size = 0;
+  if (rank == sender_rank) {
+    if (!file_exists(input_file)) {
+      std::cerr << "Error: File not found: " << input_file << "\n";
+      return 1;
+    }
+    total_file_size = file_size(input_file);
+    if (total_file_size < 0) {
+      std::cerr << "Error: Cannot read file size: " << input_file << "\n";
+      return 1;
+    }
+  }
+
+  group->all_sum(
+      &total_file_size, &total_file_size, sizeof(int64_t), jaccl::Int64);
+
+  if (!create_directory(output_dir)) {
+    std::cerr << "Error: Cannot create output directory: " << output_dir
+              << "\n";
+    return 1;
+  }
+
+  std::string output_file = output_dir == "."
+      ? basename(input_file)
+      : output_dir + "/" + basename(input_file);
+
+  if (verbose) {
+    std::printf(
+        "Rank %d of %d: Broadcasting '%s' (%ld bytes) from rank %d\n",
+        rank,
+        nranks,
+        input_file.c_str(),
+        static_cast<long>(total_file_size),
+        sender_rank);
+  }
+
+  auto t_start = std::chrono::high_resolution_clock::now();
+
+  std::int64_t num_chunks = (total_file_size + chunk_size - 1) / chunk_size;
+  if (num_chunks == 0) {
+    num_chunks = 1;
+  }
+
+  const int num_buffers = 4;
+  std::vector<std::vector<std::uint8_t>> buffers(
+      num_buffers, std::vector<std::uint8_t>(chunk_size, 0));
+
+  std::ifstream infile;
+  std::ofstream outfile;
+
+  if (rank == sender_rank) {
+    infile.open(input_file, std::ios::binary);
+    if (!infile.good()) {
+      std::cerr << "Error: Cannot open input file: " << input_file << "\n";
+      return 1;
+    }
+  }
+
+  outfile.open(output_file, std::ios::binary);
+  if (!outfile.good()) {
+    std::cerr << "Error: Cannot open output file: " << output_file << "\n";
+    return 1;
+  }
+
+  std::atomic<std::int64_t> next_read_chunk{0};
+  std::atomic<std::int64_t> next_comm_chunk{0};
+  std::atomic<std::int64_t> next_write_chunk{0};
+  std::atomic<bool> read_done{false};
+  std::atomic<bool> comm_done{false};
+
+  std::vector<std::atomic<bool>> buffer_ready(num_buffers);
+  std::vector<std::atomic<bool>> buffer_written(num_buffers);
+  for (int i = 0; i < num_buffers; i++) {
+    buffer_ready[i] = false;
+    buffer_written[i] = false;
+  }
+
+  std::vector<std::int64_t> chunk_sizes(num_chunks);
+  for (std::int64_t i = 0; i < num_chunks; i++) {
+    chunk_sizes[i] = std::min(chunk_size, total_file_size - i * chunk_size);
+  }
+
+  std::thread reader_thread;
+  if (rank == sender_rank) {
+    reader_thread = std::thread([&]() {
+      while (true) {
+        std::int64_t chunk_idx = next_read_chunk.fetch_add(1);
+        if (chunk_idx >= num_chunks) {
+          break;
+        }
+        std::int64_t offset = chunk_idx * chunk_size;
+        std::int64_t this_chunk_size = chunk_sizes[chunk_idx];
+        int buffer_idx = chunk_idx % num_buffers;
+
+        infile.seekg(offset, std::ios::beg);
+        infile.read(
+            reinterpret_cast<char*>(buffers[buffer_idx].data()),
+            this_chunk_size);
+
+        std::fill(
+            buffers[buffer_idx].begin() + this_chunk_size,
+            buffers[buffer_idx].end(),
+            0);
+
+        buffer_ready[buffer_idx] = true;
+      }
+      read_done = true;
+    });
+  } else {
+    read_done = true;
+  }
+
+  std::thread writer_thread([&]() {
+    while (true) {
+      std::int64_t chunk_idx = next_write_chunk.load();
+      if (chunk_idx >= num_chunks && comm_done) {
+        break;
+      }
+      if (chunk_idx >= num_chunks) {
+        std::this_thread::yield();
+        continue;
+      }
+
+      int buffer_idx = chunk_idx % num_buffers;
+      if (!buffer_written[buffer_idx]) {
+        std::this_thread::yield();
+        continue;
+      }
+
+      std::int64_t this_chunk_size = chunk_sizes[chunk_idx];
+      outfile.write(
+          reinterpret_cast<const char*>(buffers[buffer_idx].data()),
+          this_chunk_size);
+
+      buffer_written[buffer_idx] = false;
+      next_write_chunk.fetch_add(1);
+    }
+  });
+
+  for (std::int64_t chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) {
+    std::int64_t this_chunk_size = chunk_sizes[chunk_idx];
+    int buffer_idx = chunk_idx % num_buffers;
+
+    if (rank == sender_rank) {
+      while (!buffer_ready[buffer_idx] && !read_done) {
+        std::this_thread::yield();
+      }
+    }
+
+    std::fill(
+        buffers[buffer_idx].begin() + this_chunk_size,
+        buffers[buffer_idx].end(),
+        0);
+
+    group->all_sum(
+        buffers[buffer_idx].data(),
+        buffers[buffer_idx].data(),
+        this_chunk_size,
+        jaccl::UInt8);
+
+    buffer_written[buffer_idx] = true;
+    next_comm_chunk.fetch_add(1);
+
+    if (verbose) {
+      double progress = 100.0 * (chunk_idx + 1) / num_chunks;
+      std::printf(
+          "Rank %d: Progress %.1f%% (%ld/%ld chunks)\n",
+          rank,
+          progress,
+          static_cast<long>(chunk_idx + 1),
+          static_cast<long>(num_chunks));
+    }
+  }
+
+  comm_done = true;
+
+  if (reader_thread.joinable()) {
+    reader_thread.join();
+  }
+  writer_thread.join();
+
+  infile.close();
+  outfile.close();
+
+  auto t_end = std::chrono::high_resolution_clock::now();
+  double elapsed_ms =
+      std::chrono::duration<double, std::milli>(t_end - t_start).count();
+  double elapsed_sec = elapsed_ms / 1000.0;
+  double bandwidth_mbps = (total_file_size / (1024.0 * 1024.0)) / elapsed_sec;
+
+  if (rank == sender_rank) {
+    std::printf(
+        "Rank %d of %d: Sent %ld bytes from '%s' (%.1f MB/s)\n",
+        rank,
+        nranks,
+        static_cast<long>(total_file_size),
+        input_file.c_str(),
+        bandwidth_mbps);
+  } else {
+    std::printf(
+        "Rank %d of %d: Received %ld bytes from rank %d to '%s' (%.1f MB/s)\n",
+        rank,
+        nranks,
+        static_cast<long>(total_file_size),
+        sender_rank,
+        output_file.c_str(),
+        bandwidth_mbps);
+  }
+
+  if (verbose) {
+    std::printf(
+        "Rank %d: Total time: %.2f ms, Chunks: %ld, Chunk size: %ld bytes\n",
+        rank,
+        elapsed_ms,
+        static_cast<long>(num_chunks),
+        static_cast<long>(chunk_size));
+  }
+
+  return 0;
+}

mlx/distributed/jaccl/lib/examples/monte_carlo_pi.cpp

@@ -0,0 +1,152 @@
+// Copyright © 2025 Apple Inc.
+//
+// Monte Carlo Pi Estimation with JACCL
+//
+// This example demonstrates distributed Monte Carlo simulation using JACCL
+// to estimate the value of π. Each rank generates random points independently
+// and uses all-reduce to aggregate the results across all machines.
+//
+// The algorithm:
+// 1. Each rank generates N random points in the unit square [0,1] x [0,1]
+// 2. Count how many fall inside the quarter circle (x² + y² ≤ 1)
+// 3. Use all_sum to aggregate hits and total points across all ranks
+// 4. π ≈ 4 × (hits / total)
+//
+// Usage:
+//   Set environment variables (see README.md), then run:
+//
+//     ./jaccl_monte_carlo_pi [-n <points_per_rank>]
+//
+//   Or with mlx.launch:
+//
+//     mlx.launch --hostfile hosts.json ./jaccl_monte_carlo_pi -n 10000000
+//
+// Example output (4 ranks, 10M points each):
+//   Rank 2 of 4
+//   Local: 7854321 hits out of 10000000 points
+//   Global: 31416789 hits out of 40000000 points
+//   Estimated π = 3.141679 (error: 0.000086)
+
+#include <jaccl/jaccl.h>
+#include <jaccl/types.h>
+
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+#include <random>
+#include <string>
+#include <vector>
+
+static void usage(const char* prog) {
+  std::cerr << "Usage: " << prog << " [options]\n"
+            << "  -n <points>  Points per rank (default: 1000000)\n"
+            << "  -s <seed>    Random seed base (default: 42)\n"
+            << "  -h           Show this help\n";
+}
+
+struct MonteCarloResult {
+  int64_t hits;
+  int64_t total;
+};
+
+MonteCarloResult estimate_pi_local(int64_t num_points, unsigned int seed) {
+  std::mt19937_64 rng(seed);
+  std::uniform_real_distribution<double> dist(0.0, 1.0);
+
+  int64_t hits = 0;
+  for (int64_t i = 0; i < num_points; i++) {
+    double x = dist(rng);
+    double y = dist(rng);
+    if (x * x + y * y <= 1.0) {
+      hits++;
+    }
+  }
+
+  return {hits, num_points};
+}
+
+int main(int argc, char** argv) {
+  int64_t points_per_rank = 1000000;
+  unsigned int seed_base = 42;
+
+  for (int i = 1; i < argc; i++) {
+    std::string arg = argv[i];
+    if (arg == "-h" || arg == "--help") {
+      usage(argv[0]);
+      return 0;
+    } else if (arg == "-n" && i + 1 < argc) {
+      points_per_rank = std::atoll(argv[++i]);
+    } else if (arg == "-s" && i + 1 < argc) {
+      seed_base = static_cast<unsigned int>(std::atoi(argv[++i]));
+    } else {
+      std::cerr << "Unknown option: " << arg << "\n";
+      usage(argv[0]);
+      return 1;
+    }
+  }
+
+  auto group = jaccl::init();
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  int rank = group->rank();
+  int nranks = group->size();
+
+  std::printf("Rank %d of %d\n", rank, nranks);
+  std::printf(
+      "Generating %ld random points (seed: %u)...\n",
+      static_cast<long>(points_per_rank),
+      seed_base + static_cast<unsigned int>(rank));
+
+  auto t0 = std::chrono::high_resolution_clock::now();
+
+  MonteCarloResult local = estimate_pi_local(
+      points_per_rank, seed_base + static_cast<unsigned int>(rank));
+
+  auto t1 = std::chrono::high_resolution_clock::now();
+  double local_time =
+      std::chrono::duration<double, std::milli>(t1 - t0).count();
+
+  std::printf(
+      "Rank %d: %ld hits out of %ld points (%.2f ms)\n",
+      rank,
+      static_cast<long>(local.hits),
+      static_cast<long>(local.total),
+      local_time);
+
+  MonteCarloResult global = {0, 0};
+
+  group->all_sum(&local.hits, &global.hits, sizeof(int64_t), jaccl::Int64);
+  group->all_sum(&local.total, &global.total, sizeof(int64_t), jaccl::Int64);
+
+  if (rank == 0) {
+    double pi_estimate = 4.0 * static_cast<double>(global.hits) /
+        static_cast<double>(global.total);
+    double error = std::abs(pi_estimate - M_PI);
+
+    std::printf("\n=== Results ===\n");
+    std::printf(
+        "Global: %ld hits out of %ld points\n",
+        static_cast<long>(global.hits),
+        static_cast<long>(global.total));
+    std::printf("Estimated π = %.10f\n", pi_estimate);
+    std::printf("True π      = %.10f\n", M_PI);
+    std::printf("Error       = %.10f (%.6f%%)\n", error, 100.0 * error / M_PI);
+
+    double total_time =
+        std::chrono::duration<double, std::milli>(t1 - t0).count();
+    std::printf("\nPerformance:\n");
+    std::printf("Total points: %ld\n", static_cast<long>(global.total));
+    std::printf("Time: %.2f ms\n", total_time);
+    std::printf(
+        "Points/sec: %.0f\n",
+        static_cast<double>(global.total) / (total_time / 1000.0));
+  }
+
+  return 0;
+}

mlx/distributed/jaccl/lib/jaccl/group.h

@@ -28,9 +28,6 @@ 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,17 +176,6 @@ 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);
 }
@@ -195,18 +184,6 @@ 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,20 +44,10 @@ 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,165 +29,6 @@ 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,17 +166,6 @@ 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_;
@@ -201,29 +190,6 @@ 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,20 +45,10 @@ 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,174 +45,6 @@ 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,