JACCL: zeroed send-frame tails, a progress guard, the advertised GID and the upstream JACCL fixes

.github/workflows/fork-tests.yml

@@ -6,7 +6,8 @@ name: Fork Tests
 # The jobs:
 #   build-wheel    builds the Python wheel and uploads it.
 #   build-cpp      builds the C++ library and tests, and runs the C++ tests
-#                  with the GPU and then with the CPU as the default device.
+#                  with the GPU and then with the CPU as the default device,
+#                  and the JACCL tests with simulated verbs (tests/jaccl).
 #                  Then it runs the standalone allocation cap tests, one
 #                  process for each case.
 #   python-gpu     runs the Python tests with the GPU as the default device.
@@ -140,6 +141,19 @@ jobs:
           LLVM_PROFILE_FILE: ${{ runner.temp }}/profiles/cpu-%p-%m.profraw
         run: ./build/tests/tests
 
+      - name: Run the JACCL tests with simulated verbs
+        # tests/jaccl is its own CMake project, which the build above does
+        # not build: the real mesh, ring and RDMA headers against fake verbs,
+        # with ASan and UBSan. The step clears the coverage flags.
+        if: ${{ !cancelled() }}
+        env:
+          CFLAGS: ''
+          CXXFLAGS: ''
+        run: |
+          cmake -S tests/jaccl -B build-jaccl -DCMAKE_BUILD_TYPE=Debug
+          cmake --build build-jaccl -j $(sysctl -n hw.physicalcpu)
+          ctest --test-dir build-jaccl --output-on-failure --no-tests=error
+
       - name: Run the standalone allocation cap tests
         # The process allocation cap is install-once in each process. ctest
         # runs each entry in a fresh process, and the Metal test runs one case

docs/src/usage/distributed.rst

@@ -401,6 +401,32 @@ concatenated in rank order.
         all_gather_factory=make_side_channel,
     )
 
+A peer that does not answer
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+JACCL waits for RDMA completions in a polling loop. When the peer stops or the
+link fails, no completion arrives. A JACCL call that gets no completion for
+longer than :envvar:`MLX_JACCL_PROGRESS_TIMEOUT_MS` (30 seconds by default)
+fails. A call that gets a failed completion fails immediately.
+
+A failed call closes the communication group: JACCL destroys the queue pairs
+and the completion queues of the group and deregisters its buffers. Then, by
+default, it releases the wired GPU memory of the process (it sets the wired
+limit to zero and clears the buffer cache) and exits the process with status
+75, so that a supervisor can restart it. This teardown exit is the default
+because a rank that dies with RDMA work outstanding can leave its memory wired
+until the machine restarts, and because the call often runs on a thread where
+an exception cannot reach the program.
+
+With :envvar:`MLX_JACCL_TIMEOUT_ACTION` set to ``throw`` the call raises an
+error instead and the process continues. Each later call on the closed group
+raises an error immediately, and the group cannot be opened again. Start a new
+process to communicate again.
+
+The time that a rank waits includes the time that its peer computes, for
+example in :func:`mlx.core.distributed.recv`. Set the limit above the longest
+such wait, or set it to ``0`` to remove the limit.
+
 .. _nccl_section:
 
 Getting Started with NCCL
@@ -573,6 +599,13 @@ the following:
       ["rdma_en3", "rdma_en4", "rdma_en5", null]
    ]
 
+:envvar:`MLX_JACCL_PROGRESS_TIMEOUT_MS` is optional. It contains the longest time
+in milliseconds that a call waits for an RDMA completion before it fails and
+closes the group. The default is 30000 and 0 removes the limit.
+:envvar:`MLX_JACCL_TIMEOUT_ACTION` is optional. It selects what a failed call
+does after it closed the group: ``teardown-exit`` (the default) releases the
+GPU memory and exits the process with status 75, ``throw`` raises an error.
+
 
 NCCL
 ^^^^^

