chore: upgrade to mlx 0.32

.github/workflows/emlx.yml

@@ -80,19 +80,26 @@ jobs:
           mix test --warnings-as-errors
 
   macos:
-    name: macOS ${{ matrix.job.gpu && 'gpu' || 'cpu' }} (${{ matrix.job.elixir }}, ${{ matrix.job.otp }})
-    runs-on: macos-26
+    name: ${{ matrix.job.os }} ${{ matrix.job.gpu && 'gpu' || 'cpu' }}${{ matrix.job.compat && ' compat' || '' }} (${{ matrix.job.elixir }}, ${{ matrix.job.otp }})
+    runs-on: ${{ matrix.job.os }}
     strategy:
       fail-fast: false
       matrix:
         job:
           # Format only on Elixir 1.20 / OTP 28 so CI matches local formatter output
-          - { otp: "28.5.0.1", elixir: "1.20.1", lint: true }
-          - { otp: "25.3.2.15", elixir: "1.17.0", build: true }
-          - { otp: "25.3.2.15", elixir: "1.17.0", gpu: true }
-          - { otp: "25.3.2.15", elixir: "1.17.0", gpu: true, build: true }
+          - { os: macos-26, otp: "28.5.0.1", elixir: "1.20.1", lint: true }
+          - { os: macos-26, otp: "25.3.2.15", elixir: "1.17.0", build: true }
+          - { os: macos-26, otp: "25.3.2.15", elixir: "1.17.0", gpu: true }
+          - { os: macos-26, otp: "25.3.2.15", elixir: "1.17.0", gpu: true, build: true }
+          # #127: portable 14.0 prebuilt (LIBMLX_MACOS_COMPAT) on Sequoia
+          - { os: macos-15, otp: "28.5.0.1", elixir: "1.20.1", compat: true }
+          - { os: macos-15, otp: "28.5.0.1", elixir: "1.20.1", compat: true, gpu: true }
     env:
       MIX_ENV: test
+      # Job-level: mix.exs resolves libmlx while loading project/0 (including
+      # `mix local.hex` in Setup Mix). Per-step exports below are too late for
+      # macos-15 hosts that must use the 14.0 archive (#127).
+      LIBMLX_MACOS_COMPAT: ${{ matrix.job.compat && 'true' || '' }}
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -125,7 +132,7 @@ jobs:
         id: mix-cache # id to use in retrieve action
         with:
           path: ${{ github.workspace }}/emlx/deps
-          key: ${{ runner.os }}-Elixir-v${{ matrix.job.elixir }}-OTP-${{ matrix.job.otp }}-${{ hashFiles(format('{0}/emlx/mix.lock', github.workspace)) }}-v1
+          key: ${{ matrix.job.os }}-Elixir-v${{ matrix.job.elixir }}-OTP-${{ matrix.job.otp }}-${{ hashFiles(format('{0}/emlx/mix.lock', github.workspace)) }}-v1
 
       - name: Install dependencies
         if: ${{ steps.mix-cache.outputs.cache-hit != 'true' }}
@@ -141,6 +148,9 @@ jobs:
           if [ "${{ matrix.job.build }}" = "true" ]; then
             export LIBMLX_BUILD=true
           fi
+          if [ "${{ matrix.job.compat }}" = "true" ]; then
+            export LIBMLX_MACOS_COMPAT=true
+          fi
           mix compile --warnings-as-errors
 
       - name: Check formatting
@@ -151,6 +161,9 @@ jobs:
           if [ "${{ matrix.job.build }}" = "true" ]; then
             export LIBMLX_BUILD=true
           fi
+          if [ "${{ matrix.job.compat }}" = "true" ]; then
+            export LIBMLX_MACOS_COMPAT=true
+          fi
           mix format --check-formatted
 
       - name: Run epmd for distributed tests
@@ -167,6 +180,9 @@ jobs:
           if [ "${{ matrix.job.build }}" = "true" ]; then
             export LIBMLX_BUILD=true
           fi
+          if [ "${{ matrix.job.compat }}" = "true" ]; then
+            export LIBMLX_MACOS_COMPAT=true
+          fi
 
           if [ "${{ matrix.job.gpu }}" = "true" ]; then
             find test -name "*_test.exs" -exec bash -c 'NAME={}; echo -e "\n\n----$NAME----\n"; mix test $NAME' \;

