pool Streams, fix Device inheritance - #472
Merged
Merged
Conversation
davidkoski
commented
Sep 6, 2026
| default: die(2, "unknown --device \(CommandLine.arguments[index + 1])") | ||
| } | ||
| } | ||
| Device.setDefault(device: device) |
Member
Author
There was a problem hiding this comment.
Use task scoped calls now
davidkoski
commented
Sep 6, 2026
| } | ||
|
|
||
| /// static CPU device | ||
| /// Current CPU device. |
Member
Author
There was a problem hiding this comment.
Previously there was no distinction between these. Now (in particular for GPU devices) the distinction is real and works. If the default device is gpu/7 the .gpu will return that. Previously it would always return gpu/0.
davidkoski
commented
Sep 6, 2026
Comment on lines
-108
to
-113
| // Ask the underlying MLX C++ core for its default device rather | ||
| // than hard-coding `.gpu`. On Apple platforms with Metal this | ||
| // still resolves to GPU; on a CPU-only host (Linux without | ||
| // CUDA / no Metal) it correctly resolves to CPU. Hard-coding GPU | ||
| // here meant `defaultDevice()` / `StreamOrDevice.default` | ||
| // returned an unavailable device on those hosts. |
davidkoski
commented
Sep 6, 2026
Comment on lines
-101
to
-102
| public static let gpu = Stream(newStreamThreadUnsafe(.gpu)) | ||
| public static let cpu = Stream(newStreamThreadUnsafe(.cpu)) |
Member
Author
There was a problem hiding this comment.
Previously these were static streams, even if you were overriding the default stream. This now resolves to the default stream of the appropriate type
- see ml-explore/mlx#2118 - Streams are finite and seem to leak OS/Metal resources - implements a pool to reuse Streams -- withNewDefaultStream is fully usable - much more care is taken with GPU devices - using index != 0 now works in all cases - and inherits properly when creating new streams
davidkoski
force-pushed
the
stream-redo
branch
from
September 6, 2026 18:01
d29a02e to
3951340
Compare
This was referenced Sep 8, 2026
nwh
self-requested a review
September 16, 2026 19:57
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.
Proposed changes
see [BUG] crash while using new_stream mlx#2118
Streams are finite and seem to leak OS/Metal resources
implements a pool to reuse Streams -- withNewDefaultStream is fully usable
much more care is taken with GPU devices
using index != 0 now works in all cases
and inherits properly when creating new streams
Device.setDefaultDevice is documented as to how it works and fails -- use the task scoped calls instead
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes