[aleroot]

Proposed Changes

Follow up on ml-explore/mlx-swift-lm#575, two new public functions in MLX:

public func loadArrays(urls: [URL], stream: StreamOrDevice = .cpu) throws -> [String: MLXArray]
public func loadArraysAndMetadata(urls: [URL], stream: StreamOrDevice = .cpu) throws
    -> (arrays: [String: MLXArray], metadata: [String: String])

They load a multi-file safetensors checkpoint (e.g. a sharded model) using several threads at once, instead of one file after another.

Why

The existing loadArrays(url:) is lazy: arrays only materialize when evaluated, and evaluating a whole checkpoint in one go visits tensors in dictionary order which means the underlying reads jump around the file at effectively random offsets, on a single thread. For large models this leaves a lot of disk bandwidth on the table.

This PR moves the technique from ml-explore/mlx-swift-lm#575 down into mlx-swift, following the review feedback there that a concurrent synchronous load function "should eventually land in mlx-swift". Doing it here means every consumer (mlx-swift-lm included) can share one implementation instead of re-deriving it.

How it works

1. Each file's safetensors header is parsed directly (cheap) to learn where every tensor lives on disk.
2. Tensors are grouped into contiguous, byte-balanced ranges in file-offset order, so each worker reads sequentially rather than seeking.
3. The ranges are evaluated from concurrent work items via DispatchQueue.concurrentPerform, overlapping header parsing, I/O, and buffer copies.
4. Results are merged back in file order, so duplicate keys behave exactly like calling the serial loader file-by-file (later file wins). Files whose header can't be parsed fall back to a whole-file load.

Measured results

M4 Pro (14 cores, 24 GB), NVMe with ~6.1 GB/s ceiling:

Scenario	Serial	Concurrent	Delta
18 GB checkpoint (Muse-Glimmer-30B-4bit), cold cache	4.31 s (4.5 GB/s)	3.47 s (5.6 GB/s)	−20%
10 GB checkpoint (Qwen3.5-9B-8bit), warm cache ×3	1.25 s	1.19 s	−5–8%
10 GB checkpoint, cache-evicted	~6.3 GB/s	~6.4 GB/s	parity (disk-bound)

Honest note: current mlx core already iterates safetensors in file order at the bottom of the call stack, so the serial path is close to the disk ceiling in the evicted case. The wins are concentrated where they matter most cold loads of big models, and warm loads where copies dominate.

[aleroot]

One remaining detail: mlx-swift is pinned to MLX v0.31.1, where Load::eval_cpu uses future.wait(). MLX main now uses future.get() with stream error propagation, so the next coordinated MLX/MLX-C bump should let this checkedEval barrier surface deferred reader errors without a Swift-side workaround.

[davidkoski]

OK, I am trying to merge #450 today. I won't cut tags yet -- I want to get some testing in first, but that should unlock that part of it.

[davidkoski]

Rebase on main.

[davidkoski]

Ugh, CI won't retrigger. Well, it was clean before the rebase and I will run it locally.

[davidkoski]

Look at my comments and see what you think -- I was reading through the mlx side and I think just the sorting by position will do anything.

I rebased to make sure it still builds & fixed the test that called a test method that no longer exists.

I am not sure about the test failure (sendableCompileCapturingMaterializedArray) -- it is something from the MaterializedArray. I can't reproduce the failure here but I will investigate.

[aleroot]

Look at my comments and see what you think -- I was reading through the mlx side and I think just the sorting by position will do anything.

I rebased to make sure it still builds & fixed the test that called a test method that no longer exists.

I am not sure about the test failure (sendableCompileCapturingMaterializedArray) -- it is something from the MaterializedArray. I can't reproduce the failure here but I will investigate.

Done, for the sendableCompileCapturingMaterializedArray I could not reproduce but I have added a test coverage for cold and cached calls with different inputs, checking that the captured scalar remains evaluated, graph-free, and unchanged. This may help eventually indentifying when it fails?

[aleroot]

Found the root cause: MLX’s compilation caches are per-thread, but Swift was only clearing the current thread’s cache. If a compiled function was destroyed on another thread, its old graph survived and could be reused when a new function received the same address-derived ID. Sealed-state rebinding had the same invalidation gap.
Fixed in 8d92b24 by tracking every cache the function uses and clearing all of them on destruction or rebinding. Added two deterministic regressions that fail before the fix and pass afterward.

[davidkoski]

Change looks good. Thanks for the detective work and fix on the compile cache!!