From 4400504ad5041efbc171b350559a6c9adf8abd0c Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Tue, 14 Apr 2026 23:52:21 -0700
Subject: [PATCH] Jaccl refactor (#3412)

---
 CMakeLists.txt                                |   9 +
 mlx/distributed/jaccl/.gitignore              |   1 +
 mlx/distributed/jaccl/CMakeLists.txt          |   8 +-
 mlx/distributed/jaccl/jaccl.cpp               | 281 ++++++++-------
 mlx/distributed/jaccl/jaccl.h                 |   2 +
 mlx/distributed/jaccl/lib/CMakeLists.txt      |  83 +++++
 mlx/distributed/jaccl/lib/README.md           | 328 ++++++++++++++++++
 .../jaccl/lib/examples/CMakeLists.txt         |  40 +++
 .../jaccl/lib/examples/allreduce_bench.cpp    | 286 +++++++++++++++
 .../jaccl/lib/examples/minimal_cfg.cpp        |  36 ++
 .../jaccl/lib/examples/minimal_env.cpp        |  28 ++
 mlx/distributed/jaccl/lib/jaccl/group.h       |  58 ++++
 mlx/distributed/jaccl/lib/jaccl/jaccl.cpp     | 228 ++++++++++++
 mlx/distributed/jaccl/lib/jaccl/jaccl.h       |  85 +++++
 .../jaccl/{ => lib/jaccl}/mesh.cpp            | 135 ++++---
 mlx/distributed/jaccl/{ => lib/jaccl}/mesh.h  |  51 ++-
 .../jaccl/{ => lib/jaccl}/mesh_impl.h         | 150 +++++---
 .../jaccl/{utils.cpp => lib/jaccl/rdma.cpp}   |  25 +-
 .../jaccl/{utils.h => lib/jaccl/rdma.h}       |  17 +-
 .../jaccl/lib/jaccl/reduction_ops.h           | 117 +++++++
 .../jaccl/{ => lib/jaccl}/ring.cpp            | 134 ++++---
 mlx/distributed/jaccl/{ => lib/jaccl}/ring.h  |  49 ++-
 .../jaccl/{ => lib/jaccl}/ring_impl.h         |   6 +-
 mlx/distributed/jaccl/lib/jaccl/tcp.cpp       | 206 +++++++++++
 mlx/distributed/jaccl/lib/jaccl/tcp.h         |  67 ++++
 mlx/distributed/jaccl/lib/jaccl/types.h       | 263 ++++++++++++++
 26 files changed, 2279 insertions(+), 414 deletions(-)
 create mode 100644 mlx/distributed/jaccl/.gitignore
 create mode 100644 mlx/distributed/jaccl/lib/CMakeLists.txt
 create mode 100644 mlx/distributed/jaccl/lib/README.md
 create mode 100644 mlx/distributed/jaccl/lib/examples/CMakeLists.txt
 create mode 100644 mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp
 create mode 100644 mlx/distributed/jaccl/lib/examples/minimal_cfg.cpp
 create mode 100644 mlx/distributed/jaccl/lib/examples/minimal_env.cpp
 create mode 100644 mlx/distributed/jaccl/lib/jaccl/group.h
 create mode 100644 mlx/distributed/jaccl/lib/jaccl/jaccl.cpp
 create mode 100644 mlx/distributed/jaccl/lib/jaccl/jaccl.h
 rename mlx/distributed/jaccl/{ => lib/jaccl}/mesh.cpp (56%)
 rename mlx/distributed/jaccl/{ => lib/jaccl}/mesh.h (50%)
 rename mlx/distributed/jaccl/{ => lib/jaccl}/mesh_impl.h (69%)
 rename mlx/distributed/jaccl/{utils.cpp => lib/jaccl/rdma.cpp} (95%)
 rename mlx/distributed/jaccl/{utils.h => lib/jaccl/rdma.h} (97%)
 create mode 100644 mlx/distributed/jaccl/lib/jaccl/reduction_ops.h
 rename mlx/distributed/jaccl/{ => lib/jaccl}/ring.cpp (59%)
 rename mlx/distributed/jaccl/{ => lib/jaccl}/ring.h (52%)
 rename mlx/distributed/jaccl/{ => lib/jaccl}/ring_impl.h (99%)
 create mode 100644 mlx/distributed/jaccl/lib/jaccl/tcp.cpp
 create mode 100644 mlx/distributed/jaccl/lib/jaccl/tcp.h
 create mode 100644 mlx/distributed/jaccl/lib/jaccl/types.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index b4a2d652..a2395d02 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -321,6 +321,15 @@ FetchContent_MakeAvailable(json)
 target_include_directories(
   mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
 
+# Add standalone JACCL library (RDMA over Thunderbolt distributed backend)
+if(MLX_BUILD_CPU
+   AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
+   AND DEFINED MACOS_SDK_VERSION
+   AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
+  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx/distributed/jaccl/lib
+                   ${CMAKE_BINARY_DIR}/jaccl)
+endif()
+
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
 
 target_include_directories(
diff --git a/mlx/distributed/jaccl/.gitignore b/mlx/distributed/jaccl/.gitignore
new file mode 100644
index 00000000..ebf4281d
--- /dev/null
+++ b/mlx/distributed/jaccl/.gitignore
@@ -0,0 +1 @@
+!lib
diff --git a/mlx/distributed/jaccl/CMakeLists.txt b/mlx/distributed/jaccl/CMakeLists.txt
index 567ed672..14d9d70d 100644
--- a/mlx/distributed/jaccl/CMakeLists.txt
+++ b/mlx/distributed/jaccl/CMakeLists.txt
@@ -1,12 +1,8 @@
 if(MLX_BUILD_CPU
    AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
    AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
-  target_sources(
-    mlx
-    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp
-            ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
-            ${CMAKE_CURRENT_SOURCE_DIR}/mesh.cpp
-            ${CMAKE_CURRENT_SOURCE_DIR}/ring.cpp)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp)
+  target_link_libraries(mlx PRIVATE jaccl)
 else()
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_jaccl.cpp)
 endif()
diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp
index 5935a98f..01cc415a 100644
--- a/mlx/distributed/jaccl/jaccl.cpp
+++ b/mlx/distributed/jaccl/jaccl.cpp
@@ -1,178 +1,173 @@
 // Copyright © 2025 Apple Inc.
 
-#include <fstream>
-#include <sstream>
-
-#include <json.hpp>
-
+#include "mlx/distributed/jaccl/jaccl.h"
+#include "mlx/backend/cpu/encoder.h"
 #include "mlx/distributed/distributed_impl.h"
-#include "mlx/distributed/jaccl/mesh.h"
-#include "mlx/distributed/jaccl/ring.h"
-#include "mlx/distributed/jaccl/utils.h"
+#include "mlx/dtype_utils.h"
+
+#include <jaccl/group.h>
+#include <jaccl/jaccl.h>
 
 using GroupImpl = mlx::core::distributed::detail::GroupImpl;
-using json = nlohmann::json;
+
+namespace mlx::core::distributed::jaccl {
 
 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");
-    }
+/**
+ * Map MLX Dtype to JACCL Dtype for dispatch.
+ */
+int dtype_to_jaccl_dtype(Dtype dt) {
+  switch (dt) {
+    case bool_:
+      return ::jaccl::Dtype::Bool;
+    case int8:
+      return ::jaccl::Dtype::Int8;
+    case int16:
+      return ::jaccl::Dtype::Int16;
+    case int32:
+      return ::jaccl::Dtype::Int32;
+    case int64:
+      return ::jaccl::Dtype::Int64;
+    case uint8:
+      return ::jaccl::Dtype::UInt8;
+    case uint16:
+      return ::jaccl::Dtype::UInt16;
+    case uint32:
+      return ::jaccl::Dtype::UInt32;
+    case uint64:
+      return ::jaccl::Dtype::UInt64;
+    case float16:
+      return ::jaccl::Dtype::Float16;
+    case bfloat16:
+      return ::jaccl::Dtype::BFloat16;
+    case float32:
+      return ::jaccl::Dtype::Float32;
+    case float64:
+      return ::jaccl::Dtype::Float64;
+    case complex64:
+      return ::jaccl::Dtype::Complex64;
+    default:
+      throw std::runtime_error("[jaccl] Unsupported dtype for JACCL operation");
+  }
+}
 
-    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());
-      }
+/**
+ * Adapter that wraps a standalone jaccl::Group to implement
+ * MLX's distributed::detail::GroupImpl interface.
+ *
+ * This bridges mlx::core::array to raw pointers for JACCL operations.
+ */
+class JACCLGroup : public GroupImpl {
+ public:
+  JACCLGroup(std::shared_ptr<::jaccl::Group> group)
+      : group_(std::move(group)) {}
 
-      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.");
-        }
-      }
-    }
+  Stream communication_stream(StreamOrDevice s) override {
+    return to_stream(s, Device::cpu);
   }
 
-  int size() {
-    return devices_.size();
+  int rank() override {
+    return group_->rank();
   }
 
-  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;
+  int size() override {
+    return group_->size();
   }
 
-  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;
+  void all_sum(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_->all_sum(in_ptr, out_ptr, n_bytes, dtype);
+    });
   }
 
-  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;
+  void all_max(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_->all_max(in_ptr, out_ptr, n_bytes, dtype);
+    });
   }
 
-  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]);
+  void all_min(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_->all_min(in_ptr, out_ptr, n_bytes, dtype);
+    });
   }
 
-  std::vector<std::vector<std::vector<std::string>>> devices_;
+  void all_gather(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();
+    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]() {
+      group_->all_gather(in_ptr, out_ptr, n_bytes);
+    });
+  }
+
+  void send(const array& input, int dst, Stream stream) override {
+    auto data = input.data<char>();
+    size_t n_bytes = input.nbytes();
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_input_array(input);
+    encoder.dispatch(
+        [data, n_bytes, dst, this]() { group_->send(data, n_bytes, dst); });
+  }
+
+  void recv(array& out, int src, Stream stream) override {
+    auto data = out.data<char>();
+    size_t n_bytes = out.nbytes();
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_output_array(out);
+    encoder.dispatch(
+        [data, n_bytes, src, this]() { group_->recv(data, n_bytes, src); });
+  }
+
+  void sum_scatter(const array& input, array& output, Stream stream) override {
+    throw std::runtime_error("[jaccl] sum_scatter not supported.");
+  }
+
+  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
+    throw std::runtime_error("[jaccl] Group split not supported.");
+  }
+
+ private:
+  std::shared_ptr<::jaccl::Group> group_;
 };
 
 } // namespace
 
-namespace mlx::core::distributed::jaccl {
-
 bool is_available() {
-  return ibv().is_available();
+  return ::jaccl::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());
-    }
+  auto group = ::jaccl::init(strict);
+  if (group == nullptr) {
     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.");
-  }
+  return std::make_shared<JACCLGroup>(std::move(group));
 }
 
 } // namespace mlx::core::distributed::jaccl
diff --git a/mlx/distributed/jaccl/jaccl.h b/mlx/distributed/jaccl/jaccl.h
index d07f9ccc..cc2326b6 100644
--- a/mlx/distributed/jaccl/jaccl.h
+++ b/mlx/distributed/jaccl/jaccl.h
@@ -1,5 +1,7 @@
 // Copyright © 2025 Apple Inc.
 