docs/src/usage/environment_variables.rst

@@ -63,6 +63,28 @@ sets them.
    Prefer a ring topology for JACCL. This variable is enabled by its presence.
    ``JACCL_RING`` is accepted as a higher-priority alias.
 
+.. envvar:: MLX_JACCL_PROGRESS_TIMEOUT_MS
+
+   The longest time in milliseconds that a JACCL call waits for an RDMA
+   completion. The default is ``30000``. Set it to ``0`` or a negative value
+   to remove the limit. ``JACCL_PROGRESS_TIMEOUT_MS`` is accepted as a
+   higher-priority alias. When no completion arrives for this time, the call
+   fails and the group is closed: its queue pairs and completion queues are
+   destroyed and its buffers are deregistered. What follows is set by
+   :envvar:`MLX_JACCL_TIMEOUT_ACTION`. See the :ref:`JACCL section
+   <jaccl_section>`.
+
+.. envvar:: MLX_JACCL_TIMEOUT_ACTION
+
+   What a JACCL call does after it failed (no completion within
+   :envvar:`MLX_JACCL_PROGRESS_TIMEOUT_MS`, a failed completion, or a failed
+   poll) and closed its group. ``teardown-exit``, the default, sets the wired
+   limit to zero, clears the buffer cache and exits the process with status 75
+   without running static destructors. ``throw`` raises an error; each later
+   call on the group then raises an error immediately. Any other value keeps
+   the default. ``JACCL_TIMEOUT_ACTION`` is accepted as a higher-priority
+   alias.
+
 .. envvar:: MLX_WORLD_SIZE
 
    The total number of processes in an NCCL group.

fork.yaml

@@ -30,6 +30,13 @@ def:
     control, declared-mutable custom-kernel inputs for paged KV, and a use-after-free fix in
     Metal eval. Every change is meant to be additive and rebasable; each is gated by shape,
     device or environment checks so the upstream path stays the default elsewhere.
+
+    The fork also changes JACCL, the RDMA-over-Thunderbolt backend that runs one model across
+    two Macs: it zeroes the unused tail of every send frame, and a progress guard fails a call
+    that gets no completion, closes its group and by default releases the GPU memory and exits
+    the process with status 75. These two are not gated. The zeroing is always on, and the
+    guard is on with a 30 s limit that `JACCL_PROGRESS_TIMEOUT_MS=0` removes;
+    `JACCL_TIMEOUT_ACTION=throw` keeps the process alive and throws instead.
   sub:
     - title: "Metal eval: use-after-free with mid-eval synchronization"
       description: |
@@ -243,6 +250,74 @@ def:
       globs:
         - "python/tests/test_metal_threadgroup_bounds.py"
         - "docs/src/dev/custom_metal_kernels.rst"
