Repository navigation
Conversation
zcbenz
left a comment
There was a problem hiding this comment.
This is a statement that I would only trust someone working on Metal at Apple to make, I could try to find one in Apple to look into if it is true, but you need to prove this is not AI hallucinations first.
|
absolutely |
…arrier This adds a `__dsb(0xE)` (full system, stores only) barrier after the CPU updates the fence value. On ARM64, `std::atomic` seq_cst compiles to DMB ISH (inner-shareable), which isn't always visible to the GPU. This forces the store to the point of coherence, fixing a cluster hang bug where the CPU wrote a lock value but the GPU didn't see it.
80c634a to
de61022
Compare
|
Update — not a stealth edit:
Still working on A/B testing |
|
Updated the PR description with the full prior-art context I was missing — most importantly @vskiwi's PR #3141 (CPU-side DSB SY + GPU-side fallback, closed Feb 24 after the proponents couldn't reliably reproduce on a clean cluster), the still-open issue #3142, and @vskiwi's March 8 follow-up identifying a read-side coherence gap that this PR does not address. I should have led with that history in the original description; my apologies. The technical premise here is not new — it's the same write-side fix from #3141, originated by @rltakashige and validated by @vskiwi. I'm carrying it on a 2× M4 Max cluster running jaccl/RDMA. On reproducibility — please read the "Reproduction status" section of the updated description. I tried to A/B reproduce the hang to give you fresh evidence and could not, including a four-variant binary-verified A/B on my own cluster. My setup (2 nodes, MiniMax-M2.7) sits below the threshold where the hang has been observed by vskiwi (4× M3 Ultra, GLM-5 754B) or rltakashige. The synthetic 2-rank repro path is not the right validation method for this PR. I'm being explicit about this so you can weigh the evidence accurately, even though it weakens the case for landing. Two open questions in the updated description:
If reviving this is useful, the right move is probably to coordinate with @rltakashige and @vskiwi rather than land it under my name — they have the deeper context and the larger-cluster repro environment. |
|
Talked with the Metal team, so 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. The hack in this PR may appear to work, but it was only because it triggered some implementation quirks, not because Apple Silicon followed the ARM specs. So it is very unlikely to work for all hardwares and macOS versions, and would still break for many cases. Having said that, I'm kind of OK merging this if it helps since the spinlock fence is under a flag, but the code comment should definitely state it is a hack based on observation rather than intended API behavior. /cc @angeloskath |
|
Hmm I probably don't want to introduce any junk into the repo, since this is a hack. That was definitely the goal of upstreaming this, to help others that might be hitting it. Also appreciate the review 😄 |
|
Anyway thanks for the detailed investigation on the problem! I'm closing this PR. |
Summary
Adds a
__dsb(0xE)(DSB ST — Data Synchronization Barrier, Full System, stores only) after the CPU-side fence store inFence::update, so the value is guaranteed visible to GPU/DMA observers beforefence_waitpolls it.This is the CPU-side write-path half of the fix proposed in #3141 (which was closed without merging in February 2026). The technical premise, root-cause analysis, and reproduction work are not new — that's prior art by @rltakashige and @vskiwi. This PR resurrects the minimal write-side patch against issue #3142, which is still open.
Note on barrier scope: this PR currently uses
__dsb(0xE)(DSB ST), matching the original Feb 16, 2026 commit on @rltakashige'saddress-rdma-gpu-locksbranch. One day later that branch switched to__dsb(0xF)(DSB SY) after empirical testing showed ST insufficient on Apple Silicon — and that's what my fork (and the deployed cluster) actually carries today. I'm happy to push a follow-up commit upgrading this PR to0xFif you'd prefer the empirically-validated variant; see Why DSB SY may be needed instead of DSB ST below.Background
address-rdma-gpu-locksbranch (commit dccc541, originally__dsb(0xE)/ DSB ST)__dsb(0xF)/ DSB SY ("Use DSB SY and not ST.") after ST was found insufficient on M3 Ultrastd::atomic::store(seq_cst), GPU-side spin-loop changes ("nuclear" / "more nuclear 2")[BUG] GPU locking using METAL_FAST_SYNCH=1 and the JACCL backend. Still OPEN.Fence::waitthat this PR does not address (see Known limitation below).Root cause (per #3141 + #3142)
In
Fence::update, the fast-path CPU writes a counter that the GPU'sfence_waitMetal kernel polls:On ARM64,
std::atomicseq_cst compiles toSTLR+DMB ISH— theDMB ISHonly orders accesses within the Inner Shareable domain, which on Apple Silicon comprises the CPU cores. The integrated GPU and DMA engines live in the Full System domain and may continue to observe a stale cache line.When the underlying coherence channel is RDMA over Thunderbolt (
MLX_METAL_FAST_SYNCH=1+jacclbackend), the GPU'sfence_waitspin loop never sees the updated counter and the rank deadlocks at 100% CPU. Cross-references:__dsb(0xF)= DSB SY (full system).DSB SY(vsDMB ISH) does two things the original lacked: it ensures completion (not just ordering), and it scopes to the full system domain (not just inner-shareable).Why DSB SY may be needed instead of DSB ST
Per the ARM ARM,
DSB ST(full system, stores only —__dsb(0xE), what this PR ships) should suffice for this single-store-then-poll pattern. The original Feb 16 commit on @rltakashige's branch used0xE. One day later (Feb 17, commit3a3f1904 "Use DSB SY and not ST.") it was switched to0xF, and #3141's body documents the same finding: "DSB ST … should theoretically suffice per the ARM spec, but was found insufficient on M3 Ultra during testing — likely due to Apple Silicon's implementation-specific cache coherence behavior between CPU and GPU."My deployed fork uses
0xF, not0xE. I'm open to upgrading this PR to0xF(one extra commit on this branch — happy to push it on request).Reproduction status — please read this section
This bug cannot be easily A/B tested. I tried, and I want to be transparent about it.
What I did before submitting this rewrite (2026-04-26):
otool -tv libmlx.dylib | awk '$2 == "dsb"':__dsbinstruction emitted, upstreamfence.metalkernel)__dsb(0xE)/ DSB ST + upstreamfence.metal(this PR exactly)__dsb(0xF)/ DSB SY + upstreamfence.metal__dsb(0xF)/ DSB SY + the GPU-side spin-loop changes my fork carries (current cluster pin)mlx.launch --backend jacclwithMLX_METAL_FAST_SYNCH=1, on Thunderbolt 5 RDMA. Two tests:all_sum, 1 MB payload, 3 trials per variantall_sum+all_gatherwith cross-stream eval, 4 KB → 16 MB payloads, 60–110 seconds sustained per variantResult: all four variants pass every trial. The historical hang did not reproduce in either test on the variant-A binary that lacks the fix. I'm reporting this honestly even though it weakens the case for landing the PR — anyone trying to validate this should know the synthetic 2-rank repro path doesn't fire.
This is consistent with @vskiwi's experience reported on #3141: "We were unable to reproduce the fence deadlock on a clean cluster. Our earlier deadlock may have been caused by accumulated GPU/Metal state from prior test iterations rather than a clean-state issue." That observation is what led to #3141 being closed.
The conditions under which the hang has actually been observed (per #3142 / #3141, not by me on my cluster):
mx.evalduring exo's TP weight sharding produces ~78 rapid fence cycles immediately after JACCL init — vskiwi found this is the most reliable trigger;mlx_lm.sharded_load's bulk-eval pattern does not consistently trigger itMy cluster has 2 nodes, not 4 — and runs MiniMax-M2.7 (smaller than GLM-5) — so it sits below the threshold where the hang has been reliably observed by either vskiwi or rltakashige. That doesn't make the bug fictional; it does mean a synthetic A/B from my hardware is not the right validation method for this PR.
What I can attest to from my own cluster (2× M4 Max 128 GB, Thunderbolt 5 RDMA, jaccl backend, MiniMax-M2.7 tensor parallel + Huihui-35B-A3B pipeline): I have been running the DSB SY (
0xF) form on the CPU-side write path since Feb 16, 2026 with zero observed fence deadlocks during normal operation. That's an absence-of-symptom over ~2 months, not a controlled A/B — I want to be precise about what I'm claiming.Performance: @rltakashige measured 267 vs 269 tok/s on Llama 3.2 1B (~10K context) — within noise. My cluster has not seen any measurable decode-throughput delta either.
Verification
otool -tvon the resultinglibmlx.dylibconfirms onedsb syinstruction is emitted exactly where the patch is, and not before:Without the patch, that count is 0.
Known limitation — read-side gap (#3142, vskiwi 2026-03-08)
This PR addresses only the write side (
Fence::update). @vskiwi identified thatFence::wait(fence.cpp:67) also spins without a system barrier:When the fence counter update arrives via RDMA DMA, the polling CPU re-reads a stale L1/L2 cache line indefinitely — a separate manifestation of the same coherence gap. vskiwi's evidence (
sampleon a deadlocked process) showed 271/271 samples stuck in this read loop with the GPU completely idle.A complete fix would also add
__dsb(0xF)inside the read-side spin loop. I have intentionally not included that here so this PR matches what my cluster has actually been running. If you want, I can add it as a second commit referencing vskiwi's analysis.Authorship / coordination
The technical work behind this fix is @rltakashige's (root-cause + original patch chain on
address-rdma-gpu-locks) and @vskiwi's (independent verification + comprehensive writeup in #3141). I'm carrying the minimal write-side patch on my own fork and submitting it because (a) the upstream issue #3142 is still open, (b) the previous PR was closed without strong evidence either way, and (c) the patch has been stable in production on my cluster for ~2 months.If a maintainer wants this revived, the right move is probably to coordinate with @rltakashige and @vskiwi rather than land it under my name — they have the deeper context and the larger-cluster repro environment.
What this PR does NOT do
fence.metalfallback from fix: use DSB SY barrier and system-scope atomic load in fence for RDMA GPU visibility #3141 (system-scope atomic load after N spins). That's separate work and can be re-PR'd against this if useful.__dsbper vskiwi's Mar 8 comment (see Known limitation above).Files changed
mlx/backend/metal/fence.cpp: 1 include + 14 lines (1 new instruction + 12-line comment explaining the rationale).