[AlexCheema]

The bug

Outside prefill, every pipeline stage gathers the last stage's output from all ranks and keeps its own copy. PipelineLastLayer built the slice on top of the gather and evaluated the slice:

output = mx.distributed.all_gather(output, group=self.group)[-output.shape[0] :]
mx.eval(output)

The gather runs on the CPU, and the slice runs in a GPU command buffer that waits for it. So the GPU waits for the slowest rank. When that wait outlasts the GPU watchdog (about 5 s on Apple silicon), Metal fails the command buffer with Caused GPU Timeout Error, and the runner crashes. The send and the receive already avoid this: each is evaluated on its own, so its wait happens on the CPU, which has no watchdog (see the comment in PipelineFirstLayer). The gather was the one collective left inside a GPU command buffer.

Any pipeline instance whose later stage can fall more than a few seconds behind is exposed. An 816B-parameter MoE model (4-bit, 533 GB) on two Mac Studios (M3 Ultra, 512 GB), as 2x pipeline parallel over the ring, hit it two ways:

• exo's warmup. On the first decode step, the first stage waited in the gather while the second stage was still compiling its kernels (that machine's CPU was also busy with another job). The runner failed and the instance never became ready.
• Long chunks outside prefill. In the correctness harness (teacher forcing, is_prefill off so every rank gets logits), a 3,284-token chunk makes the first stage wait for the second stage's ~5 s of compute. It failed in 2 of 2 runs, with 4096- and with 2048-token chunks.

The fix

The gather is evaluated on its own first, then sliced:

gathered = mx.distributed.all_gather(output, group=self.group)
mx.eval(gathered)
output = gathered[-output.shape[0] :]
mx.eval(output)

Nothing else changes: the same arrays are computed, in the same order.

Evidence

On the same pair (Thunderbolt, 2x pipeline parallel over the ring), with the 816B model:

• Warmup: the instance comes up, with the other machine's CPU as busy as when it failed.
• Long chunks: the harness run that failed (teacher forcing with 4096-token chunks, after greedy decoding) passes: haystack prompts of 3.3k and 7.5k tokens.
• Serving: needle recall through /v1/chat/completions: 7/7 at 4.5k, 16k and 64k tokens, with no timeouts.

What it costs

Nothing measurable, and the outputs don't change. The same 2x pipeline-parallel runs over the ring were made with main's pipeline.py and with this change, alternating before, after, before, after. Greedy decoding was 128 tokens per prompt, plus teacher forcing over the prompts and continuations.

model	prompts	decode tok/s per prompt, before	after	greedy tokens and log-probs
the 816B MoE, 4-bit	4	10.0, 9.9, 10.0, 9.9 (both runs)	9.9, 9.9, 10.0, 9.9 / 10.0, 9.9, 10.0, 9.9	bit-identical: greedy tokens, their top-k log-probs, and teacher-forced log-probs at 1,256 positions
Qwen3.6-35B-A3B, 4-bit	5	67.7, 68.1, 67.9, 69.8, 71.1 / 67.4, 68.2, 68.0, 69.9, 71.3	67.5, 67.9, 68.0, 69.8, 71.0 (second run)	bit-identical (954 teacher-forced positions)

• The first "after" run of Qwen3.6 started with a cold kernel cache: 46.2 tok/s on its first prompt, then 65.8, 67.9, 69.9 and 70.9.
• A fifth prompt with the large model, a 109-token tool call, varied between 2 and 4.6 tok/s from run to run in both versions, so it is left out of the table. Its outputs are bit-identical too.

Tests

test_pipeline_gather.py:

• test_a_stage_evaluates_the_gather_before_slicing_it: a structural check with the gather and the send mocked. It records what mx.eval is given and requires the full gathered array before its slice. It fails before the fix: only the slice is evaluated.
• test_a_last_stage_slower_than_the_gpu_watchdog_does_not_fail_the_others (slow, Metal only): two ranks over the ring on one Mac. The last stage is 8 s late reaching the gather, longer than the GPU watchdog. Without the fix, the first rank fails with [METAL] Command buffer execution failed: Caused GPU Timeout Error; with it, both ranks finish with the right values.

Overlap with #2521

#2521 ("Pipeline parallelism: build a rank's decode graph while its input is still in flight") also changes sharding/pipeline.py. It moves the decode-step receive's wait into PipelineLastLayer (on the CPU, before any GPU work). Its PipelineLastLayer still slices the gather in a GPU command buffer, so it has this bug too. Its slow test (a sender 8 s late) doesn't cover it: the delayed rank there is the first, which hasn't built its graph yet. The two changes touch different lines of PipelineLastLayer, and whichever lands second rebases trivially.

Test plan

• uv run basedpyright (0 errors), uv run ruff check, nix fmt (no changes).
• pytest as CI runs it (pytest src -m "not slow" --import-mode=importlib) on a Mac Studio with MLX: 2504 passed, 3 skipped.
• test_pipeline_gather.py including the slow test: both pass with the fix and both fail without it (the structural one finds no evaluation of the gathered array; the slow one fails with the GPU timeout).
• The A/B runs above, and the large model's warmup, long-chunk and needle runs on two Mac Studios.

🤖 Generated with Claude Code

[AlexCheema]

Folded into #2510, which is now rebased on main and carries this change and its test, together with the same fix for task agreement, the prefix cache and warmup.