Core >= 0.31.2 thread-local streams break process-global Stream.gpu/.cpu — cross-thread eval aborts

[PowerBeef]

Summary

Once the vendored core moves past v0.31.1, Stream.gpu / Stream.cpu (and everything that resolves to them, i.e. StreamOrDevice.default on most call paths) abort at runtime when an eval happens on a different OS thread than the one that first touched MLX:

Fatal error: There is no Stream(gpu, 0) in current thread.

This is not a 0.32-only problem — it starts at core v0.31.2, so it affects any bump from today's pin, including #450 once it lands.

Root cause

Core made CommandEncoder thread-local in ml-explore/mlx#3348 (first released in v0.31.2), and default_stream() thread-local along with it: each OS thread that asks for a default stream now creates its own stream, whose command encoder is registered only in that thread's registry. gpu::eval / gpu::synchronize run inline on the calling thread and look the encoder up there (falling back to a global registry that only new_thread_unsafe_stream populates).

mlx-swift's Swift layer predates this model:

// Source/MLX/Stream.swift
public static let gpu = Stream(mlx_default_gpu_stream_new())
public static let cpu = Stream(mlx_default_cpu_stream_new())

These are process-global static lets, materialized once on whichever thread first touches MLX, and Device.defaultStream hands them to every op. Under Swift Concurrency (tasks and actors hop cooperative-pool threads; Swift Testing runs suites on worker threads) an eval of a graph bound to that stream from any other thread throws the error above, which the error handler turns into a process abort.

Reproduce (with a core ≥ 0.31.2 vendored)

Evaluate any array from two different threads:

let a = MLXArray(0 ..< 16).sum()
eval(a)                        // thread A — ok, creates Stream(gpu, 0) here
Thread.detachNewThread {
    let b = MLXArray(0 ..< 16).sum()
    eval(b)                    // thread B — fatal: no Stream(gpu, 0) in current thread
}

(Thread B's ops bind to the cached Stream.gpu, i.e. stream 0, whose encoder lives only in thread A's registry.)

Fix that worked for us

Core kept the old semantics available as new_thread_unsafe_stream(Device), which registers the encoder in the global registry — exactly the pre-0.31.2 behavior the Swift statics were built on (callers serialize per stream, which Swift users already must do). We patched the two mlx-c entry points backing the statics to return a process-wide globally-registered stream:

extern "C" mlx_stream mlx_default_gpu_stream_new(void) {
  try {
    static mlx::core::Stream s = mlx::core::new_thread_unsafe_stream(
        mlx::core::Device::DeviceType::gpu);
    return mlx_stream_new_(s);
  } ...

(same for the cpu variant). That restores any-thread eval for Stream.gpu/Stream.cpu/StreamOrDevice.default with no Swift-layer changes. Working version at PowerBeef/mlx-c branch imarello/core-0.32.1-compat (3df95f5), consumed by PowerBeef/mlx-swift branch imarello/core-0.32.1 (b0605dc) — validated on macOS against core v0.32.1 (image-generation workload + tests that eval from multiple threads).

An alternative, more idiomatic long-term shape might be making Stream.defaultStream resolution genuinely thread/task-aware in Swift, but that changes public semantics; the mlx-c patch above is the minimal compatibility restore.

Related

#446 (Package.swift jaccl excludes past 0.31.2), #450 (0.32.0 bump — will hit this at runtime), ml-explore/mlx#4350 (separate 0.32.1 JIT-path regression relevant to bumps).