Add concurrent synchronous safetensors loading across files

[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.

[davidkoski]

Suggestion from claude reviewer that I thought was reasonable: maybe use checkedEval here.

The loadArraysAndMetadata() will handle I/O errors opening the file, but since this is lazy I/O under this, might be worth doing the equivalent here so we can catch the error and report it.

[davidkoski]

And one thing I thought of: eval() has to take a lock, so I wonder if we are serializing too much here? Are we getting the full concurrent read that we want?

I wonder if asyncEval() would be better? Or async in pass 1 and sync in pass 2? If this is running on one of the scheduler queues in the back end, I wonder if we are limiting ourselves there?

I will poke around a little bit so I can better understand where this runs.

[aleroot]

I traced the complete path and you’re right:
synchronous eval holds the global evalLock across the blocking mlx_eval, so the concurrentPerform iterations serialise when they reach evaluation. The parallel reads we measured are coming from MLX’s internal I/O pools, not simultaneous Swift eval calls.
This is reinforced by my other recently merged PR: mlx#4408, which gives all ParallelFileReaders one shared, adaptive batch pool specifically so SSD queue depth is controlled across shards.
I’ll change this to use an explicit file-ordered [MLXArray], schedule the ranges with asyncEval, then perform one checkedEval barrier before returning. That should preserves the synchronous API while keeping evalLock out of the I/O wait.

[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.