mlx 0.30.7

data/mlx/mlx/distributed/jaccl/jaccl.cpp

@@ -0,0 +1,178 @@
+// Copyright © 2025 Apple Inc.
+
+#include <fstream>
+#include <sstream>
+
+#include <json.hpp>
+
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/distributed/jaccl/mesh.h"
+#include "mlx/distributed/jaccl/ring.h"
+#include "mlx/distributed/jaccl/utils.h"
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+using json = nlohmann::json;
+
+namespace {
+
+struct DeviceFile {
+  DeviceFile(const char* dev_file) {
+    std::ifstream f(dev_file);
+    json devices = json::parse(f);
+    if (!devices.is_array()) {
+      throw std::runtime_error(
+          "[jaccl] The device file should start with an array");
+    }
+
+    devices_.resize(devices.size());
+    for (int rank = 0; rank < devices.size(); rank++) {
+      auto conn = devices[rank];
+      if (!conn.is_array()) {
+        throw std::runtime_error(
+            "[jaccl] The device file should have an array of arrays");
+      }
+      if (conn.size() != devices_.size()) {
+        std::ostringstream msg;
+        msg << "[jaccl] The device file should contain the connectivity of each rank to "
+            << "all other ranks but rank " << rank << " contains only "
+            << conn.size() << " entries.";
+        throw std::runtime_error(msg.str());
+      }
+
+      devices_[rank].resize(conn.size());
+      for (int dst = 0; dst < conn.size(); dst++) {
+        auto names = conn[dst];
+        if (names.is_string()) {
+          devices_[rank][dst].push_back(names);
+        } else if (names.is_array()) {
+          for (auto name_it = names.begin(); name_it != names.end();
+               name_it++) {
+            devices_[rank][dst].push_back(*name_it);
+          }
+        } else if (!names.is_null()) {
+          throw std::runtime_error(
+              "[jaccl] Device names should be null, a string or array of strings.");
+        }
+      }
+    }
+  }
+
+  int size() {
+    return devices_.size();
+  }
+
+  bool is_valid_mesh() {
+    for (int src = 0; src < size(); src++) {
+      for (int dst = 0; dst < size(); dst++) {
+        if (devices_[src][dst].size() != static_cast<size_t>(src != dst)) {
+          return false;
+        }
+      }
+    }
+
+    return true;
+  }
+
+  bool is_valid_ring() {
+    int num_connections = devices_[0][1].size();
+    if (num_connections == 0) {
+      return false;
+    }
+
+    for (int src = 0; src < size(); src++) {
+      int left = (src + size() - 1) % size();
+      int right = (src + 1) % size();
+      for (int dst = 0; dst < size(); dst++) {
+        if (dst != left && dst != right) {
+          if (devices_[src][dst].size() != 0) {
+            return false;
+          }
+        } else {
+          if (devices_[src][dst].size() != num_connections) {
+            return false;
+          }
+        }
+      }
+    }
+
+    return true;
+  }
+
+  std::vector<std::string> extract_mesh_connectivity(int rank) {
+    std::vector<std::string> devices(size());
+    for (int dst = 0; dst < size(); dst++) {
+      if (dst != rank) {
+        devices[dst] = devices_[rank][dst][0];
+      }
+    }
+    return devices;
+  }
+
+  std::pair<std::vector<std::string>, std::vector<std::string>>
+  extract_ring_connectivity(int rank) {
+    int left = (rank + size() - 1) % size();
+    int right = (rank + 1) % size();
+
+    return std::make_pair(devices_[rank][left], devices_[rank][right]);
+  }
+
+  std::vector<std::vector<std::vector<std::string>>> devices_;
+};
+
+} // namespace
+
+namespace mlx::core::distributed::jaccl {
+
+bool is_available() {
+  return ibv().is_available();
+}
+
+std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
+  const char* dev_file = std::getenv("MLX_IBV_DEVICES");
+  const char* coordinator = std::getenv("MLX_JACCL_COORDINATOR");
+  const char* rank_str = std::getenv("MLX_RANK");
+  const char* ring = std::getenv("MLX_JACCL_RING");
+
+  if (!is_available() || !dev_file || !coordinator || !rank_str) {
+    if (strict) {
+      std::ostringstream msg;
+      msg << "[jaccl] You need to provide via environment variables a rank (MLX_RANK), "
+          << "a device file (MLX_IBV_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) "
+          << "but provided MLX_RANK=\"" << ((rank_str) ? rank_str : "")
+          << "\", MLX_IBV_DEVICES=\"" << ((dev_file) ? dev_file : "")
+          << "\" and MLX_JACCL_COORDINATOR=\""
+          << ((coordinator) ? coordinator : "");
+      throw std::runtime_error(msg.str());
+    }
+    return nullptr;
+  }
+
+  auto rank = std::atoi(rank_str);
+  bool prefer_ring = ring != nullptr;
+  DeviceFile devices(dev_file);
+
+  if (rank >= devices.size() || rank < 0) {
+    std::ostringstream msg;
+    msg << "[jaccl] Invalid rank " << rank << ". It should be between 0 and "
+        << devices.size();
+    throw std::runtime_error(msg.str());
+  }
+
+  if (prefer_ring && devices.is_valid_ring()) {
+    auto [left, right] = devices.extract_ring_connectivity(rank);
+    return std::make_shared<RingGroup>(
+        rank, devices.size(), left, right, coordinator);
+  } else if (devices.is_valid_mesh()) {
+    auto device_names = devices.extract_mesh_connectivity(rank);
+    return std::make_shared<MeshGroup>(rank, device_names, coordinator);
+  } else if (devices.is_valid_ring()) {
+    auto [left, right] = devices.extract_ring_connectivity(rank);
+    return std::make_shared<RingGroup>(
+        rank, devices.size(), left, right, coordinator);
+  } else {
+    throw std::runtime_error(
+        "[jaccl] The device file should define a valid mesh or a valid ring.");
+  }
+}
+
+} // namespace mlx::core::distributed::jaccl