+#pragma once
+
 #include "mlx/distributed/distributed.h"
 
 namespace mlx::core::distributed::jaccl {
diff --git a/mlx/distributed/jaccl/lib/CMakeLists.txt b/mlx/distributed/jaccl/lib/CMakeLists.txt
new file mode 100644
index 00000000..0cfee400
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/CMakeLists.txt
@@ -0,0 +1,83 @@
+cmake_minimum_required(VERSION 3.24)
+project(jaccl LANGUAGES CXX)
+
+# Default to Release when built standalone. Without this CMake uses an empty
+# build type (-O0) which severely impacts the reduction / memcpy hot paths.
+if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
+  set(CMAKE_BUILD_TYPE
+      Release
+      CACHE STRING "Build type" FORCE)
+endif()
+
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+include(FetchContent)
+
+# nlohmann/json for device file parsing
+message(STATUS "Downloading json for JACCL")
+FetchContent_Declare(
+  json
+  URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz)
+FetchContent_MakeAvailable(json)
+
+# Check platform and SDK version requirements
+if(NOT ${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+  message(STATUS "JACCL requires macOS (Darwin). Skipping JACCL build.")
+  return()
+endif()
+
+# Try to determine MACOS_SDK_VERSION if not set
+if(NOT DEFINED MACOS_SDK_VERSION)
+  execute_process(
+    COMMAND xcrun --sdk macosx --show-sdk-version
+    OUTPUT_VARIABLE MACOS_SDK_VERSION
+    OUTPUT_STRIP_TRAILING_WHITESPACE)
+  if(MACOS_SDK_VERSION)
+    message(STATUS "Detected macOS SDK version: ${MACOS_SDK_VERSION}")
+  endif()
+endif()
+
+if(DEFINED MACOS_SDK_VERSION AND MACOS_SDK_VERSION VERSION_LESS "26.2")
+  message(STATUS "JACCL requires macOS SDK >= 26.2. Skipping JACCL build.")
+  return()
+endif()
+
+add_library(
+  jaccl
+  ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/tcp.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/rdma.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/mesh.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/ring.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/jaccl.cpp)
+
+target_include_directories(
+  jaccl PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
+               $<INSTALL_INTERFACE:include>)
+
+target_include_directories(
+  jaccl PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
+
+target_compile_features(jaccl PUBLIC cxx_std_20)
+
+# Install targets
+install(
+  TARGETS jaccl
+  EXPORT jacclTargets
+  LIBRARY DESTINATION lib
+  ARCHIVE DESTINATION lib
+  RUNTIME DESTINATION bin
+  INCLUDES
+  DESTINATION include)
+
+install(
+  DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/
+  DESTINATION include/jaccl
+  FILES_MATCHING
+  PATTERN "*.h")
+
+install(
+  EXPORT jacclTargets
+  FILE jacclTargets.cmake
+  NAMESPACE jaccl::
+  DESTINATION lib/cmake/jaccl)
diff --git a/mlx/distributed/jaccl/lib/README.md b/mlx/distributed/jaccl/lib/README.md
new file mode 100644
index 00000000..32b2defa
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/README.md
@@ -0,0 +1,328 @@
+# JACCL
+
+**JACCL** is a low-latency distributed communication library designed for macOS
+systems with Thunderbolt 5 connectivity.
+
+## Overview
+
+JACCL leverages RDMA (Remote Direct Memory Access) over Thunderbolt to achieve
+communication latency an order of magnitude lower than traditional TCP-based
+approaches. This makes it ideal for:
+
+- Tensor parallelism in large model inference
+- High-performance distributed training
+- Low-latency collective operations between Macs
+
+JACCL was made possible by Apple's RDMA over Thunderbolt technology introduced
+in macOS 26.2.
+
+## Features
+
+- **Mesh Topology**: Fully connected communication where each node can directly
+  communicate with any other node
+- **Ring Topology**: High-bandwidth ring all-reduce for large messages
+- **Collective Operations**:
+  - `all_sum`: Sum values across all nodes
+  - `all_max`: Element-wise maximum across all nodes
+  - `all_min`: Element-wise minimum across all nodes
+  - `all_gather`: Gather data from all nodes
+- **Point-to-Point Operations**:
+  - `send`: Send data to a specific node
+  - `recv`: Receive data from a specific node
+- **Type Support**: Bool, Int8-64, UInt8-64, Float16, BFloat16, Float32,
+  Float64, Complex64
+
+## Requirements
+
+- macOS SDK >= 26.2
+- Thunderbolt 5 connectivity between nodes
+- RDMA over Thunderbolt enabled (requires macOS recovery mode setup)
+
+## Enabling RDMA over Thunderbolt
+
+RDMA over Thunderbolt must be enabled in macOS recovery mode:
+
+1. Start your Mac in [recovery mode](https://support.apple.com/en-us/102518)
+2. Open Terminal from Utilities -> Terminal
+3. Run: `rdma_ctl enable`
+4. Reboot
+
+To verify RDMA is enabled, run:
+
+```bash
+ibv_devices
+```
+
+You should see output like:
+
+```
+device          	   node GUID
+------          	----------------
+rdma_en2        	8096a9d9edbaac05
+rdma_en3        	8196a9d9edbaac05
+rdma_en5        	8396a9d9edbaac05
+```
+
+## Building
+
+JACCL can be built as a standalone library:
+
+```bash
+cd mlx/distributed/jaccl/lib
+mkdir build && cd build
+cmake ..
+make
+```
+
+You can also include it in your own project via CMake:
+
+```
+FetchContent_Declare(
+  jaccl
+  GIT_REPOSITORY https://github.com/ml-explore/mlx.git
+  GIT_TAG main
+  SOURCE_SUBDIR mlx/distributed/jaccl/lib
+)
+FetchContent_MakeAvailable(jaccl)
+```
+
+## Usage
+
+### Environment Variables
+
+The easiest way to intiialize JACCL is by using the following environment
+variables:
+
+- **JACCL_RANK** / **MLX_RANK**: The rank of this process (0-based integer)
+- **JACCL_IBV_DEVICES** / **MLX_IBV_DEVICES**: Path to a JSON file describing
+  device connectivity
+- **JACCL_COORDINATOR** / **MLX_JACCL_COORDINATOR**: IP:port of the coordinator
+  (rank 0 listener)
+- **JACCL_RING** / **MLX_JACCL_RING**: (Optional) Prefer ring topology over
+  mesh
+
+### Device File Format
+
+The device file is a JSON array where each entry describes the RDMA devices
+connecting that rank to all other ranks:
+
+```json
+[
+    [null, "rdma_en5", "rdma_en4", "rdma_en3"],
+    ["rdma_en5", null, "rdma_en3", "rdma_en4"],
+    ["rdma_en4", "rdma_en3", null, "rdma_en5"],
+    ["rdma_en3", "rdma_en4", "rdma_en5", null]
+]
+```
+
+For a valid mesh, `devices[i][j]` should contain the device name connecting
+rank `i` to rank `j`, or `null` if `i == j`.
+
+For a valid ring, only adjacent nodes should have device names (all others
+should be null).
+
+### Basic Example
+
+```cpp
+#include <iostream>
+#include <jaccl/jaccl.h>
+
+int main() {
+  // Initialize JACCL group
+  auto group = jaccl::init();
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl;
+
+  // Perform all-reduce sum
+  float input[10] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f};
+  float output[10];
+
+  group->all_sum(input, output, sizeof(input), jaccl::Float32);
+
+  std::cout << "Result: " << output[0] << std::endl;
+
+  return 0;
+}
+```
+
+You can also manually define the configuration instead of reading it from
+environment variables.
+
+```cpp
+#include <iostream>
+#include <jaccl/jaccl.h>
+
+int main() {
+  auto cfg = jaccl::Config()
+      .set_rank(0)
+      .set_coordinator("192.168.1.1:32132")
+      .set_devices({
+        {{}, {"rdma_en5"}, {"rdma_en4"}, {"rdma_en3"}},
+        {{"rdma_en5"}, {}, {"rdma_en3"}, {"rdma_en4"}},
+        {{"rdma_en4"}, {"rdma_en3"}, {}, {"rdma_en5"}},
+        {{"rdma_en3"}, {"rdma_en4"}, {"rdma_en5"}, {}}
+      });
+  auto group = jaccl::init(cfg);
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl;
+
+  // Perform all-reduce sum
+  float input[10] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f};
+  float output[10];
+
+  group->all_sum(input, output, sizeof(input), jaccl::Float32);
+
+  std::cout << "Result: " << output[0] << std::endl;
+
+  return 0;
+}
+```
+
+### Using with MLX
+
+JACCL integrates seamlessly with MLX's distributed communication:
+
+```python
+import mlx.core as mx
+
+# Initialize with JACCL backend
+world = mx.distributed.init(backend="jaccl")
+
+# Perform distributed operations
+x = mx.ones((10,))
+result = mx.distributed.all_sum(x, group=world)
+```
+
+Launch with `mlx.launch`:
+
+```bash
+mlx.launch --backend jaccl --hostfile hosts.json my_script.py
+```
+
+## Hostfile Example
+
+For use with `mlx.launch`, create a hostfile JSON:
+
+```json
+{
+    "backend": "jaccl",
+    "hosts": [
+        {
+            "ssh": "m3-ultra-1",
+            "ips": ["192.168.1.1"],
+            "rdma": [null, "rdma_en5", "rdma_en4", "rdma_en3"]
+        },
+        {
+            "ssh": "m3-ultra-2",
+            "ips": [],
+            "rdma": ["rdma_en5", null, "rdma_en3", "rdma_en4"]
+        },
+        {
+            "ssh": "m3-ultra-3",
+            "ips": [],
+            "rdma": ["rdma_en4", "rdma_en3", null, "rdma_en5"]
+        },
+        {
+            "ssh": "m3-ultra-4",
+            "ips": [],
+            "rdma": ["rdma_en3", "rdma_en4", "rdma_en5", null]
+        }
+    ]
+}
+```
+
+## Automatic Configuration
+
+MLX provides `mlx.distributed_config` to automatically discover and configure
+Thunderbolt connectivity:
+
+```bash
+# Visualize connections
+mlx.distributed_config --verbose \
+     --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \
+     --over thunderbolt --dot | dot -Tpng | open -f -a Preview
+
+# Auto-configure and generate hostfile
+mlx.distributed_config --verbose \
+     --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \
+     --over thunderbolt --backend jaccl \
+     --auto-setup --output m3-ultra-jaccl.json
+```
+
+## API
+
+The main API of JACCL is the communication group. It provides efficient
+high-level collectives.
+
+**Note: JACCL does no memory allocation. All output pointers should point to a
+location with sufficient memory allocated to hold the result.**
+
+```cpp
+class Group {
+ public:
+  virtual ~Group() {}
+
+  // Helper functions to know which process we are in the group
+  virtual int rank() = 0;
+  virtual int size() = 0;
+
+  // All reduce implementations. Input and output of the same size the
+  // reduction happens according to dtype and across the group.
+  virtual void all_sum(const void* input, void* output, size_t n_bytes, int dtype) = 0;
+  virtual void all_max(const void* input, void* output, size_t n_bytes, int dtype) = 0;
+  virtual void all_min(const void* input, void* output, size_t n_bytes, int dtype) = 0;
+
+  // All gather implementation. The output is group->size() * n_bytes.
+  virtual void all_gather(const void* input, void* output, size_t n_bytes) = 0;
+
+  // Simple send/recv primitives.
+  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;
+};
+```
+
+All that is left to use JACCL (except the communication group) is
+
+```cpp
+std::shared_ptr<Group> init(bool strict = false);
+std::shared_ptr<Group> init(const Config& cfg, bool strict = false);
+```
+
+that create the communication group from environment variables or from the
+configuration object. The latter allows one to configure JACCL using means
+other than environment variables.
+
+```cpp
+class Config {
+ public:
+
+  Config();
+
+  Config& set_rank(int rank);
+  Config& set_coordinator(std::string coordinator);
+  Config& set_devices(std::vector<std::vector<std::vector<std::string>>> devices);
+  Config& prefer_ring(bool prefer = true);
+
+  bool is_valid_mesh() const;
+  bool is_valid_ring() const;
+}
+```
+
+## License
+
+JACCL is part of MLX and is released under the same license.
+
+## Acknowledgments
+
+The name JACCL (pronounced Jackal) stands for Jack and Angelos’ Collective
+Communication Library and it is an obvious pun to Nvidia’s NCCL but also
+tribute to Jack Beasley who led the development of RDMA over Thunderbolt at
+Apple.
diff --git a/mlx/distributed/jaccl/lib/examples/CMakeLists.txt b/mlx/distributed/jaccl/lib/examples/CMakeLists.txt
new file mode 100644
index 00000000..48b7522d
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/examples/CMakeLists.txt
@@ -0,0 +1,40 @@
+cmake_minimum_required(VERSION 3.24)
+project(jaccl_examples LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+include(FetchContent)
+
+# cmake-format: off
+#
+# Treating the nearby directory like a fetched dependency to simulate what a
+# user would do in their own C++ project. Ideally you would write something like
+# the following:
+#
+# FetchContent_Declare(
+#   jaccl
+#   GIT_REPOSITORY https://github.com/ml-explore/mlx.git
+#   GIT_TAG main
+#   SOURCE_SUBDIR mlx/distributed/jaccl/lib
+# )
+#
+# cmake-format: on
+FetchContent_Declare(jaccl SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../)
+FetchContent_MakeAvailable(jaccl)
+
+# Helper function to build examples
+function(build_example SRCFILE)
+  get_filename_component(example_name ${SRCFILE} NAME_WE)
+  set(target "jaccl_${example_name}")
+  add_executable(${target} ${SRCFILE})
+  target_link_libraries(${target} PRIVATE jaccl)
+  message(STATUS "Building JACCL example: ${target}")
+endfunction()
+
+# Examples
+build_example(minimal_env.cpp)
+build_example(minimal_cfg.cpp)
+
+# Benchmarks
+build_example(allreduce_bench.cpp)
diff --git a/mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp b/mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp
new file mode 100644
index 00000000..ae1fa5a7
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp
@@ -0,0 +1,286 @@
+// Copyright © 2025 Apple Inc.
+//
+// JACCL All-Reduce Benchmark
+//
+// Measures bandwidth and latency of all_sum across a sweep of message sizes,
+// similar in spirit to the NCCL all-reduce benchmark (nccl-tests).
+//
+// Usage:
+//   Set the environment variables described in jaccl.h, then run:
+//
+//     ./jaccl_allreduce_bench [-w <warmup_iters>] [-n <iters>]
+//                             [-b <min_bytes>] [-e <max_bytes>]
+//                             [-f <step_factor>] [-d <datatype>]
+//                             [-c] [-h]
+//
+//   Or use the MLX launcher:
+//
+//     mlx.launch --hostfile hosts.json ./jaccl_allreduce_bench
+//
+//   The arguments are:
+//
+//   -w  Warmup iterations per message size  (default: 5)
+//   -n  Timed iterations per message size   (default: 20)
+//   -b  Minimum message size in bytes       (default: 1K)
+//   -e  Maximum message size in bytes       (default: 256M)
+//   -f  Multiplicative step factor          (default: 2)
+//   -d  Datatype: float32, float16, bfloat16 (default: float32)
+//   -c  Check correctness                   (default: off)
+//   -h  Print this help message
+
+#include <jaccl/jaccl.h>
+#include <jaccl/types.h>
+
+#include <chrono>
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include <iomanip>
+#include <iostream>
+#include <memory>
+#include <numeric>
+#include <string>
+#include <vector>
+
+static void usage(const char* prog) {
+  std::cerr
+      << "Usage: " << prog << " [options]\n"
+      << "  -w <warmup>    Warmup iterations          (default: 5)\n"
+      << "  -n <iters>     Timed iterations            (default: 20)\n"
+      << "  -b <min_bytes> Minimum message size        (default: 1K)\n"
+      << "  -e <max_bytes> Maximum message size        (default: 256M)\n"
+      << "  -f <factor>    Multiplicative step factor  (default: 2)\n"
+      << "  -d <dtype>     float32|float16|bfloat16    (default: float32)\n"
+      << "  -c             Check correctness\n"
+      << "  -h             Show this help\n";
+}
+
+static size_t parse_size(const char* s) {
+  char* end = nullptr;
+  double val = std::strtod(s, &end);
+  if (end && (*end == 'K' || *end == 'k'))
+    val *= 1024;
+  else if (end && (*end == 'M' || *end == 'm'))
+    val *= 1024 * 1024;
+  else if (end && (*end == 'G' || *end == 'g'))
+    val *= 1024 * 1024 * 1024;
+  return static_cast<size_t>(val);
+}
+
+static std::string fmt_bytes(size_t bytes) {
+  const char* units[] = {"B", "KB", "MB", "GB", "TB"};
+  int idx = 0;
+  double val = static_cast<double>(bytes);
+  while (val >= 1024.0 && idx < 4) {
+    val /= 1024.0;
+    idx++;
+  }
+  char buf[32];
+  if (val == static_cast<int>(val))
+    std::snprintf(buf, sizeof(buf), "%d %s", static_cast<int>(val), units[idx]);
+  else
+    std::snprintf(buf, sizeof(buf), "%.2f %s", val, units[idx]);
+  return buf;
+}
+
+// Conversion from algorithm bandwidth to bus bandwidth for a ring reduce.
+static double bus_factor(int nranks) {
+  return 2.0 * (nranks - 1) / static_cast<double>(nranks);
+}
+
+int main(int argc, char** argv) {
+  int warmup_iters = 5;
+  int timed_iters = 20;
+  size_t min_bytes = 1024;
+  size_t max_bytes = 256 * 1024 * 1024;
+  int step_factor = 2;
+  std::string dtype_str = "float32";
+  bool check = 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 == "-w" && i + 1 < argc) {
+      warmup_iters = std::atoi(argv[++i]);
+    } else if (arg == "-n" && i + 1 < argc) {
+      timed_iters = std::atoi(argv[++i]);
+    } else if (arg == "-b" && i + 1 < argc) {
+      min_bytes = parse_size(argv[++i]);
+    } else if (arg == "-e" && i + 1 < argc) {
+      max_bytes = parse_size(argv[++i]);
+    } else if (arg == "-f" && i + 1 < argc) {
+      step_factor = std::atoi(argv[++i]);
+    } else if (arg == "-d" && i + 1 < argc) {
+      dtype_str = argv[++i];
+    } else if (arg == "-c") {
+      check = true;
+    } else {
+      std::cerr << "Unknown option: " << arg << "\n";
+      usage(argv[0]);
+      return 1;
+    }
+  }
+
+  jaccl::Dtype dtype;
+  size_t elem_size;
+  if (dtype_str == "float32") {
+    dtype = jaccl::Float32;
+    elem_size = 4;
+  } else if (dtype_str == "float16") {
+    dtype = jaccl::Float16;
+    elem_size = 2;
+  } else if (dtype_str == "bfloat16") {
+    dtype = jaccl::BFloat16;
+    elem_size = 2;
+  } else {
+    std::cerr << "Unsupported dtype: " << dtype_str << "\n";
+    return 1;
+  }
+
+  auto group = jaccl::init(/* strict= */ true);
+
+  int rank = group->rank();
+  int nranks = group->size();
+
+  if (rank == 0) {
+    std::cout << "# JACCL All-Reduce Benchmark\n"
+              << "# Ranks:   " << nranks << "\n"
+              << "# Dtype:   " << dtype_str << "\n"
+              << "# Warmup:  " << warmup_iters << " iters\n"
+              << "# Timed:   " << timed_iters << " iters\n"
+              << "# Sizes:   " << fmt_bytes(min_bytes) << " .. "
+              << fmt_bytes(max_bytes) << " (x" << step_factor << ")\n"
+              << "#\n";
+
+    // Table header (NCCL-style)
+    std::cout << std::left << std::setw(14) << "#    size" << std::right
+              << std::setw(12) << "count" << std::setw(12) << "type"
+              << std::setw(14) << "time (us)" << std::setw(14) << "algo BW"
+              << std::setw(14) << "bus BW";
+    if (check)
+      std::cout << std::setw(10) << "check";
+    std::cout << "\n";
+    std::cout << std::left << std::setw(14) << "#   (bytes)" << std::right
+              << std::setw(12) << "(elems)" << std::setw(12) << ""
+              << std::setw(14) << "" << std::setw(14) << "(GB/s)"
+              << std::setw(14) << "(GB/s)";
+    if (check)
+      std::cout << std::setw(10) << "";
+    std::cout << "\n";
+  }
+
+  size_t max_elems = max_bytes / elem_size;
+  max_bytes = max_elems * elem_size;
+
+  std::vector<char> sendbuf(max_bytes);
+
+  // Fill send buffer with a simple deterministic pattern (rank + 1) casted to
+  // the target type, so correctness checks are straightforward: after all_sum
+  // every element should equal sum_{r=0}^{nranks-1} (r + 1) =
+  // nranks*(nranks+1)/2.
+  auto fill_buffer = [&](char* buf, size_t n_bytes) {
+    float val = static_cast<float>(rank + 1);
+    if (dtype == jaccl::Float32) {
+      auto* p = reinterpret_cast<float*>(buf);
+      size_t n = n_bytes / sizeof(float);
+      for (size_t i = 0; i < n; i++) {
+        p[i] = val;
+      }
+    } else if (dtype == jaccl::Float16) {
+      // Write via the library's float16_t
+      auto* p = reinterpret_cast<jaccl::float16_t*>(buf);
+      size_t n = n_bytes / sizeof(jaccl::float16_t);
+      for (size_t i = 0; i < n; i++) {
+        p[i] = jaccl::float16_t(val);
+      }
+    } else if (dtype == jaccl::BFloat16) {
+      auto* p = reinterpret_cast<jaccl::bfloat16_t*>(buf);
+      size_t n = n_bytes / sizeof(jaccl::bfloat16_t);
+      for (size_t i = 0; i < n; i++) {
+        p[i] = jaccl::bfloat16_t(val);
+      }
+    }
+  };
+
+  auto check_buffer = [&](const char* buf, size_t n_bytes) -> bool {
+    float expected = static_cast<float>(nranks) * (nranks + 1) / 2.0f;
+    float tol = (dtype == jaccl::Float32) ? 1e-5f : 1e-1f;
+    size_t n = n_bytes / elem_size;
+
+    for (size_t i = 0; i < n; i++) {
+      float val;
+      if (dtype == jaccl::Float32) {
+        val = reinterpret_cast<const float*>(buf)[i];
+      } else if (dtype == jaccl::Float16) {
+        val = static_cast<float>(
+            reinterpret_cast<const jaccl::float16_t*>(buf)[i]);
+      } else {
+        val = static_cast<float>(
+            reinterpret_cast<const jaccl::bfloat16_t*>(buf)[i]);
+      }
+      if (std::abs(val - expected) > tol) {
+        return false;
+      }
+    }
+    return true;
+  };
+
+  double bf = bus_factor(nranks);
+
+  for (size_t nbytes = min_bytes; nbytes <= max_bytes;
+       nbytes *= static_cast<size_t>(step_factor)) {
+    // Round down to element boundary
+    size_t n = std::max((nbytes / elem_size) * elem_size, elem_size);
+    size_t count = n / elem_size;
+
+    fill_buffer(sendbuf.data(), n);
+
+    // Warmup
+    for (int i = 0; i < warmup_iters; i++) {
+      group->all_sum(sendbuf.data(), sendbuf.data(), n, dtype);
+    }
+
+    // Timed iterations
+    auto t0 = std::chrono::high_resolution_clock::now();
+    for (int i = 0; i < timed_iters; i++) {
+      group->all_sum(sendbuf.data(), sendbuf.data(), n, dtype);
+    }
+    auto t1 = std::chrono::high_resolution_clock::now();
+
+    double elapsed_us =
+        std::chrono::duration<double, std::micro>(t1 - t0).count();
+    double avg_us = elapsed_us / timed_iters;
+
+    // Bandwidth in GB/s
+    double algo_bw = (static_cast<double>(n) / avg_us) / 1e3;
+    double bus_bw = algo_bw * bf;
+
+    // Correctness check
+    std::string check_result;
+    if (check) {
+      fill_buffer(sendbuf.data(), n);
+      group->all_sum(sendbuf.data(), sendbuf.data(), n, dtype);
+      check_result = check_buffer(sendbuf.data(), n) ? "OK" : "FAIL";
+    }
+
+    if (rank == 0) {
+      std::cout << std::left << std::setw(14) << n << std::right
+                << std::setw(12) << count << std::setw(12) << dtype_str
+                << std::setw(14) << std::fixed << std::setprecision(1) << avg_us
+                << std::setw(14) << std::fixed << std::setprecision(2)
+                << algo_bw << std::setw(14) << std::fixed
+                << std::setprecision(2) << bus_bw;
+      if (check)
+        std::cout << std::setw(10) << check_result;
+      std::cout << "\n";
+    }
+  }
+
+  if (rank == 0) {
+    std::cout << "# Done.\n";
+  }
+
+  return 0;
+}
diff --git a/mlx/distributed/jaccl/lib/examples/minimal_cfg.cpp b/mlx/distributed/jaccl/lib/examples/minimal_cfg.cpp
new file mode 100644
index 00000000..4aec4b2e
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/examples/minimal_cfg.cpp
@@ -0,0 +1,36 @@
+#include <jaccl/jaccl.h>
+#include <iostream>
+
+int main() {
+  auto cfg =
+      jaccl::Config()
+          .set_rank(0) // should be different per node
+          .set_coordinator("192.168.1.1:32132") // rank 0 will listen here
+          .set_devices(
+              {{{}, {"rdma_en5"}, {"rdma_en4"}, {"rdma_en3"}},
+               {{"rdma_en5"}, {}, {"rdma_en3"}, {"rdma_en4"}},
+               {{"rdma_en4"}, {"rdma_en3"}, {}, {"rdma_en5"}},
+               {{"rdma_en3"}, {"rdma_en4"}, {"rdma_en5"}, {}}});
+  auto group = jaccl::init(cfg);
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl;
+
+  // Perform all-reduce sum
+  float input[10] = {
+      1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f};
+  float output[10];
+
+  group->all_sum(input, output, sizeof(input), jaccl::Float32);
+
+  std::cout << "Result: ";
+  for (auto o : output) {
+    std::cout << o << " ";
+  }
+  std::cout << std::endl;
+
+  return 0;
+}
diff --git a/mlx/distributed/jaccl/lib/examples/minimal_env.cpp b/mlx/distributed/jaccl/lib/examples/minimal_env.cpp
new file mode 100644
index 00000000..2bc04da5
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/examples/minimal_env.cpp
@@ -0,0 +1,28 @@
+#include <jaccl/jaccl.h>
+#include <iostream>
+
+int main() {
+  // Initialize JACCL group
+  auto group = jaccl::init();
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl;
+
+  // Perform all-reduce sum
+  float input[10] = {
+      1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f};
+  float output[10];
+
+  group->all_sum(input, output, sizeof(input), jaccl::Float32);
+
+  std::cout << "Result: ";
+  for (auto o : output) {
+    std::cout << o << " ";
+  }
+  std::cout << std::endl;
+
+  return 0;
+}
diff --git a/mlx/distributed/jaccl/lib/jaccl/group.h b/mlx/distributed/jaccl/lib/jaccl/group.h
new file mode 100644
index 00000000..495458c5
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/jaccl/group.h
@@ -0,0 +1,58 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <cstddef>
+#include <memory>
+
+namespace jaccl {
+
+/**
+ * Abstract base class for a JACCL communication group.
+ */
+class Group {
+ public:
+  virtual ~Group() {}
+
+  virtual int rank() = 0;
+  virtual int size() = 0;
+
+  virtual void
+  all_sum(const void* input, void* output, size_t n_bytes, int dtype) = 0;
+
+  virtual void
+  all_max(const void* input, void* output, size_t n_bytes, int dtype) = 0;
+
+  virtual void
+  all_min(const void* input, void* output, size_t n_bytes, int dtype) = 0;
+
+  virtual void all_gather(const void* input, void* output, size_t n_bytes) = 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;
+};
+
+/**
+ * Type IDs for dispatch in the standalone JACCL library.
+ *
+ * Users pass one of these to all_sum/all_max/all_min so JACCL knows how to
+ * interpret the data for typed reduction operations.
+ */
+enum Dtype {
+  Bool = 0,
+  Int8,
+  Int16,
+  Int32,
+  Int64,
+  UInt8,
+  UInt16,
+  UInt32,
+  UInt64,
+  Float16,
+  BFloat16,
+  Float32,
+  Float64,
+  Complex64,
+};
+
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp b/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp
new file mode 100644
index 00000000..b7b68578
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp
@@ -0,0 +1,228 @@
+// Copyright © 2025 Apple Inc.
+
+#include <fstream>
+#include <sstream>
+
+#include <json.hpp>
+
+#include "jaccl/jaccl.h"
+#include "jaccl/mesh.h"
+#include "jaccl/rdma.h"
+#include "jaccl/ring.h"
+
+using json = nlohmann::json;
+
+namespace {
+
+std::vector<std::vector<std::vector<std::string>>> parse_devices_json(
+    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");
+  }
+
+  std::vector<std::vector<std::vector<std::string>>> result(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());
+    }
+
+    result[rank].resize(conn.size());
+    for (int dst = 0; dst < conn.size(); dst++) {
+      auto names = conn[dst];
+      if (names.is_string()) {
+        result[rank][dst].push_back(names);
+      } else if (names.is_array()) {
+        for (auto name_it = names.begin(); name_it != names.end(); name_it++) {
+          result[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.");
+      }
+    }
+  }
+
+  return result;
+}
+
+template <typename First, typename... Rest>
+const char* getenv(First first, Rest... rest) {
+  const char* rs = std::getenv(first);
+  if (rs != nullptr) {
+    return rs;
+  }
+  if constexpr (sizeof...(rest) > 0) {
+    return getenv(rest...);
+  }
+  return rs;
+}
+
+} // namespace
+
+namespace jaccl {
+
+Config::Config() : rank_(0), size_(0) {}
+
+Config& Config::set_rank(int rank) {
+  rank_ = rank;
+  return *this;
+}
+
+Config& Config::set_coordinator(std::string coordinator) {
+  coordinator_ = std::move(coordinator);
+  return *this;
+}
+
+Config& Config::set_devices(
+    std::vector<std::vector<std::vector<std::string>>> devices) {
+  devices_ = std::move(devices);
+  size_ = devices_.size();
+  for (int r = 0; r < size_; r++) {
+    if (size_ != devices_[r].size()) {
+      std::ostringstream msg;
+      msg << "[jaccl] The full connectivity matrix should be provided but we have "
+          << size_ << " rows and row " << r << " has " << devices_[r].size()
+          << " columns.";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  return *this;
+}
+
+Config& Config::prefer_ring(bool prefer /* = true */) {
+  prefer_ring_ = prefer;
+  return *this;
+}
+
+bool Config::is_valid_mesh() const {
+  if (size_ < 2) {
+    return false;
+  }
+
+  for (int src = 0; src < size_; src++) {
+    for (int dst = 0; dst < size_; dst++) {
+      if ((src == dst && devices_[src][dst].size() != 0) ||
+          (src != dst && devices_[src][dst].size() == 0)) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+bool Config::is_valid_ring() const {
+  if (size_ < 2) {
+    return false;
+  }
+
+  int num_connections = devices_[0][1].size();
+  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() != num_connections) {
+          return false;
+        }
+      }
+    }
+  }
+  return true;
+}
+
+std::vector<std::string> Config::get_mesh_connectivity() const {
+  if (!is_valid_mesh()) {
+    throw std::runtime_error("[jaccl] The devices do not form a valid mesh.");
+  }
+  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>>
+Config::get_ring_connectivity() const {
+  if (!is_valid_ring()) {
+    throw std::runtime_error("[jaccl] The devices do not form a valid ring.");
+  }
+  int left = (rank_ + size_ - 1) % size_;
+  int right = (rank_ + 1) % size_;
+
+  return std::make_pair(devices_[rank_][left], devices_[rank_][right]);
+}
+
+std::optional<Config> Config::from_env() {
+  const char* dev_file = getenv("JACCL_IBV_DEVICES", "MLX_IBV_DEVICES");
+  const char* coordinator =
+      getenv("JACCL_COORDINATOR", "MLX_JACCL_COORDINATOR");
+  const char* rank_str = getenv("JACCL_RANK", "MLX_RANK");
+  const char* ring = getenv("JACCL_RING", "MLX_JACCL_RING");
+
+  if (!dev_file || !coordinator || !rank_str) {
+    return std::nullopt;
+  }
+
+  return Config()
+      .set_rank(std::atoi(rank_str))
+      .set_coordinator(coordinator)
+      .set_devices(parse_devices_json(dev_file))
+      .prefer_ring(ring != nullptr);
+}
+
+bool is_available() {
+  return ibv().is_available();
+}
+
+std::shared_ptr<Group> init(bool strict /* = false */) {
+  auto cfg = Config::from_env();
+  if (!cfg.has_value() && strict) {
+    std::ostringstream msg;
+    msg << "[jaccl] You need to provide via environment variables a rank "
+        << "(JACCL_RANK/MLX_RANK), a device file (JACCL_IBV_DEVICES/"
+        << "MLX_IBV_DEVICES) and a coordinator ip/port (JACCL_COORDINATOR/"
+        << "MLX_JACCL_COORDINATOR).";
+    throw std::runtime_error(msg.str());
+  }
+
+  return init(*cfg, strict);
+}
+
+std::shared_ptr<Group> init(const Config& cfg, bool strict /* = false */) {
+  if (cfg.get_prefer_ring() && cfg.is_valid_ring()) {
+    auto [left, right] = cfg.get_ring_connectivity();
+    return std::make_shared<RingGroup>(
+        cfg.get_rank(), cfg.get_size(), left, right, cfg.get_coordinator());
+  } else if (cfg.is_valid_mesh()) {
+    auto mesh = cfg.get_mesh_connectivity();
+    return std::make_shared<MeshGroup>(
+        cfg.get_rank(), mesh, cfg.get_coordinator());
+  } else if (cfg.is_valid_ring()) {
+    auto [left, right] = cfg.get_ring_connectivity();
+    return std::make_shared<RingGroup>(
+        cfg.get_rank(), cfg.get_size(), left, right, cfg.get_coordinator());
+  } else {
+    if (!strict) {
+      return nullptr;
+    }
+
+    throw std::runtime_error(
+        "[jaccl] The configuration should define a valid mesh or a valid ring.");
+  }
+}
+
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/lib/jaccl/jaccl.h b/mlx/distributed/jaccl/lib/jaccl/jaccl.h
new file mode 100644
index 00000000..a4df21ae
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/jaccl/jaccl.h
@@ -0,0 +1,85 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "jaccl/group.h"
+
+namespace jaccl {
+
+class Config {
+ public:
+  Config();
+
+  Config& set_rank(int rank);
+  Config& set_coordinator(std::string coordinator);
+  Config& set_devices(
+      std::vector<std::vector<std::vector<std::string>>> devices);
+  Config& prefer_ring(bool prefer = true);
+
+  bool is_valid_mesh() const;
+  bool is_valid_ring() const;
+
+  int get_rank() const {
+    return rank_;
+  }
+
+  int get_size() const {
+    return size_;
+  }
+
+  std::string get_coordinator() const {
+    return coordinator_;
+  }
+
+  bool get_prefer_ring() const {
+    return prefer_ring_;
+  }
+
+  static std::optional<Config> from_env();
+
+  friend std::shared_ptr<Group> init(const Config& cfg, bool strict);
+
+ private:
+  std::vector<std::string> get_mesh_connectivity() const;
+  std::pair<std::vector<std::string>, std::vector<std::string>>
+  get_ring_connectivity() const;
+
+  int rank_;
+  int size_;
+  std::string coordinator_;
+  std::vector<std::vector<std::vector<std::string>>> devices_;
+  bool prefer_ring_;
+};
+
+/**
+ * Check if JACCL (RDMA over Thunderbolt) is available on this system.
+ */
+bool is_available();
+
+/**
+ * Initialize a JACCL communication group from environment variables.
+ *
+ * Reads configuration from environment variables:
+ *   - JACCL_RANK / MLX_RANK: The rank of this process
+ *   - JACCL_IBV_DEVICES / MLX_IBV_DEVICES: Path to the device connectivity
+ *     JSON file
+ *   - JACCL_COORDINATOR / MLX_JACCL_COORDINATOR: IP:port of the coordinator
+ *   - JACCL_RING / MLX_JACCL_RING: If set, prefer ring topology
+ *
+ * Args:
+ *   strict: If true, throw on failure. If false, return nullptr.
+ *
+ * Returns:
+ *   A shared_ptr to the Group, or nullptr on failure.
+ */
+std::shared_ptr<Group> init(bool strict = false);
+
+/**
+ * Initialize a JACCL communication group from an explicit Config object.
+ */
+std::shared_ptr<Group> init(const Config& cfg, bool strict = false);
+
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/mesh.cpp b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp
similarity index 56%
rename from mlx/distributed/jaccl/mesh.cpp
rename to mlx/distributed/jaccl/lib/jaccl/mesh.cpp
index 52f43b94..6ea1da58 100644
--- a/mlx/distributed/jaccl/mesh.cpp
+++ b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp
@@ -1,19 +1,18 @@
 // 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"
+#include "jaccl/mesh.h"
+#include "jaccl/reduction_ops.h"
+#include "jaccl/types.h"
 
-namespace mlx::core::distributed::jaccl {
+namespace jaccl {
 
 MeshGroup::MeshGroup(
     int rank,
     const std::vector<std::string>& device_names,
-    const char* coordinator_addr)
+    const std::string& coordinator_addr)
     : rank_(rank),
       size_(device_names.size()),
-      side_channel_(rank_, size_, coordinator_addr),
+      side_channel_(rank_, size_, coordinator_addr.c_str()),
       connections_(create_connections(device_names)) {
   if (size_ > MESH_MAX_PEERS) {
     std::ostringstream msg;
@@ -26,7 +25,7 @@ MeshGroup::MeshGroup(
   initialize();
 
   // Make sure every node has reached here before continuing
-  side_channel_.all_gather<int>(0);
+  side_channel_.barrier();
 
   // Create the mesh implementation object
   mesh_ = MeshImpl(rank_, size_, connections_, buffers_);
@@ -61,9 +60,9 @@ void MeshGroup::initialize() {
     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.
+  // 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());
@@ -106,29 +105,27 @@ void MeshGroup::allocate_buffers() {
     for (int i = 0; i < NUM_BUFFERS; i++) {
       // Mesh buffers
       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_) {
+          // This is our send buffer so register it with all pds so we can
+          // send it to all connected devices.
           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 {
+        } else {
+          // This is the recv buffer from rank j so register it to rank j's
+          // protection domain.
           buffers_[k * NUM_BUFFERS * size_ + i * size_ + j]
               .register_to_protection_domain(connections_[j].protection_domain);
         }
       }
 
       // Ring buffers (see ring group for the logic below)
-      // We register send buffers to both the right and the left.
       int left = (rank_ + size_ - 1) % size_;
       int right = (rank_ + 1) % size_;
+      // We register send buffers to both the right and the left.
       ring_send_buffers_[k * NUM_BUFFERS * 2 + i * 2 + 0]
           .register_to_protection_domain(connections_[right].protection_domain);
       ring_recv_buffers_[k * NUM_BUFFERS * 2 + i * 2 + 0]
@@ -141,78 +138,68 @@ void MeshGroup::allocate_buffers() {
   }
 }
 
-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_sum(
+    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);
+    all_reduce<T>(input, output, n_bytes, 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_max(
+    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);
+    all_reduce<T>(input, output, n_bytes, 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_min(
+    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);
+    all_reduce<T>(input, output, n_bytes, 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]() {
-    mesh_.all_gather(in_ptr, out_ptr, n_bytes);
-  });
+void MeshGroup::all_gather(const void* input, void* output, size_t n_bytes) {
+  mesh_.all_gather(
+      static_cast<const char*>(input), static_cast<char*>(output), n_bytes);
 }
 
-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]() { mesh_.send(data, n_bytes, dst); });
+void MeshGroup::send(const void* input, size_t n_bytes, int dst) {
+  mesh_.send(static_cast<const char*>(input), n_bytes, dst);
 }
 
-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]() { mesh_.recv(data, n_bytes, src); });
+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::all_reduce(
-    const array& input,
-    array& output,
-    Stream stream,
+    const void* input,
+    void* output,
+    size_t n_bytes,
     ReduceOp reduce_op) {
-  auto in_ptr = input.data<T>();
-  auto out_ptr = output.data<T>();
-  int64_t size = input.size();
-  auto& encoder = cpu::get_command_encoder(stream);
-  encoder.set_input_array(input);
-  encoder.set_output_array(output);
-  encoder.dispatch([in_ptr, out_ptr, size, this, reduce_op]() {
-    if (size_ > 2 &&
-        ((std::is_same_v<T, bfloat16_t> && size > 65536) ||
-         size >= 8 * 1024 * 1024 / sizeof(T))) {
-      ring_.all_reduce<2>(in_ptr, out_ptr, size, 1, reduce_op);
-    } else {
-      mesh_.all_reduce(in_ptr, out_ptr, size, 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 (size_ > 2 &&
+      ((std::is_same_v<T, bfloat16_t> && count > 256 * 1024) ||
+       count >= 8 * 1024 * 1024 / static_cast<int64_t>(sizeof(T)))) {
+    ring_.all_reduce<2>(in_ptr, out_ptr, count, 1, reduce_op);
+  } else {
+    mesh_.all_reduce(in_ptr, out_ptr, count, reduce_op);
+  }
 }
 
-} // namespace mlx::core::distributed::jaccl
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/mesh.h b/mlx/distributed/jaccl/lib/jaccl/mesh.h
similarity index 50%
rename from mlx/distributed/jaccl/mesh.h
rename to mlx/distributed/jaccl/lib/jaccl/mesh.h
index ed51361a..3184f1fb 100644
--- a/mlx/distributed/jaccl/mesh.h
+++ b/mlx/distributed/jaccl/lib/jaccl/mesh.h
@@ -2,14 +2,12 @@
 
 #pragma once
 
-#include "mlx/distributed/distributed_impl.h"
-#include "mlx/distributed/jaccl/mesh_impl.h"
-#include "mlx/distributed/jaccl/ring_impl.h"
-#include "mlx/distributed/jaccl/utils.h"
+#include "jaccl/group.h"
+#include "jaccl/mesh_impl.h"
+#include "jaccl/rdma.h"
+#include "jaccl/ring_impl.h"
 
-using GroupImpl = mlx::core::distributed::detail::GroupImpl;
-
-namespace mlx::core::distributed::jaccl {
+namespace jaccl {
 
 /**
  * The JACCL communication group for a fully connected mesh. We expect one
@@ -20,16 +18,12 @@ namespace mlx::core::distributed::jaccl {
  * information and then configure the connections to be ready for RDMA
  * operations.
  */
-class MeshGroup : public GroupImpl {
+class MeshGroup : public Group {
  public:
   MeshGroup(
       int rank,
       const std::vector<std::string>& device_names,
-      const char* coordinator_addr);
-
-  Stream communication_stream(StreamOrDevice s) override {
-    return to_stream(s, Device::cpu);
-  }
+      const std::string& coordinator_addr);
 
   int rank() override {
     return rank_;
@@ -39,27 +33,26 @@ class MeshGroup : public GroupImpl {
     return size_;
   }
 
-  void all_sum(const array& input, array& output, Stream stream) override;
-  void all_max(const array& input, array& output, Stream stream) override;
-  void all_min(const array& input, array& output, Stream stream) override;
-  void all_gather(const array& input, array& output, Stream stream) override;
-  void send(const array& input, int dst, Stream stream) override;
-  void recv(array& out, int src, Stream stream) override;
+  void all_sum(const void* input, void* output, size_t n_bytes, int dtype)
+      override;
 
-  void sum_scatter(const array& input, array& output, Stream stream) override {
-    throw std::runtime_error("[jaccl] sum_scatter not supported.");
-  }
+  void all_max(const void* input, void* output, size_t n_bytes, int dtype)
+      override;
 
-  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
-    throw std::runtime_error("[jaccl] Group split not supported.");
-  }
+  void all_min(const void* input, void* output, size_t n_bytes, int dtype)
+      override;
+
+  void all_gather(const void* input, void* output, size_t n_bytes) 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 all_reduce(
-      const array& input,
-      array& output,
-      Stream stream,
+      const void* input,
+      void* output,
+      size_t n_bytes,
       ReduceOp reduce_op);
 
   /**
@@ -86,4 +79,4 @@ class MeshGroup : public GroupImpl {
   RingImpl ring_;
 };
 
-} // namespace mlx::core::distributed::jaccl
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h
similarity index 69%
rename from mlx/distributed/jaccl/mesh_impl.h
rename to mlx/distributed/jaccl/lib/jaccl/mesh_impl.h
index fc486a39..0201f7ff 100644
--- a/mlx/distributed/jaccl/mesh_impl.h
+++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h
@@ -2,13 +2,16 @@
 
 #pragma once
 
+#include <memory>
 #include <span>
 
-#include "mlx/distributed/jaccl/utils.h"
+#include "jaccl/rdma.h"
 
 constexpr int MESH_MAX_PEERS = 8;
+constexpr int MESH_PIPELINE = 2;
+constexpr int64_t MAX_BUFFER_SIZE = FRAME_SIZE * (1 << (BUFFER_SIZES - 1));
 
-namespace mlx::core::distributed::jaccl {
+namespace jaccl {
 
 class MeshImpl {
  public:
@@ -17,20 +20,25 @@ class MeshImpl {
       int size,
       std::vector<Connection>& conns,
       std::vector<SharedBuffer>& buffers)
-      : rank_(rank), size_(size), connections_(conns), buffers_(buffers) {}
+      : rank_(rank),
+        size_(size),
+        connections_(conns),
+        buffers_(buffers),
+        staging_mem_(
+            std::make_unique<char[]>(MESH_PIPELINE * MAX_BUFFER_SIZE)) {}
 
   MeshImpl() : rank_(0), size_(1) {}
 
   template <typename T, typename ReduceOp>
-  void
-  all_reduce(const T* in_ptr, T* out_ptr, int64_t size, ReduceOp 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));
-    }
+  void all_reduce(const T* in, T* out, int64_t size, ReduceOp reduce_op) {
+    // Fully connected all reduce with deterministic reduction order.
+    //
+    // We copy rank 0's data to the output buffer and then we reduce every
+    // subsequent rank in-place in the output.
+    //
+    // Our own data is copied to a staging buffer to ensure we can reduce it in
+    // the output when needed.
 
-    // 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;
@@ -38,21 +46,36 @@ class MeshImpl {
     int64_t total = static_cast<int64_t>(size);
     int num_peers = size_ - 1;
 
+    // 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 completed_recv_begin[MESH_MAX_PEERS] = {0};
-    int completed_recv_end[MESH_MAX_PEERS] = {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(
-          data + read_offset,
-          data + std::min(read_offset + N, total),
+          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++;
@@ -61,14 +84,14 @@ class MeshImpl {
     }
 
     // Main loop
-    //
-    // Keep going until we have no longer data in flight.
-    while (in_flight > 0) {
+    while (reduce_chunk < total_chunks) {
       // 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.
+      // Also copy the next chunk into the staging area and advance our
+      // completed "receives".
       //
       // If a receive is completed then advance the pointer of completed
       // receives.
@@ -84,10 +107,16 @@ class MeshImpl {
         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(
-                data + read_offset,
-                data + std::min(read_offset + N, total),
+                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;
@@ -97,38 +126,70 @@ class MeshImpl {
         }
 
         else if (work_type == RECV_WR) {
-          completed_recv_end[rank]++;
+          recv_end[rank]++;
         }
       }
 
-      // Process the completed recv
+      // Process the received chunks in order.
       //
-      // 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);
+      // Rank 0 is always copied as is. Our rank is always read from the
+      // staging area.
+      while (reduce_chunk < total_chunks) {
+        // w is our write location so break if it is ahead of the read location.
+        int64_t w = static_cast<int64_t>(reduce_chunk) * N;
+        if (w >= read_offset) {
+          break;
+        }
+        // We want to reduce the 'reduce_chunk' chunk but it hasn't arrived
+        // yet.
+        if (recv_end[reduce_rank] <= reduce_chunk) {
+          break;
+        }
+        int b = reduce_chunk % PIPELINE;
+        int64_t elems = std::min(N, total - w);
+
+        // Data is read from the staging area
+        if (reduce_rank == rank_) {
+          if (reduce_rank == 0) {
+            std::copy_n(local_staging(b), elems, out + w);
+          } else {
+            reduce_op(local_staging(b), out + w, elems);
+          }
+        }
+
+        // Data is read from the recv buffers
+        else {
+          if (reduce_rank == 0) {
+            std::copy_n(
+                recv_buffer(sz, b, reduce_rank).begin<T>(), elems, out + w);
+          } else {
+            reduce_op(
+                recv_buffer(sz, b, reduce_rank).begin<T>(), out + w, elems);
+          }
+
+          // 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++;
           }
         }
-        completed_recv_begin[r] = s;
+
+        // Means we processed that chunk so move to the next one
+        reduce_rank++;
+        if (reduce_rank >= size_) {
+          reduce_rank = 0;
+          reduce_chunk++;
+        }
       }
     }
+
+    // Drain remaining in-flight completions (outstanding sends).
+    while (in_flight > 0) {
+      ibv_wc wc[WC_NUM];
+      int n = poll(connections_, WC_NUM, wc);
+      in_flight -= n;
+    }
   }
 
   void all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes) {
@@ -353,6 +414,7 @@ class MeshImpl {
   int size_;
   std::span<Connection> connections_;
   std::span<SharedBuffer> buffers_;
+  std::unique_ptr<char[]> staging_mem_;
 };
 
-} // namespace mlx::core::distributed::jaccl
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/utils.cpp b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp
similarity index 95%
rename from mlx/distributed/jaccl/utils.cpp
rename to mlx/distributed/jaccl/lib/jaccl/rdma.cpp
index dc87d428..7dc23043 100644
--- a/mlx/distributed/jaccl/utils.cpp
+++ b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp
@@ -5,7 +5,7 @@
 #include <iostream>
 #include <sstream>
 
-#include "mlx/distributed/jaccl/utils.h"
+#include "jaccl/rdma.h"
 
 #define LOAD_SYMBOL(symbol, variable)                               \
   {                                                                 \
@@ -31,7 +31,7 @@ void* page_aligned_alloc(size_t num_bytes) {
 
 } // namespace
 
-namespace mlx::core::distributed::jaccl {
+namespace jaccl {
 
 IBVWrapper::IBVWrapper() {
   librdma_handle_ = dlopen("librdma.dylib", RTLD_NOW | RTLD_GLOBAL);
@@ -182,11 +182,20 @@ const Destination& Connection::info() {
   ibv_port_attr port_attr;
   ibv().query_port(ctx, 1, &port_attr);
   ibv_gid gid;
-  ibv().query_gid(ctx, 1, 1, &gid);
+  for (int i = 0; i < port_attr.gid_tbl_len; i++) {
+    ibv_gid tmp;
+    if (ibv().query_gid(ctx, 1, i, &tmp) == 0) {
+      if (*(uint64_t*)&tmp.raw[0] == 0 && *(uint16_t*)&tmp.raw[8] == 0 &&
+          *(uint16_t*)&tmp.raw[10] == 0xffff) {
+        gid = tmp;
+        break;
+      }
+    }
+  }
 
   src.local_id = port_attr.lid;
   src.queue_pair_number = queue_pair->qp_num;
-  src.packet_sequence_number = 7; // TODO: Change to sth random
+  src.packet_sequence_number = 7;
   src.global_identifier = gid;
 
   return src;
@@ -287,10 +296,10 @@ std::vector<Connection> create_connections(
 
 SideChannel::SideChannel(int rank, int size, const char* addr)
     : rank_(rank), size_(size) {
-  auto address = detail::parse_address(addr);
+  auto address = parse_address(addr);
 
   if (rank_ == 0) {
-    detail::TCPSocket server(IBV_TAG);
+    TCPSocket server(IBV_TAG);
     server.listen(IBV_TAG, address);
 
     for (int i = 0; i < size - 1; i++) {
@@ -311,7 +320,7 @@ SideChannel::SideChannel(int rank, int size, const char* addr)
     }
   } else {
     sockets_.push_back(
-        detail::TCPSocket::connect(
+        TCPSocket::connect(
             IBV_TAG, address, 4, 1000, [](int attempt, int wait) {
               std::cerr << IBV_TAG << " Connection attempt " << attempt
                         << " waiting " << wait << " ms" << std::endl;
@@ -326,4 +335,4 @@ SideChannel::SideChannel(SideChannel&& sc)
   sc.size_ = -1;
 }
 
-} // namespace mlx::core::distributed::jaccl
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/utils.h b/mlx/distributed/jaccl/lib/jaccl/rdma.h
similarity index 97%
rename from mlx/distributed/jaccl/utils.h
rename to mlx/distributed/jaccl/lib/jaccl/rdma.h
index 8faa7740..b28c4ee1 100644
--- a/mlx/distributed/jaccl/utils.h
+++ b/mlx/distributed/jaccl/lib/jaccl/rdma.h
@@ -5,10 +5,11 @@
 #include <infiniband/verbs.h>
 
 #include <span>
+#include <sstream>
 #include <unordered_map>
 #include <vector>
 
-#include "mlx/distributed/utils.h"
+#include "jaccl/tcp.h"
 
 constexpr const char* IBV_TAG = "[jaccl]";
 constexpr int SEND_WR = 1;
@@ -19,8 +20,6 @@ constexpr int BUFFER_SIZES = 8;
 constexpr int NUM_BUFFERS = 2;
 constexpr int FRAME_SIZE = 4096;
 
-namespace detail = mlx::core::distributed::detail;
-
 namespace {
 
 template <typename T, typename = void>
@@ -45,7 +44,7 @@ inline std::pair<int, int64_t> buffer_size_from_message(int64_t msg) {
 
 } // namespace
 
-namespace mlx::core::distributed::jaccl {
+namespace jaccl {
 
 /**
  * Wrapper for the ibverbs API.
@@ -334,10 +333,16 @@ class SideChannel {
     return result;
   }
 
+  void barrier() {
+    // Twice has proven to be more robust to initialization issues.
+    all_gather<int>(0);
+    all_gather<int>(0);
+  }
+
  private:
   int rank_;
   int size_;
-  std::vector<detail::TCPSocket> sockets_;
+  std::vector<TCPSocket> sockets_;
 };
 
-} // namespace mlx::core::distributed::jaccl
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h b/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h
new file mode 100644
index 00000000..a9b3e5e7
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h
@@ -0,0 +1,117 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <cstddef>
+
+#include "jaccl/types.h"
+
+namespace jaccl {
+
+template <typename T>
+struct SumOp {
+  void operator()(const T* input, T* output, size_t N) const {
+    for (size_t i = 0; i < N; i++) {
+      output[i] = output[i] + input[i];
+    }
+  }
+};
+
+template <typename T>
+struct MaxOp {
+  void operator()(const T* input, T* output, size_t N) const {
+    for (size_t i = 0; i < N; i++) {
+      output[i] = (output[i] > input[i]) ? output[i] : input[i];
+    }
+  }
+};
+
+template <typename T>
+struct MinOp {
+  void operator()(const T* input, T* output, size_t N) const {
+    for (size_t i = 0; i < N; i++) {
+      output[i] = (output[i] < input[i]) ? output[i] : input[i];
+    }
+  }
+};
+
+//
+// The last piece of the puzzle to use the native bf16 while compiling a single
+// binary for all Macs is to compile these functions with
+// target("arch=armv8.6-a").
+//
+// Now we can simply check in runtime and call them only when they are
+// supported.
+//
+
+#if defined(__aarch64__)
+
+__attribute__((target("arch=armv8.6-a"))) inline void
+native_bf16_sum(const void* input, void* output, size_t N) {
+  auto in = reinterpret_cast<const __bf16*>(input);
+  auto out = reinterpret_cast<__bf16*>(output);
+  for (size_t i = 0; i < N; i++) {
+    out[i] = out[i] + in[i];
+  }
+}
+
+__attribute__((target("arch=armv8.6-a"))) inline void
+native_bf16_max(const void* input, void* output, size_t N) {
+  auto in = reinterpret_cast<const __bf16*>(input);
+  auto out = reinterpret_cast<__bf16*>(output);
+  for (size_t i = 0; i < N; i++) {
+    out[i] = (out[i] > in[i]) ? out[i] : in[i];
+  }
+}
+
+__attribute__((target("arch=armv8.6-a"))) inline void
+native_bf16_min(const void* input, void* output, size_t N) {
+  auto in = reinterpret_cast<const __bf16*>(input);
+  auto out = reinterpret_cast<__bf16*>(output);
+  for (size_t i = 0; i < N; i++) {
+    out[i] = (out[i] < in[i]) ? out[i] : in[i];
+  }
+}
+
+template <>
+struct SumOp<bfloat16_t> {
+  void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const {
+    if (has_native_bf16_support()) {
+      native_bf16_sum(input, output, N);
+    } else {
+      for (size_t i = 0; i < N; i++) {
+        output[i] = output[i] + input[i];
+      }
+    }
+  }
+};
+
+template <>
+struct MaxOp<bfloat16_t> {
+  void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const {
+    if (has_native_bf16_support()) {
+      native_bf16_max(input, output, N);
+    } else {
+      for (size_t i = 0; i < N; i++) {
+        output[i] = (output[i] > input[i]) ? output[i] : input[i];
+      }
+    }
+  }
+};
+
+template <>
+struct MinOp<bfloat16_t> {
+  void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const {
+    if (has_native_bf16_support()) {
+      native_bf16_min(input, output, N);
+    } else {
+      for (size_t i = 0; i < N; i++) {
+        output[i] = (output[i] < input[i]) ? output[i] : input[i];
+      }
+    }
+  }
+};
+
+#endif // defined(__aarch64__)
+
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/ring.cpp b/mlx/distributed/jaccl/lib/jaccl/ring.cpp
similarity index 59%
rename from mlx/distributed/jaccl/ring.cpp
rename to mlx/distributed/jaccl/lib/jaccl/ring.cpp
index 4a09d85c..96ddb0e0 100644
--- a/mlx/distributed/jaccl/ring.cpp
+++ b/mlx/distributed/jaccl/lib/jaccl/ring.cpp
@@ -1,22 +1,21 @@
 // Copyright © 2026 Apple Inc.
 
-#include "mlx/distributed/jaccl/ring.h"
-#include "mlx/backend/cpu/encoder.h"
-#include "mlx/distributed/reduction_ops.h"
-#include "mlx/dtype_utils.h"
+#include "jaccl/ring.h"
+#include "jaccl/reduction_ops.h"
+#include "jaccl/types.h"
 
-namespace mlx::core::distributed::jaccl {
+namespace jaccl {
 
 RingGroup::RingGroup(
     int rank,
     int size,
     const std::vector<std::string>& left_devices,
     const std::vector<std::string>& right_devices,
-    const char* coordinator_addr)
+    const std::string& coordinator_addr)
     : rank_(rank),
       size_(size),
       n_conns_(left_devices.size()),
-      side_channel_(rank_, size_, coordinator_addr),
+      side_channel_(rank_, size_, coordinator_addr.c_str()),
       left_(create_connections(left_devices)),
       right_(create_connections(right_devices)) {
   if (left_.size() > RING_MAX_CONNS || right_.size() > RING_MAX_CONNS) {
@@ -29,8 +28,8 @@ RingGroup::RingGroup(
   // Initialize all the connections and allocate buffers
   initialize();
 
-  // Make sure every node has reached here before continuing
-  side_channel_.all_gather<int>(0);
+  // Make sure every node has reached here before continuing.
+  side_channel_.barrier();
 
   // Create the ring implementation object
   ring_ = RingImpl(rank_, size_, left_, right_, send_buffers_, recv_buffers_);
@@ -60,9 +59,9 @@ void RingGroup::initialize() {
     conn.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.
+  // 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> left_info;
   for (auto& conn : left_) {
     left_info.emplace_back(conn.info());
@@ -126,40 +125,48 @@ void RingGroup::allocate_buffers() {
   }
 }
 
-void RingGroup::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 RingGroup::all_sum(
+    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);
+    all_reduce<T>(input, output, n_bytes, SumOp<T>{});
   });
 }
 
-void RingGroup::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 RingGroup::all_max(
+    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);
+    all_reduce<T>(input, output, n_bytes, MaxOp<T>{});
   });
 }
 
-void RingGroup::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 RingGroup::all_min(
+    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);
+    all_reduce<T>(input, output, n_bytes, MinOp<T>{});
   });
 }
 
-void RingGroup::all_gather(const array& input, array& output, Stream stream) {
-  auto in_ptr = input.data<char>();
-  auto out_ptr = output.data<char>();
-  int64_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]() {
-    ring_.all_gather(in_ptr, out_ptr, n_bytes, n_conns_);
-  });
+void RingGroup::all_gather(const void* input, void* output, size_t n_bytes) {
+  ring_.all_gather(
+      static_cast<const char*>(input),
+      static_cast<char*>(output),
+      n_bytes,
+      n_conns_);
 }
 
-void RingGroup::send(const array& input, int dst, Stream stream) {
+void RingGroup::send(const void* input, size_t n_bytes, int dst) {
   int right = (rank_ + 1) % size_;
   int left = (rank_ + size_ - 1) % size_;
   if (dst != right && dst != left) {
@@ -168,16 +175,10 @@ void RingGroup::send(const array& input, int dst, Stream stream) {
         << "but tried to send to " << dst << " from " << rank_ << std::endl;
     throw std::runtime_error(msg.str());
   }
-  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]() {
-    ring_.send(data, n_bytes, dst, n_conns_);
-  });
+  ring_.send(static_cast<const char*>(input), n_bytes, dst, n_conns_);
 }
 
-void RingGroup::recv(array& out, int src, Stream stream) {
+void RingGroup::recv(void* output, size_t n_bytes, int src) {
   int right = (rank_ + 1) % size_;
   int left = (rank_ + size_ - 1) % size_;
   if (src != right && src != left) {
@@ -186,42 +187,29 @@ void RingGroup::recv(array& out, int src, Stream stream) {
         << "but tried to recv from " << src << " to " << rank_ << std::endl;
     throw std::runtime_error(msg.str());
   }
-  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]() {
-    ring_.recv(data, n_bytes, src, n_conns_);
-  });
+  ring_.recv(static_cast<char*>(output), n_bytes, src, n_conns_);
 }
 
 template <typename T, typename ReduceOp>
 void RingGroup::all_reduce(
-    const array& input,
-    array& output,
-    Stream stream,
+    const void* input,
+    void* output,
+    size_t n_bytes,
     ReduceOp reduce_op) {
-  auto in_ptr = input.data<T>();
-  auto out_ptr = output.data<T>();
-  int64_t size = input.size();
-  int64_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, size, n_bytes, this, reduce_op]() {
-    if (size < size_ * 2 * n_conns_) {
-      ring_.all_reduce<1, T, ReduceOp>(in_ptr, out_ptr, size, 1, reduce_op);
-      return;
-    }
+  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_.all_reduce<1, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op);
+    return;
+  }
 
-    if (n_bytes <= 65536) {
-      ring_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, size, 1, reduce_op);
-      return;
-    }
+  if (n_bytes <= 65536) {
+    ring_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op);
+    return;
+  }
 
-    ring_.all_reduce<2, T, ReduceOp>(
-        in_ptr, out_ptr, size, n_conns_, reduce_op);
-  });
+  ring_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, count, n_conns_, reduce_op);
 }
 
-} // namespace mlx::core::distributed::jaccl
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/ring.h b/mlx/distributed/jaccl/lib/jaccl/ring.h
similarity index 52%
rename from mlx/distributed/jaccl/ring.h
rename to mlx/distributed/jaccl/lib/jaccl/ring.h
index b3ce2f7b..f0bc1ba9 100644
--- a/mlx/distributed/jaccl/ring.h
+++ b/mlx/distributed/jaccl/lib/jaccl/ring.h
@@ -2,13 +2,11 @@
 
 #pragma once
 
-#include "mlx/distributed/distributed_impl.h"
-#include "mlx/distributed/jaccl/ring_impl.h"
-#include "mlx/distributed/jaccl/utils.h"
+#include "jaccl/group.h"
+#include "jaccl/rdma.h"
+#include "jaccl/ring_impl.h"
 
-using GroupImpl = mlx::core::distributed::detail::GroupImpl;
-
-namespace mlx::core::distributed::jaccl {
+namespace jaccl {
 
 /**
  * The JACCL communication group for a ring where each node is connected to its
@@ -19,18 +17,14 @@ namespace mlx::core::distributed::jaccl {
  * information and then configure the connections to be ready for RDMA
  * operations.
  */
-class RingGroup : public GroupImpl {
+class RingGroup : public Group {
  public:
   RingGroup(
       int rank,
       int size,
       const std::vector<std::string>& left_devices,
       const std::vector<std::string>& right_devices,
-      const char* coordinator_addr);
-
-  Stream communication_stream(StreamOrDevice s) override {
-    return to_stream(s, Device::cpu);
-  }
+      const std::string& coordinator_addr);
 
   int rank() override {
     return rank_;
@@ -40,27 +34,26 @@ class RingGroup : public GroupImpl {
     return size_;
   }
 
-  void all_sum(const array& input, array& output, Stream stream) override;
-  void all_max(const array& input, array& output, Stream stream) override;
-  void all_min(const array& input, array& output, Stream stream) override;
-  void all_gather(const array& input, array& output, Stream stream) override;
-  void send(const array& input, int dst, Stream stream) override;
-  void recv(array& out, int src, Stream stream) override;
+  void all_sum(const void* input, void* output, size_t n_bytes, int dtype)
+      override;
 
-  void sum_scatter(const array& input, array& output, Stream stream) override {
-    throw std::runtime_error("[jaccl] sum_scatter not supported.");
-  }
+  void all_max(const void* input, void* output, size_t n_bytes, int dtype)
+      override;
 
-  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
-    throw std::runtime_error("[jaccl] Group split not supported.");
-  }
+  void all_min(const void* input, void* output, size_t n_bytes, int dtype)
+      override;
+
+  void all_gather(const void* input, void* output, size_t n_bytes) 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 all_reduce(
-      const array& input,
-      array& output,
-      Stream stream,
+      const void* input,
+      void* output,
+      size_t n_bytes,
       ReduceOp reduce_op);
 
   /**
@@ -86,4 +79,4 @@ class RingGroup : public GroupImpl {
   RingImpl ring_;
 };
 
-} // namespace mlx::core::distributed::jaccl
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h
similarity index 99%
rename from mlx/distributed/jaccl/ring_impl.h
rename to mlx/distributed/jaccl/lib/jaccl/ring_impl.h
index ce883d1f..3ee91b37 100644
--- a/mlx/distributed/jaccl/ring_impl.h
+++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h
@@ -4,11 +4,11 @@
 
 #include <span>
 
-#include "mlx/distributed/jaccl/utils.h"
+#include "jaccl/rdma.h"
 
 constexpr int RING_MAX_CONNS = 4;
 
-namespace mlx::core::distributed::jaccl {
+namespace jaccl {
 
 class RingImpl {
  public:
@@ -628,4 +628,4 @@ class RingImpl {
   std::span<SharedBuffer> recv_buffers_;
 };
 
-} // namespace mlx::core::distributed::jaccl
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/lib/jaccl/tcp.cpp b/mlx/distributed/jaccl/lib/jaccl/tcp.cpp
new file mode 100644
index 00000000..08f2338c
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/jaccl/tcp.cpp
@@ -0,0 +1,206 @@
+// Copyright © 2025 Apple Inc.
+
+#include <netdb.h>
+#include <unistd.h>
+#include <cstring>
+#include <sstream>
+#include <thread>
+
+#include "jaccl/tcp.h"
+
+namespace jaccl {
+
+/**
+ * Parse a sockaddr from an ip and port provided as strings.
+ */
+address_t parse_address(const std::string& ip, const std::string& port) {
+  struct addrinfo hints, *res;
+  std::memset(&hints, 0, sizeof(hints));
+  hints.ai_family = AF_UNSPEC;
+  hints.ai_socktype = SOCK_STREAM;
+
+  int status = getaddrinfo(ip.c_str(), port.c_str(), &hints, &res);
+  if (status != 0) {
+    std::ostringstream msg;
+    msg << "Can't parse address " << ip << ":" << port;
+    throw std::runtime_error(msg.str());
+  }
+
+  address_t result;
+  memcpy(&result.addr, res->ai_addr, res->ai_addrlen);
+  result.len = res->ai_addrlen;
+  freeaddrinfo(res);
+
+  return result;
+}
+
+/**
+ * Parse a sockaddr provided as an <ip>:<port> string.
+ */
+address_t parse_address(const std::string& ip_port) {
+  auto colon = ip_port.find(":");
+  if (colon == std::string::npos) {
+    std::ostringstream msg;
+    msg << "Can't parse address " << ip_port;
+    throw std::runtime_error(msg.str());
+  }
+  std::string ip(ip_port.begin(), ip_port.begin() + colon);
+  std::string port(ip_port.begin() + colon + 1, ip_port.end());
+
+  return parse_address(ip, port);
+}
+
+TCPSocket::TCPSocket(const char* tag) {
+  sock_ = socket(AF_INET, SOCK_STREAM, 0);
+  if (sock_ < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't create socket (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+}
+
+TCPSocket::TCPSocket(TCPSocket&& s) {
+  sock_ = s.sock_;
+  s.sock_ = -1;
+}
+
+TCPSocket& TCPSocket::operator=(TCPSocket&& s) {
+  if (this != &s) {
+    sock_ = s.sock_;
+    s.sock_ = -1;
+  }
+  return *this;
+}
+
+TCPSocket::TCPSocket(int s) : sock_(s) {}
+
+TCPSocket::~TCPSocket() {
+  if (sock_ > 0) {
+    shutdown(sock_, 2);
+    close(sock_);
+  }
+}
+
+int TCPSocket::detach() {
+  int s = sock_;
+  sock_ = -1;
+  return s;
+}
+
+void TCPSocket::listen(const char* tag, const address_t& addr) {
+  int success;
+
+  // Make sure we can launch immediately after shutdown by setting the
+  // reuseaddr option so that we don't get address already in use errors
+  int enable = 1;
+  success = setsockopt(sock_, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int));
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't enable reuseaddr (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+  success = setsockopt(sock_, SOL_SOCKET, SO_REUSEPORT, &enable, sizeof(int));
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't enable reuseport (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Bind the socket to the address and port
+  success = bind(sock_, addr.get(), addr.len);
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't bind socket (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Prepare waiting for connections
+  success = ::listen(sock_, 0);
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't listen (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+}
+
+TCPSocket TCPSocket::accept(const char* tag) {
+  int peer = ::accept(sock_, nullptr, nullptr);
+  if (peer < 0) {
+    std::ostringstream msg;
+    msg << tag << " Accept failed (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  return TCPSocket(peer);
+}
+
+void TCPSocket::send(const char* tag, const void* data, size_t len) {
+  while (len > 0) {
+    auto n = ::send(sock_, data, len, 0);
+    if (n <= 0) {
+      std::ostringstream msg;
+      msg << tag << " Send failed with errno=" << errno;
+      throw std::runtime_error(msg.str());
+    }
+    len -= n;
+    data = static_cast<const char*>(data) + n;
+  }
+}
+
+void TCPSocket::recv(const char* tag, void* data, size_t len) {
+  while (len > 0) {
+    auto n = ::recv(sock_, data, len, 0);
+    if (n <= 0) {
+      std::ostringstream msg;
+      msg << tag << " Recv failed with errno=" << errno;
+      throw std::runtime_error(msg.str());
+    }
+    len -= n;
+    data = static_cast<char*>(data) + n;
+  }
+}
+
+TCPSocket TCPSocket::connect(
+    const char* tag,
+    const address_t& addr,
+    int num_retries,
+    int wait,
+    std::function<void(int, int)> cb) {
+  int sock, success;
+
+  // Attempt to connect `num_retries` times with exponential backoff.
+  for (int attempt = 0; attempt < num_retries; attempt++) {
+    // Create the socket
+    sock = socket(AF_INET, SOCK_STREAM, 0);
+    if (sock < 0) {
+      std::ostringstream msg;
+      msg << tag << " Couldn't create socket to connect (error: " << errno
+          << ")";
+      throw std::runtime_error(msg.str());
+    }
+
+    success = ::connect(sock, addr.get(), addr.len);
+    if (success == 0) {
+      break;
+    }
+
+    if (cb != nullptr) {
+      cb(attempt, wait);
+    }
+    if (wait > 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(wait));
+    }
+
+    wait <<= 1;
+  }
+
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't connect (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  return TCPSocket(sock);
+}
+
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/lib/jaccl/tcp.h b/mlx/distributed/jaccl/lib/jaccl/tcp.h
new file mode 100644
index 00000000..8c879135
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/jaccl/tcp.h
@@ -0,0 +1,67 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <sys/socket.h>
+#include <functional>
+#include <string>
+
+namespace jaccl {
+
+struct address_t {
+  sockaddr_storage addr;
+  socklen_t len;
+
+  const sockaddr* get() const {
+    return (struct sockaddr*)&addr;
+  }
+};
+
+/**
+ * Parse a sockaddr from an ip and port provided as strings.
+ */
+address_t parse_address(const std::string& ip, const std::string& port);
+
+/**
+ * Parse a sockaddr provided as an <ip>:<port> string.
+ */
+address_t parse_address(const std::string& ip_port);
+
+/**
+ * Small wrapper over a TCP socket to simplify initiating connections.
+ */
+class TCPSocket {
+ public:
+  TCPSocket(const char* tag);
+  TCPSocket(const TCPSocket&) = delete;
+  TCPSocket& operator=(const TCPSocket&) = delete;
+  TCPSocket(TCPSocket&& s);
+  TCPSocket& operator=(TCPSocket&&);
+  ~TCPSocket();
+
+  void listen(const char* tag, const address_t& addr);
+  TCPSocket accept(const char* tag);
+
+  void send(const char* tag, const void* data, size_t len);
+  void recv(const char* tag, void* data, size_t len);
+
+  int detach();
+
+  operator int() const {
+    return sock_;
+  }
+
+  static TCPSocket connect(
+      const char* tag,
+      const address_t& addr,
+      int num_retries = 1,
+      int wait = 0,
+      std::function<void(int, int)> cb = nullptr);
+
+ private:
+  TCPSocket(int sock);
+
+  int sock_;
+};
+
+} // namespace jaccl
diff --git a/mlx/distributed/jaccl/lib/jaccl/types.h b/mlx/distributed/jaccl/lib/jaccl/types.h
new file mode 100644
index 00000000..defae731
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/jaccl/types.h
@@ -0,0 +1,263 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <complex>
+#include <cstdint>
+#include <type_traits>
+
+#if defined(__aarch64__) && defined(__APPLE__)
+#include <sys/sysctl.h>
+#endif
+
+namespace jaccl {
+
+namespace {
+
+union float_bits {
+  float f;
+  uint32_t u;
+};
+
+} // namespace
+
+#ifdef __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
+
+#include <arm_fp16.h>
+
+#else
+
+#define __JACCL_HALF_NAN__ 0x7D00
+
+//
+// The MLX float16 compatibility fallback. Redefined here to keep JACCL
+// standalone. Also it doesn't really need all the ops MXL defines.
+//
+struct float16_t {
+  uint16_t bits_;
+
+  float16_t(const float& x) : bits_(0) {
+    float_bits in;
+    in.f = x;
+
+    // Extract sign
+    uint32_t x_sign_32 = in.u & uint32_t(0x80000000);
+    uint16_t x_sign_16 = (x_sign_32 >> 16);
+
+    if (std::isnan(x)) {
+      bits_ = x_sign_16 | uint16_t(__JACCL_HALF_NAN__);
+    } else {
+      // Union
+      float_bits inf_scale, zero_scale, magic_bits;
+
+      // Find exponent bits and take the max supported by half
+      uint32_t x_expo_32 = in.u & uint32_t(0x7f800000);
+      uint32_t max_expo_32 = uint32_t(0x38800000);
+      x_expo_32 = x_expo_32 < max_expo_32 ? max_expo_32 : x_expo_32;
+      x_expo_32 += uint32_t(15) << 23;
+
+      // Handle scaling to inf as needed
+      inf_scale.u = uint32_t(0x77800000);
+      zero_scale.u = uint32_t(0x08800000);
+
+      // Combine with magic and let addition do rounding
+      magic_bits.u = x_expo_32;
+      magic_bits.f += (std::abs(x) * inf_scale.f) * zero_scale.f;
+
+      // Take the lower 5 bits of the exponent
+      uint32_t x_expo_16 = ((magic_bits.u >> 13) & uint32_t(0x7c00));
+
+      // Collect the lower 12 bits which have the mantissa
+      uint32_t x_mant_16 = magic_bits.u & uint32_t(0x0fff);
+
+      // Combine sign, exp and mantissa
+      bits_ = (x_sign_16 | uint16_t(x_expo_16 + x_mant_16));
+    }
+  }
+
+  operator float() const {
+    float_bits out;
+
+    uint32_t x_sign_32 = (bits_ << 16) & uint32_t(0x80000000);
+    uint32_t base = (bits_ << 16);
+    uint32_t two_base = base + base;
+
+    uint32_t denorm_max = 1u << 27;
+    if (two_base < denorm_max) {
+      out.u = uint32_t(126) << 23; // magic mask
+      out.u |= (two_base >> 17); // Bits from fp16
+      out.f -= 0.5f; // magic bias
+    } else {
+      out.u = uint32_t(0xE0) << 23; // exponent offset
+      out.u += (two_base >> 4); // Bits from fp16
+      float out_unscaled = out.f; // Store value
+      out.u = uint32_t(0x7800000); // exponent scale
+      out.f *= out_unscaled;
+    }
+
+    // Add sign
+    out.u |= x_sign_32;
+
+    return out.f;
+  }
+
+  bool operator<(float16_t x) {
+    return static_cast<float>(*this) < static_cast<float>(x);
+  }
+
+  bool operator>(float16_t x) {
+    return static_cast<float>(*this) > static_cast<float>(x);
+  }
+
+  float16_t operator+(float16_t x) {
+    return static_cast<float>(*this) + static_cast<float>(x);
+  }
+
+  float16_t& operator+=(float16_t x) {
+    *this = *this + x;
+    return *this;
+  }
+};
+
+#endif
+
+//
+// Check at runtime if the CPU supports native bfloat16 (FEAT_BF16).
+//
+// This allows us to compile once for all Macs but still enable the feature if
+// it is supported.
+//
+inline bool has_native_bf16_support() {
+#if defined(__aarch64__) && defined(__APPLE__)
+  static bool has_support = []() {
+    int value = 0;
+    size_t value_size = sizeof(value);
+    int success = sysctlbyname(
+        "hw.optional.arm.FEAT_BF16", &value, &value_size, nullptr, 0);
+    return success == 0 & value != 0;
+  }();
+  return has_support;
+#else
+  return false;
+#endif
+}
+
+//
+// The MLX bfloat16 compatibility fallback.
+//
+// Contrary to float16 we always define it and we 'll use
+// has_native_bf16_support to decide if we are going to use __bf16 instead
+// during runtime.
+//
+
+#define __JACCL_BFLOAT_NAN__ 0x7FC0
+
+struct bfloat16_t {
+  uint16_t bits_;
+
+  bfloat16_t(const float& x) {
+    if (std::isnan(x)) {
+      bits_ = __JACCL_BFLOAT_NAN__;
+    } else {
+      float_bits in;
+      in.f = x;
+      in.u += (in.u >> 16 & 1) + uint32_t(0x7FFF);
+      bits_ = in.u >> 16;
+    }
+  }
+
+  operator float() const {
+    float_bits out;
+    out.u = ((uint32_t)bits_) << 16;
+    return out.f;
+  }
+
+  bool operator<(bfloat16_t x) {
+    return static_cast<float>(*this) < static_cast<float>(x);
+  }
+
+  bool operator>(bfloat16_t x) {
+    return static_cast<float>(*this) > static_cast<float>(x);
+  }
+
+  bfloat16_t operator+(bfloat16_t x) {
+    return static_cast<float>(*this) + static_cast<float>(x);
+  }
+
+  bfloat16_t& operator+=(bfloat16_t x) {
+    *this = *this + x;
+    return *this;
+  }
+};
+
+using complex64_t = std::complex<float>;
+
+inline bool operator<(complex64_t lhs, complex64_t rhs) {
+  return lhs.real() < rhs.real() ||
+      (lhs.real() == rhs.real() && lhs.imag() < rhs.imag());
+}
+
+inline bool operator>(complex64_t lhs, complex64_t rhs) {
+  return lhs.real() > rhs.real() ||
+      (lhs.real() == rhs.real() && lhs.imag() > rhs.imag());
+}
+
+template <typename T>
+struct type_identity {
+  using type = T;
+};
+
+#define JACCL_GET_TYPE(x) typename decltype(x)::type
+
+/**
+ * Dispatch a function for all supported types based on a Dtype.
+ */
+template <typename F>
+void dispatch_all_types(int dtype, F&& f) {
+  switch (dtype) {
+    case Dtype::Bool:
+      f(type_identity<bool>{});
+      break;
+    case Dtype::Int8:
+      f(type_identity<int8_t>{});
+      break;
+    case Dtype::Int16:
+      f(type_identity<int16_t>{});
+      break;
+    case Dtype::Int32:
+      f(type_identity<int32_t>{});
+      break;
+    case Dtype::Int64:
+      f(type_identity<int64_t>{});
+      break;
+    case Dtype::UInt8:
+      f(type_identity<uint8_t>{});
+      break;
+    case Dtype::UInt16:
+      f(type_identity<uint16_t>{});
+      break;
+    case Dtype::UInt32:
+      f(type_identity<uint32_t>{});
+      break;
+    case Dtype::UInt64:
+      f(type_identity<uint64_t>{});
+      break;
+    case Dtype::Float16:
+      f(type_identity<float16_t>{});
+      break;
+    case Dtype::BFloat16:
+      f(type_identity<bfloat16_t>{});
+      break;
+    case Dtype::Float32:
+      f(type_identity<float>{});
+      break;
+    case Dtype::Float64:
+      f(type_identity<double>{});
+      break;
+    case Dtype::Complex64:
+      f(type_identity<complex64_t>{});
+      break;
+  }
+}
+
+} // namespace jaccl