emlx/Makefile

@@ -47,6 +47,12 @@ ifeq ($(UNAME_S),Darwin)
     MLX_SO ?= $(MLX_LIB_DIR)/libmlx.dylib
     LDFLAGS += -undefined dynamic_lookup -flat_namespace -rpath @loader_path/mlx/lib
     LDFLAGS += -framework Metal -framework Foundation -framework Accelerate
+    # Match the precompiled libmlx deployment target so ld does not warn about
+    # linking a newer minos dylib (mlx-build Darwin variants).
+    ifneq ($(MACOSX_DEPLOYMENT_TARGET),)
+      CFLAGS += -mmacosx-version-min=$(MACOSX_DEPLOYMENT_TARGET)
+      LDFLAGS += -mmacosx-version-min=$(MACOSX_DEPLOYMENT_TARGET)
+    endif
 		MAKE_DEFAULT_JOBS = $(shell sysctl -n hw.ncpu)
 else
     MLX_SO ?= $(MLX_LIB_DIR)/libmlx.so

emlx/README.md

@@ -81,6 +81,25 @@ The binaries are always downloaded to match the current configuration, so you sh
 
 The version of the MLX binary to download. By default EMLX will always use the latest version possible.
 
+##### `LIBMLX_MACOS_COMPAT`
+
+Defaults to `false`.
+
+On Apple Silicon macOS, precompiled libmlx archives are keyed by deployment target ([mlx-build](https://github.com/cocoa-xu/mlx-build#macos)):
+
+| deployment target | runs on     | AOT NAX kernels |
+|-------------------|-------------|-----------------|
+| `26.2` (default)  | macOS 26.2+ | yes             |
+| `14.0`            | macOS 14+   | no              |
+
+The reduced featureset is only the missing **ahead-of-time NAX kernels** — MLX's fast GEMM/attention paths on Apple's `MetalPerformancePrimitives` (Metal 4). The `14.0` archive still provides normal Metal GPU ops. On macOS 26.2+, `LIBMLX_ENABLE_JIT=true` can still JIT-compile NAX at runtime even with the `14.0` archive.
+
+Set `LIBMLX_MACOS_COMPAT=true` to download the `14.0` archive. Use this on macOS 15 (and any host older than 26.2); the default `26.2` build will refuse to compile there instead of crashing the VM at runtime.
+
+##### `LIBMLX_DEPLOYMENT_TARGET`
+
+Optional explicit override of the macOS deployment target segment. Accepted values: `26.2`, `14.0`. When set, this wins over `LIBMLX_MACOS_COMPAT`.
+
 ##### `LIBMLX_ENABLE_JIT`
 
 Defaults to `false`.
@@ -102,3 +121,14 @@ The directory to store the downloaded and built archives in. Defaults to the sta
 If you want to compile MLX from source, you can do so by setting the `LIBMLX_BUILD` environment variable to `true`.
 
 Environment variables listed in the previous section will still apply.
+
+#### Testing on macOS 15
+
+On a macOS 15 machine locally:
+
+```bash
+export LIBMLX_MACOS_COMPAT=true
+cd emlx && mix deps.get && mix test
+```
+
+Without that flag, compilation fails with a message naming `LIBMLX_MACOS_COMPAT` instead of crashing the VM at runtime.

emlx/c_src/emlx_async.hpp

@@ -1,21 +1,22 @@
 // Async NIF dispatch built on top of emlx::Worker.
 //
-// MLX 0.31.2 makes both Metal CommandEncoders (mlx/backend/metal/device.cpp:
-// `static thread_local std::unordered_map<int, CommandEncoder> encoders;`)
-// and the per-device default Stream (mlx/stream.cpp: `static thread_local
-// auto default_streams = ...`) thread-local. Because `mlx::core::eval` walks
-// the tape and calls `gpu::eval(arr)` *directly* on the calling thread (it
-// is NOT trampolined to a `scheduler::StreamThread`; see
-// mlx/transforms.cpp:eval_impl), every op for a given GPU stream — both
-// graph construction AND eval — must happen on the OS thread that called
-// `mlx::core::new_stream(d)` for that stream. Otherwise the eval thread's
-// thread-local encoder map will not contain an entry for the stream's
-// index, producing the runtime error
-//   "There is no Stream(gpu, N) in current thread."
+// MLX makes Metal CommandEncoders thread-local (0.31+) and, as of 0.32
+// (ml-explore/mlx#3537), CPU CommandEncoders too. The per-device default
+// Stream is also thread-local (mlx/stream.cpp). Because `mlx::core::eval`
+// walks the tape and calls `gpu::eval` / `cpu::eval` *directly* on the
+// calling thread (it is NOT trampolined to a `scheduler::StreamThread`),
+// every op for a given stream — both graph construction AND eval — must
+// happen on the OS thread that called `mlx::core::new_stream(d)` for that
+// stream. Otherwise the eval thread's encoder map misses the stream
+// index and raises
+//   "There is no Stream(gpu|cpu, N) in current thread."
 //
 // Consequence for EMLX: every NIF that touches the MLX graph must run on
-// the worker thread that owns the stream. We achieve this without
-// rewriting each NIF body by:
+// the worker thread that owns the stream. Mixed-device operand lists are
+// normalized in Elixir (`prepare_tensors!/1`) so a CPU scalar created via
+// `Nx.to_tensor/1` (default backend) is moved onto the op's device before
+// the graph is built. We achieve worker routing without rewriting each
+// NIF body by:
 //
 //   1. Defining each "sync" NIF (e.g. `add`, `reshape`, `eval`, ...) as
 //      a normal C++ function with the ERTS NIF signature.

emlx/c_src/emlx_worker.hpp

@@ -48,13 +48,17 @@ class Worker {
   using Job = std::function<void()>;
 
   // Spawns the worker thread. The fresh mlx::core::Stream is allocated
-  // *inside* the worker thread — MLX 0.31.2 has thread-local state for
-  // GPU streams (mlx-lm#1090, mlx-lm#1179: "There is no Stream(gpu, N)
-  // in current thread"), so a stream created on thread A cannot be
-  // synchronized or dispatched to from thread B. We block here until
-  // the worker thread has created the stream and signalled ready, so
-  // that callers can rely on stream() / device() being valid the
-  // moment the constructor returns.
+  // *inside* the worker thread — MLX has thread-local command encoders
+  // for GPU (0.31+) and CPU (0.32+, ml-explore/mlx#3537), so a stream
+  // created on thread A cannot be synchronized or dispatched to from
+  // thread B ("There is no Stream(gpu|cpu, N) in current thread"). We
+  // block here until the worker thread has created the stream and
+  // signalled ready, so callers can rely on stream() / device() being
+  // valid the moment the constructor returns.
+  //
+  // Mixed-device graphs (e.g. GPU tensor + Nx scalar on the default CPU
+  // backend) are handled in Elixir by promoting operands onto the op's
+  // target device in `prepare_tensors!/1` before the NIF runs.
   explicit Worker(mlx::core::Device device)
       : device_(device), stream_(/*placeholder index*/ -1, device) {
     std::promise<mlx::core::Stream> stream_promise;
@@ -165,13 +169,10 @@ class Worker {
 
 private:
   void thread_main(std::promise<mlx::core::Stream> stream_promise) {
-    // Allocate the stream on THIS thread (MLX 0.31.2 thread-locality
-    // requirement — see constructor comment). Pin all MLX ops issued
-    // from this thread to our stream by making it the per-thread
-    // default. Graph-construction NIFs continue to run on BEAM
-    // scheduler threads (and use those threads' defaults); only ops
-    // invoked *inside* a posted job (currently mx::eval and
-    // mx::synchronize) inherit this binding.
+    // Allocate the stream on THIS thread (MLX thread-locality requirement —
+    // see constructor comment). Pin all MLX ops issued from this thread to
+    // our stream by making it the per-thread default. Only ops invoked
+    // *inside* a posted job inherit this binding.
     try {
       mlx::core::Stream stream = mlx::core::new_stream(device_);
       mlx::core::set_default_stream(stream);

emlx/lib/emlx.ex

@@ -1403,29 +1403,80 @@ defmodule EMLX do
   defp wrap_tensor(tuple, device) when is_tuple(tuple),
     do: tuple |> Tuple.to_list() |> Enum.map(&{device, &1}) |> List.to_tuple()
 
-  defp prepare_tensors_list!(tensors_list, device) do
-    Enum.map_reduce(tensors_list, device, fn
-      {dev, ref}, device when is_tensor(dev, ref) ->
-        {ref, merge_device(device, dev)}
+  # Unwrap EMLX tensors to raw refs for the NIF, after moving every operand
+  # onto the op's merged device. Required since MLX 0.32: CPU streams are
+  # thread-local, so a GPU op cannot eval a CPU scalar that `Nx.to_tensor/1`
+  # allocated on the default CPU backend (`Stream(cpu, N)` lives on the CPU
+  # worker). Promoting here keeps the graph single-device / single-worker.
+  defp prepare_tensors!(tensors) do
+    target = merge_devices(tensors)
 
-      bad_tensor, _device ->
-        raise ArgumentError, "expected a EMLX tensor, got: #{inspect(bad_tensor)}"
-    end)
+    prepared =
+      Enum.map(tensors, fn
+        {dev, ref} = tensor when is_tensor(dev, ref) ->
+          elem(ensure_on_device!(tensor, target), 1)
+
+        [{dev, ref} | _] = list when is_tensor(dev, ref) ->
+          Enum.map(list, fn tensor -> elem(ensure_on_device!(tensor, target), 1) end)
+
+        bad_tensor ->
+          raise ArgumentError, "expected a EMLX tensor, got: #{inspect(bad_tensor)}"
+      end)
+
+    {prepared, target}
   end
 
-  defp prepare_tensors!(tensors) do
-    Enum.map_reduce(tensors, :cpu, fn
+  defp merge_devices(tensors) do
+    Enum.reduce(tensors, :cpu, fn
       {dev, ref}, device when is_tensor(dev, ref) ->
-        {ref, merge_device(device, dev)}
+        merge_device(device, dev)
 
-      [{dev, ref} | _] = tensors, device when is_tensor(dev, ref) ->
-        prepare_tensors_list!(tensors, device)
+      list, device when is_list(list) ->
+        Enum.reduce(list, device, fn
+          {dev, ref}, acc when is_tensor(dev, ref) -> merge_device(acc, dev)
+          bad, _ -> raise ArgumentError, "expected a EMLX tensor, got: #{inspect(bad)}"
+        end)
 
-      bad_tensor, _device ->
-        raise ArgumentError, "expected a EMLX tensor, got: #{inspect(bad_tensor)}"
+      bad, _ ->
+        raise ArgumentError, "expected a EMLX tensor, got: #{inspect(bad)}"
     end)
   end
 
+  defp ensure_on_device!({dev, _} = tensor, dev), do: tensor
+
+  defp ensure_on_device!(tensor, target_device), do: to_device(tensor, target_device)
+
+  # Passthrough EMLX tensors into a runtime_call still live on the
+  # eval_program / default worker's stream. Same device atom (`:cpu`) as the
+  # bound `runtime_call_worker`, so `to_device/2` alone is a no-op — eval on
+  # the home worker, then contiguous onto the bound stream.
+  defp adopt_passthrough_tensor(
+         %Nx.Tensor{data: %EMLX.Backend{ref: {device, ref} = tensor_ref}} = tensor
+       )
+       when is_tensor(device, ref) do
+    case Process.get(:emlx_command_queue) do
+      {worker, ^device} = bound ->
+        Process.delete(:emlx_command_queue)
+
+        try do
+          eval(tensor_ref)
+        after
+          Process.put(:emlx_command_queue, bound)
+        end
+
+        new_ref =
+          EMLX.NIF.to_device(worker, ref, device)
+          |> unwrap!()
+          |> await_worker()
+          |> wrap_tensor(device)
+
+        %{tensor | data: %{tensor.data | ref: new_ref}}
+
+      _ ->
+        tensor
+    end
+  end
+
   defp merge_device(:gpu, _), do: :gpu
   defp merge_device(_, :gpu), do: :gpu
   defp merge_device(_, _), do: :cpu
@@ -1498,10 +1549,11 @@ defmodule EMLX do
     end
   end
 
-  # CPU and GPU operations do not share thread-local Metal encoder state, so
-  # routing a CPU tensor through a GPU queue (or vice-versa) is safe — MLX
-  # inserts the necessary cross-stream synchronization internally. We therefore
-  # do NOT force an intermediate eval; we let MLX manage the graph dependency.
+  # CPU and GPU workers each own thread-local stream encoders (MLX 0.32+).
+  # Cross-device operand lists are normalized in `prepare_tensors!/1` before
+  # the NIF runs; when a process-bound queue's device differs from the
+  # requested tensor device we fall back to that device's default worker
+  # unless `:cross_device_promotion` is enabled.
   defp resolve_cross_device(requested, worker, bound) do
     if Application.get_env(:emlx, :cross_device_promotion, false) do
       if Application.get_env(:emlx, :warn_cross_device, false) do
@@ -1585,25 +1637,6 @@ defmodule EMLX do
       opts: opts
     } = Enum.at(runtime_calls, callback_index)
 
-    {args_container, {[], []}} =
-      Nx.Defn.Composite.traverse(
-        args_template,
-        {args_binaries, positions},
-        fn leaf, {[bin | bins_rest], [pos | pos_rest]} ->
-          value =
-            case pos && Enum.at(tensors, pos) do
-              %Nx.Tensor{data: %EMLX.Backend{quantization_config: %EMLX.Quantization.Config{}}} =
-                  t ->
-                t
-
-              _ ->
-                bin |> Nx.from_binary(leaf.type) |> Nx.reshape(leaf.shape)
-            end
-
-          {value, {bins_rest, pos_rest}}
-        end
-      )
-
     callback_queue = %EMLX.CommandQueue{
       ref: EMLX.Application.runtime_call_worker(dev),
       device: dev
@@ -1613,6 +1646,30 @@ defmodule EMLX do
       try do
         binaries =
           EMLX.CommandQueue.with_queue(callback_queue, fn ->
+            # Allocate reconstructed args on this bound worker (MLX 0.32:
+            # CPU streams are thread-local). Only quantized passthrough
+            # tensors still need an explicit adopt from the eval_program
+            # worker — they keep their original refs.
+            {args_container, {[], []}} =
+              Nx.Defn.Composite.traverse(
+                args_template,
+                {args_binaries, positions},
+                fn leaf, {[bin | bins_rest], [pos | pos_rest]} ->
+                  value =
+                    case pos && Enum.at(tensors, pos) do
+                      %Nx.Tensor{
+                        data: %EMLX.Backend{quantization_config: %EMLX.Quantization.Config{}}
+                      } = t ->
+                        adopt_passthrough_tensor(t)
+
+                      _ ->
+                        bin |> Nx.from_binary(leaf.type) |> Nx.reshape(leaf.shape)
+                    end
+
+                  {value, {bins_rest, pos_rest}}
+                end
+              )
+
             result = callback.(args_container, opts)
 
             [result]
@@ -1990,12 +2047,21 @@ defmodule EMLX do
   end
 
   # Materialises defn input lazy refs to real bound %Nx.Tensor{} values on
-  # `dev` (copying any non-EMLX-backed tensor).
+  # `dev`. Already-EMLX tensors on another device are moved with `to_device/2`
+  # — required since MLX 0.32: `eval_program` runs on `dev`'s worker, and a
+  # lazy CPU graph (e.g. `Nx.broadcast`) still tagged `Stream(cpu, N)` cannot
+  # be eval'd on the GPU worker when `opts[:device]`/`default_device()` is `:gpu`.
   defp materialise_input_tensors(params, dev) do
     Enum.map(params, fn lazy ->
       case lazy.() do
-        %Nx.Tensor{data: %EMLX.Backend{}} = t -> t
-        %Nx.Tensor{} = t -> Nx.backend_copy(t, {EMLX.Backend, device: dev})
+        %Nx.Tensor{data: %EMLX.Backend{ref: {^dev, _}}} = t ->
+          t
+
+        %Nx.Tensor{data: %EMLX.Backend{ref: ref}} = t ->
+          %{t | data: %{t.data | ref: to_device(ref, dev)}}
+
+        %Nx.Tensor{} = t ->
+          Nx.backend_copy(t, {EMLX.Backend, device: dev})
       end
     end)
   end