Metal: distributed collectives run on the CPU stream, and with the default fence each crossing costs ~180 µs (MLX_METAL_FAST_SYNCH: ~4 µs)

Summary

On the Metal backend AllReduce, AllGather, Send, Recv and ReduceScatter
have no eval_gpu and throw (mlx/backend/metal/distributed.cpp:17-18, v0.32.2),
while the CUDA backend implements them (mlx/backend/cuda/distributed.cu:12). On
Metal, then, every collective in a sharded forward becomes a CPU<->GPU stream
crossing, and eval synchronizes those with a Fence (mlx/transforms.cpp, the
fences[...].wait(...) / Fence{stream} paths).

With the default fence a crossing costs ~180 µs on my setup; with
MLX_METAL_FAST_SYNCH=1 it costs ~4 µs. I wanted to write the size of that gap
down somewhere, since the flag is easy to miss and isn't mentioned in the
distributed docs.

Single-process repro

M3 Ultra, macOS 26, mlx 0.32.2. Chain of 101 GPU blocks with one trivial op
between each; the two arms are interleaved and I take the min of 40, because
measuring them sequentially makes the difference very sensitive to background
load.

# run twice: once plain, once with MLX_METAL_FAST_SYNCH=1
import time, os, mlx.core as mx

N, REPS = 101, 40
x0 = mx.ones((1, 1, 4096), dtype=mx.bfloat16)
W = mx.random.normal((4096, 2048)).astype(mx.bfloat16)
V = mx.random.normal((2048, 4096)).astype(mx.bfloat16)
one = mx.array(1.0, dtype=mx.bfloat16)
mx.eval(x0, W, V, one)

def run(stream):
    x = x0
    for _ in range(N):
        x = mx.multiply(((x @ W) @ V).astype(mx.bfloat16), one, stream=stream)
    mx.eval(x)

def timed(stream):
    t0 = time.perf_counter()
    run(stream)
    return (time.perf_counter() - t0) * 1e6 / N

for _ in range(5):
    run(mx.gpu); run(mx.cpu)
G, C = [], []
for _ in range(REPS):
    G.append(timed(mx.gpu)); C.append(timed(mx.cpu))
g, c = min(G), min(C)
print(f"FAST_SYNCH={os.environ.get('MLX_METAL_FAST_SYNCH', '0')} "
      f"gpu {g:.1f} us/op, cpu {c:.1f} us/op, crossing {c - g:.1f} us")

Three runs of each:

Two processes, real collectives

2 ranks over a Thunderbolt-RDMA (jaccl) group, all_sum on [1,1,4096] bf16
interleaved with GPU compute, 101 chained reduces per step, 1000 steps:

An isolated all_sum of that shape is 20.4 µs, and 101 chained with no GPU work
between them are 8.2 µs each, so the transport isn't the limiting factor here.

Where the cost seems to go

A standalone Swift probe (no MLX), N=101, min of 60, µs/op:

So closing and reopening an encoder is close to free, and the event round trip is
~1.8 µs; the remaining ~127 µs behaves like the GPU draining and refilling because
encodeWaitForEvent sits mid-command-buffer — it grows with work in flight
(115.8 / 152.4 / 182.1 µs as the dispatch grows). The fast fence avoids that by
dispatching a one-thread kernel that spins on a shared-memory counter inside the
already-open encoder (mlx/backend/metal/fence.cpp, use_fast), whereas the
default path goes through CommandEncoder::signal_event / wait_event
(mlx/backend/metal/device.cpp:499-509), both of which call end_encoding()
first. Happy to post the probe source if it would be useful.

Questions

1. Would it be worth documenting MLX_METAL_FAST_SYNCH in the distributed docs?
   Someone benchmarking distributed MLX on Metal today will measure the default
   path without necessarily realising there is another one.
2. Is selecting the fast fence automatically on Metal 3.2+ something you would
   consider, or does #3142 (GPU stuck in the fence_wait loop with fast synch and
   the JACCL backend) keep it opt-in deliberately? I ran 1000 steps x 101
   collectives on a 2-rank group without a hang, but that is well short of the
   4-node case in that report, so I may simply not have hit it. One small
   thought, offered tentatively: a max-iteration bound in fence_wait would turn
   that failure into an error rather than a wedged GPU — though I realise it
   would not address the coherence issue itself.
3. Longer term, is a Metal eval_gpu for collectives something on the roadmap, so
   they can sit in the GPU timeline the way they do on CUDA?

Environment: Apple M3 Ultra (512 GB), macOS 26, mlx 0.32.2, jaccl backend over
Thunderbolt RDMA. Line numbers are against v0.32.2.

Start by reading mlx/backend/metal/distributed.cpp alongside mlx/backend/cuda/distributed.cu, then trace the fence paths in mlx/transforms.cpp, mlx/backend/metal/fence.cpp, and mlx/backend/metal/device.cpp:499-509. Review the distributed documentation and issue #3142 before choosing a scope. Done requires a maintainer-approved decision on documentation, fence selection, or GPU collectives, with corresponding tests or benchmarks.