Request: Support for concurrent inference of independent models from separate threads

[1-ashraful-islam]

Summary

I'm building an application that runs two independent ML models concurrently on Apple Silicon (both macOS and iOS). Each model processes different inputs and produces different outputs - no shared arrays or dependencies between them. Currently this crashes due to thread-safety issues in the Metal backend.

I've researched the existing issues (#2133, #2067, #2086) and PR #2104. I'd like to discuss whether this use case can be supported and what the best path forward might be.

Environment

• MLX version: 0.30.3
• Platform: macOS and iOS (Apple Silicon)
• API: C++ (direct MLX C++ API)
• Build: Built from source

Use Case

Thread A: model_a → inference(input_a) → output_a
Thread B: model_b → inference(input_b) → output_b
(concurrent execution, completely independent)

Both models are:

• Loaded once at startup via mx::import_function()
• Wrapped with mx::compile() for kernel fusion
• Called from their own dedicated threads

Current Behavior

Concurrent inference produces Metal assertion failures such as:

-[_MTLCommandBuffer addScheduledHandler:]: failed assertion
'Scheduled handler provided after commit call'

A command encoder is already encoding to this command buffer

I can provide more detailed stack traces if helpful.

What I've Tried

1. Dedicated streams per model using StreamContext

// Construction
dedicated_stream_ = mx::new_stream(mx::Device::gpu);

// Inference
mx::StreamContext ctx(dedicated_stream_);
auto outputs = compiled_model_({input_array});
mx::eval(outputs[0]);
mx::synchronize();

Result: Crashes. Two threads using StreamContext concurrently race on the global default_streams_ map in Scheduler::set_default_stream().

2. Default stream (no StreamContext)

Both threads use the default GPU stream.

Result: Crashes. Both threads race on the shared DeviceStream::buffer and DeviceStream::encoder fields in get_command_buffer() / commit_command_buffer().

3. Mutex serialization

This works but negates the benefit of running models in parallel.

Root Cause Analysis

I traced this to two issues:

A. StreamContext modifies global state

set_default_stream() writes to Scheduler::default_streams_, a non-thread-safe unordered_map. Concurrent StreamContext usage corrupts this or causes incorrect restore on destruction.

B. DeviceStream lacks synchronization

Even with unique stream indices, DeviceStream::buffer and encoder are accessed without locks in the eval path, causing races when operations interleave.

Potential Solutions

I'd appreciate guidance on which (if any) of these aligns with MLX's design:

Option 1: Thread-local default streams

Store per-thread stream overrides in thread-local storage. get_default_stream() checks TLS first, set_default_stream() writes to TLS. This makes StreamContext thread-safe without locks.

Option 2: Per-DeviceStream synchronization

