[BUG] GPU locking using METAL_FAST_SYNCH=1 and the JACCL backend

[rltakashige]

Describe the bug
In certain occasions, the GPU gets stuck in the while(1) loop in fence_wait and can never escape, leaving the GPU "locked" at 100% usage until the machine is restarted.

We've looked into the issue for a while and can't seem to pin down the specific cause. It seems as though, from my investigation, that the GPU in fence_wait does not see the timestamp when it is updated by the CPU. It may be a timing issue or some sort of coherence issue.

To Reproduce
This is fairly difficult to reproduce (happens non-deterministically as far as I can tell). The most consistent way I have found is to run two large models in parallel using RDMA. It is also more likely on 4 nodes vs 2.

The context that we encounter this issue is as follows:

• Using the standard stream_generate with either a tensor/pipeline sharded model in Exo, we see this issue occasionally.
• I would like to see if this occurs in MLX LM, but I have not gotten around to this.

Additional context
#3141 attempts to resolve this, and the number of GPU locks seems to be drastically reduced. However, this is not an exhaustive fix at the very least.

[vskiwi]

We've been investigating this issue on a 4× Mac Studio M3 Ultra 512 GB cluster (macOS 26.3, Thunderbolt 5 full-mesh RDMA, JACCL backend) running GLM-5 754B tensor parallel. Here's what we found.

TL;DR

Changing HazardTrackingModeUntracked to HazardTrackingModeDefault in allocator.cpp:15 completely eliminates the GPU locking / SIGABRT with zero measurable performance impact on our distributed workload.

Root cause

There are two independent issues at play:

1. Fence visibility (addressed by #3141)

DSB SY in fence.cpp fixes the CPU→GPU memory ordering problem. Without it, DMB ISH (inner shareable) doesn't propagate fence counter writes to the GPU. This causes fence_wait to spin on a stale value. #3141 resolves this.

2. Buffer lifetime during in-flight command buffers (this issue)

Even with #3141 applied, SIGABRT crashes persist on longer contexts (50K+ tokens). The crash stack is always:

Triggered thread (com.Metal.CompletionQueueDispatch):
  mlx::core::gpu::check_error(MTL::CommandBuffer*)
  MTL::CommandBuffer::addCompletedHandler(...)

The root cause: fence.cpp completion handlers capture only fence_ but not the data buffers used by the dispatched kernels. Compare with eval.cpp (lines 45-64) which captures shared_ptr<Data> for all input buffers.

With ResourceHazardTrackingModeUntracked (allocator.cpp:15), Metal does not automatically track buffer usage by command buffers. When MetalAllocator::clear_cache() calls buf->release() or BufferCache::reuse_from_cache() recycles a buffer, in-flight command buffers may still reference it — causing GPU page faults → SIGABRT.

What we tested

We ran 7 experiments on 4× M3 Ultra 512 GB (GLM-5 754B, tensor parallel, JACCL RDMA):

Approach	50K tokens ctx	Notes
Clean (DSB SY + GPU fix from #3141)	SIGABRT ~3 min	clear_cache in crash stack
Remove residency_set_.erase from free_ callback	SIGABRT ~4.5 min	Different crash path, no clear_cache in stack
Capture x.data_shared_ptr() in fence completion handlers	DEADLOCK	~2 GB memory headroom → allocator starvation
Metal buf->retain() / buf->release() in fence handlers	SIGABRT ~5 min	Reduces frequency, doesn't eliminate
Combo: residency fix + Metal retain	SIGABRT ~5 min	Still crashes on repeated 50K runs
HazardTrackingModeDefault	PASS (all tests)	0 crashes, 0% perf impact

Why shared_ptr capture doesn't work

Capturing data_shared_ptr() in fence completion handlers prevents the buffer from entering the cache. With ~190 GB model on 192 GB RAM (~2 GB headroom), even one held buffer starves the allocator. The completion handler fires after CB completion, but allocation pressure occurs before CB commit → circular deadlock.

Performance comparison (HazardTrackingModeDefault vs baseline)

Benchmarked on 4× M3 Ultra, tensor parallel, JACCL RDMA:

Context	Baseline (Untracked) decode	Tracked decode	Delta
0 tokens	23.8 tok/s	23.6 tok/s	-0.8%
10K tokens	21.3 tok/s	21.2 tok/s	-0.5%
50K tokens	16.1 tok/s	16.1 tok/s	-0.2%

Prefill speeds also identical (365 vs 365 tok/s at 10K, 220 vs 218 tok/s at 50K). Peak memory unchanged (228 GB/node at 50K).

Suggested fix

The minimal fix is changing one line in mlx/backend/metal/allocator.cpp:

-constexpr size_t resource_options =
-    MTL::ResourceStorageModeShared | MTL::ResourceHazardTrackingModeUntracked;
+constexpr size_t resource_options =
+    MTL::ResourceStorageModeShared | MTL::ResourceHazardTrackingModeDefault;

We understand Untracked was chosen for performance, but our benchmarks show no measurable impact for distributed RDMA workloads. Single-node performance may differ — we defer to the MLX team's judgment on whether this should be unconditional or conditional on FAST_SYNCH=1.

An alternative would be to properly retain data buffers in fence.cpp completion handlers (matching what eval.cpp does), but our experiments show this causes deadlocks when memory headroom is tight. A more surgical fix might be possible with deeper knowledge of the allocator/scheduler internals.

Environment

• macOS 26.3 (25D125)
• Apple M3 Ultra × 4 (512 GB each)
• Thunderbolt 5 full-mesh RDMA (JACCL)
• MLX 0.30.7.dev (rltakashige fork, address-rdma-gpu-locks branch + our HazardTrackingModeDefault patch)
• Model: GLM-5-8bit-MXFP8 (767 GB, 78 layers, tensor parallel)

[vskiwi]

We tested #3144 on our 4× Mac Studio M3 Ultra 512 GB cluster (macOS 26.3, Thunderbolt 5 full-mesh RDMA, JACCL backend, GLM-5 754B tensor parallel).

Upstream main with #3144 only (no DSB SY from #3141)

Deadlock on first inference request — all 4 runners at 100% CPU, GPU idle (~40W), classic fence_wait spin. Required cluster reboot to recover.

Upstream main with #3144 + DSB SY (__dsb(0xF) after fence store in fence.cpp)

Test	Tokens (actual)	Result	Prefill	Time
Smoke ("Say OK")	7	PASS	—	2s
~36K target	50,406	PASS	222 tok/s	3m47s
~50K target	70,006	PASS	421 tok/s	2m46s
~67K target	93,806	PASS	358 tok/s	4m22s

All tests passed on first attempt.

Conclusion

#3144 fixes inter-command-buffer synchronization for fence buffers, but does not address CPU→GPU memory visibility through RDMA. On ARM64, std::memory_order_seq_cst compiles to DMB ISH (Inner Shareable domain), which is not visible to the GPU. Adding DSB SY (full system barrier) after the fence store in Fence::update() resolves the remaining deadlocks.

The DSB SY patch from #3141 is still needed for JACCL/RDMA workloads, in addition to #3144.

[rltakashige]

I don't think the issue is resolved yet. Will still keep this issue open.

[vskiwi]

@rltakashige We continued investigating after #3141 was closed and found the read-side half of the fence visibility problem.

Your fork and #3141 added __dsb(0xF) to Fence::update (the write path). But Fence::wait (the read path) still spins without a system barrier:

// fence.cpp, fast path (line 65)
while (f.cpu_value()[0] < count) {
    // compiles to LDAR + DMB ISH — invisible to GPU/DMA writes
}

DMB ISH is scoped to CPU cores only. When the fence counter is updated via RDMA DMA, the CPU re-reads a stale L1/L2 cache line indefinitely. One-line fix:

 while (f.cpu_value()[0] < count) {
+    __dsb(0xF);  // full system barrier — observe GPU/DMA writes
 }

Evidence — sample on a deadlocked process (4× M3 Ultra 512GB, MLX 0.31.0, JACCL RDMA, GLM-5 tensor parallel):

Thread (StreamThread):
  271 Fence::wait::$_2::operator()()
    101 IOGPUMetalBuffer::contents
    170 IOGPUMetalResource::virtualAddress

All 271 samples stuck in the read loop. GPU completely idle. Classic stale-cache-line symptom.

Reproduces 100% through exo (per-layer mx.eval during tensor parallel sharding creates ~78 rapid fence cycles shortly after JACCL init). Also reproduced with a standalone script on vanilla 0.31.0, though at lower frequency (~3%).

Are you still seeing GPU locks with your current fork? If so, could you try adding this one line to your Fence::wait and check if it helps? Happy to collaborate on a PR — the read-side fix naturally complements your write-side work.

[Harperbot]

Sharing a related data point from production.

We run a multi-model MLX pipeline on M1 Ultra (64GB) that loads/unloads 8+ models sequentially. We hit a similar class of cross-process Metal state corruption — in our case, IOGPUMemory.cpp:492 completeMemory() prepare count underflow kernel panics rather than fence_wait hangs, but the root cause overlaps: concurrent or racing Metal command buffer operations across process boundaries.

We built MetalGuard (v0.4.0) to address this. Relevant features for this issue:

1. Cross-process mutual exclusion — file-based lock (acquire_mlx_lock() / mlx_exclusive_lock()) prevents two MLX processes from touching Metal simultaneously. Self-healing stale locks after crashes.

2. Hardware-aware auto-configuration — MetalGuard.recommended_config() detects your chip/memory tier and returns appropriate safety thresholds.

3. Pre-generate health probe — probe_metal_health() verifies the Metal command queue is alive before starting inference. If the GPU is in a bad state, this fails at a controlled point (~1ms) instead of hanging indefinitely.

For the fence_wait infinite loop specifically, the process lock prevents the concurrent execution pattern that triggers it. While it doesn't fix the underlying Metal driver coherence issue, it eliminates the trigger path.

from metal_guard import mlx_exclusive_lock

with mlx_exclusive_lock("node_1_inference"):
    result = mlx_lm.generate(model, tokenizer, prompt=prompt)

Single file, zero dependencies, MIT licensed. Would be interested to hear if it helps with the RDMA multi-node scenario.