+    - title: "JACCL: zeroed send-frame tails, a progress guard and the advertised GID"
+      description: |
+        JACCL is the RDMA-over-Thunderbolt backend. Upstream has changed none of the three
+        things below.
+
+        Send frames: JACCL posts every frame at the full size of its scratch buffer, so a
+        payload shorter than the buffer went to the peer followed by the bytes of earlier
+        sends. `stage_send_frame` (`send_frame.h`) copies the payload and zeroes the rest of
+        the frame, and every mesh and ring send path stages through
+        `SharedBuffer::stage_send`. The posted sizes and the receive paths are unchanged.
+
+        Progress guard: the polling loops waited for completions without a limit and never
+        read `wc.status`, so a rank whose peer had died spun on a full core until it was
+        killed, and a failed completion was used as data. `ProgressGuard`
+        (`progress_guard.h`) is called after every poll of every mesh and ring loop. A
+        completion with a failure status, or a failed `ibv_poll_cq`, is an error at once; no
+        completion for longer than `JACCL_PROGRESS_TIMEOUT_MS` /
+        `MLX_JACCL_PROGRESS_TIMEOUT_MS` (default 30000 ms; `0` or a negative value removes
+        the limit) is an error. The clock is read once per 1024 polls that return nothing, so
+        the polling fast path gains a counter and no clock read. The failed call first
+        destroys the group's queue pairs and completion queues and deregisters its buffers;
+        the ring does so only after all wires have returned. Then, by default
+        (`JACCL_TIMEOUT_ACTION` / `MLX_JACCL_TIMEOUT_ACTION` = `teardown-exit`), it prints the
+        reason, runs the memory release hook that MLX's `jaccl.cpp` registers
+        (`set_wired_limit(0)`, then `clear_cache()`, at most 10 s) and leaves through
+        `std::_Exit(75)`, a restartable status. With `throw` it throws a `[jaccl]`
+        `std::runtime_error` instead; every later call on that group throws immediately and
+        the group cannot be reopened in the process. The ring keeps its own guarded
+        `dispatch_wires`, which also reports the errors of the ring pool wires that it used
+        to drop. A mesh with more than one wire (upstream's second cable) has one `MeshImpl`,
+        and so one guard, for each wire, and the wires of a `MeshGroup` share one failure
+        flag: a failed wire closes its own queue pairs and buffers, a wire that waits in the
+        same call stops with `WireStopped` and closes its own, and every later call on the
+        group is refused. The mesh runs its wires through the
+        shared `dispatch_wires` in `threadpool.h`, which upstream lets drop the error of a
+        pool wire; the fork throws that error after all wires returned. A region that the
+        verbs refuse to deregister is reported and kept for the destructor, as a refused
+        queue pair or completion queue is. Each completion
+        queue has four times the entries its queue pair can have outstanding
+        (`COMPLETION_QUEUE_DEPTH`, used by `mesh.cpp` and `ring.cpp` when they create the
+        queues). MLX's `jaccl.cpp` registers the release hook each time it wraps a new JACCL
+        group.
+
+        Advertised GID: `Connection::info()` advertises the IPv4-mapped GID of the port, but
+        `queue_pair_rtr` sent from GID index 1 whatever the index of that GID. The connection
+        now remembers the index where the advertised GID was found and sends from it;
+        behaviour is unchanged where that index is 1.
+
+        `tests/jaccl` is a separate CMake project that the MLX build does not build; the
+        `build-cpp` job of Fork Tests builds it and runs it. It
+        compiles the real mesh, ring and RDMA headers against simulated verbs and checks the
+        frame tails, the guard and the release path under ASan, UBSan and TSan.
+      globs:
+        - "mlx/distributed/jaccl/jaccl.cpp"
+        - "mlx/distributed/jaccl/lib/README.md"
+        - "mlx/distributed/jaccl/lib/jaccl/mesh.cpp"
+        - "mlx/distributed/jaccl/lib/jaccl/mesh_impl.h"
+        - "mlx/distributed/jaccl/lib/jaccl/progress_guard.h"
+        - "mlx/distributed/jaccl/lib/jaccl/rdma.cpp"
+        - "mlx/distributed/jaccl/lib/jaccl/rdma.h"
+        - "mlx/distributed/jaccl/lib/jaccl/ring.cpp"
+        - "mlx/distributed/jaccl/lib/jaccl/ring_impl.h"
+        - "mlx/distributed/jaccl/lib/jaccl/send_frame.h"
+        - "mlx/distributed/jaccl/lib/jaccl/threadpool.h"
+        - "docs/src/usage/distributed.rst"
+        - "docs/src/usage/environment_variables.rst"
+        - "tests/jaccl/*"
+        - "tests/jaccl/fake_verbs/infiniband/verbs.h"
     - title: "Quantized tile bounds and distinct loader symbols"
       description: |
         Sorted gathered matrix products use the remaining reduction width for a
@@ -290,6 +365,20 @@ def:
       globs:
         - "mlx/backend/metal/kernels/quantized_nax.h"
         - "mlx/backend/metal/kernels/fp_quantized_nax.h"
+    - title: "JIT instantiation of the NAX head-dim split attention kernel"
+      description: |
+        `get_steel_attention_nax_kernel` builds the JIT source of both `attention_nax` and
+        `attention_nax_dsplit`, and upstream passes `bv` to both. `attention_nax_dsplit` has
+        no `BV` template parameter (its value head dim is `BD`), so every later argument
+        shifts by one, an integer lands in the `MaskType` slot, and the library does not
+        compile. JIT builds such as mlx-swift's compile this kernel at runtime for full
+        attention with head dim 256 or 512 and at least 1,024 query tokens on NAX GPUs, and
+        the caller's process exited with "Unable to build metal library from source". The
+        fork passes the split kernel the same arguments as its instantiation in
+        `steel_attention_nax.metal`; the non-split kernel is unchanged. Builds that use the
+        precompiled metallib were not affected. This is a latent upstream bug (PR #48).
+      globs:
+        - "mlx/backend/metal/jit_kernels.cpp"
     - title: "64-bit products in the JACCL ring and the vector SDPA layout check"
       description: |
         The JACCL ring computes each wire's byte offset as `lr * n_wires * size`. Upstream
@@ -415,7 +504,8 @@ def:
         `fork-tests.yml` builds the fork with Metal and runs its C++ and Python tests on
         GitHub-hosted macOS runners. Two jobs build in parallel. One job builds the Python
         wheel. The other job builds the C++ library and tests, and runs the C++ tests with the
-        GPU and then with the CPU as the default device. Then it runs the standalone allocation
+        GPU and then with the CPU as the default device, and builds and runs the JACCL tests
+        with simulated verbs (`tests/jaccl`). Then it runs the standalone allocation
         cap tests with ctest, one process for each case. The Python tests run in shard jobs:
         one set with the GPU as the default device, and one set with the CPU.
         `.github/scripts/fork-tests/python-shards.json` gives the shards.

mlx/distributed/jaccl/jaccl.cpp

@@ -4,9 +4,11 @@
 #include "mlx/backend/cpu/encoder.h"
 #include "mlx/distributed/distributed_impl.h"
 #include "mlx/dtype_utils.h"
+#include "mlx/memory.h"
 
 #include <jaccl/group.h>
 #include <jaccl/jaccl.h>
+#include <jaccl/progress_guard.h>
 
 #include <optional>
 #include <sstream>
@@ -168,26 +170,43 @@ class JACCLGroup : public GroupImpl {
   std::shared_ptr<::jaccl::Group> group_;
 };
 
+/**
+ * Run by a JACCL teardown exit (see progress_guard.h) after the group released
+ * its queue pairs. The exit skips static destructors, so release the wired GPU
+ * memory here: unwire first (the residency set shrinks to nothing), then drop
+ * the buffer cache. Best effort; it never throws into the exit path.
+ */
+void release_memory_before_exit() {
+  try {
+    set_wired_limit(0);
+  } catch (...) {
+  }
+  try {
+    clear_cache();
+  } catch (...) {
+  }
+}
+
+std::shared_ptr<GroupImpl> make_group(std::shared_ptr<::jaccl::Group> group) {
+  if (group == nullptr) {
+    return nullptr;
+  }
+  ::jaccl::set_memory_release(&release_memory_before_exit);
+  return std::make_shared<JACCLGroup>(std::move(group));
+}
+
 } // namespace
 
 bool is_available() {
   return ::jaccl::is_available();
 }
 
 std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
-  auto group = ::jaccl::init(strict);
-  if (group == nullptr) {
-    return nullptr;
-  }
-  return std::make_shared<JACCLGroup>(std::move(group));
+  return make_group(::jaccl::init(strict));
 }
 
 std::shared_ptr<GroupImpl> init(bool strict, AllGatherFactory factory) {
-  auto group = ::jaccl::init(strict, factory);
-  if (group == nullptr) {
-    return nullptr;
-  }
-  return std::make_shared<JACCLGroup>(std::move(group));
+  return make_group(::jaccl::init(strict, factory));
 }
 
 } // namespace mlx::core::distributed::jaccl

