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

[jonathan308]

Summary

Three fixes to JACCL, the RDMA-over-Thunderbolt backend, on current main, with the upstream JACCL fixes that the fork did not have in front of them.

1. A send frame no longer carries bytes of earlier sends after a short payload.
2. A rank whose peer stops no longer polls forever. The call fails with an error, the group releases its RDMA resources, and every later call on that group fails at once.
3. The queue pair sends from the GID that was advertised to the peer.

This PR contains and extends #26. #26 carries the send-frame commit alone on an older base. Here the same commit (authorship kept) is on current main, formatted with clang-format and described in fork.yaml, which were the two open items of #26. #26 can be closed when this one is accepted.

Status: draft. The mesh path was run on two Macs over Thunderbolt RDMA (results below). The ring path was run only with simulated verbs.

Commits

Commit	Author	Subject
d92347ba	Drifter4242 (upstream)	Fix jaccl ring all_gather: direction 1 slice was not mirrored (backport of ml-explore/mlx ml-explore#4443)
d3739307	Ronald Mannak (upstream)	chore: Fix assertion in ReduceScatter::eval_cpu (backport of ml-explore/mlx ml-explore#4557)
46149688	Cheng (upstream)	Add a scoped_env utility for tests (partial backport of ml-explore/mlx ml-explore#4474)
7112cbf8	Ronald Mannak (upstream)	Do not cache failed distributed backend initialization (backport of ml-explore/mlx ml-explore#4558)
95e89385	Gajesh Naik	fix(jaccl): clear the unused tail of each send frame
cbe63d14	jonathan308	fix(jaccl): bound the polling loops and release a group that fails
8785a85e	jonathan308	fix(jaccl): send from the GID that was advertised to the peer
c9dcfb50	jonathan308	docs(jaccl): describe the progress timeout
fc71f7fb	jonathan308	test(jaccl): add tests with simulated verbs
d0f8b0a8	jonathan308	docs(forkdiff): describe the JACCL changes and the upstream JACCL fixes
4241bc45	jonathan308	style(jaccl): format tests/jaccl/CMakeLists.txt with cmake-format
b7658964	jonathan308	fix(jaccl): do the ring all_gather slice arithmetic in 64 bits

Upstream backports

Upstream main was read at 0d8236fa (2026-10-09). Five upstream commits after the fork base (77a0c1e8) touch mlx/distributed/, and one more touches the CPU code of a distributed op.

Upstream	Commit	Date	What it fixes	In this PR
ml-explore/mlx#4210	adf21dea	2026-08-14	MPI backend is used only for a world size above 1	Already in the fork
ml-explore/mlx#4281	a44fc8c0	2026-08-16	TCP ring backend: a transfer of 2 GiB or more	Already in the fork
ml-explore/mlx#4443	7241f12e	2026-09-15	JACCL ring all_gather: both directions sent the same slice, so the upper half of each peer's region was not transferred	Backported unchanged; a later commit here (b7658964) casts the first factor of its slice arithmetic to 64 bits, because the code scanner flags the int product
ml-explore/mlx#4557	42987b66	2026-09-28	ReduceScatter::eval_cpu asserted that it has no inputs, so a build with assertions stopped at each sum_scatter	Backported, no changes
ml-explore/mlx#4558	8803e524	2026-10-01	A failed, non-strict backend initialization registered an empty group, and a later strict init returned it	Backported, no changes
ml-explore/mlx#4607	f4232c7f	2026-10-02	Feature: a second cable for the JACCL mesh	Not backported

• The test of Do not cache failed distributed backend initialization ml-explore/mlx#4558 uses the scoped_env test helper of ml-explore/mlx#4474. Only that helper is taken from Add a scoped_env utility for tests ml-explore/mlx#4474. The other changes of Add a scoped_env utility for tests ml-explore/mlx#4474 (test refactors) are not included.
• [jaccl] second cable for the mesh ml-explore/mlx#4607 is a feature, not a fix. It makes one MeshImpl for each wire and moves dispatch_wires out of RingImpl. The progress guard of this PR closes one MeshImpl and joins the ring wires inside RingImpl::dispatch_wires, so the guard must be reworked for that structure and tested again on hardware. That is a separate PR.
• Upstream has none of the three fixes of this PR: at 0d8236fa the JACCL loops do not read wc.status, have no time limit, post full frames without clearing, and use sgid_index = 1.

The three fixes

1. Send-frame tail (95e89385)

JACCL posts each frame at the full size of its scratch buffer. When the payload is shorter than the buffer, the bytes of earlier sends stay after it and go to the peer. stage_send_frame (new send_frame.h) copies the payload and sets the rest of the frame to zero. All mesh and ring send paths stage through SharedBuffer::stage_send. The posted sizes and the receive paths do not change.

2. Progress guard (cbe63d14)

Before: each mesh and ring polling loop waited for completions without a limit and did not read the status of a work completion. A rank whose peer had stopped used a full core until it was killed. A failed completion was used as data.

After: ProgressGuard::check() (new progress_guard.h) runs after each poll of each loop (6 loops in mesh_impl.h, 4 in ring_impl.h).

• A completion with a status other than IBV_WC_SUCCESS is an error at once. A negative return of ibv_poll_cq is an error at once.
• No completion for longer than the limit is an error.
• The clock is read once per 1024 polls that return nothing, not on each poll. Measured without sanitizers, the guard adds 0.7 ns per poll. A clock read on each poll adds 12.7 ns.
• The call throws a std::runtime_error that starts with [jaccl] and names the call and the cause. It does not exit the process. Scheduler::enqueue catches the exception on the CPU stream and reports it to the caller, so the caller can do its own cleanup.
• The errors of the ring pool wires now reach the caller. dispatch_wires dropped them before, because it used future::wait().
• The completion-queue size does not change. Each connection has at most 2 sends and 2 receives posted and its queue has 64 entries. The tests measured a maximum of 4 queued completions.

New environment variable. JACCL_PROGRESS_TIMEOUT_MS, with the alias MLX_JACCL_PROGRESS_TIMEOUT_MS, is the longest time in milliseconds that a call waits for a completion. The default is 120000. 0 or a negative value removes the limit. A value that is not a number keeps the default. The value is read once in a process. The wait of a rank includes the time that its peer computes (for example in recv), so the limit must be above the longest such wait.

The group is closed for good after a failure. This is the same for each exception that leaves a call:

• The queue pair and the completion queue of each connection are destroyed, which cancels the posted work, and all buffers are deregistered. The ring does this after all wires have returned.
• The protection domains and the device contexts stay open until the group is destroyed. MLX keeps the group until the process exits.
• Each later call on the group throws immediately. The group cannot be opened again in the process. A new process is necessary.
• The peer gets no message. Its own guard ends its wait.

3. Advertised GID (8785a85e)

Connection::info() finds the IPv4-mapped GID of the port and sends it to the peer. queue_pair_rtr then set the source GID by the fixed index 1. The two agree only when the mapped GID is at index 1. The connection now keeps the index where it found the advertised GID and sends from that index. Nothing changes where that index is 1.

On the two test Macs the mapped GID is at index 2. The commit passed normal traffic on that pair. The code without it also passes on that pair, so the hardware runs do not show that the fix is necessary.

Tests with simulated verbs

tests/jaccl (new) compiles the real mesh_impl.h, ring_impl.h and rdma.h with a stand-in for <infiniband/verbs.h>. A fake fabric completes a send when it meets a posted receive of the peer. It can lose or fail one completion and it can make the poll fail. The directory is a separate CMake project, in the same way as mlx/distributed/jaccl/lib/examples, so the MLX build does not build it and no header of it is installed.

cmake -S tests/jaccl -B build/jaccl-tests            # add -DJACCL_TESTS_TSAN=ON for TSan
cmake --build build/jaccl-tests
ctest --test-dir build/jaccl-tests --output-on-failure

At d0f8b0a8: 8 of 8 CTest tests pass with ASan and UBSan, and 8 of 8 pass with TSan. At the head (4241bc45, the format change): 8 of 8 pass with ASan and UBSan.

The same test sources were also built against the headers of each commit of this branch (a private harness that exports one revision at a time):

Headers	Send-frame groups (6)	Guard cases (38)
main (cb77239b)	0 pass, 6 fail: 4 find old bytes after the payload, 1 has no staging helper to test, 1 is the ring all_gather defect	not run
after the ml-explore#4443 backport (d92347ba to 7112cbf8)	1 pass (ring all_gather), 5 fail	not run
send-frame commit (95e89385)	6 pass	4 pass, 34 fail
guard commit (cbe63d14) and later	6 pass	38 pass with ASan and UBSan, 38 pass with TSan

The 34 guard cases that fail without the guard:

• A peer that does not answer, in each mesh call and each ring call with 1 and 2 wires (16 cases): the loop does not stop.
• A lost completion (4 cases): the mesh all-reduce drain loop, a mesh receive in the middle of a five-frame transfer, and a ring send on 2 wires for each wire.
• A failed completion or a failed poll (12 cases): the failed completion is used as data, or the loop does not stop.
• A destroy call that the verbs refuse, and the read of the limit (2 cases).

The 4 cases that pass without the guard are three transfers with waits below the limit and the count of the fake objects. With the guard, each stalled call fails between its limit (250 ms in the tests) and one second after it, the queue pairs and completion queues of the rank are destroyed exactly once, each region is deregistered exactly once, and each later call is refused in less than 100 ms. With the limit removed, a stalled call continues to wait and a failed completion is still an error.

Runs on two Macs

Two Apple silicon Macs (an M3 Ultra and an M5 Max, macOS 27) with a Thunderbolt RDMA link and a two-rank JACCL mesh. These runs used the three fix commits before the rebase onto the backports. The rebase changes the ring all_gather code, one assertion and the registration of a failed initialization. None of these is on the mesh path that was run.

Run	Result
Normal traffic, guard on: 1.38 GiB and 48,000 small reductions	Zero mismatches. The driver reports a success status on good completions, so the status check does not fail good traffic.
Peer ended, JACCL_PROGRESS_TIMEOUT_MS=4000	The other rank reported [jaccl] mesh all_reduce: no completion for 4001 ms … The group is closed and exited about 4 s after its peer ended.
Peer ended with the model loaded, limit 60 s	The other rank reported mesh recv: no completion for 60001 ms and exited by itself. Neither Mac kept a process or wired memory. The next session gave exact results.
Peer ended, code without the guard	The other rank stayed in jaccl::MeshImpl::all_reduce → tbt_poll_cq at 100% CPU for more than nine minutes.
Normal traffic with the GID commit	Pass. The code without it also passes on this pair.

CI on this pull request

At 4241bc45:

• metal-build-and-test (macOS, deployment target 26.2): the JACCL library (tcp, rdma, mesh, ring, jaccl) and mlx/distributed/jaccl/jaccl.cpp are compiled and linked. 292 C++ test cases pass and 1,869 Python tests run without a failure.
• Linux Fedora (aarch64 and x86_64) and the Linux ASan and UBSan test jobs pass.
• Check Lint and the fork-diff check pass.

No CI job runs JACCL traffic: the runners have no RDMA device.

Not shown

• Three or more ranks. The simulated tests and the hardware runs use two ranks, so the mesh scatter-gather all-reduce for more than two ranks was not run.
• Two CPU streams that use one group at the same time. This was not safe before this PR and is not made safe by it.
• The ring path on real hardware. The hardware runs used the mesh.
• A failed work completion on real hardware. None occurred in the runs. The status check was run only with simulated verbs.
• Whether the destruction of a queue pair with posted work can block or fail on the real driver. It returned in the two runs above. When it fails, the group prints [jaccl] Could not destroy a queue pair (N) and keeps the handle for its destructor.
• The two upstream Python tests that come with the backports (test_all_gather_rank_data, test_strict_init_after_fallback) were not run locally. ring_test_distributed.py in the Fork Tests workflow runs them; that workflow did not run on this pull request.
• tests/jaccl in CI. No workflow builds it.

Fork bookkeeping

• fork.yaml has a section for the JACCL fork work and a section that lists the cherry-picks as upstream work. python3 scripts/check_forkdiff.py --upstream-ref refs/remotes/upstream/main passes: 202 changed files, 202 described. base.hash does not change, because a cherry-pick does not move the merge-base.
• The fork-diff page was not rendered locally.
• The pre-commit hooks were not run locally. The Check Lint job ran them: clang-format, black and isort passed on the first push, cmake-format wanted one argument per line in one foreach of tests/jaccl/CMakeLists.txt, and 4241bc45 applies that. The job passes at the head.
• docs/src/usage/environment_variables.rst and docs/src/usage/distributed.rst describe the new variable and what occurs when a peer does not answer.
• mlx-swift vendors MLX. Its pin must move before the Swift stack gets these changes. The draft that moves it is build: pin mlx and mlx-c to the JACCL fixes and the owned bootstrap bridge mlx-swift#54.

How AI was used

Claude Code (Anthropic) wrote the two fork fixes after the send-frame commit, the tests, the documentation changes, the backport selection, the commit messages and this description, and ran the simulated-verbs tests and the two-Mac sessions, at the direction of the account owner, who asked for this draft. The send-frame commit is by Gajesh Naik and is unchanged apart from formatting. The upstream commits are cherry-picks with their original authors. AGENTS.md and CLAUDE.md in this repository are upstream's files and say that an agent must not push or open a pull request; this draft was opened on the explicit request of the account owner, in this fork only. Nothing is proposed to ml-explore/mlx.

Update 2026-10-10: guard v2 and the NAX MoE simdgroup-skip revert

Two commits on top of b7658964:

• ccf517a0 reverts upstream c7ff35d97 (the NAX quantised-MoE simdgroup skip, ml-explore/mlx PR 4352), which reached this fork through the v0.32.2 squash (734241bb), not by ancestry. On a two-Mac JACCL pair with one NAX GPU, MoE serving with the skip lost all-reduce completions and left 92 to 96 GiB wired until reboot. Output is unchanged (stores stay bounded by sgp_sm).
• fdae3174: default progress limit 30 s (was 120 s); after the group is released the default action is a teardown exit (memory release hook: set_wired_limit(0) then clear_cache(), 10 s budget, then exit status 75); JACCL_TIMEOUT_ACTION=throw keeps the old behaviour; completion queues sized 4x the outstanding work. Queue pair attributes untouched; no global collective lock.

Checked:

• Simulated verbs: ASan/UBSan 8/8 at ccf517a0; ASan/UBSan 10/10 and TSan 10/10 at fdae3174 (new teardown-exit and 30 s default-exit cases).
• Two Macs over Thunderbolt 5 RDMA (M3 Ultra + M5 Max): raw collective check 1 GiB exact; dead peer with no variable set: survivor exits 75 by itself 30.3 to 33.3 s after the peer's SIGTERM, wired memory back on both Macs (three runs, plus once with a 35B MoE stage loaded); Qwen3.6 35B A3B and GPT-OSS 20B across the pair: outputs exact against the pair runs before the revert, speeds within run-to-run spread, 56 of 56 sustained requests each with no lost completion; on the M5 alone the revert is bit-identical in output and speed.

Note for consumers: the d-inference libs/mlx pin (cb77239b) and the MLX nested in its libs/mlx-swift also carry the skip today; they need to move to this head (with a rebuilt metallib) for any distributed build.

🤖 Generated with Claude Code