[CUDA] distributed all_sum barrier hangs in cu::AtomicEvent::wait on Blackwell (sm_121/sm_110) during warmup — ring sockets healthy

Summary

On a 5-node pipeline-parallel ring that mixes Apple Metal nodes with Blackwell CUDA nodes (NVIDIA GB10 sm_121 and Jetson Thor Tegra sm_110), the mx.distributed.all_sum-based barrier hangs deterministically during warmup. All ranks freeze; the CUDA ranks busy-poll forever in mlx::core::cu::AtomicEvent::wait on a CUDA event that is never signaled.

This is distinct from #3862: there the theory was a dead ring SocketThread after a transient TCP reset. Here the ring sockets are healthy — ESTABLISHED, Recv-Q=0, Send-Q=0, no data in flight, no reset. The hang is not in socket send/recv; it is a never-signaled CUDA completion event inside mx.eval.

Exact location (py-spy --native, identical on both CUDA nodes)

The barrier is exo's mx_barrier:

mx.eval(mx.distributed.all_sum(mx.array(1.0), group=group,
                               stream=mx.default_stream(mx.Device(mx.cpu))))

Stack of the hung MainThread (GB10 sm_121; Thor sm_110 identical):

clock_nanosleep (libc)
mlx::core::cu::event_wait (libmlx)
mlx::core::cu::AtomicEvent::wait (libmlx)
mlx::core::Event::wait (libmlx)
mlx::core::array::wait (libmlx)
mlx::core::eval (libmlx)
<python> mx_barrier(...)

Key observations

• Two barriers exist in the load path. The barrier called during weight loading (shard_and_load) passes; the barrier called during warmup (after the first forward pass, i.e. with pending GPU work) hangs. The hang correlates with GPU work being in flight when the barrier's mx.eval runs.
• Sockets are clean. ss -tnp on the hung process: all ring connections ESTABLISHED with empty queues.
• Threads: 58 threads, nearly all in futex_do_wait, one spinning (the nanosleep poll inside cu::AtomicEvent::wait). CPU ~200–300%.
• Flaky, not always fatal: on some runs the warmup barrier does complete, which suggests a race in CUDA event signaling rather than a hard deadlock.

Hypothesis

Because mx.array(1.0) is created on the default device (CUDA), the CPU-stream all_sum needs a device→host dependency, and the resulting completion CUDA event is not reliably signaled on Blackwell (sm_121/sm_110) when GPU work is already queued. cu::AtomicEvent::wait then polls forever.

What we tried (does not help — please don't suggest these)

Two application-level rewrites of the barrier both regressed (they broke the previously-working loading barrier), which points at the CUDA backend event mechanism rather than array placement:

1. Wrapping array creation + all_sum in with mx.stream(cpu):.
2. Temporarily mx.set_default_device(mx.cpu) around the barrier.

Reproduction

Pipeline-shard mlx-community/Qwen3-0.6B-4bit across 5 nodes (mix of Metal + Blackwell CUDA) via a ring backend; run a warmup forward. The warmup barrier hangs within seconds and stays frozen. A tiny model reproduces it, so it is not memory-related.

Environment

• mlx 0.32.0.dev20260717+7a1d4f5 (main is 5 commits ahead, none CUDA-event related), source-built.
• CUDA nodes: NVIDIA GB10 (sm_121, CUDA 13.0) and Jetson Thor Tegra Blackwell (sm_110, JetPack 7.1), each built with -DCMAKE_CUDA_ARCHITECTURES=121 / =110.
• Metal nodes: MacBook Pro M5 Max / M4 Max.

Related: #3862 (same "wedge in Event::wait" symptom, different — socket — root-cause framing).

Start with exo's mx_barrier entry point and the mx.eval call using mx.distributed.all_sum, then trace the CUDA-side wait shown at cu::AtomicEvent::wait. Reproduce the warmup barrier on the listed Blackwell configurations while comparing it with the loading barrier and healthy ring sockets. Done means the warmup barrier completes reliably without regressing the loading barrier.