Implement Group.split() for the JACCL mesh backend

mlx/distributed/jaccl/jaccl.cpp

@@ -161,7 +161,14 @@ class JACCLGroup : public GroupImpl {
   }
 
   std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
-    throw std::runtime_error("[jaccl] Group split not supported.");
+    // MPI's convention, which mlx.distributed follows: a negative key means
+    // order the child by the parent's rank.
+    auto child = group_->split(color, (key < 0) ? group_->rank() : key);
+    if (child == nullptr) {
+      throw std::runtime_error(
+          "[jaccl] A negative color leaves this rank in no group.");
+    }
+    return std::make_shared<JACCLGroup>(std::move(child));
   }
 
  private:

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

@@ -4,19 +4,36 @@
 
 #include <cstddef>
 #include <memory>
+#include <stdexcept>
 
 namespace jaccl {
 
 /**
  * Abstract base class for a JACCL communication group.
+ *
+ * Groups are created and held through shared_ptr. A group made by `split`
+ * bootstraps over its parent's side channel and so keeps the parent alive,
+ * which is why `shared_from_this` has to be available here.
  */
-class Group {
+class Group : public std::enable_shared_from_this<Group> {
  public:
   virtual ~Group() {}
 
   virtual int rank() = 0;
   virtual int size() = 0;
 
+  /**
+   * Build a new group from the members of this one that pass the same color,
+   * ordered by key and then by rank in the parent, which is the rule MPI uses.
+   *
+   * Collective over this group: every member has to call it, including a
+   * member that ends up in no child, because the colors are exchanged over
+   * this group. Pass a negative color to take part without joining a child.
+   */
+  virtual std::shared_ptr<Group> split(int color, int key) {
+    throw std::runtime_error("[jaccl] Group split not supported.");
+  }
+
   virtual void
   all_sum(const void* input, void* output, size_t n_bytes, int dtype) = 0;
 

mlx/distributed/jaccl/lib/jaccl/mesh.cpp

@@ -4,14 +4,93 @@
 #include "jaccl/reduction_ops.h"
 #include "jaccl/types.h"
 
+#include <algorithm>
+#include <sstream>
+
 namespace jaccl {
 
+std::shared_ptr<Group> MeshGroup::split(int color, int key) {
+  // Everybody learns everybody's colour and key, so each rank can work out the
+  // membership of its own child with no further agreement.
+  auto colors = side_channel_.all_gather<int>(color);
+  auto keys = side_channel_.all_gather<int>(key);
+
+  std::vector<int> members;
+  for (int r = 0; r < size_; r++) {
+    if (color >= 0 && colors[r] == color) {
+      members.push_back(r);
+    }
+  }
+  // MPI's ordering: by key, ties broken by rank in the parent.
+  std::stable_sort(members.begin(), members.end(), [&](int a, int b) {
+    return keys[a] != keys[b] ? keys[a] < keys[b] : a < b;
+  });
+
+  // A negative colour means take part but join no child. The two all_gathers
+  // above have already run, which is the point: this rank has done its share
+  // of the collective and can leave.
+  if (color < 0) {
+    return nullptr;
+  }
+
+  int child_rank =
+      std::find(members.begin(), members.end(), rank_) - members.begin();
+
+  // A child can only be built from members this rank can already reach. The
+  // hostfile is the only description of connectivity the runtime ever sees, so
+  // a pair missing there cannot be recovered here.
+  std::vector<std::string> child_devices(members.size());
+  for (size_t i = 0; i < members.size(); i++) {
+    int m = members[i];
+    if (m == rank_) {
+      continue;
+    }
+    if (device_names_[m].empty()) {
+      std::ostringstream msg;
+      msg << "[jaccl] Cannot split: rank " << rank_ << " has no device to rank "
+          << m
+          << ", which shares its colour. A child group can only contain "
+             "directly connected members.";
+      throw std::runtime_error(msg.str());
+    }
+    child_devices[i] = device_names_[m];
+  }
+
+  // The child bootstraps over the parent, so the parent has to outlive it, and
+  // holding a shared_ptr in the closure is what guarantees that.
+  //
+  // Every rank of the parent runs this on every child round, including ranks
+  // in a different child or in none, because the all_gather below is
+  // parent-wide. That is what keeps differently sized children in step: the
+  // length negotiation inside the parent's container all_gather covers the
+  // whole parent at once, so two children can never disagree about how many
+  // bytes are in flight.
+  auto self = std::static_pointer_cast<MeshGroup>(shared_from_this());
+  AllGatherFn child_all_gather =
+      [self, members](const char* src, char* dst, size_t n_bytes) {
+        auto all = self->side_channel_.all_gather(
+            std::vector<char>(src, src + n_bytes));
+        for (size_t i = 0; i < members.size(); i++) {
+          std::copy(
+              all[members[i]].begin(),
+              all[members[i]].begin() + n_bytes,
+              dst + i * n_bytes);
+        }
+      };
+
+  return std::make_shared<MeshGroup>(
+      child_rank,
+      child_devices,
+      SideChannel(child_rank, members.size(), std::move(child_all_gather)));
+}
+
 MeshGroup::MeshGroup(
     int rank,
     const std::vector<std::string>& device_names,
     SideChannel sc)
     : rank_(rank),
       size_(device_names.size()),
+      device_names_(device_names),
       side_channel_(std::move(sc)),
       connections_(create_connections(device_names)) {
   if (size_ > MESH_MAX_PEERS) {

mlx/distributed/jaccl/lib/jaccl/mesh.h

@@ -51,6 +51,8 @@ class MeshGroup : public Group {
 
   void barrier() override;
 
+  std::shared_ptr<Group> split(int color, int key) override;
+
  private:
   template <typename T, typename ReduceOp>
   void all_reduce(
@@ -80,6 +82,9 @@ class MeshGroup : public Group {
 
   int rank_;
   int size_;
+  // Kept so that `split` can build a child's row from it. Entry i is the
+  // device reaching rank i, empty for this rank itself.
+  std::vector<std::string> device_names_;
   SideChannel side_channel_;
   std::vector<Connection> connections_;
   std::vector<SharedBuffer> buffers_;

mlx/distributed/jaccl/lib/jaccl/ring.cpp

@@ -1,6 +1,8 @@
 // Copyright © 2026 Apple Inc.
 
 #include "jaccl/ring.h"
+
+#include <stdexcept>
 #include "jaccl/reduction_ops.h"
 #include "jaccl/types.h"
 
@@ -255,4 +257,12 @@ void RingGroup::reduce_scatter(
       in_ptr, out_ptr, total, n_conns_, reduce_op);
 }
 
+std::shared_ptr<Group> RingGroup::split(int color, int key) {
+  throw std::runtime_error(
+      "[jaccl] Group split is not supported for a ring. A subset of a ring is "
+      "only a ring when its members are contiguous in ring order, and there is "
+      "no path around the gap when they are not. Configure the group as a mesh "
+      "to split it.");
+}
+
 } // namespace jaccl

mlx/distributed/jaccl/lib/jaccl/ring.h

@@ -53,6 +53,14 @@ class RingGroup : public Group {
 
   void barrier() override;
 
+  /**
+   * A ring child is only a ring when its members are contiguous in ring order,
+   * and JACCL has no way to route around the gap when they are not, so this
+   * stays unimplemented rather than silently building something that is not a
+   * ring. Configure the group as a mesh if you need to split it.
+   */
+  std::shared_ptr<Group> split(int color, int key) override;
+
  private:
   template <typename T, typename ReduceOp>
   void all_reduce(