Add a mutex to DeviceStream so different streams can execute in parallel while same-stream access is serialized. (I saw concerns about deadlock in PR #2104 discussion.)

Option 3: Explicit stream passing for compiled functions

Allow callers to pass a stream directly when invoking compiled functions, bypassing the default stream mechanism entirely.

Option 4: Document as unsupported

If this isn't a priority use case, documenting the limitation clearly would also help users plan accordingly.

Questions

1. Is concurrent multi-model inference something MLX aims to support?
2. Is PR Metal thread safety #2104 the intended fix? Should discussion continue there instead?
3. Are there workarounds I've missed that avoid serialization?
4. Would any of the proposed solutions be welcome as a contribution?

Related

• Thread safety: Ongoing issue for thread safety in MLX #2133 - Thread safety tracking issue
• [BUG] thread issues with evaluation #2067 - Thread issues with evaluation
• [BUG] c++ compile cache should be thread safe #2086 - Compiler cache thread safety
• Metal thread safety #2104 - Metal thread safety PR
• Discussion Question about running `mx.eval` in separate threads #1448 - mx.eval in separate threads

Thanks for MLX - the performance on Apple Silicon is excellent. Happy to provide more details or test proposed fixes.

[awni]

I think it will be a while before we can provide true thread safety. It's something I think we should prioritize more. But it will be a spectrum to get there (some cases may not be supported others may) and doing it well without taking performance hits will take some time.

On the other hand you can get effective multi-threading by using streams in MLX:

with mx.stream(stream_1):
  output_1 = first_model(input_1)

with mx.stream(stream_2):
  output_2 = first_model(input_2)

mx.eval(output_1, output_2)

Something like that will in theory run both models in parallel up to the capacity of the GPU.

The other option is to use multiple processes.

[1-ashraful-islam]

Thanks! I dug into the MLX source to understand this better.

I see that separate stream indices map to different DeviceStream objects with separate command buffers - so no race there. The issue is StreamContext modifying the global default_streams_ map concurrently, which can cause arrays to get assigned the wrong stream, leading both threads to the same DeviceStream.

More context on my setup: I'm using MediaPipe, where each model is a separate Calculator node. MediaPipe's executor schedules calculators on a thread pool based on input availability - I don't control when each runs. So my architecture looks like:

MediaPipe Graph:
  ┌─────────────────────────┐       ┌─────────────────────────┐
  │ Calculator A            │       │ Calculator B            │
  │ (Thread from pool)      │       │ (Thread from pool)      │
  │                         │       │                         │
  │ StreamContext(stream_1) │       │ StreamContext(stream_2) │
  │ model_a(input)          │       │ model_b(input)          │
  │ mx::eval()              │       │ mx::eval()              │
  └─────────────────────────┘       └─────────────────────────┘
           ↓                                   ↓
     Triggered by                        Triggered by
     upstream data                       upstream data
     (independent)                       (independent)

I can't batch into a single mx::eval(out_a, out_b) because the calculators don't know about each other and fire at different times based on their input streams.

Questions based on source analysis:

1. Would thread-local default streams work? If get/set_default_stream() used TLS instead of the global map, each thread's StreamContext would be isolated. Each calculator would reliably get its own stream, and since stream indices map to separate DeviceStream objects, they wouldn't race. Is there a reason this wouldn't be viable?

2. Minimal serialization workaround: If I serialize just the StreamContext + array creation (few microseconds), then let eval() run concurrently on different streams - would that work? Or does the eval path also read default_streams_?

3. Explicit stream for ImportedFunction: I see most ops accept StreamOrDevice, but ImportedFunction::operator() uses default_stream() internally (export.cpp:1015). Adding a stream parameter there would let me bypass StreamContext entirely. Would that be a reasonable addition?

Happy to prototype any of these if you think they're viable paths.

[zcbenz]

Thread safety means a lot of things, for example there are a few levels of thread safety that users may expect:

1. Multiple threads can share the same stream.
2. Each thread can only use its own stream.
   1. Streams in different threads can wait for each other.
   2. Array in different threads can not depend on each other, but you can safely eval it in any thread.
   3. Array can only be evaled in the thread that it was created.

To simply put, level-1 would be disastrous to implement without rethinking the design of command encoder: image 2 threads running eval_gpu at the same time, there are just so many race conditions.

Level 2-i is very tricky that, we have to solve race conditions in many corner cases and performance-sensitive code routes, and we are unlikely to get any kind of performance gain and much more likely to get higher overhead with slower inference.

Level 2-ii is the ideal thread safety in my opinion, it is practical to implement and maintain, and actually useful.

Level 2-iii would be easier to implement, we just need to add locks in proper places, the tricky part would be avoiding locks in certain performance critical functions though.

There are also devils in details that we only get to know when started to implement thread safely, for example the default stream problem mentioned in this issue, and I think making it thread local would be a good design — it is basically how CUDA backend handles current device.

Anyway I think this is something that we can gradually achieve with small steps, if anyone is trying to work on this: please create small PRs that tackles one problem at one time, do not create PRs that no one is going to review.

[1-ashraful-islam]

Level 2 - ii with thread local stream - sounds like would solve the problem we have. Each Mediapipe calculator running in its own thread just need to be able to run independent eval without clobbering each other.

appreciate any further pointers you may have for that approach.

[zcbenz]

I would try to achieve 2-ii with following steps (in separate PRs):

1. Move some of Device methods like end_encoding and add_temporary to DeviceStream - we want to make it thread-safe to get CommandEncoder, but not using it, so a simple design would be making Device thread-safe and moving the unsafe methods out.
2. (optionally) Merge DeviceStream into CommandEncoder - after previous step the extra abstraction feels unnecessary, and we don't have it in CUDA backend.
3. Make it thread-safe to get CommandEncoder - we probably want to have a metal::get_command_encoder function that does not go through Device, so we can simply store them in a thread local static variable.
4. (At this point we should have achieved 2-iii)
5. Make it thread-safe to create Event and Fence, and make sure their methods do not share global states (we don't want to use locks as the classes are performance critical).
6. Make suer eval do not share global state.

[zcbenz]

I plan to work on this soon, which I think would be a faster route to fix #3216 than going back and forth trying to locate whether it is a Metal bug or user code bug.