[BUG] Distributed pipeline generation: Metal fence handoff deadlocks under MLX_METAL_FAST_SYNCH=1 (TCP ring; orphaned fence_wait kernel locks GPU until reboot) and hits the ~5s GPU watchdog when unset (kIOGPUCommandBufferCallbackErrorTimeout at ~7.3k tokens)

[geekybraindev]

Summary

Running pipeline-style distributed inference over the plain TCP ring backend silently deadlocks when large activation tensors cross the GPU-stream ↔ CPU-communication-stream fence handoff with MLX_METAL_FAST_SYNCH=1. Both ranks sit at ~100% CPU forever, memory flat, no error. Killing the wedged processes releases their memory but leaves the GPU compute engine occupied by the orphaned fence_wait spin kernel — after that, any Metal eval from any process on the machine hangs, and only a reboot recovers (in one instance the machine escalated to a full userspace freeze). We believe this is the same family as #3142 / #3141, but note the fence layer is transport-independent (ring.cpp and jaccl.cpp declare the identical CPU communication stream), so the "JACCL" framing of #3142 undersells the blast radius: plain-Ethernet TCP-ring users hit it too. This matches silent multi-node hangs long reported in the wild (e.g. mlx_lm.server --pipeline over 10 GbE wedging mid-generation after ~2k tokens on ~60-layer models; exo-explore/exo#1666 / #1853 / #1726).

With the env var unset, the deadlock disappears — but long generations instead die at the macOS GPU watchdog (details below). Distributed pipeline users currently choose between an untimed spin that deadlocks and a driver watchdog that kills their process.

Environment

• Mac Studio M2 Ultra 192 GB + MacBook Pro M5 Max 128 GB, macOS 26.5.1 (25F80)
• mlx 0.31.2, mlx-lm 0.31.3 (PyPI), Python 3.13
• Reproduces on 2 ranks on ONE machine over a loopback TCP ring, and across two Macs over 10 GbE

Minimal repro (no model needed, ~30 s)

fence_stress.py (~90 lines, happy to attach): 2-rank loopback ring; per iteration it replays mlx-lm's per-token pipeline pattern — GPU compute → CPU-stream send/recv_like/all_gather → GPU compute, with mx.async_eval double-buffering exactly like mlx_lm/generate.py.

echo '[["127.0.0.1:19100"],["127.0.0.1:19101"]]' > /tmp/hostfile.json
# rank 1 then rank 0, both with:
MLX_HOSTFILE=/tmp/hostfile.json MLX_RANK=<r> MLX_METAL_FAST_SYNCH=1 \
  python fence_stress.py --log /tmp/rank<r>.log --rows 256 --dim 5120 --layers 4 --scalar-allsum

