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

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.

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)

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.

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

fix: DSB SY barrier + system-scope atomic load for GPU fence visibility

@rltakashige

@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.

Note on METAL_FAST_SYNCH : according to the Metal team, there is no support for any kind of atomic memory coherence between CPU/GPU, nor any public way to implement a spinlock in the GPU that is guaranteed to work for all GPUs. So this flag would not be guaranteed to work, and I'm marking this issue as won't fix.

@zcbenz

Thanks @zcbenz ! Shame, but understandable.

(As an fyi, for the past several months, we've gotten around this by exiting the spin loop after a fixed number of iterations. While this can cause slightly incorrect outputs, I feel that this is better for our use case than the alternative which is irrecoverable. So we can make do)

Document MLX_METAL_FAST_SYNCH as unreliable instead of warning at run…

Warn once when a fast fence is created with MLX_METAL_FAST_SYNCH=1

Document MLX_METAL_FAST_SYNCH as unreliable instead of warning at run…

@rltakashige

A data point on one sub-class of these hangs that does not need the CPU/GPU spinlock at all, with a fix.

Shape. With MLX_METAL_FAST_SYNCH=1, a GPU stream consuming an array that another GPU stream produced in the same eval is ordered with a fence_wait spin kernel on the consumer queue polling a fence_update kernel on the producer queue (fence.cpp, cross_device=false). Two command buffers on two queues are not guaranteed to overlap. When the consumer's runs first, the spin never sees the update: on main fence_wait is while (1), so that is a hang; on @rltakashige's branch it is a 10 s stall per handoff. Resources are allocated HazardTrackingModeUntracked, so the driver does not order the queues either.

Where it showed up. 2× Mac Studio M3 Ultra, macOS 26.6, JACCL tensor parallel, mlx-lm with a vision model: a fixed 9–14 s pause on about one image turn in seven, identical on both ranks, never on text turns. Image turns add mx.eval calls on the default stream (vision embeddings, KV restore/flush) whose arrays mlx-lm's generation stream consumes inside the same eval. An instrumented build tied every stall to a GPU fence_wait exhausting its bound on a fence carrying one layer's 44.7 MB KV cache, both sides GPU-encoded. Coherence-style changes (reload cadence, one page per fence, explicit cache clean) did nothing; ordering did.

Fix. Use a per-fence MTLSharedEvent for same-device cross-stream handoffs (encodeSignalEvent on the producer queue, encodeWait on the consumer queue), which is what the non-fast path already does, and keep the spin only across the CPU/GPU boundary, where the flag pays off. One file, mlx/backend/metal/fence.cpp. Details, the production numbers (two days, ~650 requests, three models rotating, zero give-ups; 2/7 → 0/23 stalls in the reproduction harness; warm prefill +8%, decode unchanged) and a standalone stress script are in rltakashige#4.

A branch against main is at aidiffuser/mlx-jaccl-fix-small-recv:metal-fast-synch-cross-stream-events-v2 (clang-format clean, test_eval.py passes with the flag off). I have not run it with the flag on against main on this machine: the stress also exercises the CPU→GPU direction, which on main hangs in the class this issue is about, and an orphaned unbounded spin kernel would take out a production GPU. I understand the flag is unsupported; this only narrows what it relies on the spinlock for. Happy to open the PR if that is useful.

Co-authored with Claude Fable 5.1.

Navigation Menu

Repository navigation

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

TL;DR

Root cause

1. Fence visibility (addressed by #3141)

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

What we tested

Why `shared_ptr` capture doesn't work

Performance comparison (`HazardTrackingModeDefault` vs baseline)

Suggested fix

Environment

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

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

Conclusion

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

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

Description

Activity

vskiwi commented on Feb 18, 2026

TL;DR

Root cause

1. Fence visibility (addressed by #3141)

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

What we tested

Why shared_ptr capture doesn't work

Performance comparison (HazardTrackingModeDefault vs baseline)

Suggested fix

Environment

vskiwi commented on Feb 19, 2026

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

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

Conclusion

rltakashige commented on Feb 20, 2026

vskiwi commented on Mar 8, 2026

Harperbot commented on Apr 12, 2026

zcbenz commented on Apr 28, 2026

rltakashige commented on Apr 28, 2026

aidiffuser commented on Sep 10, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Why `shared_ptr` capture doesn't work

Performance comparison (`HazardTrackingModeDefault` vs baseline)

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