Make CommandEncoder thread local#3348
Merged
zcbenz merged 1 commit intoml-explore:mainfrom Apr 1, 2026
Merged
Conversation
zcbenz
force-pushed
the
encoder-refactor-4
branch
4 times, most recently
from
d27cbb7 to
a58624e
Compare
angeloskath
approved these changes
Apr 1, 2026
Member
angeloskath
left a comment
angeloskath
left a comment
There was a problem hiding this comment.
Awesome! And really nice broken up in several PRs the past few days. Kudos
zcbenz
force-pushed
the
encoder-refactor-4
branch
2 times, most recently
from
247c227 to
0aefd88
Compare
zcbenz
force-pushed
the
encoder-refactor-4
branch
from
0aefd88 to
8a40a06
Compare
Thump604
added a commit
to Thump604/mlx-lm
that referenced
this pull request
Apr 1, 2026
MLX ml-explore/mlx#3348 (Make CommandEncoder thread local) makes each thread maintain its own Metal command encoder. The module-level generation_stream created at import time belongs to the main thread's command encoder. When generation runs in a thread pool (e.g., via vllm-mlx BatchedEngine), mx.eval and mx.synchronize on that stream fail with: RuntimeError: There is no Stream(gpu, 0) in current thread. Fix: replace the module-level mx.Stream with a thread-local factory function. Each thread gets its own stream on first use. The public name `generation_stream` becomes a function returning the current thread's stream, keeping the call-site diff minimal. Backward compatible: on MLX versions without thread-local command encoders, per-thread streams work identically to a shared stream.
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Refs #3078, #3216.
~CommandEncoderinstead of~Scheduler: so we can clean up when each thread exits, instead of when process exits, and avoid issues about thread safety and destruction order.CommandEncoderin thread local storage and throw error when accessing a stream not created from current thread.