feat: sequence-parallel Ring Attention for MLX

[abendrothj]

Related to/Closes #39

Sequence-parallel Ring Attention for MLX

Implements distributed Ring Attention prefill for verified MLX Llama and Qwen3 architectures, with replicated decode caches, capability-gated placement, per-node memory admission, failure-safe cleanup, dashboard selection, integration coverage, and reproducible benchmarks.

Rescoped: this PR previously bundled four fixes that were found while validating ring on a heterogeneous cluster but do not depend on it. They are now separate: #2289 (test suite is broken on main), #2290 (loguru diagnose), #2291 (CUDA_HOME auto-detection). Accelerator-memory reporting was dropped entirely in favour of #2216 — see Memory admission below. This branch is now 33 files / +2521, ring only.

What's included

• RingAttentionLayer (src/exo/worker/engines/mlx/ring_attention.py): wraps only the attention module (never a full decoder block), splits the sequence across ranks, and rotates KV blocks around the ring while accumulating attention with a numerically stable online-softmax merge.
• Overlapped communication: dedicated CPU send/receive streams (MLX distributed send/recv are CPU ops; unified memory lets them consume Metal-produced KV without copies). Receives are posted before the current block's attention is scheduled, and the lazy attention recurrence is explicitly mx.async_eval-ed to create the computation window the transfer overlaps with. Transport peer (always the ring neighbor) is correctly separated from KV origin rank.
• Batching: decode on a ring instance is replicated with no communication, so it forwards to the wrapped attention module and inherits every cache type the model supports, including BatchKVCache. --no-batch is not required.
• Deadlock avoidance: KV payloads are materialised before sends are posted and receive templates are allocated on the CPU receive stream, so transport never waits on unfinished GPU work. This removes a cross-rank circular wait seen on long-sequence prefill under MLX_METAL_FAST_SYNCH, where comm ops posted with unfinished GPU inputs block on Metal shared events inside a bounded command-buffer queue.
• Stall containment: ring prefill runs under a progress watchdog (EXO_PREFILL_STALL_TIMEOUT, default 300s, kicked per chunk) that dumps all thread stacks and exits the runner, converting any residual hang into a clean instance failure rather than a stuck request.
• Placement: Sharding.Ring is gated on MlxRing transport, ≥2 nodes, model-card supports_ring, and per-node admission (see below).
• Model cards: ring support declared for verified Llama 3.x cards; supports_ring defaults off for everything else.
• Tests: unit/integration coverage including scheduling-order regression (communication posted before compute is waited on), online-merge equivalence vs full attention, and real 2-rank and 3-rank distributed Metal tests over the MLX ring backend.
• Docs/benchmarks: README instructions (MLX_METAL_FAST_SYNCH=1), exo-bench --sharding ring support.

Memory admission

Every ring rank replicates the full weights and must hold the long-context prefill working set, so admitting on weights alone over-admits: a 4 GB rank passed the old check for a model whose 16K prefill peaks well past 4 GB. estimate_ring_node_memory adds an estimated KV working set (16K-capped context, 4x multiplier derived from observed Metal/CUDA peaks) to the replicated model size.

This is a pure function of the model card and uses whatever memory figures the node already reports. An earlier revision of this branch also added NVML-based VRAM reporting, which overlaps #2216 (filed before this branch, using nvidia-smi). That half has been dropped — #2216 should own accelerator-memory reporting, and this gate picks up VRAM-aware numbers automatically once it lands, with no change here.

Benchmarks

Collected with exo-bench against a live 2-node exo cluster running this branch (--warmup 1 --repeat 3, prefix caching disabled, MLX_METAL_FAST_SYNCH=1 uv run exo on both nodes).

Environment: both nodes on a single MacBook Pro (M4 Pro, 14-core, 24 GB, macOS 26.5.2, MLX 0.32.0.dev20260709). Model: mlx-community/Llama-3.2-1B-Instruct-4bit.

⚠️ Because both ranks share one GPU, these numbers demonstrate functional correctness and stability, not multi-device scaling. Multi-machine numbers will differ qualitatively; treat this as a smoke-test baseline, and prompt_tps as the primary ring prefill metric.

Ring (sequence-parallel prefill), 2 nodes

pp	tg	prompt_tps (3 runs)	gen_tps (mean)	peak mem/node
2048	64	2137 / 2114 / 2069 (mean 2114)	125	1.12 GB
4096	64	1857 / 1849 / 1712 (mean 1806)	104	1.55 GB
8192	64	1283 / 1281 / 1212 (mean 1258)	58	2.36 GB

Pipeline baseline (same nodes, transport, model, params)

pp	tg	prompt_tps (3 runs)	gen_tps (mean)	peak mem/node
2048	64	240¹ / 2012 / 1993 (median 1993)	165	1.24 GB
4096	64	1998 / 2005 / 1993 (median 1998)	159	1.56 GB
8192	64	1990 / 2008 / 719² (median 1990)	138	1.88 GB

¹ cold first iteration (kernel compilation) · ² transient slowdown, likely thermal/contention

Reproduce with:

uv run bench/exo_bench.py --model Llama-3.2-1B-Instruct-4bit \
  --pp 2048,4096,8192 --tg 64,64,64 --min-nodes 2 --max-nodes 2 \
  --instance-meta ring --sharding ring --warmup 1 --repeat 3
# baseline: same command with --sharding pipeline

Real two-node validation — heterogeneous Metal + CUDA over a physical network

• Node A: MacBook Pro M4 Pro, 24 GB (Metal, MLX_METAL_FAST_SYNCH=1)
• Node B: Debian 13 LXC, RTX 3050 Laptop 4 GB VRAM (CUDA, mlx-cuda-12 0.32.0, driver 610.43.02), ~6 ms RTT over Tailscale LAN

Result: 20/20 consecutive 4096-token ring prefills, zero hangs, identical deterministic output, ~468 prompt tok/s (428–508), ~42 tok/s decode. 8K/16K OOM the 4 GB rank cleanly (surfaced as HTTP 500, no wedged runners). Concurrent requests batch correctly, and a 3570-token ring prefill produced output identical to the pipeline baseline at temperature 0. The 16K same-host deadlock could not be probed on this hardware due to VRAM constraints.

Note that reproducing the CUDA half of this setup currently requires #2290 and #2291.