feat: sequence-parallel Ring Attention for MLX - #2255
Open
abendrothj wants to merge 6 commits into
Open
Conversation
Implement 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, UI selection, integration coverage, and reproducible benchmarks.
Ring decode re-implemented attention by hand and hard-required the plain KVCache, forcing --no-batch cluster-wide. Decode on a ring instance is replicated with no communication, so forward to the wrapped attention module instead: decode stays byte-for-byte identical to the non-ring path and inherits support for every cache type the model handles, including BatchKVCache. Verified against a live 2-node batching-enabled cluster: concurrent requests batch correctly and a 3570-token ring prefill produces output identical to the pipeline baseline at temperature 0.
loguru's diagnose mode calls repr() on every local in a traceback; repr of MLX CUDA objects can segfault mid-log, replacing the actual runner exception with an opaque signal 11. Observed masking a missing-CUDA_HOME error as a crash during heterogeneous ring validation.
MLX's CUDA backend JIT-compiles kernels with NVRTC at runtime (first triggered by the distributed send in ring prefill) and fails with 'Can not find locations of CUDA headers' unless CUDA_HOME/CUDA_PATH is set. The pip nvidia-cuda-runtime package ships the headers; resolve them from the nvidia namespace package in the runner bootstrap when nothing is configured.
…ll watchdog The intermittent long-sequence prefill hang was a cross-rank circular wait: with MLX_METAL_FAST_SYNCH, comm ops posted with unfinished GPU inputs block on Metal shared events inside a bounded command-buffer queue, and two ranks can each be waiting on GPU work that is queued behind an event only the other rank's transport would signal. Materialise KV payloads before posting sends and allocate receive templates on the CPU receive stream so transport never waits on the GPU and always drains. As containment for any residual stall mode, ring prefill now runs under a progress watchdog (EXO_PREFILL_STALL_TIMEOUT, default 300s, kicked per chunk) that dumps all thread stacks and exits the runner, converting an indefinite hang into a clean instance failure.
… set Workers now report dedicated accelerator memory via NVML and placement admits against min(RAM, VRAM); weight-only RAM admission let a 4GB-VRAM rank into 16K-context Ring placements it could never serve. Ring admission additionally requires an estimated prefill working set (KV cache at a 16K-capped context with a 4x multiplier derived from observed Metal/CUDA peaks) on every rank, since Ring replicates weights and exists specifically for long-context prefill.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
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.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.Sharding.Ringis gated onMlxRingtransport, ≥2 nodes, model-cardsupports_ring, and per-node admission of the full model size (weights are replicated; only prefill sequence is sharded).supports_ringdefaults off for everything else.MLX_METAL_FAST_SYNCH=1),exo-bench --sharding ringsupport.Benchmarks
Collected with
exo-benchagainst a live 2-node exo cluster running this branch (--warmup 1 --repeat 3, prefix caching disabled,MLX_METAL_FAST_SYNCH=1 uv run exoon 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.Ring (sequence-parallel prefill), 2 nodes
Pipeline baseline (same nodes, transport, model, params)
¹ 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 pipelineUpdates since original submission
Continuous batching (
--no-batchno longer required, fixed in0a0a6c4f): the incompatibility was decode-only —_decode_stepre-implemented attention by hand and hard-gated onisinstance(cache, KVCache). Fix: forward decode to the wrapped attention module (return self.original_layer(x, mask=mask, cache=cache)), inheriting support forBatchKVCacheand all other cache types. Verified on a live 2-node cluster: 3 concurrent requests batched and answered correctly; a 3570-token ring prefill produced output identical to pipeline baseline at temperature 0. Net diff: −40 lines of duplicated attention code, +2 regression tests.Real two-node validation — heterogeneous Metal + CUDA ring over a physical network:
MLX_METAL_FAST_SYNCH=1)mlx-cuda-120.32.0, driver 610.43.02), ~6 ms RTT over Tailscale LANResult: 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). The 16K same-host deadlock could not be probed on this hardware due to VRAM constraints.
Follow-up issues identified during CUDA validation (pre-existing, not regressions — all addressed in this branch):
CUDA_HOMEis auto-detected from bundlednvidia-cuda-runtime-cu12packages when present.diagnose=False) so CUDA errors aren't masked by a segfault in_format_value.Stress validation: 12/12 consecutive 8,192-token same-host Metal prefills clean after the above fixes; 543 tests passing including watchdog, CUDA environment, accelerator admission, and memory-filtering coverage.
Known issues / follow-ups
mx.evalandmx.async_eval. Did not reproduce at ≤8192 (9/9 runs clean) or in the distributed test suite. Needs investigation on genuine multi-node hardware with adequate VRAM headroom.Testing
uv run basedpyright— 0 errorsuv run ruff check— cleanuv run pytest— 543 passed, 3 skipped (includes 2-rank and 3-rank distributed Metal ring tests, BatchKVCache decode regression, watchdog, and admission tests)Refiled from #2217, which was closed unintentionally. Follow-up comments from that PR are incorporated above.