From 49bc3fe6c6c87adc94d1ec48d013f75c15e87889 Mon Sep 17 00:00:00 2001
From: Daniel Kristofik <daniel.kristofik@gmail.com>
Date: Wed, 29 Apr 2026 22:40:58 +0200
Subject: [PATCH] Fix JACCL GID selection on Apple Thunderbolt RDMA
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Connection::info() in rdma.cpp scans the GID table for IPv4-mapped IPv6
GIDs (::ffff:x.x.x.x, RoCE v2 format). Apple Thunderbolt RDMA exposes only
link-local IPv6 GIDs (fe80::...) — the filter never matches and gid is left
uninitialized. The garbage value is then sent to the peer via the side
channel, causing the kernel to reject the QP RTR transition with EINVAL
(errno 22).

Initialize gid to zero, factor the GID selection into a try_gid() helper,
and add a fallback that prefers index 1 (the actual rdma_enX port GID on
Apple TB; index 0 is typically derived from a non-RDMA interface and
routes elsewhere, surfacing as errno 60 ETIMEDOUT). The IPv4-mapped path
remains the preferred match, so RoCE v2 setups are unaffected.

Tested on 2x Mac Studio M4 Max + Thunderbolt 5 mesh, macOS 26.4.1.
Distributed init succeeds, sustained tensor-parallel inference works.
---
 mlx/distributed/jaccl/lib/jaccl/rdma.cpp | 60 +++++++++++++++++++++---
 1 file changed, 54 insertions(+), 6 deletions(-)

diff --git a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp
index 7dc23043d2..62cd43b99b 100644
--- a/mlx/distributed/jaccl/lib/jaccl/rdma.cpp
+++ b/mlx/distributed/jaccl/lib/jaccl/rdma.cpp
@@ -181,13 +181,61 @@ const Destination& Connection::info() {
 
   ibv_port_attr port_attr;
   ibv().query_port(ctx, 1, &port_attr);
-  ibv_gid gid;
-  for (int i = 0; i < port_attr.gid_tbl_len; i++) {
+  ibv_gid gid = {};
+
+  auto try_gid = [&](int i, bool require_ipv4_mapped) -> bool {
+    if (i < 0 || i >= port_attr.gid_tbl_len) {
+      return false;
+    }
     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;
+    if (ibv().query_gid(ctx, 1, i, &tmp) != 0) {
+      return false;
+    }
+    if (require_ipv4_mapped) {
+      if (*(uint64_t*)&tmp.raw[0] != 0 || *(uint16_t*)&tmp.raw[8] != 0 ||
+          *(uint16_t*)&tmp.raw[10] != 0xffff) {
+        return false;
+      }
+    } else {
+      bool is_zero = true;
+      for (int j = 0; j < 16; j++) {
+        if (tmp.raw[j] != 0) {
+          is_zero = false;
+          break;
+        }
+      }
+      if (is_zero) {
+        return false;
+      }
+    }
+    gid = tmp;
+    return true;
+  };
+
+  // 1. Prefer IPv4-mapped IPv6 GID (RoCE v2 standard).
+  bool found = false;
+  for (int i = 0; i < port_attr.gid_tbl_len; i++) {
+    if (try_gid(i, /*require_ipv4_mapped=*/true)) {
+      found = true;
+      break;
+    }
+  }
+
+  // 2. Fallback for Apple Thunderbolt RDMA, which exposes only link-local
+  //    IPv6 (fe80::...) GIDs. Prefer index 1 (the actual rdma_enX port GID;
+  //    index 0 on Apple TB is typically derived from a non-RDMA interface
+  //    and routes elsewhere, surfacing as RTR errno 60 ETIMEDOUT).
+  if (!found) {
+    for (int i : {1, 0}) {
+      if (try_gid(i, /*require_ipv4_mapped=*/false)) {
+        found = true;
+        break;
+      }
+    }
+  }
+  if (!found) {
+    for (int i = 2; i < port_attr.gid_tbl_len; i++) {
+      if (try_gid(i, /*require_ipv4_mapped=*/false)) {
         break;
       }
     }