data/mlx/mlx/distributed/jaccl/jaccl.h

@@ -0,0 +1,12 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/distributed/distributed.h"
+
+namespace mlx::core::distributed::jaccl {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available();
+std::shared_ptr<GroupImpl> init(bool strict = false);
+
+} // namespace mlx::core::distributed::jaccl

data/mlx/mlx/distributed/jaccl/mesh.cpp

@@ -0,0 +1,451 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/distributed/jaccl/mesh.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/distributed/reduction_ops.h"
+#include "mlx/dtype_utils.h"
+
+constexpr int MAX_PEERS = 8;
+
+namespace mlx::core::distributed::jaccl {
+
+MeshGroup::MeshGroup(
+    int rank,
+    const std::vector<std::string>& device_names,
+    const char* coordinator_addr)
+    : rank_(rank),
+      size_(device_names.size()),
+      side_channel_(rank_, size_, coordinator_addr),
+      connections_(create_connections(device_names)) {
+  if (size_ > MAX_PEERS) {
+    std::ostringstream msg;
+    msg << "[jaccl] The JACCL mesh supports up to " << MAX_PEERS
+        << " peers but " << size_ << " were provided.";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Initialize all the connections and allocate buffers
+  initialize();
+
+  // Make sure every node has reached here before continuing
+  side_channel_.all_gather<int>(0);
+}
+
+void MeshGroup::initialize() {
+  // Create the queue pairs
+  for (auto& conn : connections_) {
+    if (conn.ctx == nullptr) {
+      continue;
+    }
+    conn.allocate_protection_domain();
+    conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR);
+    conn.create_queue_pair();
+  }
+
+  allocate_buffers();
+
+  // First init all connections
+  for (int peer = 0; peer < size_; peer++) {
+    if (peer == rank_) {
+      continue;
+    }
+    connections_[peer].queue_pair_init();
+  }
+
+  // Gather the information to be exchanged, this also serves as a barrier so
+  // that all peers have initialized their connections before attempting to
+  // transition to RTS.
+  std::vector<Destination> info;
+  for (auto& conn : connections_) {
+    info.emplace_back(conn.info());
+  }
+  auto all_infos = side_channel_.all_gather(info);
+
+  // Transition queue pairs to RTS
+  for (int peer = 0; peer < size_; peer++) {
+    if (peer == rank_) {
+      continue;
+    }
+    auto peer_info = all_infos[peer][rank_];
+    connections_[peer].queue_pair_rtr(peer_info);
+    connections_[peer].queue_pair_rts();
+  }
+}
+
+void MeshGroup::allocate_buffers() {
+  // Deregister any buffers and free the memory
+  buffers_.clear();
+
+  // Allocate the memory
+  for (int k = 0; k < BUFFER_SIZES; k++) {
+    for (int i = 0; i < NUM_BUFFERS; i++) {
+      for (int j = 0; j < size_; j++) {
+        buffers_.emplace_back(FRAME_SIZE * (1 << k));
+      }
+    }
+  }
+
+  for (int k = 0; k < BUFFER_SIZES; k++) {
+    for (int i = 0; i < NUM_BUFFERS; i++) {
+      for (int j = 0; j < size_; j++) {
+        // This is our send buffer so register it with all pds so we can send
+        // it to all connected devices.
+        if (j == rank_) {
+          for (auto& conn : connections_) {
+            if (conn.ctx != nullptr) {
+              buffers_[k * NUM_BUFFERS * size_ + i * size_ + j]
+                  .register_to_protection_domain(conn.protection_domain);
+            }
+          }
+        }
+
+        // This is the recv buffer from rank j so register it to rank j's
+        // protection domain.
+        else {
+          buffers_[k * NUM_BUFFERS * size_ + i * size_ + j]
+              .register_to_protection_domain(connections_[j].protection_domain);
+        }
+      }
+    }
+  }
+}
+
+void MeshGroup::all_sum(const array& input, array& output, Stream stream) {
+  dispatch_all_types(output.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    all_reduce<T>(input, output, stream, detail::SumOp<T>{});
+  });
+}
+
+void MeshGroup::all_max(const array& input, array& output, Stream stream) {
+  dispatch_all_types(output.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    all_reduce<T>(input, output, stream, detail::MaxOp<T>{});
+  });
+}
+
+void MeshGroup::all_min(const array& input, array& output, Stream stream) {
+  dispatch_all_types(output.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    all_reduce<T>(input, output, stream, detail::MinOp<T>{});
+  });
+}
+
+void MeshGroup::all_gather(const array& input, array& output, Stream stream) {
+  auto in_ptr = input.data<char>();
+  auto out_ptr = output.data<char>();
+  size_t n_bytes = input.nbytes();
+  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, this]() {
+    // Copy our data to the appropriate place
+    std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes);
+
+    // Fully connected all gather
+    char* data = out_ptr;
+    char* our_data = out_ptr + rank_ * n_bytes;
+    auto [sz, N] = buffer_size_from_message(n_bytes);
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2;
+    int64_t total = static_cast<int64_t>(n_bytes);
+    int num_peers = size_ - 1;
+
+    // Counters to maintain the state of transfers
+    int in_flight = 0;
+    int read_offset = 0;
+    int completed_send_count[PIPELINE] = {0};
+    int write_offset[MAX_PEERS] = {0};
+
+    // Prefill the pipeline
+    int buff = 0;
+    while (read_offset < total && buff < PIPELINE) {
+      post_recv_all(sz, buff);
+      std::copy(
+          our_data + read_offset,
+          our_data + std::min(read_offset + N, total),
+          send_buffer(sz, buff).begin<char>());
+      post_send_all(sz, buff);
+
+      buff++;
+      in_flight += 2 * num_peers;
+      read_offset += N;
+    }
+
+    // Main loop
+    //
+    // Keep going until we have no longer data in flight.
+    while (in_flight > 0) {
+      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--;
+
+        // Send completed. If all sends completed then send the next chunk.
+        if (work_type == SEND_WR && read_offset < total) {
+          completed_send_count[buff]++;
+          if (completed_send_count[buff] == num_peers) {
+            std::copy(
+                our_data + read_offset,
+                our_data + std::min(read_offset + N, total),
+                send_buffer(sz, buff).begin<char>());
+            post_send_all(sz, buff);
+
+            completed_send_count[buff] = 0;
+            in_flight += num_peers;
+            read_offset += N;
+          }
+        }
+
+        // Recv completed. If we have more chunks then post another recv.
+        else if (work_type == RECV_WR) {
+          std::copy(
+              recv_buffer(sz, buff, rank).begin<char>(),
+              recv_buffer(sz, buff, rank).begin<char>() +
+                  std::min(N, total - write_offset[rank]),
+              data + rank * n_bytes + write_offset[rank]);
+          write_offset[rank] += N;
+          if (write_offset[rank] + N * (PIPELINE - 1) < total) {
+            recv_from(sz, rank, buff);
+            in_flight++;
+          }
+        }
+      }
+    }
+  });
+}
+
+void MeshGroup::send(const array& input, int dst, Stream stream) {
+  auto data = input.data<char>();
+  int64_t n_bytes = input.nbytes();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(input);
+  encoder.dispatch([data, n_bytes, dst, this]() {
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE;
+    auto [sz, N] = buffer_size_from_message(n_bytes);
+
+    int in_flight = 0;
+    int64_t read_offset = 0;
+
+    // Prefill the pipeline
+    int buff = 0;
+    while (read_offset < n_bytes && buff < PIPELINE) {
+      std::copy(
+          data + read_offset,
+          data + std::min(read_offset + N, n_bytes),
+          send_buffer(sz, buff).begin<char>());
+      send_to(sz, dst, buff);
+
+      buff++;
+      read_offset += N;
+      in_flight++;
+    }
+
+    // Main loop
+    while (in_flight > 0) {
+      // Poll the hardware for completions.
+      //
+      // If a send was completed and we have more data to send then go ahead
+      // and send them.
+      ibv_wc wc[WC_NUM];
+      int n = connections_[dst].poll(WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int rank = wc[i].wr_id & 0xff;
+
+        in_flight--;
+
+        if (read_offset < n_bytes) {
+          std::copy(
+              data + read_offset,
+              data + std::min(read_offset + N, n_bytes),
+              send_buffer(sz, buff).begin<char>());
+          send_to(sz, dst, buff);
+
+          read_offset += N;
+          in_flight++;
+        }
+      }
+    }
+  });
+}
+
+void MeshGroup::recv(array& out, int src, Stream stream) {
+  auto data = out.data<char>();
+  int64_t n_bytes = out.nbytes();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  encoder.dispatch([data, n_bytes, src, this]() {
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE;
+    auto [sz, N] = buffer_size_from_message(n_bytes);
+
+    int in_flight = 0;
+    int64_t write_offset = 0;
+
+    // Prefill the pipeline
+    int buff = 0;
+    while (N * buff < n_bytes && buff < PIPELINE) {
+      recv_from(sz, src, buff);
+
+      in_flight++;
+      buff++;
+    }
+
+    // Main loop
+    while (in_flight > 0) {
+      // Poll the hardware for completions.
+      //
+      // If a recv was completed copy it to the output and if we have more
+      // data to fetch post another recv.
+      ibv_wc wc[WC_NUM];
+      int n = connections_[src].poll(WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int rank = wc[i].wr_id & 0xff;
+
+        in_flight--;
+
+        std::copy(
+            recv_buffer(sz, buff, src).begin<char>(),
+            recv_buffer(sz, buff, src).begin<char>() +
+                std::min(n_bytes - write_offset, static_cast<int64_t>(N)),
+            data + write_offset);
+        write_offset += N;
+
+        if (write_offset + (PIPELINE - 1) * N < n_bytes) {
+          recv_from(sz, src, buff);
+
+          in_flight++;
+        }
+      }
+    }
+  });
+}
+
+template <typename T, typename ReduceOp>
+void MeshGroup::all_reduce(
+    const array& input,
+    array& output,
+    Stream stream,
+    ReduceOp reduce_op) {
+  auto in_ptr = input.data<T>();
+  auto out_ptr = output.data<T>();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(input);
+  encoder.set_output_array(output);
+  encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() {
+    // If not inplace all reduce then copy the input to the output first
+    if (in_ptr != out_ptr) {
+      std::memcpy(out_ptr, in_ptr, size * sizeof(T));
+    }
+
+    // Fully connected all reduce
+    T* data = out_ptr;
+    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 * MAX_PEERS * 2;
+    int64_t total = static_cast<int64_t>(size);
+    int num_peers = size_ - 1;
+
+    // Counters to maintain the state of transfers
+    int in_flight = 0;
+    int64_t read_offset = 0;
+    int completed_send_count[PIPELINE] = {0};
+    int completed_recv_begin[MAX_PEERS] = {0};
+    int completed_recv_end[MAX_PEERS] = {0};
+
+    // Prefill the pipeline
+    int buff = 0;
+    while (read_offset < total && buff < PIPELINE) {
+      post_recv_all(sz, buff);
+      std::copy(
+          data + read_offset,
+          data + std::min(read_offset + N, total),
+          send_buffer(sz, buff).begin<T>());
+      post_send_all(sz, buff);
+
+      buff++;
+      in_flight += 2 * num_peers;
+      read_offset += N;
+    }
+
+    // Main loop
+    //
+    // Keep going until we have no longer data in flight.
+    while (in_flight > 0) {
+      // Poll the hardware for completions.
+      //
+      // If a send was completed mark how many completions we have received
+      // for that buffer. If we have sent the buffer to all peers we can
+      // reuse the buffer so copy the next chunk of data and send it to all.
+      //
+      // If a receive is completed then advance the pointer of completed
+      // receives.
+      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) {
+            std::copy(
+                data + read_offset,
+                data + std::min(read_offset + N, total),
+                send_buffer(sz, buff).begin<T>());
+            post_send_all(sz, buff);
+
+            completed_send_count[buff] = 0;
+            in_flight += num_peers;
+            read_offset += N;
+          }
+        }
+
+        else if (work_type == RECV_WR) {
+          completed_recv_end[rank]++;
+        }
+      }
+
+      // Process the completed recv
+      //
+      // For each rank we have a range of completed recv defined by a begin
+      // and end inclusive and exlusive in standard C++ fashion.
+      //
+      // When there is an unprocessed receive we first check if we have
+      // finished sending the write location. If so then we reduce in-place
+      // and then check if there is more to be received and post a recv.
+      for (int r = 0; r < size_; r++) {
+        int s = completed_recv_begin[r];
+        int e = completed_recv_end[r];
+        int w = s * N;
+        while (w < read_offset && e - s > 0) {
+          int buff = s % PIPELINE;
+          reduce_op(
+              recv_buffer(sz, buff, r).begin<T>(),
+              data + w,
+              std::min(N, total - w));
+          w += N;
+          s++;
+          if (w + (PIPELINE - 1) * N < total) {
+            recv_from(sz, r, buff);
+            in_flight++;
+          }
+        }
+        completed_recv_begin[r] = s;
+      }
+    }
+  });
+}
+
+} // namespace mlx::core::distributed::jaccl