Use MaterializedArray for Sendable conformance by davidkoski · Pull Request #335 · ml-explore/mlx-swift-lm

davidkoski · 2026-06-09T23:21:14Z

Proposed changes

adopt changes from MaterializedArray is a Sendable MLXArray mlx-swift#418
we don't need private box types -- the technique becomes general
it also opens up some potential for synchronous evaluation

Checklist

Put an x in the boxes that apply.

I have read the CONTRIBUTING document
I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
I have added tests that prove my fix is effective or that my feature works
I have updated the necessary documentation (if needed)

davidkoski · 2026-06-09T23:21:46Z

-    private let context: SerialAccessContainer<EmbedderModelContext>
+    private let context: EmbedderModelContext


We don't need this as the Context is now Sendable

davidkoski · 2026-06-09T23:22:02Z

-        get async {
-            await context.read { $0.configuration }
-        }
+        context.configuration


These become synchronous accessors.

davidkoski · 2026-06-09T23:22:27Z

+    public func perform<R>(
+        _ action: @Sendable (EmbedderModelContext) throws -> R
+    ) rethrows -> R {


And we can have a synchronous perform()

davidkoski · 2026-06-09T23:24:19Z

@@ -27,52 +29,52 @@ import MLXLMCommon
 /// }
 /// ```
 public final class EmbedderModelContainer: Sendable {


Is this type still needed? Maybe. EmbedderModelContext is a struct, so this gives reference semantics -- users share the instance. Potentially EmbedderModelContext could become a class and we remove this? It would need to be immutable to do so. Keeping it for now.

davidkoski · 2026-06-09T23:25:11Z

    }
 }
+
+extension MaterializedModule: EmbeddingModel, BaseLanguageModel where LayerType: EmbeddingModel {


This is how we make MaterializedModule usable as a EmbeddingModel

davidkoski · 2026-06-09T23:26:01Z

-    public var model: any EmbeddingModel
+    public var model: any EmbeddingModel & Sendable
    public var tokenizer: any Tokenizer
    public let pooling: Pooling


Potentially we want this to be var

davidkoski · 2026-06-09T23:26:55Z

 /// `Pooling` takes the sequence of hidden states from a transformer model and collapses them
 /// into a single vector using strategies like mean, max, or token selection.
-open class Pooling: Module {
+public struct Pooling: Sendable {


I don't see why this should be Module. It isn't attached to the model itself (so e.g. update() semantics are not needed). I changed it like this and everything built cleanly.

davidkoski · 2026-06-09T23:27:09Z

    ],
    dependencies: [
-        .package(url: "https://github.com/ml-explore/mlx-swift", .upToNextMinor(from: "0.31.4")),
+        .package(url: "https://github.com/ml-explore/mlx-swift", branch: "materialized-array"),


Pick up that branch for now.

aleroot · 2026-06-10T17:42:25Z

-        try await context.read {
-            try await action($0)
-        }
+        try await action(context)


This seems to drop the serialization guarantee previously provided. The old SerialAccessContainer.read held an async mutex for the full duration of the closure, including suspension points, so concurrent perform calls could not overlap on the same model/tokenizer/pooling context.

With this direct call, two tasks can now enter perform concurrently and use the same underlying context values at the same time, while the type-level documentation still says the container “guarantees single threaded access.” Even if MaterializedModule makes the wrapped model sendable, that does not by itself preserve the previous exclusive-access contract for all work reachable through the context.

This may not be a real issue, except a documentation kind of inconsistency

I think that is a good point -- it won't matter as much for the Embedders because they don't have state, but for example the KVCache is state. It has to have exclusive access or be some kind of copy-in/copy-out setup.

For example in ChatSession it had to play some games with private classes and knowledge of what was safe:

private func streamMap<R: Sendable>() { try await cache.update { cache in let model = await model.perform { context in SendableBox(context.model) }.consume()

This would get exclusive access to the KVCache and then "borrow" the model/weights -- it treated them as Sendable, but they were not represented that way in the type system. This had to be internal implementation because it required care to make sure it was thread safe.

My hope is that providing Sendable pieces would allow anyone to write this code and do it safely.

But your point about what is serial access vs not is important.

aleroot · 2026-06-10T18:06:29Z

    ) {
        self.configuration = configuration
-        self.model = model
+        self.model = MaterializedModule(model)


context.model is still exposed as any EmbeddingModel & Sendable; EmbeddingModel inherits BaseLanguageModel, which inherits the Module mutation APIs. This now stores a MaterializedModule, whose update/apply/train/freeze overrides trap with fatalError. Existing callers mutating the model through perform { try $0.model.update(...) } can still compile through the existential API but now crash at runtime.

Yes, I would prefer it to be typed as MaterializedModule and then these functions are unavailable. That would require either the ModelContext to be generic OR MaterializedModule to not be generic (probably subclasses).

I need to play around with that and see if one works better than the other. Just getting this to compile took some effort -- perhaps it can be done better. This change has to be done in the context of integrating it!

The fact that the Materialized variants violate Liskov substitution isn't great, but I think the benefits of documenting and obtaining Sendable in a way that is interoperable is probably worth it, but I would appreciate feedback from people who use it!

I should note that MaterializedModule has these methods marked as unavailable so the compile time warning would be available if it were typed like that.

A non-generic MaterializedModule won't work -- it has to be final for the Sendable to work on it. I can't make it a struct because you can't subclass a struct.

I will look at making ModelContext generic, but that will change things all over the place.

We may have to live with some limitations on MaterializedModule, but perhaps if it is encapsulated inside things like ModelContext it won't be an issue.

- adopt changes from ml-explore/mlx-swift#418 - we don't need private box types -- the technique becomes general - it also opens up some potential for synchronous evaluation

davidkoski commented Jun 9, 2026

View reviewed changes

aleroot reviewed Jun 10, 2026

View reviewed changes

Use MaterializedArray for Sendable conformance

8b5c7b7

- adopt changes from ml-explore/mlx-swift#418 - we don't need private box types -- the technique becomes general - it also opens up some potential for synchronous evaluation

davidkoski force-pushed the materialized-array branch from da4f7d5 to 8b5c7b7 Compare June 18, 2026 20:12

		private let context: SerialAccessContainer<EmbedderModelContext>
		private let context: EmbedderModelContext

Uh oh!

Conversation

davidkoski commented Jun 9, 2026

Proposed changes

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aleroot Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

davidkoski Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

aleroot Jun 10, 2026 •

edited

Loading

davidkoski Jun 10, 2026 •

edited

Loading