[rltakashige]

Motivation

There is an issue on Macs that means that an explicit synchronization is necessary for memory to be updated from L1 cache. This means that GPU locks can occur when a spin wait does not see the updated timestamp.

Changes

Updated in my own personal fork.

Why It Works

https://github.com/ARM-software/acle/releases

Test Plan

Manual Testing

Tested manually that no GPU locks occur (even with multiple simultaneous instances running) and that the performance differential is negligible (267 vs 269 tps on Llama 3.2 1B at an approx 10k context.)

---

I have seen a GPU lock, specifically when sending a particularly large chat completion while the model was loading. However, I have since been unable to reproduce and this may be something I did wrong. Please do create an issue and tag me if any GPU locks do occur.

[vskiwi]

Hey @rltakashige — the fence.cpp fix in your address-rdma-gpu-locks branch is excellent. We've been hitting exactly this Fence::wait deadlock on our 4× Mac Studio M3 Ultra cluster for weeks and independently traced it down to the same root cause (RDMA DMA writes not visible to GPU spin-wait due to DMB ISH only covering the inner shareable domain).

Have you considered submitting the DSB SY fix as a PR to upstream ml-explore/mlx? It's a very targeted 1-line change in fence.cpp with a clear rationale — should be an easy review for the MLX team.

The change in the exo PR bundles it with the hanning() function and other unrelated bits — splitting just the fence fix into a clean upstream PR would:

1. Get the fix into the next MLX release (benefiting everyone using JACCL/RDMA)
2. Remove the need for exo to depend on a fork
3. Let us pin to a standard MLX version instead of a git ref

We have detailed crash data and stack traces from our testing that could support the PR description if helpful (SIGABRT in check_error(MTL::CommandBuffer*), Fence::wait deadlocks on 4-node tensor parallel, AMCC interrupt correlation from macOS 26.2, etc.)

Your prior PR #3033 (jaccl small recv) went through cleanly — this one should too.

[AlexCheema]

LGTM — ARM64 DSB SY barrier fix is correct, pinned in uv.lock + nix by commit SHA. Must upstream to ml-explore/mlx. CI green.

[AlexCheema]

PR #1489 Review: "Use custom fork that resolves GPU locks"

Summary

This PR switches the macOS MLX dependency from the upstream PyPI release (v0.30.6) to a custom fork (rltakashige/mlx-jaccl-fix-small-recv @ 50487b41) to resolve GPU locks caused by ARM L1 cache synchronization issues on Apple Silicon. It also fixes Nix build issues for mlx-cuda-13 on Linux.

CI: All 6 checks passing.

Note: This appears to be an updated version of PR #1453, with a newer fork commit and additional Nix fixes.

---

Files Changed (6)

1. README.md — Add Nix run instructions, add Xcode prerequisite, whitespace fixes
2. flake.nix — Filter MLX package by p.source ? git to disambiguate two MLX entries in uv.lock
3. nix/mlx.nix — Point MLX build at custom fork (version 0.30.7.dev20260217+50487b41)
4. pyproject.toml — Pin macOS MLX to custom git fork; remove version pin for darwin
5. python/parts.nix — Fix mlx-cuda-13 RPATH issues on Linux; add venvIgnoreCollisions for file conflicts
6. uv.lock — Updated lock file reflecting fork switch; removes mlx-metal package for macOS

---

Detailed Analysis

1. MLX Fork Switch (pyproject.toml, uv.lock, nix/mlx.nix, flake.nix)

• macOS MLX changed from PyPI 0.30.6 to git fork at 50487b41
• Linux MLX remains pinned to 0.30.6 from PyPI
• mlx-metal package removed from lock file (fork bundles metal shaders)
• flake.nix disambiguates MLX packages by checking for git source

Positives:

• Good PR description explaining the ARM L1 cache synchronization root cause
• Manual testing documented (267 vs 269 tps — negligible perf difference)
• Honest disclosure about one observed GPU lock during model loading

Concerns:

• Personal fork dependency — Same concern as Address rdma gpu locks #1453. Is there an upstream MLX PR for the JACCL fix? Long-term maintenance depends on getting this merged upstream.
• Removed version pin for darwin MLX ("mlx; sys_platform == 'darwin'") — if the git source becomes unavailable, uv could resolve to any version.
• The fork is at a dev version (0.30.7.dev...) — this is unreleased code. What's the delta from upstream 0.30.6? Just the JACCL fix, or other changes too?

2. Nix Linux Fixes (python/parts.nix)

• mlx-cuda-13 override adds NVIDIA build inputs and addAutoPatchelfSearchPath for CUDA libraries
• autoPatchelfIgnoreMissingDeps for libcuda.so.1 (driver stub, only on GPU hosts)
• venvIgnoreCollisions for mlx* and nvidia* site-packages paths to resolve file collisions between mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13

Assessment: This is solid Nix packaging work. The commit message (by JakeHillion/Claude) clearly explains both problems and solutions. The collision handling via fnmatch patterns is reasonable — keeping the first path encountered is fine for type-checking environments.

3. README.md

• Adds Nix run shortcut (nix run .#exo)
• Adds Xcode as a prerequisite (needed for Metal toolchain)
• Minor whitespace normalization

Assessment: Good quality-of-life improvements.

---

Comparison with PR #1453

This PR (#1489) supersedes #1453 with:

• A newer fork commit (50487b41 vs d94f81a2)
• Nix fixes for mlx-cuda-13 on Linux (not in Address rdma gpu locks #1453)
• No code changes to generate.py, bootstrap.py, or channels.py (those were in Address rdma gpu locks #1453)
• Better PR description with motivation and manual test results

---

Overall Assessment

Positives:

• Well-motivated: GPU locks are a real problem, and the ARM cache sync explanation is sound
• Manual testing shows negligible performance impact
• Nix Linux fixes are well-explained and correct
• CI all green
• README improvements are welcome

Concerns:

1. Personal fork as a production dependency — What's the plan for upstreaming the JACCL fix to ml-explore/mlx? This should not stay on a personal fork long-term.
2. Scope of fork changes — What's the exact delta between upstream 0.30.6 and this fork? Is it just the cache sync fix, or are there other changes?
3. Removed version pin on darwin — Could lead to unexpected resolution if git source is unavailable
4. One observed GPU lock during model loading was noted but not reproduced — worth monitoring

Recommendation: The core change (fork with GPU lock fix) addresses a real problem and has been manually validated. The Nix fixes are solid. Main risk is long-term maintenance of the personal fork — would be good to clarify the upstream plan.