mlx/distributed/jaccl/lib/README.md

@@ -102,6 +102,27 @@ variables:
   (rank 0 listener)
 - **JACCL_RING** / **MLX_JACCL_RING**: (Optional) Prefer ring topology over
   mesh
+- **JACCL_PROGRESS_TIMEOUT_MS** / **MLX_JACCL_PROGRESS_TIMEOUT_MS**: (Optional)
+  The longest time in milliseconds that a call waits for an RDMA completion.
+  The default is 30000. Zero or a negative value removes the limit.
+- **JACCL_TIMEOUT_ACTION** / **MLX_JACCL_TIMEOUT_ACTION**: (Optional) What a
+  failed call does: `teardown-exit` (the default) or `throw`. Any other value
+  keeps the default.
+
+A call fails when it gets no completion for this time, a failed work
+completion or a failed poll. It first destroys the queue pairs and completion
+queues of the group and deregisters its buffers (the ring waits until all of
+its wires have returned). With `teardown-exit` it then prints the reason, runs
+the memory release hook registered with `jaccl::set_memory_release` (MLX
+registers one that sets the wired limit to zero and clears the buffer cache;
+it gets at most 10 seconds) and exits the process with status 75 (EX_TEMPFAIL,
+restartable) through `std::_Exit`. With `throw` the caller gets the exception;
+every later call on that group throws immediately, and the group cannot be
+opened again in the same process. Set the limit above the longest time that a
+rank can wait for a peer, for example in `recv` while the peer computes.
+
+Each completion queue has `4 * (MAX_SEND_WR + MAX_RECV_WR)` entries, four times
+what its queue pair can have outstanding.
 
 ### Device File Format
 
