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
ModelContext is now Sendable and ModelContainer is deprecated
the model in ModelContext is now immutable -- suitable for inference, but not training/fine tuning
TrainableLanguageModel and TrainableModelContext provide mutable models
but require separate API (loadTrainable) to load

Breaking API Changes

BaseLanguageModel is not a Module any more -- not all LanguageModels are mutable
see TrainableLanguageModel if mutation is needed (past load)
UserInput.Image/Audio, MLXArray -> MaterializedArray
ModelContainer is deprecated (ModelContext is now Sendable) and contains a TrainableModelContext rather than a ModelContext (now immutable)
- only code that references the types directly will need to change

See ml-explore/mlx-swift-examples#488 for adoption of the new API.

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.

OK, MaterializedModule is no longer a Module and I think that removes a lot of API surface that would e confusing.

It does mean if you want to mutate the model after load you have to indicate that when you load it (loadTrainable vs load). There is a TrainableModelContext and a ModelContext to keep these separate.

Everything should work as it did with the exception of LoRA style fine tuning -- the fact that the model is immutable on load breaks that path. I will update the PR description to reflect the breaks.

davidkoski · 2026-07-08T19:47:45Z

-//
-// `Gemma4AssistantDraftModel` is deliberately non-Sendable (see the
-// design note at Gemma4Assistant.swift:153–155 — Embedding is a
-// reference type and cross-domain access must go through
-// `MTPDrafterContainer.perform`). `nonisolated(unsafe)` is appropriate
-// here because the suite is serialized — only one test runs at a time,
-// and the cache becomes read-only after first population.


This comment is obsolete, though the cache here looks like it probably needs at least a lock -- it isn't thread safe.

davidkoski · 2026-07-08T19:48:29Z

+        let warmInput = try await context.processor.prepare(
+            input: UserInput(prompt: benchmarkPrompt))
+        let warmStream = try generate(input: warmInput, parameters: parameters, context: context)


Use ModelContext instead of ModelContainer. The latter still works, but is deprecated.

davidkoski · 2026-07-08T19:49:22Z

 /// ```
