update for mlx v0.32.2

[CharlieTLe]

Proposed changes

Updates the vendored mlx to v0.32.0 and mlx-c to fba4470 (its latest; those bindings were generated against 0.31.2, so 0.32.0's new ops aren't exposed yet). MLX_VERSION is bumped in Package.swift and xcode/xcconfig/Cmlx.xcconfig.

Adaptations the bump requires:

• FFT gained a norm parameter. Every mlx_fft_* signature changed, so this adds a public FFTNorm enum (.backward default, preserving current behaviour) and threads it through Source/MLX/FFT.swift. Also wraps the new fftfreq/rfftfreq.

• SEEK_END in the in-memory IO stream was wrong — it computed from the read position rather than the end of the buffer. Latent until v0.32.0's load_safetensors added a size check via seek(0, end) + tell(), which made saveToData/loadArrays round-trips fail. SaveTests covers it.

• cpu/jit_compiler.cpp now shells out via std::system(), which is unavailable on iOS. Routed through mlx-conditional the same way compiled.cpp already is, and excluded from the direct build so it isn't compiled twice.

• gemv is JIT compiled as of v0.32.0 (jit_kernels.cpp gained get_gemv_kernel using metal::gemv()), so gemv.metal no longer needs compiling into the metallib and comes out of fix-metal-includes.sh. The jit source lists themselves need no change — Derive the jit-source lists in update-mlx.sh instead of hand-listing them #445 derives them, and the derived sets match what a hand bump would produce (48 Metal targets, 22 CUDA sources).

• Package.swift excludes and xcodeproj target membership updated for added/removed upstream files (jaccl moved under lib/, nccl_stub gone, new cuda/dirs.cpp and qmm sources, new backend/common/metal_kernel.{h,cpp}), plus regenerated include/, include-framework/ and mlx-generated/.

New API now reachable from mlx-c fba4470:

• the reducing slice-updates, wired into ArrayAtIndices so pure-slice at[] updates use mlx_slice_update_add/prod/min/max instead of a scatter (mirroring upstream's mlx_add_item), falling back to scatter otherwise
• exportToDot / graphDescription / NodeNamer in a new Source/MLX/GraphUtils.swift

CMakeLists.txt declares mlx ahead of mlx-c so both build paths use the same mlx — mlx-c currently pins v0.31.2. That can be dropped once mlx-c pins 0.32.0 itself, e.g. via ml-explore/mlx-c#123.

Per review, the mlx.patch scaffolding and the local new_thread_unsafe_stream shim were both removed in 2d27240 — see the two threads below. tools/update-mlx.sh, README.md and MAINTENANCE.md are byte-identical to main again.

Known blockers

Two failures remain, both from upstream changes rather than from this change. Neither is fixed here.

1. The ordering abort. Device::operator< and Stream::operator< are not strict weak orderings, and 0.31.2 started running them on every eval() via std::set<Stream> in eval_impl. libc++ validates the comparator under the Xcode 27 SDK and aborts wherever cpu and gpu streams are both open, so the test run aborts rather than reporting failures. Reported as ml-explore/mlx#4083 and fixed by #4086, but that is on main and unreleased — v0.32.0 is still the newest tag. This needs an mlx release containing #4086.

2. Thread-affine streams. Streams became thread affine in 0.31.2 (#3348, #3281) and upstream treats it as intended (#3727 was closed pointing at a consumer-side fix in mlx-lm). Because Stream.gpu/.cpu are cached process-wide, MLXRandomTests/testRandomThreadsDifferent, testRandomThreadsSame and TransformTests/testVmapThreadSafety fail with There is no Stream(gpu, 0) in current thread. They pass on 0.31.1. Making mlx-swift's default streams resolve per thread would fix them and needs no new API, but the stream model wants its own design rather than being settled inside a version bump, so it's deferred to follow-up work — see the Stream.swift thread. It doesn't block this bump.

Verification

On macOS 27.0 (Apple Silicon), Xcode 27.0:

swift build	succeeds
xcodebuild -scheme mlx-swift-Package -destination 'generic/platform=iOS'	succeeds
xcode/MLX.xcodeproj (-scheme MLX, macOS)	succeeds
cmake .. -G Ninja && ninja	succeeds, libs and examples link, ./tutorial runs
pre-commit run --all-files	clean
swift test	aborts — blocker 1 above
scripts/verify-docs.sh	fails — see below

With the patch still applied locally, the full suite was 554 tests / 0 failures, so the two blockers above are the only outstanding failures.

verify-docs.sh fails in ExtractAPI, which can't find any C++ stdlib header (algorithm, cmath, iostream) while parsing the Cmlx target. This is not from this change: it fails identically with my new headers removed, and it failed the same way before the bump on unmodified metal-cpp. It looks like Xcode 27 parsing C++ headers in C mode. I couldn't get a clean baseline run to prove it — SwiftPM rejected the manifest in a baseline worktree under the docc plugin — so worth confirming on CI's Xcode.

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)

On tests: FFTTests (norm scaling per mode, round-trips, fftfreq/rfftfreq), GraphUtilsTests, slice-update cases in ArrayAtTests including a slice-vs-scatter agreement check, and a SaveTests regression for the SEEK_END fix. On docs: the new public API carries DocC comments; no article changes are needed now that the stream work is out.

[davidkoski]

There is an incoming 0.32.0 mlx-c in ml-explore/mlx-c#123. #445 should take care of the JIT issue.

[davidkoski]

We should not need this -- we should rely on mlx and mlx-c tags. Anything that needs special handling should be done statically via tools at the time when we make the new release.

In particular SwiftPM can't apply patches (or make any source mods) so this won't really work.

[CharlieTLe]

You're right, and the SwiftPM point settles it — removed in 2d27240. I'd wired the patch into the places this repo controls (update-mlx.sh, both workflows, the CMake PATCH_COMMAND, a README step), but a package consumer gets the submodule pristine and never runs any of it, so they'd build with the bug regardless. That's the path that actually matters and I hadn't covered it. There's no static option either: the patch touches headers inside the submodule, so there's nothing checked in for tools/ to bake in at release time. update-mlx.sh, README.md and MAINTENANCE.md are byte-identical to main again.

What that leaves: ml-explore/mlx#4086 (the ordering fix, from #4083) is on main but unreleased — v0.32.0 is still the newest tag and git tag --contains for it is empty. So on the Xcode 27 SDK, libc++ validates the comparator and eval() aborts out of std::set<Stream> in eval_impl. It isn't confined to a couple of tests; it fires wherever cpu and gpu streams are both open, so the run aborts rather than reporting failures. Builds are all fine — swift build, iOS, and the xcodeproj all succeed.

So this needs an mlx release containing #4086 before it can go green. Happy to hold it until then and re-pin, or leave it open as-is if you'd rather have the diff visible in the meantime.

Separately, I rebased onto #445 and dropped my hand-added jit-source entries — I checked the derived sets match what the hand bump produced (48 Metal targets, 22 CUDA sources), so the committed mlx-generated output is unchanged. One leftover I'd flag: CMakeLists.txt still declares mlx ahead of mlx-c to pin v0.32.0, because mlx-c currently pins v0.31.2 and the two build paths would otherwise disagree. That comes out once mlx-c pins 0.32.0 itself, e.g. via #123.

[davidkoski]

We have been pretty loose about streams until now. I think we may need to:

• manage a per task/thread stream, e.g. if you create a new task you get an implicit new stream unless you use some kind of withNewStream the of API
• because of [BUG] crash while using new_stream mlx#2118 I think we may need to manage a pool of streams -- we can't create and destroy them at will

[CharlieTLe]

Understood — reading this as direction rather than something to solve here. Leaving it out of this PR.

Recording where I got to, in case it's useful later. The minimal step looks like it needs no new API: mlx_default_gpu_stream_new() already calls default_stream(DeviceType::gpu), and default_stream_storage is static thread_local as of 0.31.2, so each call already returns the calling thread's stream. The only reason it breaks is that Stream.gpu/.cpu are static let, so the first thread to touch them wins. Making those (and StreamOrDevice.cpu/.gpu, and Device.defaultStream, which is a stored property) resolve per thread would fix the three failing tests without a shim or an mlx-c binding.

It wouldn't be a complete answer, which I think is your point: withNewDefaultStream installs one new_stream as a @TaskLocal, so a task that hops threads across an await still breaks, and it does nothing for #2118 — Stream(_ device:) still creates a stream per call. Pooling plus per-task streams is the real fix.

So for this PR the three thread tests fail (MLXRandomTests/testRandomThreadsDifferent, testRandomThreadsSame, TransformTests/testVmapThreadSafety — all pass on 0.31.1). I've left them failing rather than skipping them, so the behaviour change stays visible; happy to add XCTSkips pointing at this thread if you'd rather.

One correction to what I wrote earlier: I said this had to land with or before the bump. That was wrong — the suite can't be green anyway until mlx releases the operator< fix, so the two are independent.

[PowerBeef]

Sharing notes from a working bump of the vendored core to v0.32.1 on a fork, in case they save this PR a cycle when it moves forward (or forward to 0.32.1):

• JIT abort at 0.32.1 (does not affect 0.32.0): core 0.32.1's qmv() passes a 6th template argument (results_per_simdgroup, from Optimize large NVFP4 QMV on M5 Max mlx#3961) while the [[kernel]] qmv templates still declare 5 params — with nojit_kernels.cpp excluded (this repo's current configuration), the first quantized matvec aborts at runtime with a Metal front-end "too many template arguments" error. Reported upstream at [BUG] Metal JIT build: qmv template instantiation passes 6 args to 5-param kernel templates (runtime abort) mlx#4350. We switched Cmlx to the nojit path (exclude jit_kernels.cpp instead and resolve kernels from the prebuilt metallib), which sidesteps it.
• jaccl excludes: the minimal fix from Package.swift exclude list breaks when the mlx pin moves past v0.31.2 (jaccl restructured) #446 (- mesh/ring/utils.cpp, + mlx/mlx/distributed/jaccl/lib) is confirmed working at 0.32.x.
• mlx-c: the vendored mlx-c needs its compile-cache calls updated for the 0.32.x core API (mlx::core::detail::compile_cache() handle); mlx-c's latest tag (v0.6.0) predates it.
• tools/update-mlx.sh: the metal kernel make-list needs gemv and searchsorted added at 0.32.x or the regenerated sources fail to link.
• NAX kernels: with a deployment target below 26.2, the generated *_nax.cpp sources reference symbols the header gates out — we compiled the base build with a MLX_METAL_NO_NAX define and excluded the seven mlx-generated/*_nax.cpp files (a 26.2+ variant can re-enable them).
• FFT: the C API grew a norm parameter; the Swift FFT.swift call sites need MLX_FFT_NORM_BACKWARD threaded through (~20 sites; the fftshift/ifftshift helpers are unaffected).

Everything above is public at PowerBeef/mlx-swift branch imarello/core-0.32.1 (f2459b6) + PowerBeef/mlx-c branch imarello/core-0.32.1-compat — feel free to lift anything useful.

[davidkoski]

cmake builds will fail until mlx-c gets a tag

[davidkoski]

OK, not making it today. Next step, debug CUDA build:

terminate called after throwing an instance of 'std::runtime_error'
  what():  cudaStreamSynchronize(stream_) failed: driver shutting down

claude says:

This error means a CUDA API call (cudaStreamSynchronize) was attempted after the CUDA runtime (libcudart) or driver had already begun unloading/shutting down for the process.

It is typically an exit-time teardown issue rather than a failure of the GPU kernel itself.

This is running example1 and indeed it looks like a shutdown error:

  9 [ra] 0x0000729cae8bb384 <unknown> in libstdc++.so.6.0.33
 10 [ra] 0x00005a46c5ced323 mlx::core::check_cuda_error(char const*, cudaError) + 352 in example1
 11 [ra] 0x00005a46c582e09d mlx::core::cu::CommandEncoder::synchronize() + 72 in example1
 12 [ra] 0x00005a46c582c7e8 mlx::core::cu::CommandEncoder::~CommandEncoder() + 27 in example1
 13 [ra] 0x00005a46c583aa3c std::pair<int const, mlx::core::cu::CommandEncoder>::~pair() + 31 in example1
 14 [ra] 0x00005a46c583aa5b void std::destroy_at<std::pair<int const, mlx::core::cu::CommandEncoder> >(std::pair<int const, mlx::core::cu::CommandEncoder>*) + 27 in example1

[davidkoski]

Status: waiting for feedback from mlx folks. It may be a lifecycle issue:

let deviceCtx = mlx_device_new_type(MLX_GPU, 0)
let streamCtx = mlx_stream_new_thread_unsafe(deviceCtx)
mlx_stream_free(streamCtx)
mlx_device_free(deviceCtx)

is a minimal repro. The program crashes when the static std::unordered_map<int, CommandEncoder> encoders gets destroyed on exit.