[michaeltrefry]

sc-12937: fix "There is no Stream(gpu, N) in current thread" at the root — share stream encoders across threads

Shortcut story sc-12937

Root cause (verified, corrects the story's hypothesis)

MLX registers each stream's CommandEncoder in a thread_local map (metal::get_command_encoders() in device.cpp, same shape for cpu in encoder.cpp) and gpu::eval() runs inline on the calling thread — so a stream is only usable from the OS thread that created it. Any graph node whose primitive stream was created on another thread throws There is no Stream(gpu, N) in current thread.

Rust makes cross-thread graphs unavoidable: Array is Send, libtest runs every #[test] on its own spawned thread, and pmetal moves work across tokio workers. This is not a 0.31.2→0.32.0 regression: the thread_local map + throw already existed at v0.31.2; the full suite just never ran to completion before sc-12940, so all 20 failures were latent. What 0.32.0 added is the escape hatch this PR uses: new_thread_unsafe_stream (upstream #3578) registers encoders in a global fallback map that get_command_encoder consults from any thread.

Relationship to sc-12787 (PR #17, merged while this was in flight)

PR #17's random.rs change (eagerly evaluating the global RNG state) fixed the RNG-mediated subset of this class — which happens to cover all 20 of the story's original failures (they all crossed threads via the global random key). Verified on pristine current main: the original nn::value_and_grad repro passes, but the general class still fails — a lazy graph built on one thread and evaluated on another dies with the exact story error (2 of the 4 new regression tests here are red on current main without this patch). The Rust-side RNG workaround stays as belt-and-braces; this PR fixes the class at the MLX level.

The fix

mlx-sys/patches/thread-shared-streams.patch — 4th fork patch in the per-file apply_patches.sh flow (sc-12746/sc-12780 conventions, required=true): gpu::new_stream and cpu::new_stream delegate to the global-map registration used by upstream's own new_thread_unsafe_stream. Every stream (including per-thread defaults) becomes visible from every thread.

• Sequential cross-thread eval (the failing pattern everywhere above) is safe.
• Concurrent eval on the same stream from two threads remains the caller's responsibility — identical contract to upstream's thread-unsafe API. Follow-ups for pmetal's eval-serialization discipline, the lock-free global-map read, and encoder lifetime are tracked in sc-12959.

Regression tests

mlx-tests/tests/test_cross_thread_stream.rs — 4 deterministic tests (the original 20 were test-order-dependent), serialized behind a static Mutex so plain parallel cargo test can't race MLX's non-thread-safe eval:

test	on current main (unpatched)	patched
test_eval_graph_built_on_other_thread	FAILED (Stream(gpu, 0))	ok
test_eval_graph_on_spawned_thread	FAILED (Stream(gpu, 2))	ok
test_global_random_state_across_threads	ok (covered by PR #17)	ok
test_evaluated_array_used_from_other_thread	ok (guard)	ok

Verification

• Original story repro (cargo test -p pmetal-mlx-rs --lib nn::value_and_grad -- --test-threads=1): was 3/5 failing on 932beb4 → 5/5 pass.
• Full suite cargo test --all --no-fail-fast -- --test-threads=1: 964 passed / 0 failed (pre-PR-sc-12787: wire up CI on GitHub-hosted macos-15 (validate.yml never executed) + fix everything it catches #17 baseline was 939 passed / 20 failed).
• cargo fmt --all -- --check clean; cargo clippy -- -D warnings clean.
• Patch verified to apply through a from-scratch FetchContent rebuild (fresh clone + all 4 patches via apply_patches.sh, idempotency guard intact).
• Adversarial review (fresh-context subagent) verdict SHIP after fixes: tests serialized against parallel-run SIGABRT, red-without-patch coverage added for the spawned-eval direction, stale build.rs comment fixed; residual design risks filed as sc-12959.

🤖 Generated with Claude Code