-public final class EmbedderModelContainer: Sendable {
-    private let context: SerialAccessContainer<EmbedderModelContext>
+@available(*, deprecated, message: "use EmbedderModelContext instead")


We can start removing this. Everything should still work but it is no longer needed for Sendable conformance.

davidkoski · 2026-07-08T19:50:08Z

+    private var _context: EmbedderModelContext
+    private let lock = NSLock()
+    private var context: EmbedderModelContext {


Because the API allows mutation (of a value type) we do it under lock so the mutation itself is thread safe.

davidkoski · 2026-07-08T19:52:21Z

@@ -40,8 +40,8 @@ import MLX
 public struct SpeculativeDecodingConfig: Sendable {


I will file an issue for this, but I think this needs to be a reference type -- as it is written, if you have multiple ChatSessions with deferred loads on the draft model, each of them will load it into memory. It works but is probably not the intent (since the whole idea around it is being sensitive to memory pressure).

davidkoski · 2026-07-08T19:53:03Z

-        images: consuming [UserInput.Image],
-        videos: consuming [UserInput.Video],
-        audios: consuming [UserInput.Audio]
+        images: [UserInput.Image],
+        videos: [UserInput.Video],
+        audios: [UserInput.Audio]


These types are all Sendable now so this doesn't have to be consuming.

davidkoski · 2026-07-08T19:54:03Z

-                    // prepare the cache, if needed.  note:
-                    // this is using the LanguageModel (not Sendable) outside
-                    // the protective lock.  Assuming the weights are not
-                    // being mutated behind the scenes, this will obey the MLXArray
-                    // contract that they be evaluated if used across threads.
-                    // This is internal to the implementation and this technique
-                    // should not be used in calling code.
-                    //
-                    // The benefit is that callers can be running multiple
-                    // ChatSessions in parallel, as long as the instances
-                    // are distinct.  In particular the KVCache cannot
-                    // be shared and that is the lock that is held here.
-
-                    let model = await model.perform { context in
-                        SendableBox(context.model)
-                    }.consume()


This fancy dance is no longer needed since the model itself is now Sendable. I think this is one of the biggest gains -- you no longer need tricky code and long explanations about why you think it is safe.

davidkoski · 2026-07-08T19:57:18Z

-    return generate(
-        input: input, context: context, iterator: iterator,
-        didGenerate: didGenerate)
+    fatalError("not implemented")


This has been deprecated for a while. I think it could be ported, but ModelContext would have to accept (I think) MaterializedModule also. Easier to mark as unavailable, but if it becomes an issue then we look at bringing it back.

davidkoski · 2026-07-08T19:59:53Z

+    }
+}
+
+public typealias TrainableBaseLanguageModel = BaseLanguageModel & Module


BaseLanguageModel (LanguageModel) is no longer a Module. That looks fine except in the LoRA case, where it is a breaking change. We want this because we want to wrap the trainable (Module) based models with MaterializedModule to make them Sendable (for inference cases).

davidkoski · 2026-07-08T20:01:15Z

        case ciImage(CIImage)
        case url(URL)
-        case array(MLXArray)
+        case array(MaterializedArray)


Potential source break

davidkoski · 2026-07-09T17:18:48Z

@ronaldmannak @aleroot @manojmahapatra FYI -- I think this is ready to go, along with the equivalent ml-explore/mlx-swift#418

There are a couple of breaking changes -- we can't have both mutable and Sendable (which we did before -- it was fine unless you used it!)

manojmahapatra · 2026-07-10T05:44:22Z

+/// ) { progres in ... }
+/// ```
+@freestanding(expression)
+public macro huggingFaceLoadTrainabledModel(


Question: huggingFaceLoadTrainabledModel looks like this is a typo?

manojmahapatra · 2026-07-10T06:03:24Z

 ) -> ModelContext =
    #externalMacro(module: "MLXHuggingFaceMacros", type: "LoadContextMacro")

+/// Load a `TrainableModelContext` using default hub client and tokenizer loader with progress.


Question: Do we also need to add a no progressHandler overload of huggingFaceLoadTrainableModel?

For consistency we probably should, but let me see if I can make it work with a default on the progress handler (or maybe nullable). The macro side actually supplies a default if not present.

The macro side actually supplies a default if not present.

yea, I see that.

manojmahapatra · 2026-07-10T06:11:18Z

I think overall the changes look good to me. 👍

davidkoski · 2026-07-10T06:13:00Z

@ronaldmannak @aleroot @manojmahapatra FYI -- I think this is ready to go, along with the equivalent ml-explore/mlx-swift#418

There are a couple of breaking changes -- we can't have both mutable and Sendable (which we did before -- it was fine unless you used it!)

I was thinking about the breaking change for ModelContainer. I think I might be able to hold a mutable model (behind a lock) and callers can update() it. Then when done and you try to run inference it can switch it to a MaterializedModule. It can't switch back (you could reload it) but this would work for the examples -- they train or augment, then do inference.

I will give that a try in the morning. This API break is the only thing that gives me pause. I think the rest of it cleans up a lot of sensibility issues nicely.

manojmahapatra · 2026-07-10T06:25:11Z

I was thinking about the breaking change for ModelContainer. I think I might be able to hold a mutable model (behind a lock) and callers can update() it. Then when done and you try to run inference it can switch it to a MaterializedModule. It can't switch back (you could reload it) but this would work for the examples -- they train or augment, then do inference.

That sounds reasonable to me. one question - would the state machine be explicit here? In short, the container stays mutable until the first inference call, then transitions one way into a materialized state afterward?

davidkoski · 2026-07-10T06:49:03Z

That sounds reasonable to me. one question - would the state machine be explicit here? In short, the container stays mutable until the first inference call, then transitions one way into a materialized state afterward?

Yes, exactly.

Though looking at mlx-swift-examples, I am not sure it will work. It does the LoRAContainer augmentation in perform(). This "works" because it doesn't mutate the ModelContext, it mutates the Module itself.

In fact there are no calls to update() at all because it mutate the reference type.

Maybe better just to take the API hit and fix this correctly. 😭

manojmahapatra · 2026-07-13T05:22:54Z

 /// applies optional quantization, and
 /// updates the model with the weights.
 public func loadWeights(
-    modelDirectory: URL, model: BaseLanguageModel,


Question: I pulled the PR locally and it fails to build because ModelConversion.convert(...) still takes model as BaseLanguageModel, while the implementation uses Module APIs like parameters(), leafModules(), and update(...). Should this be changed to TrainableBaseLanguageModel / BaseLanguageModel & Module?

Yes, that sounds right -- it it curious that it was building for me. I will take a look today. Thanks!

Rebased on main & this issue is fixed.

manojmahapatra · 2026-07-13T05:33:11Z

@@ -6,6 +6,8 @@ import MLXNN

 /// Container for models that guarantees single threaded access.


Suggestion: The “guarantees single threaded access” wording feels misleading now. perform only reads _context under the lock; the closure itself runs after that access is complete, so multiple perform calls can execute concurrently. Should we update this doc comment to describe ModelContainer as a deprecated lightweight wrapper around ModelContext instead?

Yes, good idea. I can document how it used to work vs how it works now. In practice things like ChatSession actually used it this way (the current way) but had to do so very very carefully.

Updated the documentation on ModelContainer.

davidkoski · 2026-07-13T16:07:29Z

investigate ModelConversion.convert
update docs on ModelContainer

davidkoski · 2026-07-20T20:17:01Z

OK, I thought about ModelContainer and came up with a new approach that is mostly not source breaking. Details are in the class docs for it.

llm-tool from mlx-swift-examples requires two changes:

--- a/Tools/llm-tool/LLMTool.swift
+++ b/Tools/llm-tool/LLMTool.swift
@@ -196,7 +196,7 @@ struct GenerateArguments: ParsableArguments, Sendable {
     }
 
     func prepare(
-        _ context: inout ModelContext
+        _ context: inout TrainableModelContext
     ) {
         if let extraEosToken {
             context.configuration.extraEOSTokens.insert(extraEosToken)
@@ -204,7 +204,7 @@ struct GenerateArguments: ParsableArguments, Sendable {
     }
 
     func generate(
-        input: LMInput, context: ModelContext
+        input: LMInput, context: ModelContextProviding
     ) async throws -> (GenerateCompletionInfo, String) {
         var output = ""
         for await item in try MLXLMCommon.generate(

The LoRATrainingExample (app) requires none.

- 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

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

davidkoski mentioned this pull request Jul 8, 2026

MaterializedArray is a Sendable MLXArray ml-explore/mlx-swift#418

Open

4 tasks

davidkoski commented Jul 8, 2026

View reviewed changes

manojmahapatra reviewed Jul 10, 2026

View reviewed changes

aleroot approved these changes Jul 10, 2026

View reviewed changes

davidkoski force-pushed the materialized-array branch from c311dda to 78dec31 Compare July 10, 2026 17:05

davidkoski requested a review from angeloskath July 10, 2026 21:35

manojmahapatra reviewed Jul 13, 2026

View reviewed changes

davidkoski force-pushed the materialized-array branch from 78dec31 to 9c04a10 Compare July 14, 2026 20:23

This was referenced Jul 17, 2026

cb2-7 (chain 7/8): scheduler-aware loadModelContainer overloads (batching front door) PicoMLX/mlx-swift-lm#30

Open

[DO NOT MERGE] Continuous batching: integrated stack (whole-diff CI + review) PicoMLX/mlx-swift-lm#23

Draft

davidkoski force-pushed the materialized-array branch from 9c04a10 to c3ccd0c Compare July 20, 2026 20:13

This was referenced Jul 20, 2026

Minimal adoption of materialized-array from mlx-swift-lm ml-explore/mlx-swift-examples#488

Open

update for MaterializedArray ml-explore/mlx-swift-examples#489

Open

davidkoski added 4 commits July 23, 2026 09:47

Use MaterializedArray for Sendable conformance

fa4d70b

- 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

redo ModelContainer -- mostly do not break source compatibility

e1f1e5e

more Sendable

4ef365c

update docs, items for examples

ee189e3

davidkoski force-pushed the materialized-array branch from 2bedcf0 to ee189e3 Compare July 23, 2026 16:47

davidkoski mentioned this pull request Jul 31, 2026

update documentation and skills #489

Open

4 tasks

thechriswebb added the cat-misc label Aug 25, 2026

davidkoski mentioned this pull request Sep 1, 2026

Declare module weights as compile state in compiled decode paths #589

Merged

5 tasks

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

		@@ -40,8 +40,8 @@ import MLX
		public struct SpeculativeDecodingConfig: Sendable {

		@@ -6,6 +6,8 @@ import MLXNN

		/// Container for models that guarantees single threaded access.

Conversation

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

Uh oh!

Proposed changes

Breaking API 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!

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!

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!

davidkoski commented Jul 9, 2026

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!

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

Uh oh!

Uh oh!

davidkoski commented Jul 10, 2026

Uh oh!

manojmahapatra commented Jul 10, 2026

Uh oh!

davidkoski commented Jun 9, 2026 •

edited

Loading

aleroot Jun 10, 2026 •

edited

Loading

davidkoski Jun 10, 2026 •

edited

Loading

manojmahapatra commented Jul 10, 2026 •

edited

Loading

davidkoski commented Jul 13, 2026 •

edited

Loading