fix(mlx): run all of a runner's GPU work on one stream (fixes tensor-parallel stalls with fast synch)

[AlexCheema]

Problem

Tensor parallel over RDMA (MlxJaccl) stalls under concurrent load with fast synch on, which is exo's default and what JACCL needs for low latency. Within minutes of a few clients sending requests to Qwen3.8-27B on two Mac minis, a runner stops making progress. The stuck-runner watchdog then kills it after 90 s, and every request in flight fails. With fast synch off, the same load runs clean.

Stack samples of both ranks during a stall:

• One rank:
  • the main thread waits in mlx::core::eval_impl;
  • its communication stream spins in mlx::core::Fence::wait for its own GPU;
  • its GPU is at 100%.
• The other rank: its communication stream waits in jaccl::MeshImpl::all_reduce for the first rank's contribution.
• Where (py-spy): both main threads are evaluating the same prefill chunk (mx.eval([c.state for c in prompt_cache]) in mlx-lm's generate_step).

Cause

With MLX_METAL_FAST_SYNCH, a GPU op that needs the result of an op on another GPU stream spins on the GPU until it is there. A tensor-parallel forward's GPU work also spins, waiting for each all_sum on the communication stream.

mlx-lm runs the model on its own generation stream. exo's own ops ran on the default stream: prefix cache copies and trims, snapshot restores, snapshots. So graphs went from the generation stream to the default stream and back into the next tensor-parallel eval. Two GPU streams then wait on each other inside an eval that is also waiting on the communication stream, and nothing moves.

Fix

Run all of a runner's GPU work on one stream:

• MlxBuilder makes mlx-lm's generation stream the runner thread's default GPU stream;
• exo's generation code uses that same stream instead of creating its own.

Collectives still run on the CPU (communication) stream, as before.

Evidence

Reproduction without exo (two Mac minis, M4 Pro 64 GB, Thunderbolt 5 RDMA, JACCL, MLX main at 0e3ff364). A 64-layer tensor-parallel prefill (matmul + all_sum per layer) whose caches are built by an op on the default stream from work on the generation stream:

	Fast synch on	Fast synch off
Two GPU streams (exo before this PR)	hangs within the first 9 steps, every run	300/300
One GPU stream (this PR)	500/500, also with a decode step in flight	—

Regression test: test_tensor_parallel_prefill_from_exo_built_caches_does_not_deadlock (marked slow, Metal only). It runs the same shape as two ranks on one Mac (ring backend). With this PR it passes 3/3 in about 7 s; without it, it hangs 3/3.

exo under load (Qwen3.8-27B tensor parallel over RDMA on 2 Mac minis, 6 concurrent clients, fast synch on):

	Requests OK	Failed	Stuck-runner restarts
Before (3 × 10 min)	337	36	6
This PR (6 × 10 min, two pairs of minis)	1,116	0	0
This PR (3 × 30 min)	1,614	0	0

With the old build, even sequential single requests stalled once in four rounds of the speed test below.

Speed. Single requests (731-token prompt, 256 tokens out, no prefix cache, medians of 20, two rounds, fast synch on):

	Before	This PR
Qwen3.8-27B tensor parallel over RDMA, generation tok/s	25.9	25.9, 25.9
Qwen3.8-27B tensor parallel over RDMA, prompt tok/s	195.0	195.0, 195.0
Llama-3.2-1B on one node, generation tok/s	253.7, 267.3	266.7, 267.0
Llama-3.2-1B on one node, prompt tok/s	2,867, 2,638	2,923, 2,933

(The other "before" Qwen round stalled.)

For reference, on this PR's build with fast synch off:

• single requests generate the tensor-parallel Qwen at 8.6–9.0 tok/s, against 25.9–26.0 with it on;
• under the 6-client load, 156 requests complete per 10 min (2 runs), against 186 with it on.

So fast synch is worth keeping on for JACCL, and with this PR it no longer stalls.

Test plan

• uv run basedpyright (0 errors), uv run ruff check, ruff format --check on the changed files
• uv run pytest src -m "not slow" (CI's command): 471 passed, 3 skipped
• New tests: exo's generation stream is mlx-lm's; the helper makes it the default; slow two-rank deadlock test (above)
• Hardware: reproduction, load A/B and speed above

🤖 Generated with Claude Code

[AlexCheema]

Follow-up: #2393 fixes the second cause of the tensor-parallel fast-synch stalls. A prompt eval that started while the previous decode step's collectives were still running could deadlock, mostly when another model's runner shared one of the Macs. With both PRs, RDMA tensor-parallel chaos testing (3 runs × 3 h, Qwen3.8-27B over 2 of 4 Mac minis, fast synch on, node crashes, freezes and partitions) had 0 stuck runners.