@@ -321,6 +342,26 @@ class Config {
 }
 ```
 
+## Tests
+
+`tests/jaccl` in the MLX repository has tests that do not use RDMA hardware.
+They compile the mesh, ring and RDMA headers with a small stand-in for
+`<infiniband/verbs.h>` and simulate the verbs. They check that each send frame
+holds its payload and then only zero bytes, and that a call fails and the group
+is closed when a completion does not arrive or reports a failure. The directory
+is a separate CMake project, so the MLX build does not build it:
+
+```bash
+cmake -S tests/jaccl -B build/jaccl-tests
+cmake --build build/jaccl-tests
+ctest --test-dir build/jaccl-tests --output-on-failure
+```
+
+The tests use AddressSanitizer and UndefinedBehaviorSanitizer. Configure with
+`-DJACCL_TESTS_TSAN=ON` to use ThreadSanitizer. The guard tests measure time: a
+stalled call must fail between its limit of 250 ms and one second after it, so
+a machine under heavy load can make them fail.
+
 ## License
 
 JACCL is part of MLX and is released under the same license.

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

@@ -32,10 +32,16 @@ MeshGroup::MeshGroup(
   // Make sure every node has reached here before continuing
   side_channel_.barrier();
 
-  // Create the mesh implementation objects
+  // Create the mesh implementation objects. The wires share one failure flag.
+  auto failed = std::make_shared<std::atomic<bool>>(false);
   for (int w = 0; w < n_wires_; w++) {
     meshes_.emplace_back(
-        rank_, size_, connections_[w], buffers_[w], scatter_buffers_[w]);
+        rank_,
+        size_,
+        connections_[w],
+        buffers_[w],
+        scatter_buffers_[w],
+        failed);
   }
 }
 
@@ -47,7 +53,7 @@ void MeshGroup::initialize() {
         continue;
       }
       conn.allocate_protection_domain();
-      conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR);
+      conn.create_completion_queue(COMPLETION_QUEUE_DEPTH);
       conn.create_queue_pair();
     }
   }