Repository navigation
Sync upstream ml-explore/mlx-swift (mlx v0.32.2, mlx-c ebc88f1) via real merge - #20
Merged
Merged
Conversation
…-explore#398) mlx_detail_compile and mlx_closure_apply both return int (0=success, 1=failure) but their return values were silently ignored. When an error fires inside a withError scope the MLX error handler stores the error in an ErrorBox instead of calling fatalError; execution then continues past the failed call, innerCall returns an empty result vector, and the single/two/three-array compile overloads crash with a Swift 'Index out of range' trap — bypassing withError entirely. Fix: capture both return values and early-return [] from innerCall on failure. The placeholder return from the compile overloads is never observed by the caller because withError throws before the value is used. Adds three regression tests covering the single-array, two-array, and [MLXArray]->[MLXArray] compile overloads.
* hold evalLock while computing tostring - tostring internally calls eval, which is not thread safe * remove show-sdk-version -- ci breakage
…coded GPU (ml-explore#422) Closes ml-explore#395. `_resolveGlobalDefaultDevice()` returned `.gpu` whenever no explicit `setDefault(...)` had been called, so on a CPU-only host (Linux with no Metal / no CUDA) `Device.defaultDevice()` and `StreamOrDevice.default` both returned a device that simply doesn't exist, causing every subsequent op routed to the default stream to fail. The deprecated `Device.init()` already shows the right pattern: call `mlx_get_default_device()` and let the C++ core pick the device its build supports. Use that path instead of the hard-coded `.gpu`.
* Support CUDA compilation with SPM * Enable CUDA by default on Linux * Add code generation capabilities in CudaBuild plugin * Use `tools/update-mlx.sh` to generate `cuda_jit_sources.h`
Port `mlx.optimizers.MultiOptimizer` to Swift: a meta-optimizer that routes each model parameter to the first sub-optimizer whose filter predicate matches its flattened key path, with the last optimizer acting as the unconditioned fallback. This allows using different optimizers or hyperparameters for different subsets of a model's weights. `update(model:gradients:)` partitions the flattened gradient tree by first-matching filter and delegates each subset to the corresponding child optimizer; `innerState()` surfaces the children's state. It conforms to the `Optimizer` protocol directly since it holds no per-parameter state of its own. Closes ml-explore#228 Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Add MLXArray and DType extensions * Simplify magic numbers * Return MLXArray instead of Double
Port `mlx.optimizers.Muon` (MomentUm Orthogonalized by Newton-Schulz): a momentum optimizer that orthogonalizes the update for >=2D parameters via a 5-step quintic Newton-Schulz iteration, scaling the learning rate by sqrt(max(1, rows/cols)). 0D/1D parameters fall back to a plain momentum/Nesterov update. Mirrors the reference (weight decay, Nesterov default on, ns_steps=5). Tests pin the 1D momentum/Nesterov math exactly and exercise the 2D orthogonalization path (shape preserved, finite, parameter moved). Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…rs) (ml-explore#433) Port the mlx.optimizers schedule functions to Swift: exponentialDecay, stepDecay, cosineDecay, linearSchedule, and joinSchedules. Each returns a (Int) -> Float mapping a step to a learning rate, composing with the mutable `learningRate` exposed by the optimizers (assign optimizer.learningRate = schedule(step) each step). mlx-swift previously had no LR schedules. Tests pin exact values against the reference docstring examples. Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- does not build on iOS, etc. - fix ml-explore#436
* Add FP8 conversion utilities.
- add an iOS build in CI to catch iOS-only build breaks - move the docs build into the same macOS runner
* Add global scale support to quantized layers Store optional globalScale on QuantizedLinear and QuantizedEmbedding so nvfp4 weights can preserve the scale needed by lower-level MLX quantize/dequantize operations. Forward the scale when creating or dequantizing weights, add a direct pre-quantized QuantizedEmbedding initializer to match QuantizedLinear, and guard global-scale execution paths on Metal because MLX does not support globalScale dequantization there. Add focused tests for the new layer state and parameter exposure without broadening the generic quantization API surface. * Keep nvfp4 global scale on the quantized path for Metal Since MLX encodes nvfp4 group scales pre-multiplied by (448 * 6) / globalScale, the global scale factors out of dequantization linearly. Exploit this by always running quantizedMM (supported on Metal without globalScale) and applying globalScale / (448 * 6) to the output, instead of dequantizing the full weight and falling back to a plain matmul. This keeps weights quantized end-to-end on all backends, removes the dequantizedWeight fallback from QuantizedLinear and QuantizedEmbedding, and applies the scale before the bias add so the bias is unaffected. Add Metal tests that validate the GPU quantized path against CPU dequantized(globalScale:) references for QuantizedLinear (matmul + bias) and QuantizedEmbedding (asLinear and index lookup).
…er:) (ml-explore#448) The initializer retains a box holding the finalizer closure and hands it to mlx as an opaque payload, but `finalizerTrampoline` read it back with `takeUnretainedValue`, which never consumes the `passRetained` +1. The local `free(ptr:)` that would have balanced it was never passed anywhere — only the trampoline is — so it was dead code and the box leaked. The box owns the closure, so the closure's captures leak with it. The initializer's own documented example captures an `IOSurface` and comments that it "implicitly releases it when it returns"; it did not, and the surface was pinned for the lifetime of the process. `testIOSurface` covers the path but only prints, so nothing caught it. `takeRetainedValue` consumes the +1 exactly once. mlx-c documents the dtor as "Callback for when the buffer is no longer needed" and calls it once per buffer, and the trampoline has exactly one caller, so a single consume is correct. Removes the dead `free(ptr:)`. Adds a differential regression test: with the bug the finalizer runs but its captures are never released (witness releases 0); with the fix they are (1). The test drains mlx's allocator cache first, since a freed buffer's dtor is not guaranteed to have run at scope exit, and uses lock-guarded counters because mlx may invoke the dtor from its own thread. Full suite: 534 tests, 0 failures. Co-authored-by: Fred Jura <digital@afya.ai>
…them (ml-explore#445) `tools/update-mlx.sh` hand-tracked two lists that mlx's own CMake derives, so each one silently went stale whenever the vendored mlx submodule gained a jit source. The CUDA `MLX_JIT_SOURCES` argument named 16 `device/` headers explicitly, where `mlx/backend/cuda/CMakeLists.txt` uses `file(GLOB)` over `device/*.h` and `device/*.cuh`. A bump that adds one regenerated a `cuda_jit_sources.h` missing its array, failing later as `use of undeclared identifier 'jit_source_*'`: mlx v0.31.2 needs 17 and v0.32.0 needs 22, against the 16 the frozen list produced. The Metal step named 46 `make` targets explicitly, where `mlx/backend/metal/CMakeLists.txt` declares them via `make_jit_source()`. A bump that adds one left `jit/includes.h` declaring a `mlx::core::metal::<name>()` that no generated source defined, i.e. an undefined symbol at link time on macOS and iOS: v0.31.2 declares 47 and v0.32.0 declares 48, against the same 46. Both lists are now derived from the vendored tree at run time. Against the currently pinned mlx (v0.31.1) the derived values are identical to the removed literals, and the regenerated `cuda_jit_sources.h` is byte-for-byte identical to the committed one, so this is a no-op at the current pin. Each derivation is guarded: with `set -e` an empty result would not fail on its own, and bare `make` with no targets would quietly build `all` instead.
* Reap encuda's child processes with waitpid instead of polling isRunning `waitUntilExitWorkaround()` polled `Process.isRunning` in a 50ms sleep loop, added because `waitUntilExit()` "never returns" on aarch64 Linux. It did not fix that hang -- `isRunning` reads the same Foundation bookkeeping that stalls, so the workaround only changed a blocked call into an infinite sleep loop. Observed during a parallel SwiftPM build of the CUDA backend for linux/arm64 (aarch64-unknown-linux-gnu, CUDA 13, sm_121): `nvcc` finishes and writes its output in full, then becomes a zombie that is never reaped, while `encuda-tool` spins in `hrtimer_nanosleep` holding a completed 15MB transpiled .cpp on disk. Invoking encuda directly on the same file with identical arguments completes in 2-3 seconds, so the work itself is fine -- the process simply never notices the child finished. It is load-dependent: `--jobs 2` transpiled all 96 .cu files without incident, while an unconstrained build (SwiftPM defaults to every core) deadlocked. That is consistent with pipe file descriptors leaking between concurrently spawned siblings -- the hung `encuda-tool` held both the read and the write end of the same pipes, so no reader could ever see EOF. `waitpid` answers the only question this code has: did the child exit, and with what status. It does not depend on Foundation's reaper, and it collects the zombie. On ECHILD -- Foundation's own reaper winning the race -- it falls back to `terminationStatus`, which is authoritative in exactly that case. Exit status is decoded without the C macros, which Swift does not import: the low 7 bits carry the terminating signal and the next 8 the exit code, an encoding shared by Linux and Darwin. Signals are reported as `128 + signal`, following the shell convention, so a killed compiler is still a non-zero failure rather than a silent success. Verified on macOS: exit codes 0, 7 and 42 round-trip exactly; SIGKILL and SIGTERM report 137 and 143; no reapable child remains afterwards. The aarch64 Linux deadlock itself has not been reproduced locally -- it needs a CUDA toolchain and a parallel build to trigger -- so this is a fix for a diagnosed root cause rather than one observed to clear the symptom. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01K1rzxACJqmxVqgRsa4xj3R * Don't read terminationStatus before Foundation has published it The waitpid change cleared the deadlock it was written for -- an unconstrained SwiftPM build of mlx's CUDA backend for aarch64 Linux (CUDA 13.2, sm_121, 16 jobs) now runs through the transpile step at full speed instead of stalling forever on a zombie nvcc. But it replaced the hang with a crash. `waitpid` and Foundation's own reaper thread race for every child. When Foundation wins, our waitpid returns ECHILD and the fallback read `terminationStatus` immediately. That is not safe: Foundation publishes the result in two steps, and reading the status inside the window before `isRunning` goes false traps the process. *** Program crashed: System trap at 0x0000ffff90b6fb50 *** Thread 0 "encuda-tool" crashed: 0 Process.terminationStatus.getter + 80 in libFoundation.so 1 Process.waitForExitStatus() + 107 in encuda-tool at Source/Encuda/encuda-utils.swift:73:24 2 Encuda.Compile.run() + 1883 in encuda-tool Roughly one transpile in 250 died this way, at a different .cu file each run. The build failure it produces is unusually opaque: the crash report goes to the transpiler's own output, so `swift build` exits 1 with no diagnostic at all, and the last thing in the log is an ordinary "Compiling scan.cu to scan.cpp" progress line. So wait for `isRunning` to clear before reading the status -- but on a deadline, since an unbounded wait is the hang this function exists to avoid. On expiry it reports 254 rather than trapping or spinning: the caller turns any non-zero status into a build error, which is diagnosable, and neither a trap nor a hang is. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RdegWrNm5PKBqeXdkgmn9Q --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
* nested initializers
Three argument lists in the CUDA build are assembled as long `+` chains that the
type checker struggles with. On a Swift 6.2 toolchain they exceed its budget
outright:
encuda-compile.swift:66:17: error: the compiler is unable to type-check this
expression in reasonable time; try breaking up the expression into distinct
sub-expressions
plugin.swift:172:13: error: the compiler is unable to type-check this
expression in reasonable time; ...
Appends into local arrays instead, which is behaviourally identical and is what
the diagnostic suggests. The plugin's two chains move into `compileArguments` and
`linkArguments` helpers to keep `createBuildCommands` readable, and adding a
further argument to either list no longer risks tipping the type checker over.
Note this does *not* make the package buildable on Swift 6.2: SwiftPM 6.2 ignores
non-Swift files produced by a build-tool plugin ("Only Swift is supported for
generated plugin source files at this time"), so every generated .cpp is dropped
and the CUDA symbols go undefined at link time. The 6.3 floor is load-bearing.
This is purely a type-check-cost and readability change.
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
…data race on the compiler cache (ml-explore#461) * add a deterministic repro for the compile/eval lock-order inversion CompiledFunction.call takes the per-instance lock (ml-explore#226) and then, in innerCall, the process-global evalLock (ml-explore#339), holding it across mlx_closure_apply. On a cold cache entry that runs compile_trace, which calls back into the Swift body being compiled; if that body calls another compiled function -- every MLXNN activation does, e.g. silu calls the file-scope compiledSilu -- the nested call takes a second per-instance lock while evalLock is held. thread H: L_cf(outer) -> evalLock -> wants L_cf(shared) thread W: L_cf(shared) -> wants evalLock The orders are inverted and the two threads deadlock permanently. CompileLockRepro owns its threads, uses fresh compile() results per round so the compiler cache is cold, signals from inside the traced body so W is ordered after H enters the trace, and holds the window open with a usleep in that body. A per-round deadline turns the hang into a hard failure. It is its own executable because a wedged round cannot be recovered and would take a shared test process down with it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * take evalLock outermost in CompiledFunction.call Establishes one global lock order, evalLock -> CompiledFunction.lock, so the inversion described in the previous commit cannot form. The reverse order is not avoidable on its own: the trace has to run under evalLock (the compiler cache is an unguarded global and the tracing flag is a plain static, not thread_local), the trace calls back into Swift, and the traced body calls other compiled functions -- every MLXNN activation is a file-scope CompiledFunction. So evalLock -> lock happens regardless. The only fix available is to make it the *only* order, which is what taking evalLock at the top of call(_:) does. evalLock is recursive, so the nested acquisition from inside a trace is free. innerCall already took evalLock before doing any real work and held it through mlx_closure_apply, so this widens the held region only by the state gathering and closure creation at the top of innerCall. lock becomes an NSRecursiveLock so that a compiled function whose body re-enters that same compiled function on one thread no longer self- deadlocks on a non-recursive NSLock. CompileLockRepro: 20/20 rounds pass; without this it wedges on round 1. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * skip the per-instance lock when a compiled function has no state ml-explore#226 added CompiledFunction.lock to guard the observed state. The _updateInternal swaps in innerCall are the only mutation of anything reachable from self -- every other stored property is let, and id is written once at init -- so when inputs and outputs are both empty the lock has nothing to guard. With the previous commit evalLock is held for the whole call, which already serializes the compiler cache and every other caller, so dropping the acquisition changes nothing observable. Optional on top of the previous commit; the deadlock is fixed without it. It removes a mutex round-trip from the hot path of the common case: every compiled function in MLXNN is stateless. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * check the evalLock -> CompiledFunction.lock ordering in debug builds A passing CompileLockRepro only shows that the race did not fire, so assert the ordering itself. evalLock is recursive, and a successful try() on a recursive lock cannot distinguish 'this thread owns it' from 'nobody owns it', so ownership has to be recorded explicitly. withEvalLock is now the single choke point for evalLock; in debug builds it keeps a per-thread depth (pthread-specific, so the bookkeeping itself needs no lock -- it runs while holding evalLock). CompiledFunction acquires its per-instance lock through withInstanceLock, which requires that depth to be non-zero. Moving the per-instance lock back outside evalLock -- the shape that deadlocks -- trips it on the first compiled call, with no race to lose. All of it compiles out of release builds. CompileLockOrderingTests covers: that the bookkeeping can report 'not held' and is per thread (otherwise a violation count of zero would prove nothing), that a stateful compiled function really reaches the checked acquisition, and the nested case where one compiled function is called from inside another's trace. Verified by inverting the order locally: the assertion fires immediately with 'CompiledFunction.lock was acquired without holding evalLock'. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * erase the compiler cache under evalLock in CompiledFunction.deinit `mlx_detail_compile_erase` is `compiler_cache().erase(fun_id)` on the process-global `mlx::core::detail::CompilerCache`, which has no synchronization of its own. `CompiledFunction.deinit` called it with no lock held, so it could run concurrently with another thread's `CompilerCache::find` -- `cache_[fun_id]`, which can itself insert and rehash, followed by `entries.push_back` -- and with the whole trace that fills in the entry `find` returned a reference to. Unlike the lock-order inversion fixed in the preceding commits, this is a data race rather than a deadlock, and it predates them. Every other use of that cache already happens under `evalLock`, so the erase takes it too. `evalLock` is recursive, so a deinit that fires on a thread already inside a compiled call or an `eval` -- a traced body dropping the last reference to some other compiled function, for instance -- re-enters for free, and a thread holding a `CompiledFunction.lock` holds `evalLock` already, so no new edge is added to the lock graph. A deinit on a thread with no MLX lock simply waits for the call in flight. Nothing in the library hands a `CompiledFunction` reference to a back-end worker thread, so this erase can never be what an `evalLock` holder is waiting for. The erase stays synchronous rather than deferred: `id` is the object's own address, so a later `CompiledFunction` can be allocated at the same address and receive the same id, and a deferred erase could delete that new function's entries. Two tests. The first is the interleaving itself: one thread is parked inside `compile_trace` holding `evalLock` (it signals a semaphore from inside the traced body), the other drops the last reference to a warmed compiled function at that instant, and the erase is required to complete only after the first thread has left the traced body -- which is what serialization against `find` means. The second runs a stream of erases against a stream of cold compiles, covering the insert-and-rehash window too, and checks the results are still correct. Both assert the debug counters so a pass cannot be vacuous.
StreamOrDevice.stream(_:) discarded the stream it was given and returned the default stream of the default device instead, so operations that passed an explicit stream silently ran somewhere else. Return the given stream and add a regression test.
* test Module parameters and compile
* update for mlx v0.32.0 mlx v0.32.2 - carry `Source/Cmlx/mlx.patch` for mlx's `Device::operator<` / `Stream::operator<`, which are not strict weak orderings and now run on every `eval()` via `std::set<Stream>`. libc++ aborts on them under the Xcode 27 SDK. Still broken upstream in v0.32.0 and on main, no issue filed. Applied by `tools/patch-mlx.sh` from `update-mlx.sh` and CI, since `git submodule update` restores a pristine checkout. - streams became thread affine in v0.31.2, which broke the multi-threaded tests. Adopt mlx's `new_thread_unsafe_stream` (added upstream for language bindings that lock themselves -- mlx-swift serializes on `evalLock`) via a small C shim, as mlx-c does not bind it. - FFT ops gained a `norm` parameter: add `FFTNorm` and thread it through, plus `fftfreq`/`rfftfreq`. - wrap the new reducing slice-updates in `ArrayAtIndices` and add `exportToDot`/`graphDescription`/`NodeNamer`. - fix `SEEK_END` in the in-memory IO stream -- it computed from the read position rather than the end of the buffer, which v0.32.0's `load_safetensors` size check exposed. - route `cpu/jit_compiler.cpp` through `mlx-conditional`: it shells out via `std::system()`, unavailable on iOS. - `gemv` is JIT compiled as of v0.32.0, so drop `gemv.metal` from the metallib in `fix-metal-includes.sh`. The jit source lists themselves no longer need touching -- ml-explore#445 derives them from the vendored tree, and the derived sets match what a hand-bump would have produced here (48 Metal targets, 22 CUDA sources). * disable cuda runs for now -- need clear_streams exposed in mlx-c --------- Co-authored-by: David Koski <dkoski@apple.com>
* Add countNonzero + docs * Add tests
Add distributed communication and sharded neural network layers
Port mlx.core.distributed and distributed layer primitives to Swift:
- Backend & Primitives:
- Enable the ring distributed backend (CPU socket-based communication).
- Add MLXDistributed.Group with rank/size/split and 8 communication
primitives (allSum, allGather, send, recv, etc.).
- Default distributed operations to the CPU stream.
- Make initialize and split throw on failure.
- Layers & Optimization:
- Add AllToShardedLinear, ShardedToAllLinear, and quantized counterparts
(subclassing Linear / QuantizedLinear for seamless model replacement).
- Implement shardLinear, shardInPlace (supporting predicate sharding),
sumGradients, averageGradients, and clipGradNormSharded.
- Ensure correct boundary checks and integer scaling for quantized layers.
- Testing & Infrastructure:
- Add a reusable multi-process test harness (DistributedHarness) for
spawning local ranks via loopback sockets.
- Port Python ring tests (collectives, dtypes, donation, send/recv).
- Port layer, predicate sharding, and gradient tests; enable in CI.
* add thin logging shim
* patch for Cuda builds - local apply of ml-explore/mlx#4480 + other shutdown items from ml-explore/mlx#4506 - fixes crash on exit for cuda builds when using cmake
* pool Streams, fix Device inheritance - see ml-explore/mlx#2118 - Streams are finite and seem to leak OS/Metal resources - implements a pool to reuse Streams -- withNewDefaultStream is fully usable - much more care is taken with GPU devices - using index != 0 now works in all cases - and inherits properly when creating new streams Co-authored-by: Nick Henderson <525247+nwh@users.noreply.github.com>
* replace and improve integration tests - previously we used a python script to generate swift code and check output against known values from mlx (python) - this is replaced with a new and more powerful script (claude rewrite) - much more integration test coverage and tools to examine that - fixes per problems found by new integration tests - linspace output type, add dtype override - linspace default output type is float32 (matches python) - convolve .same padding side matches python (padRight) - nanToNum default for +/- infinity honors optionals (mlx defaults) - tensorDot axes default 1 -> 2 (matches python) - MLXNN Pooling padWidths matches python - MLXNN Alibi slop function matches python - MLXNN GRU apply bias to match python - MLXNN Lion optimizer default beta matches python - MLXNN Adafactor moving average matches python - MLXNN Muon zero power newton shulz matches python * fix linspace disambiguation * split integration tests into their own test target -- we need to run them with TF32 disabled so we can guarantee consistent results on NAX vs non NAX hardware Co-authored-by: Nick Henderson <525247+nwh@users.noreply.github.com>
* pick up fix-custom-io from mlx-c - I/O callbacks now return number of bytes read/written to indicate failure or success
* fix MLXLogger - fix ml-explore#491 - OSLog.Logger.isEnabled(type:) is not actually available since macOS 11 (and friends) - this causes crashes on iOS/macOS before 26.4 - use different API
…lx-swift Real merge of upstream/main (b4f7684, "Fix attention mask for float16 (ml-explore#496)") instead of a snapshot copy, so upstream history is preserved and future syncs only have to deal with new commits. - Re-vendor mlx v0.32.2 (1f8e74e) and mlx-c ebc88f1 as plain directories and re-apply the SharpAI patch series on top (SSD expert streaming, TurboKV, fast loaders, mlx-c fast.* wrappers, quieter Metal compile logging). - Drop the old lazy per-thread Metal command encoder patch: upstream's Swift StreamPool creates streams with mlx_stream_new_thread_unsafe, so they already work across threads, and a second encoder for the same stream index would silently break ordering. - mlx-c: forward the Swift stream in mlx_fast_streamed_gather_mm (it used only the device, which crashes when the graph is evaluated on another thread) and make the prefetch cache's try_take check the stored size before copying. - Regenerate Source/Cmlx/mlx-generated with tools/update-mlx.sh; keep moe_stream.metal. - Swift: keep the fork's MLXFast bindings next to upstream's crossEntropy and the compile scoped-default-device fix. - Package.swift stays on swift-tools-version 5.12 (upstream's 6.3;(experimentalCGen) is only needed for the Linux CUDA plugin, which this fork does not ship). - Restore upstream fixes that an earlier sync (877d992) had silently reverted: Module.update dictionary keys (ml-explore#356), ellipsis indexing copy (ml-explore#358/ml-explore#363), init order (ml-explore#367), WiredMemory cancellation guard (ml-explore#369), release.yml (ml-explore#471). - Add SSDStreamingTests (streamedGatherMM, preadInto, preadIntoOffset, and evaluation on a different thread). - Document the layout and the merge procedure in UPSTREAM_MERGE_PLAN.md. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Commit 8540be8 copied the root of the mlx repository (docs, python bindings, C++ tests, benchmarks, CMake files, examples, ...) into Source/. None of it is part of any SwiftPM target; it only bloats the repository and the stray Examples/extensions/axpby/axpby.metal breaks builds that compile every .metal file under Source/. No file under Source/Cmlx, Source/MLX, Source/MLXNN, etc. is touched. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This was referenced Oct 5, 2026
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
20 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Proposed changes
Syncs the fork with upstream
ml-explore/mlx-swift(b4f7684, i.e. 0.32.3+) as a real merge (--no-ff, no cherry-picks), moving the vendoredmlxto v0.32.2 andmlx-ctoebc88f1and re-applying the SharpAI patch series on top. Replaces the bot branchsync/upstream-latest, which cannot be merged cleanly.AI disclosure: this PR was prepared with Claude Code (Claude Sonnet 5.5) and reviewed by a human before it was opened.
What changed
mlx_stream_new_thread_unsafe, logging shim add thin logging shim ml-explore/mlx-swift#484, Cuda manifest patch patch for Cuda builds ml-explore/mlx-swift#478, float16 mask fix Fix attention mask forfloat16ml-explore/mlx-swift#496, …).Source/Cmlx/mlx= upstream mlx v0.32.2 + carried patches (SSD streamer /moe_stream_op,moe_stream.metal, TurboQuantsdpa_vector.h,fast.{h,cpp}hooks, I/O fast loaders).Source/Cmlx/mlx-c= upstreamebc88f1+fast.{h,cpp}wrappers.Source/Cmlx/mlx-generatedregenerated with upstreamtools/update-mlx.sh;moe_stream.metalkept.Package.swiftstays onswift-tools-version: 5.12with the CUDA plugin stubbed (unchanged from the fork's current setup).UPSTREAM_MERGE_PLAN.md(carried patches + how to merge upstream safely).Things worth reviewing
877d992had reverted upstream add gh for the release build ml-explore/mlx-swift#356, Fix wired mem race condition ml-explore/mlx-swift#358, Fix dictionary key lookup in Module.update(parameters:) ml-explore/mlx-swift#363, fix #364 -- x[.ellipsis] returns self instead of a "copy" ml-explore/mlx-swift#367, break init order dependency ml-explore/mlx-swift#369 and Fix wired-memory ticket lifecycle during cancellation ml-explore/mlx-swift#471. They are back;git diff upstream/mainover the Swift sources now only shows the intentional fork changes (MLXFast.swiftbindings,Transforms+Compile.swiftscoped default device, tests,Package.swift).StreamPoolcreates every Swift stream withmlx_stream_new_thread_unsafe(global encoder map), so streams already work across threads. Keeping a second lazily-created encoder per thread would break ordering between threads.mlx_fast_streamed_gather_mmnow forwards the stream it was given (it only forwarded the device, so evaluating on another thread broke); the prefetch cachetry_takeonly returns an entry whose stored size equals the requested length.d7a5ec7): commit8540be8had dumped ~300 mlx repo-root files (docs, CI, tests, …) intoSource/.preadinLoad) are now stored per stream and re-thrown on a later eval of that stream; Swift streams are shared process-wide, so a stale error can surface in an unrelated test. Found inmlx-swift-lmtests that callsave(arrays:url:)on the same file their lazy arrays were loaded from (save_safetensorstruncates the file before evaluating). Fixed on themlx-swift-lmside (the tests nowevalfirst); no change needed here.Tests
xcodebuild test -scheme mlx-swift-Package -destination 'platform=macOS' -skipPackagePluginValidation -parallel-testing-enabled NO: 392 tests, 2 skipped, 0 failures.Tests/MLXTests/SSDStreamingTests.swift(against a real safetensors file):testStreamedGatherMMReadsExpertSlab,testStreamedGatherMMEvaluatesOnAnotherThread(fails without the stream-forwarding fix),testPreadIntoOverwritesWholeArray,testPreadIntoOffsetFillsSelectedSlots.mlx-swift-lmmain(+ a 12-line test fix) built against this branch: 676 XCTest (6 skipped) + 973 Swift Testing, all passing, same counts as against the previous forkmain.pre-commit.Known leftovers (not changed here)
mlx_fast_submit_prefetchhas nomlx_error;mlx_fast_set_prefetch_enabledhas no try/catch; the mlx-c PAPPS header/fd caches are unbounded; the custom mlx sources are wired into SwiftPM only, not the mlx CMake build.Checklist
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes🤖 Generated with Claude Code