• --rows 1 (10 KB messages): clean for 2,000,000 iterations (also clean at 500k with heavier compute, and at 217 ms/iter GPU residency for 7k iterations).
• --rows 256 (2.6 MB messages ≈ a ~256-token prompt's prefill activation): deadlocks within seconds.

What the wedge looks like inside (sample of both ranks)

• Main thread parked: mlx::core::eval → array::wait → Event::wait → -[IOSurfaceSharedEvent waitUntilSignaledValue:timeoutMS:]
• Both SocketThread::worker threads idle in condition_variable::wait — zero pending socket tasks; the network layer is NOT stalled.
• A CPU-stream thread burning 100% in the fast-fence spin (fence.cpp:65 while (f.cpu_value()[0] < count) {} — hot frames show -[IOGPUMetalBuffer contents]), waiting on a timestamp the GPU side never bumps.
• GPU utilization pinned at 100% — the fence_wait kernel (fence.metal while(1) on a non-atomic volatile load, no timeout) occupying the queue. Operator-visible tell: GPU gauges read pegged while producing nothing.

Second bug: GPU poisoning after process death

After kill -9 of both wedged ranks (their wired memory IS released): a trivial mx.eval from a fresh process hangs forever; killing every remaining GPU-client process does not recover; only reboot restores the GPU. On a poisoned machine every subsequent distributed run "wedges" instantly at startup — we suspect this causes serial misdiagnosis in the wild.

With MLX_METAL_FAST_SYNCH unset: deterministic watchdog crash on long generations

The deadlock is gone (the synthetic killer config ran 100,000/100,000 iterations clean; real mlx_lm.server --pipeline runs complete normally). But long single generations die deterministically — three runs at 7,325 / 7,391 / 7,280 tokens (DeepSeek-V2.5-1210-4bit, 2 ranks over 10 GbE) with:

libc++abi: terminating due to uncaught exception of type std::runtime_error:
[METAL] Command buffer execution failed: Caused GPU Timeout Error
(00000002:kIOGPUCommandBufferCallbackErrorTimeout)

thrown from mlx::core::gpu::check_error in the completion handler (SIGABRT), rank 1 first. Mechanism: the slow path parks committed command buffers on encodeWait for cross-rank fence data; the macOS ~5 s GPU watchdog counts waiting as running and kills the buffer — the GPU is actually idle when it dies. (Same class as the stream=mx.cpu workaround discussion in #3207, but here the waiting buffers are the compute graph's fence waits, which users can't move off the GPU stream.)

We also confirmed the FAST_SYNCH deadlock is network-independent: it reproduces on loopback and on a verified single-interface 10 GbE link (after we eliminated an ARP-flux problem from dual-homed hosts — which the ring's missing SO_KEEPALIVE had been converting into separate silent stalls).

Static observations that may help localize

1. fence.metal fence_wait spins on a plain volatile load (not a Metal atomic) with thread_scope_system fences only; the CPU signals via std::atomic_uint store (fence.cpp:112). No atomic RMW forces cache-line ownership transfer; a missed invalidation is an unrecoverable spin — and there is no timeout/error path in either fast-path wait (the slow path's waitUntilSignaledValue at least dies loudly).
2. The fast GPU wait is dispatched into the current open command encoder without end_encoding() (fence.cpp:87-94) and buffers batch many ops (MLX_MAX_OPS_PER_BUFFER), so a spinning wait can block a later-encoded fence_update the peer rank transitively needs through the ring round-trip — a cross-rank ordering deadlock whose probability rises with tensor size / ops-per-buffer, matching the observed size dependence.
3. Unrelated but found while auditing ring.cpp (happy to file separately): the SocketThread worker reads/mutates sends_/recvs_ and task.buffer/size outside queue_mutex_ (ring.cpp:233-259 — a data race); recv()==0 peer-close is misclassified as EAGAIN via stale errno; after 10 errors the worker exits and orphans all pending promises; no timeout or SO_KEEPALIVE anywhere — so ANY distributed stall is silent and permanent for mlx_lm.server users.

Asks

• Do not leave GPU command buffers occupied or committed across fence waits: split/commit before the wait on the slow path (fixes the watchdog kills), and end_encoding() before the fast GPU wait (breaks the cross-rank ordering cycle).
• A bounded wait + error path in both fast-path spins (GPU kernel and CPU loop) — converts a silent machine-killer into a diagnosable exception.
• We can provide: the stress harness, full sample captures of both ranks mid-wedge (synthetic and mlx_lm.server), macOS crash reports for the watchdog kills, and the two-machine 10 GbE reproduction recipe.

[katlun-lgtm]

Confirming an independent reproduction on another machine class — M3 Max (64 GB), macOS 26, mlx 0.31.2 (PyPI) — with a small synthetic harness that replays mlx-lm's PipelineMixin per-step pattern (recv_like → GPU compute → send → all_gather, with mx.async_eval double-buffering), 2-rank loopback ring.

The env-var dependence you describe is exact here:

• MLX_METAL_FAST_SYNCH=1 → wedges on the first iteration; both ranks pinned at 100% CPU, memory flat, no error.
• unset → runs clean (completed a 300-iteration loop at ~325 it/s).

sample of both wedged ranks matches your diagnosis — the main thread is parked on the GPU event while the socket workers sit idle, so the network layer is not the stall:

-[IOSurfaceSharedEvent waitUntilSignaledValue:timeoutMS:]   (IOSurface)
mlx::core::Event::wait()
mlx::core::array::wait()
  ... mlx::core::Fence::wait(mlx::core::Stream, mlx::core::array const&)::$_2 ...
mlx::core::distributed::ring::(anonymous)::SocketThread::worker()   // in condition_variable::wait, no pending tasks

Two notes that may help narrow it:

• This is the plain TCP ring (not JACCL), so I can confirm the transport-independence point first-hand.
• In my harness the wedge appears even at rows=1 (~20 KB messages) once a per-step scalar all_sum is in the loop alongside the pipeline send/recv_like/all_gather — i.e. it reproduces without large activations, consistent with the fence handoff (rather than message size) being the trigger. My harness differs from yours, so treat this as an extra data point rather than a contradiction of your rows=1 clean run — it might implicate the scalar-reduction fence crossing specifically.

On the second (GPU-poisoning) bug: on this config, after kill -9 of both wedged ranks a fresh GPU mx.eval did recover without a reboot — so the poisoning may not be fully deterministic across setups, even though your reports show it clearly can (and does) happen.

Happy to share the harness and the full sample captures if they'd help.

[geekybraindev]

Thanks @katlun-lgtm — a third machine class (M3 Max) with an independently written harness is exactly the confirmation this needed. Your sample matches ours frame-for-frame (main thread in Event::wait → IOSurfaceSharedEvent, socket workers idle on their CVs, the CPU-stream thread hot in the fence.cpp:65 spin).

On the scalar-all_sum point, our data agrees with your "extra data point" framing, with one nuance:

• In our synthetic harness, --rows 1 --scalar-allsum (with 8 matmul layers per iteration) ran clean for 500,000 iterations on the M2 Ultra — so on this chip/OS the per-step scalar reduction alone didn't trigger at small message sizes.
• But the real mlx_lm.server --pipeline wedges at small per-token messages regardless: DeepSeek-V2-Lite (~4 KB/token) froze ~30 tokens into generation after a clean multi-MB prefill, and our two-machine DeepSeek-V2.5 runs wedged both with a per-iteration scalar all_sum patched into the generation loop (8,548 tokens) and without it (1,920 tokens) — n=1 each, huge variance.

So in our observations, message size and the scalar-reduction crossing both raise the per-iteration trigger probability, but neither is necessary — consistent with the fence handoff itself being the race. That your harness triggers at rows=1 once the scalar reduction joins the loop (where ours didn't) may just be chip/OS timing differences, which would also fit a nondeterministic race.

On post-kill -9 recovery: on our M2 Ultra (macOS 26.5.1) the GPU never came back — a trivial mx.eval of a matmul from a fresh process hangs forever, killing every remaining GPU client doesn't help, and once the poisoned state escalated to a full userspace freeze (SSH dead, power cycle required). Your M3 Max recovering suggests the poisoning severity varies across chips/OS builds; agreed it isn't fully deterministic.

Our harness is inline below so everything stays in the thread (launch recipe as in the issue body; --rows / --layers / --scalar-allsum map to the result columns above). Full mid-wedge sample captures (synthetic and mlx_lm.server) and the watchdog crash reports are available on request.

For anyone hitting this in production meanwhile: our operating workaround is MLX_METAL_FAST_SYNCH unset, distributed generations capped at ~6,500 tokens (under the ~7.3k-token GPU-watchdog kill we measured on this model/config), and longer outputs delivered as continuation requests over the preserved prompt cache.

fence_stress.py (~90 lines, no model needed)

#!/usr/bin/env python3
"""fence_stress.py — minimal 2-rank TCP-ring fence stress (no model).

Replays the per-token cross-stream fence pattern of a 2-rank mlx-lm pipeline
(deepseek_v2.py:394-409 shape): GPU compute -> CPU-stream ring comm -> GPU
compute, with async_eval double-buffering like generate.py. With
MLX_METAL_FAST_SYNCH=1 each iteration crosses the Metal fast-fence path
(~6 wait/update pairs per iter across both ranks); unset, it uses SharedEvents.

Wedge signature: the log stops growing while the process stays alive.
Launch: MLX_HOSTFILE=<hostfile> MLX_RANK=<0|1> python fence_stress.py --log <file>
"""
import argparse
import os
import time

import mlx.core as mx

p = argparse.ArgumentParser()
p.add_argument("--iters", type=int, default=200000)
p.add_argument("--dim", type=int, default=5120)  # DeepSeek-V2.5 hidden size
p.add_argument("--log", required=True)
p.add_argument("--max-seconds", type=int, default=1800)
p.add_argument("--layers", type=int, default=1)  # GPU matmuls per iter (compute depth)
p.add_argument("--scalar-allsum", action="store_true")  # mimic 0035 patch: 1-elem all_sum/iter
p.add_argument("--rows", type=int, default=1)  # activation rows; >socket-buffer forces partial I/O
p.add_argument("--idle-after-barrier", type=int, default=0)  # secs; mimics mlx_lm.server's model-load ring idle
args = p.parse_args()

group = mx.distributed.init(backend="ring")
rank, size = group.rank(), group.size()
assert size == 2, f"need exactly 2 ranks, got {size}"

log = open(args.log, "a", buffering=1)
fs = os.environ.get("MLX_METAL_FAST_SYNCH", "<unset>")
log.write(f"START rank={rank} pid={os.getpid()} fastsynch={fs} "
          f"dim={args.dim} iters={args.iters} t={time.time():.3f}\n")

# upstream-recommended barrier after init (mlx#3149)
mx.eval(mx.distributed.all_sum(mx.ones(10), group=group, stream=mx.cpu))
log.write(f"BARRIER-OK rank={rank} t={time.time():.3f}\n")

if args.idle_after_barrier:
    log.write(f"IDLE {args.idle_after_barrier}s (ring quiet, like a model load) t={time.time():.3f}\n")
    time.sleep(args.idle_after_barrier)
    log.write(f"IDLE-OVER t={time.time():.3f}\n")

w = mx.random.normal((args.dim, args.dim)).astype(mx.float16)   # GPU weight
x = mx.random.normal((args.rows, args.dim)).astype(mx.float16)  # rows*dim*2 bytes per message
w2 = mx.random.normal((args.dim, args.dim)).astype(mx.float16)
mx.eval(w, x, w2)

def compute(h):
    for _ in range(args.layers):
        h = mx.tanh(h @ w)                                     # GPU
    return h

def step(xin):
    if rank == 1:
        h = compute(xin)                                       # GPU
        s = mx.distributed.send(h, 0, group=group)             # CPU stream
        h2 = mx.depends(h, s)                                  # GPU node <- CPU dep
    else:
        r = mx.distributed.recv_like(xin, 1, group=group)      # CPU stream
        h2 = compute(r)                                        # GPU <- CPU->GPU fence
    g = mx.distributed.all_gather(h2, group=group)             # GPU->CPU + CPU op
    out = mx.tanh(g[:1] @ w2)                                  # GPU <- CPU->GPU fence
    if args.scalar_allsum:
        # 0035-patch pattern: scalar all_sum per iteration (size<size_ ring path,
        # direction -1 => shares the p2p send socket-direction)
        stop = mx.distributed.all_sum(mx.array([0], dtype=mx.int32), group=group)
        out = mx.depends(out, stop)
    return out

t0 = time.time()
prev = None
for i in range(args.iters):
    out = step(x)
    mx.async_eval(out)          # double-buffer like generate.py GenerationBatch
    if prev is not None:
        mx.eval(prev)
    prev = out
    if i % 100 == 0:
        log.write(f"i={i} t={time.time():.3f}\n")
    if time.time() - t0 > args.max_seconds:
        log.write(f"TIME-CAP i={i} t={time.time():.3f}\n")
        break
if prev is not None:
    mx.eval(prev)
log.write(f"DONE rank={rank} iters_run={i + 1} elapsed={time.time() - t0:.1f}s t={time.time():.3f}\n")