From 9cfd8999b80338cfcff6fe9bb10a96bb618964d9 Mon Sep 17 00:00:00 2001 From: Paulo Valente <16843419+polvalente@users.noreply.github.com> Date: Tue, 21 Jul 2026 22:00:10 -0300 Subject: [PATCH 1/4] chore: upgrade to mlx 0.32 --- .github/workflows/emlx.yml | 26 +- emlx/Makefile | 6 + emlx/README.md | 30 ++ emlx/c_src/emlx_async.hpp | 29 +- emlx/c_src/emlx_worker.hpp | 29 +- emlx/lib/emlx.ex | 146 +++++++--- emlx/lib/emlx/native/expr.ex | 8 +- emlx/mix.exs | 150 +++++++++- emlx_axon/Makefile | 10 +- emlx_axon/c_src/llama_plugin.cpp | 146 +++++----- emlx_axon/c_src/qwen3_plugin.cpp | 484 +++++++++++++++---------------- 11 files changed, 656 insertions(+), 408 deletions(-) diff --git a/.github/workflows/emlx.yml b/.github/workflows/emlx.yml index b47b8e4..ac7203b 100644 --- a/.github/workflows/emlx.yml +++ b/.github/workflows/emlx.yml @@ -80,17 +80,20 @@ 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 steps: @@ -125,7 +128,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 +144,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 +157,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 +176,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' \; diff --git a/emlx/Makefile b/emlx/Makefile index ab2f8ad..4da7d74 100644 --- a/emlx/Makefile +++ b/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 diff --git a/emlx/README.md b/emlx/README.md index 053ea8c..b7839a8 100644 --- a/emlx/README.md +++ b/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. diff --git a/emlx/c_src/emlx_async.hpp b/emlx/c_src/emlx_async.hpp index 0447adc..4a9be70 100644 --- a/emlx/c_src/emlx_async.hpp +++ b/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 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. diff --git a/emlx/c_src/emlx_worker.hpp b/emlx/c_src/emlx_worker.hpp index 5b4b714..e9c6a21 100644 --- a/emlx/c_src/emlx_worker.hpp +++ b/emlx/c_src/emlx_worker.hpp @@ -48,13 +48,17 @@ class Worker { using Job = std::function; // 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 stream_promise; @@ -165,13 +169,10 @@ class Worker { private: void thread_main(std::promise 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); diff --git a/emlx/lib/emlx.ex b/emlx/lib/emlx.ex index 46b80cf..7bfb358 100644 --- a/emlx/lib/emlx.ex +++ b/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 diff --git a/emlx/lib/emlx/native/expr.ex b/emlx/lib/emlx/native/expr.ex index fda642a..de99382 100644 --- a/emlx/lib/emlx/native/expr.ex +++ b/emlx/lib/emlx/native/expr.ex @@ -2272,7 +2272,13 @@ defmodule EMLX.Native.Expr do wrapped_callback = fn {value, chain_in}, _opts -> callback.(value) - {value, Nx.add(chain_in, 1)} + # Keep the keepalive counter on BinaryBackend. `Nx.to_tensor(1)` would + # otherwise allocate an EMLX scalar on whichever worker + # `resolve_worker/1` picks inside the runtime_call callback (often a + # different OS thread than the one that owns Stream(cpu, 0)), and + # MLX 0.32's thread-local CPU encoders then fail at `to_binary`. + one = Nx.tensor(1, type: @hook_chain_type, backend: Nx.BinaryBackend) + {value, Nx.add(chain_in, one)} end {result_refs, state} = diff --git a/emlx/mix.exs b/emlx/mix.exs index 45bfc07..b82293f 100644 --- a/emlx/mix.exs +++ b/emlx/mix.exs @@ -3,8 +3,13 @@ defmodule EMLX.MixProject do @app :emlx @version "0.4.0" - @mlx_version "0.31.2" + @mlx_version "0.32.0" @source_url "https://github.com/elixir-nx/emlx" + # mlx-build Darwin archives from v0.31.0+ encode the deployment target in the + # filename. 26.2 ships AOT NAX (Metal 4); 14.0 is the portable reduced set. + @macos_deployment_target_latest "26.2" + @macos_deployment_target_compat "14.0" + @macos_deployment_targets [@macos_deployment_target_compat, @macos_deployment_target_latest] require Logger @@ -41,7 +46,8 @@ defmodule EMLX.MixProject do "EMLX_CACHE_DIR" => libmlx_config.cache_dir, "EMLX_VERSION" => @version, "LIBMLX_ENABLE_DEBUG" => to_string(libmlx_config.features.debug?), - "FINE_INCLUDE_DIR" => Fine.include_dir() + "FINE_INCLUDE_DIR" => Fine.include_dir(), + "MACOSX_DEPLOYMENT_TARGET" => libmlx_config.macos_deployment_target } end, @@ -70,7 +76,7 @@ defmodule EMLX.MixProject do [ {:elixir_make, "~> 0.6"}, {:fine, "~> 0.1", runtime: false}, - {:nx, "~> 0.12"}, + {:nx, "~> 0.12.0"}, {:telemetry, "~> 1.0"}, {:ex_doc, "~> 0.34", only: :docs} ] @@ -219,10 +225,17 @@ defmodule EMLX.MixProject do build?: to_boolean(System.get_env("LIBMLX_BUILD")) } - variant = to_variant(features) - current_target = current_target!() + macos_deployment_target = + unless features.build? do + resolve_macos_deployment_target(current_target) + end + + # Deployment target is a Darwin archive-name segment (mlx-build), not a + # sorted feature flag — keep it ahead of `-debug`/`-jit`. + variant = to_variant(features, macos_deployment_target) + cache_dir = if dir = System.get_env("LIBMLX_CACHE") do Path.expand(dir) @@ -259,6 +272,7 @@ defmodule EMLX.MixProject do dir: Path.join(cache_dir, "libmlx-#{version}-#{current_target}#{variant}"), features: features, variant: variant, + macos_deployment_target: macos_deployment_target, cache_dir: cache_dir } end @@ -269,18 +283,124 @@ defmodule EMLX.MixProject do String.downcase(to_string(var)) in ["1", "true", "on", "yes", "y"] end - defp to_variant(features) do - [ - if(features.build?, do: "build", else: nil), - if(features.debug?, do: "debug", else: nil), - if(features.jit?, do: "jit", else: nil) - ] - |> Enum.filter(&(&1 != nil)) - |> Enum.sort() - |> Enum.map(&"-#{&1}") - |> Enum.join("") + defp to_variant(features, macos_deployment_target) do + feature_variant = + [ + if(features.build?, do: "build"), + if(features.debug?, do: "debug"), + if(features.jit?, do: "jit") + ] + |> Enum.filter(&(&1 != nil)) + |> Enum.sort() + |> Enum.join("-") + + feature_variant = if feature_variant != "", do: "-#{feature_variant}", else: "" + + case macos_deployment_target do + nil -> feature_variant + deployment_target -> "-#{deployment_target}#{feature_variant}" + end + end + + defp resolve_macos_deployment_target(target) do + if String.contains?(target, "apple-darwin") do + selected = selected_macos_deployment_target() + assert_macos_deployment_target_compatible!(selected) + selected + end + end + + defp selected_macos_deployment_target do + case System.get_env("LIBMLX_DEPLOYMENT_TARGET") do + nil -> + if to_boolean(System.get_env("LIBMLX_MACOS_COMPAT")) do + @macos_deployment_target_compat + else + @macos_deployment_target_latest + end + + target when target in @macos_deployment_targets -> + target + + other -> + Mix.raise(""" + Invalid LIBMLX_DEPLOYMENT_TARGET=#{inspect(other)}. + + Supported values: #{Enum.join(@macos_deployment_targets, ", ")}. + + Use #{@macos_deployment_target_latest} (default) for the full macOS feature set \ + including AOT NAX kernels, or #{@macos_deployment_target_compat} / \ + LIBMLX_MACOS_COMPAT=true for the reduced-featureset build that runs on macOS 14+. + """) + end + end + + defp assert_macos_deployment_target_compatible!(deployment_target) do + # Cross-compiling via TARGET_* — host sw_vers is not the runtime OS. + if current_target_from_env() do + :ok + else + case macos_product_version() do + nil -> + :ok + + host_version -> + min_version = parse_dotted_version!(deployment_target) + + if version_lt?(host_version, min_version) do + Mix.raise(""" + Precompiled libmlx for macOS deployment target #{deployment_target} cannot run \ + on this host (macOS #{format_dotted_version(host_version)}). + + This is the failure mode behind https://github.com/elixir-nx/emlx/issues/127 \ + (VM crash: mutex lock failed). + + For macOS 15 and other hosts older than #{@macos_deployment_target_latest}, use \ + the reduced-featureset archive: + + export LIBMLX_MACOS_COMPAT=true + + or set LIBMLX_DEPLOYMENT_TARGET=#{@macos_deployment_target_compat} explicitly. + """) + else + :ok + end + end + end end + defp macos_product_version do + case :os.type() do + {:unix, :darwin} -> + case System.cmd("sw_vers", ["-productVersion"], stderr_to_stdout: true) do + {version, 0} -> parse_dotted_version!(String.trim(version)) + _ -> nil + end + + _ -> + nil + end + end + + defp parse_dotted_version!(version) do + case String.split(version, ".", trim: true) do + [major, minor | _rest] -> + {String.to_integer(major), String.to_integer(minor)} + + [major] -> + {String.to_integer(major), 0} + end + rescue + ArgumentError -> + Mix.raise("Could not parse macOS version #{inspect(version)}") + end + + defp version_lt?({maj_a, min_a}, {maj_b, min_b}) do + maj_a < maj_b or (maj_a == maj_b and min_a < min_b) + end + + defp format_dotted_version({major, minor}), do: "#{major}.#{minor}" + defp download_and_unarchive(args) do libmlx_config = libmlx_config() diff --git a/emlx_axon/Makefile b/emlx_axon/Makefile index cab9c7b..723a03f 100644 --- a/emlx_axon/Makefile +++ b/emlx_axon/Makefile @@ -12,14 +12,20 @@ LLAMA_OBJECT = $(OBJ_DIR)/llama_plugin.o QWEN3_SOURCE = c_src/qwen3_plugin.cpp LLAMA_SOURCE = c_src/llama_plugin.cpp PLUGIN_ABI = $(EMLX_PLUGIN_INCLUDE_DIR)/emlx/plugin/abi.hpp +# Rebuild plugin objects when MLX's ABI surface changes — not only when our +# sources change. A stale .o compiled against an older ops.h keeps undefined +# refs like the pre-0.32 3-arg `astype` and fails dlopen against newer libmlx. +MLX_OPS_H = $(MLX_INCLUDE_DIR)/mlx/ops.h UNAME_S := $(shell uname -s) ifeq ($(UNAME_S),Darwin) PLATFORM_CFLAGS = -Wglobal-constructors -Wexit-time-destructors -Werror=global-constructors -Werror=exit-time-destructors PLATFORM_LDFLAGS = -undefined dynamic_lookup -flat_namespace -rpath @loader_path/mlx/lib -framework Metal -framework Foundation -framework Accelerate + MLX_SHARED_LIB = $(MLX_LIB_DIR)/libmlx.dylib else PLATFORM_CFLAGS = PLATFORM_LDFLAGS = -Wl,-rpath,'$$ORIGIN/mlx/lib' + MLX_SHARED_LIB = $(MLX_LIB_DIR)/libmlx.so endif CFLAGS = -fPIC -isystem $(call esc,$(MLX_INCLUDE_DIR)) -I$(call esc,$(EMLX_PLUGIN_INCLUDE_DIR)) -Wall -std=c++20 -O3 -fvisibility=hidden $(PLATFORM_CFLAGS) @@ -28,11 +34,11 @@ LDFLAGS = -L$(call esc,$(MLX_LIB_DIR)) -lmlx -shared $(PLATFORM_LDFLAGS) all: $(call esc,$(QWEN3_SO)) $(call esc,$(LLAMA_SO)) @ echo > /dev/null -$(call esc,$(QWEN3_OBJECT)): $(QWEN3_SOURCE) $(call esc,$(PLUGIN_ABI)) +$(call esc,$(QWEN3_OBJECT)): $(QWEN3_SOURCE) $(call esc,$(PLUGIN_ABI)) $(call esc,$(MLX_OPS_H)) $(call esc,$(MLX_SHARED_LIB)) @ mkdir -p "$(OBJ_DIR)" $(CXX) $(CFLAGS) -c "$<" -o "$@" -$(call esc,$(LLAMA_OBJECT)): $(LLAMA_SOURCE) $(call esc,$(PLUGIN_ABI)) +$(call esc,$(LLAMA_OBJECT)): $(LLAMA_SOURCE) $(call esc,$(PLUGIN_ABI)) $(call esc,$(MLX_OPS_H)) $(call esc,$(MLX_SHARED_LIB)) @ mkdir -p "$(OBJ_DIR)" $(CXX) $(CFLAGS) -c "$<" -o "$@" diff --git a/emlx_axon/c_src/llama_plugin.cpp b/emlx_axon/c_src/llama_plugin.cpp index 812935b..e0bbda2 100644 --- a/emlx_axon/c_src/llama_plugin.cpp +++ b/emlx_axon/c_src/llama_plugin.cpp @@ -193,26 +193,26 @@ bool validate_kv_cache_bn(const mlx::core::array &k_cache, mlx::core::array apply_rope(const mlx::core::array &a, int head_dim, int offset, const mlx::core::array &rope_freqs, - const mlx::core::Device &device) { - auto offsets = mlx::core::full({a.shape(0)}, offset, mlx::core::int32, device); + mlx::core::StreamOrDevice stream) { + auto offsets = mlx::core::full({a.shape(0)}, offset, mlx::core::int32, stream); return mlx::core::fast::rope( - a, head_dim, false, std::nullopt, 1.0f, offsets, rope_freqs, device); + a, head_dim, false, std::nullopt, 1.0f, offsets, rope_freqs, stream); } mlx::core::array linear_in_out(const mlx::core::array &x, const mlx::core::array &weight, - const mlx::core::Device &device) { + mlx::core::StreamOrDevice stream) { if (x.ndim() == 3 && x.shape(1) == 1) { - auto x_2d = mlx::core::reshape(x, {x.shape(0), x.shape(2)}, device); - auto out = mlx::core::matmul(x_2d, weight, device); - return mlx::core::reshape(out, {x.shape(0), 1, weight.shape(1)}, device); + auto x_2d = mlx::core::reshape(x, {x.shape(0), x.shape(2)}, stream); + auto out = mlx::core::matmul(x_2d, weight, stream); + return mlx::core::reshape(out, {x.shape(0), 1, weight.shape(1)}, stream); } - return mlx::core::matmul(x, weight, device); + return mlx::core::matmul(x, weight, stream); } mlx::core::array linear_out_in(const mlx::core::array &x, const mlx::core::array &weight, - const mlx::core::Device &device) { + mlx::core::StreamOrDevice stream) { return mlx::core::tensordot( - x, weight, std::vector{static_cast(x.ndim()) - 1}, std::vector{1}, device); + x, weight, std::vector{static_cast(x.ndim()) - 1}, std::vector{1}, stream); } bool validate_dense_layer(const mlx::core::array &hidden, const LayerParams &layer, @@ -266,35 +266,35 @@ bool validate_dense_layer(const mlx::core::array &hidden, const LayerParams &lay } mlx::core::array build_prefill_mask(const mlx::core::array &q, int T_new, int valid_len, - const mlx::core::Device &device) { + mlx::core::StreamOrDevice stream) { auto mask_dtype = q.dtype(); - auto zero_val = mlx::core::zeros({}, mask_dtype, device); - auto neginf_val = mlx::core::full({}, -std::numeric_limits::infinity(), mask_dtype, device); + auto zero_val = mlx::core::zeros({}, mask_dtype, stream); + auto neginf_val = mlx::core::full({}, -std::numeric_limits::infinity(), mask_dtype, stream); int kv_offset = valid_len - T_new; auto row = mlx::core::reshape( - mlx::core::arange(T_new, mlx::core::int32, device), {1, 1, T_new, 1}, device); + mlx::core::arange(T_new, mlx::core::int32, stream), {1, 1, T_new, 1}, stream); auto col = mlx::core::reshape( - mlx::core::arange(valid_len, mlx::core::int32, device), {1, 1, 1, valid_len}, device); + mlx::core::arange(valid_len, mlx::core::int32, stream), {1, 1, 1, valid_len}, stream); auto causal_bool = mlx::core::less_equal( - col, mlx::core::add(row, mlx::core::array(kv_offset, mlx::core::int32), device), device); - return mlx::core::where(causal_bool, zero_val, neginf_val, device); + col, mlx::core::add(row, mlx::core::full({}, kv_offset, mlx::core::int32, stream), stream), stream); + return mlx::core::where(causal_bool, zero_val, neginf_val, stream); } mlx::core::array sdpa(const mlx::core::array &q_rope, const mlx::core::array &k_valid, const mlx::core::array &v_valid, float scale, int T_new, - int valid_len, const mlx::core::Device &device) { + int valid_len, mlx::core::StreamOrDevice stream) { return (T_new == 1) ? mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, device) + q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, stream) : mlx::core::fast::scaled_dot_product_attention( q_rope, k_valid, v_valid, scale, "array", - build_prefill_mask(q_rope, T_new, valid_len, device), std::nullopt, device); + build_prefill_mask(q_rope, T_new, valid_len, stream), std::nullopt, stream); } mlx::core::array layer_dense_impl(const mlx::core::array &hidden, const LayerParams &layer, KVCache &kv, int offset, float scale, int head_dim, const mlx::core::array &rope_freqs, float eps, - const mlx::core::Device &device, mlx::core::array *k_out, + mlx::core::StreamOrDevice stream, mlx::core::array *k_out, mlx::core::array *v_out) { int B = hidden.shape(0); int T_new = hidden.shape(1); @@ -304,45 +304,45 @@ mlx::core::array layer_dense_impl(const mlx::core::array &hidden, const LayerPar int attn_width = N_q * D; int valid_len = offset + T_new; - auto xn = mlx::core::fast::rms_norm(hidden, *layer.norm1, eps, device); - auto q_flat = linear_in_out(xn, *layer.q_proj, device); - auto k_flat = linear_in_out(xn, *layer.k_proj, device); - auto v_flat = linear_in_out(xn, *layer.v_proj, device); + auto xn = mlx::core::fast::rms_norm(hidden, *layer.norm1, eps, stream); + auto q_flat = linear_in_out(xn, *layer.q_proj, stream); + auto k_flat = linear_in_out(xn, *layer.k_proj, stream); + auto v_flat = linear_in_out(xn, *layer.v_proj, stream); - auto q = mlx::core::reshape(q_flat, {B, T_new, N_q, D}, device); - auto k = mlx::core::reshape(k_flat, {B, T_new, N_kv, D}, device); - auto v = mlx::core::reshape(v_flat, {B, T_new, N_kv, D}, device); + auto q = mlx::core::reshape(q_flat, {B, T_new, N_q, D}, stream); + auto k = mlx::core::reshape(k_flat, {B, T_new, N_kv, D}, stream); + auto v = mlx::core::reshape(v_flat, {B, T_new, N_kv, D}, stream); - auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, device); - auto k_bn = mlx::core::transpose(k, {0, 2, 1, 3}, device); - auto v_bn = mlx::core::transpose(v, {0, 2, 1, 3}, device); + auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, stream); + auto k_bn = mlx::core::transpose(k, {0, 2, 1, 3}, stream); + auto v_bn = mlx::core::transpose(v, {0, 2, 1, 3}, stream); - auto q_rope = apply_rope(q_bn, D, offset, rope_freqs, device); - auto k_rope = apply_rope(k_bn, D, offset, rope_freqs, device); + auto q_rope = apply_rope(q_bn, D, offset, rope_freqs, stream); + auto k_rope = apply_rope(k_bn, D, offset, rope_freqs, stream); auto k_cache_owned = std::move(*kv.k); auto v_cache_owned = std::move(*kv.v); auto k_upd = mlx::core::slice_update( - k_cache_owned, k_rope, to_shape({0, 0, offset, 0}), to_shape({B, N_kv, valid_len, D}), device); + k_cache_owned, k_rope, to_shape({0, 0, offset, 0}), to_shape({B, N_kv, valid_len, D}), stream); auto v_upd = mlx::core::slice_update( - v_cache_owned, v_bn, to_shape({0, 0, offset, 0}), to_shape({B, N_kv, valid_len, D}), device); + v_cache_owned, v_bn, to_shape({0, 0, offset, 0}), to_shape({B, N_kv, valid_len, D}), stream); - auto k_valid = mlx::core::slice(k_upd, to_shape({0, 0, 0, 0}), to_shape({B, N_kv, valid_len, D}), device); - auto v_valid = mlx::core::slice(v_upd, to_shape({0, 0, 0, 0}), to_shape({B, N_kv, valid_len, D}), device); + auto k_valid = mlx::core::slice(k_upd, to_shape({0, 0, 0, 0}), to_shape({B, N_kv, valid_len, D}), stream); + auto v_valid = mlx::core::slice(v_upd, to_shape({0, 0, 0, 0}), to_shape({B, N_kv, valid_len, D}), stream); - auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, scale, T_new, valid_len, device); - auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, device); - auto attn_out = mlx::core::reshape(attn_out_bthd, {B, T_new, attn_width}, device); - auto attn_projected = linear_in_out(attn_out, *layer.o_proj, device); - auto attn_hidden = mlx::core::add(hidden, attn_projected, device); + auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, scale, T_new, valid_len, stream); + auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, stream); + auto attn_out = mlx::core::reshape(attn_out_bthd, {B, T_new, attn_width}, stream); + auto attn_projected = linear_in_out(attn_out, *layer.o_proj, stream); + auto attn_hidden = mlx::core::add(hidden, attn_projected, stream); - auto xn2 = mlx::core::fast::rms_norm(attn_hidden, *layer.norm2, eps, device); - auto gate = linear_in_out(xn2, *layer.gate_proj, device); - auto up = linear_in_out(xn2, *layer.up_proj, device); + auto xn2 = mlx::core::fast::rms_norm(attn_hidden, *layer.norm2, eps, stream); + auto gate = linear_in_out(xn2, *layer.gate_proj, stream); + auto up = linear_in_out(xn2, *layer.up_proj, stream); auto mlp = mlx::core::multiply( - mlx::core::multiply(gate, mlx::core::sigmoid(gate, device), device), up, device); - auto mlp_out = linear_in_out(mlp, *layer.down_proj, device); + mlx::core::multiply(gate, mlx::core::sigmoid(gate, stream), stream), up, stream); + auto mlp_out = linear_in_out(mlp, *layer.down_proj, stream); if (k_out != nullptr) { *k_out = k_upd; @@ -351,12 +351,12 @@ mlx::core::array layer_dense_impl(const mlx::core::array &hidden, const LayerPar *v_out = v_upd; } - return mlx::core::add(attn_hidden, mlp_out, device); + return mlx::core::add(attn_hidden, mlp_out, stream); } bool v_layer_dense(const mlx::core::array &hidden, const LayerParams &layer, KVCache &kv, int offset, double scale, int head_dim, const mlx::core::array &rope_freqs, - double eps, const mlx::core::Device &device, mlx::core::array &out, + double eps, mlx::core::StreamOrDevice stream, mlx::core::array &out, mlx::core::array &k_upd, mlx::core::array &v_upd, std::string &error) { try { if (!validate_dense_layer(hidden, layer, kv, rope_freqs, offset, head_dim, error)) { @@ -364,7 +364,7 @@ bool v_layer_dense(const mlx::core::array &hidden, const LayerParams &layer, KVC } out = layer_dense_impl( hidden, layer, kv, offset, static_cast(scale), head_dim, rope_freqs, - static_cast(eps), device, &k_upd, &v_upd); + static_cast(eps), stream, &k_upd, &v_upd); return true; } catch (const std::exception &e) { error = e.what(); @@ -380,7 +380,7 @@ bool v_forward_greedy_from_hidden( std::vector &kv, const mlx::core::array &norm, const mlx::core::array &lm_head, int offset, double scale, int head_dim, const mlx::core::array &rope_freqs, double eps, - const mlx::core::Device &device, mlx::core::array &token_out, + mlx::core::StreamOrDevice stream, mlx::core::array &token_out, std::vector &k_out, std::vector &v_out, std::string &error) { try { @@ -408,7 +408,7 @@ bool v_forward_greedy_from_hidden( mlx::core::array v_new = *kv[i].v; current = layer_dense_impl( current, layers[i], kv[i], offset, static_cast(scale), head_dim, rope_freqs, - static_cast(eps), device, &k_new, &v_new); + static_cast(eps), stream, &k_new, &v_new); k_out.push_back(k_new); v_out.push_back(v_new); } @@ -423,14 +423,14 @@ bool v_forward_greedy_from_hidden( } auto last = (T == 1) - ? mlx::core::reshape(current, {B, H}, device) + ? mlx::core::reshape(current, {B, H}, stream) : mlx::core::reshape( - mlx::core::slice(current, to_shape({0, T - 1, 0}), to_shape({B, T, H}), device), - {B, H}, device); + mlx::core::slice(current, to_shape({0, T - 1, 0}), to_shape({B, T, H}), stream), + {B, H}, stream); - auto normed = mlx::core::fast::rms_norm(last, norm, static_cast(eps), device); - auto logits = linear_out_in(normed, lm_head, device); - auto token = mlx::core::argmax(logits, 1, false, device); + auto normed = mlx::core::fast::rms_norm(last, norm, static_cast(eps), stream); + auto logits = linear_out_in(normed, lm_head, stream); + auto token = mlx::core::argmax(logits, 1, false, stream); token_out = token; return true; @@ -449,7 +449,7 @@ bool v_forward_greedy_ids_chunk( const mlx::core::array &norm, const mlx::core::array &lm_head, int offset, int count, double scale, int head_dim, const mlx::core::array &rope_freqs, double eps, bool submit_each_step, - const mlx::core::Device &device, + mlx::core::StreamOrDevice stream, std::vector &token_out, std::vector &k_out, std::vector &v_out, std::string &error) { try { @@ -496,9 +496,9 @@ bool v_forward_greedy_ids_chunk( for (int step = 0; step < count; ++step) { int B = current_ids.shape(0); int T = current_ids.shape(1); - auto ids = mlx::core::reshape(current_ids, {B * T}, device); + auto ids = mlx::core::reshape(current_ids, {B * T}, stream); auto current = mlx::core::reshape( - mlx::core::take(embed_tokens, ids, 0, device), {B, T, embed_tokens.shape(1)}, device); + mlx::core::take(embed_tokens, ids, 0, stream), {B, T, embed_tokens.shape(1)}, stream); next_k_cache.clear(); next_v_cache.clear(); @@ -515,7 +515,7 @@ bool v_forward_greedy_ids_chunk( mlx::core::array v_new = *kv.v; current = layer_dense_impl( current, layers[layer_idx], kv, current_offset, static_cast(scale), head_dim, - rope_freqs, static_cast(eps), device, &k_new, &v_new); + rope_freqs, static_cast(eps), stream, &k_new, &v_new); next_k_cache.push_back(k_new); next_v_cache.push_back(v_new); } @@ -525,14 +525,14 @@ bool v_forward_greedy_ids_chunk( int H_out = current.shape(2); auto last = (T_out == 1) - ? mlx::core::reshape(current, {B_out, H_out}, device) + ? mlx::core::reshape(current, {B_out, H_out}, stream) : mlx::core::reshape( - mlx::core::slice(current, to_shape({0, T_out - 1, 0}), to_shape({B_out, T_out, H_out}), device), - {B_out, H_out}, device); + mlx::core::slice(current, to_shape({0, T_out - 1, 0}), to_shape({B_out, T_out, H_out}), stream), + {B_out, H_out}, stream); - auto normed = mlx::core::fast::rms_norm(last, norm, static_cast(eps), device); - auto logits = linear_out_in(normed, lm_head, device); - auto token = mlx::core::argmax(logits, 1, false, device); + auto normed = mlx::core::fast::rms_norm(last, norm, static_cast(eps), stream); + auto logits = linear_out_in(normed, lm_head, stream); + auto token = mlx::core::argmax(logits, 1, false, stream); if (submit_each_step) { std::vector eval_arrays; @@ -546,7 +546,7 @@ bool v_forward_greedy_ids_chunk( } token_arrays.push_back(token); - current_ids = mlx::core::reshape(token, {B_out, 1}, device); + current_ids = mlx::core::reshape(token, {B_out, 1}, stream); k_cache.swap(next_k_cache); v_cache.swap(next_v_cache); current_offset += 1; @@ -591,7 +591,7 @@ plugin_layer_dense(const emlx::plugin::call_t &call, auto v_updated = operands[7]; if (!v_layer_dense( operands[0], layer, cache, offset, f64_from_bits(call.attrs[1]), - head_dim, operands[12], f64_from_bits(call.attrs[3]), call.device, + head_dim, operands[12], f64_from_bits(call.attrs[3]), call.stream, output, k_updated, v_updated, error)) { return error; } @@ -696,7 +696,7 @@ plugin_forward_dense(const emlx::plugin::call_t &call, if (!v_forward_greedy_from_hidden( operands[0], layers, caches, operands[tail], operands[tail + 1], offset, f64_from_bits(call.attrs[2]), head_dim, - operands[tail + 2], f64_from_bits(call.attrs[4]), call.device, token, + operands[tail + 2], f64_from_bits(call.attrs[4]), call.stream, token, keys, values, error)) { return error; } @@ -748,13 +748,13 @@ plugin_chunk_dense(const emlx::plugin::call_t &call, operands[0], operands[1], layers, caches, operands[tail], operands[tail + 1], offset, count, f64_from_bits(call.attrs[3]), head_dim, operands[tail + 2], f64_from_bits(call.attrs[5]), - call.attrs[6] == 1, call.device, tokens, keys, values, error)) { + call.attrs[6] == 1, call.stream, tokens, keys, values, error)) { return error; } outputs.reserve(expected_outputs); outputs.push_back(mlx::core::reshape( - mlx::core::stack(tokens, 0, call.device), {count}, call.device)); + mlx::core::stack(tokens, 0, call.stream), {count}, call.stream)); for (size_t index = 0; index < keys.size(); ++index) { outputs.push_back(std::move(keys[index])); outputs.push_back(std::move(values[index])); diff --git a/emlx_axon/c_src/qwen3_plugin.cpp b/emlx_axon/c_src/qwen3_plugin.cpp index 24a2bc4..f4d7545 100644 --- a/emlx_axon/c_src/qwen3_plugin.cpp +++ b/emlx_axon/c_src/qwen3_plugin.cpp @@ -234,31 +234,31 @@ bool validate_qkv_cache_attention(const mlx::core::array &q, // ── Dense/quantized linear projection helpers ──────────────────────────── mlx::core::array linear_in_out(const mlx::core::array &x, const mlx::core::array &weight, - const mlx::core::Device &device) { + mlx::core::StreamOrDevice stream) { if (x.ndim() == 3 && x.shape(1) == 1) { - auto x_2d = mlx::core::reshape(x, {x.shape(0), x.shape(2)}, device); - auto out = mlx::core::matmul(x_2d, weight, device); - return mlx::core::reshape(out, {x.shape(0), 1, weight.shape(1)}, device); + auto x_2d = mlx::core::reshape(x, {x.shape(0), x.shape(2)}, stream); + auto out = mlx::core::matmul(x_2d, weight, stream); + return mlx::core::reshape(out, {x.shape(0), 1, weight.shape(1)}, stream); } - return mlx::core::matmul(x, weight, device); + return mlx::core::matmul(x, weight, stream); } mlx::core::array linear_out_in(const mlx::core::array &x, const mlx::core::array &weight, - const mlx::core::Device &device) { + mlx::core::StreamOrDevice stream) { return mlx::core::tensordot(x, weight, std::vector{static_cast(x.ndim()) - 1}, - std::vector{1}, device); + std::vector{1}, stream); } mlx::core::array apply_linear(const mlx::core::array &x, const LinearWeight &w, - const mlx::core::Device &device) { + mlx::core::StreamOrDevice stream) { if (w.quantized) { std::optional biases_opt = w.biases != nullptr ? std::make_optional(*w.biases) : std::nullopt; return mlx::core::quantized_matmul(x, *w.weight, *w.scales, biases_opt, w.transpose, - w.group_size, w.bits, w.mode, device); + w.group_size, w.bits, w.mode, stream); } - return w.transpose ? linear_out_in(x, *w.weight, device) - : linear_in_out(x, *w.weight, device); + return w.transpose ? linear_out_in(x, *w.weight, stream) + : linear_in_out(x, *w.weight, stream); } int linear_weight_out_features(const LinearWeight &w) { @@ -443,19 +443,19 @@ bool validate_dense_layer(const mlx::core::array &hidden, const LayerParams &lay // Shared causal/prefill mask builder used by every attention path below. mlx::core::array build_prefill_mask(const mlx::core::array &q, int T_new, int valid_len, - const mlx::core::Device &device) { + mlx::core::StreamOrDevice stream) { auto mask_dtype = q.dtype(); - auto zero_val = mlx::core::zeros({}, mask_dtype, device); + auto zero_val = mlx::core::zeros({}, mask_dtype, stream); auto neginf_val = - mlx::core::full({}, -std::numeric_limits::infinity(), mask_dtype, device); + mlx::core::full({}, -std::numeric_limits::infinity(), mask_dtype, stream); int kv_offset = valid_len - T_new; - auto row = mlx::core::reshape(mlx::core::arange(T_new, mlx::core::int32, device), - {1, 1, T_new, 1}, device); - auto col = mlx::core::reshape(mlx::core::arange(valid_len, mlx::core::int32, device), - {1, 1, 1, valid_len}, device); + auto row = mlx::core::reshape(mlx::core::arange(T_new, mlx::core::int32, stream), + {1, 1, T_new, 1}, stream); + auto col = mlx::core::reshape(mlx::core::arange(valid_len, mlx::core::int32, stream), + {1, 1, 1, valid_len}, stream); auto causal_bool = mlx::core::less_equal( - col, mlx::core::add(row, mlx::core::array(kv_offset, mlx::core::int32), device), device); - return mlx::core::where(causal_bool, zero_val, neginf_val, device); + col, mlx::core::add(row, mlx::core::full({}, kv_offset, mlx::core::int32, stream), stream), stream); + return mlx::core::where(causal_bool, zero_val, neginf_val, stream); } bool validate_tensor_offset(const mlx::core::array &offset, int capacity, @@ -473,68 +473,68 @@ bool validate_tensor_offset(const mlx::core::array &offset, int capacity, } mlx::core::array clamp_offset(const mlx::core::array &offset, int maximum, - const mlx::core::Device &device) { - auto offset_i32 = mlx::core::astype(offset, mlx::core::int32, device); - auto zero = mlx::core::array(0, mlx::core::int32); - auto upper = mlx::core::array(maximum, mlx::core::int32); - return mlx::core::minimum(mlx::core::maximum(offset_i32, zero, device), upper, - device); + mlx::core::StreamOrDevice stream) { + auto offset_i32 = mlx::core::astype(offset, mlx::core::int32, stream); + auto zero = mlx::core::full({}, 0, mlx::core::int32, stream); + auto upper = mlx::core::full({}, maximum, mlx::core::int32, stream); + return mlx::core::minimum(mlx::core::maximum(offset_i32, zero, stream), upper, + stream); } mlx::core::array rope_with_positions(const mlx::core::array &input, const mlx::core::array &offset, int dims, float theta, - const mlx::core::Device &device) { + mlx::core::StreamOrDevice stream) { const int batch = input.shape(0); const int tokens = input.shape(1); const int heads = input.shape(2); const int half = dims / 2; auto positions = mlx::core::add( - mlx::core::arange(tokens, mlx::core::int32, device), offset, device); - positions = mlx::core::reshape(positions, {1, tokens, 1}, device); + mlx::core::arange(tokens, mlx::core::int32, stream), offset, stream); + positions = mlx::core::reshape(positions, {1, tokens, 1}, stream); if (batch != 1) - positions = mlx::core::broadcast_to(positions, {batch, tokens, 1}, device); - auto frequency_index = mlx::core::arange(0, dims, 2, mlx::core::float32, device); + positions = mlx::core::broadcast_to(positions, {batch, tokens, 1}, stream); + auto frequency_index = mlx::core::arange(0, dims, 2, mlx::core::float32, stream); auto exponent = mlx::core::divide( - frequency_index, mlx::core::array(static_cast(dims)), device); + frequency_index, mlx::core::full({}, static_cast(dims), mlx::core::float32, stream), stream); auto inverse_frequency = mlx::core::exp( mlx::core::multiply( - exponent, mlx::core::array(-std::log(theta), mlx::core::float32), - device), - device); + exponent, mlx::core::full({}, static_cast(-std::log(theta)), mlx::core::float32, stream), + stream), + stream); auto angles = mlx::core::multiply( - mlx::core::astype(positions, mlx::core::float32, device), - mlx::core::reshape(inverse_frequency, {1, 1, half}, device), device); + mlx::core::astype(positions, mlx::core::float32, stream), + mlx::core::reshape(inverse_frequency, {1, 1, half}, stream), stream); auto cosine = mlx::core::astype( - mlx::core::reshape(mlx::core::cos(angles, device), {batch, tokens, 1, half}, - device), - input.dtype(), device); + mlx::core::reshape(mlx::core::cos(angles, stream), {batch, tokens, 1, half}, + stream), + input.dtype(), stream); auto sine = mlx::core::astype( - mlx::core::reshape(mlx::core::sin(angles, device), {batch, tokens, 1, half}, - device), - input.dtype(), device); + mlx::core::reshape(mlx::core::sin(angles, stream), {batch, tokens, 1, half}, + stream), + input.dtype(), stream); auto cosine_full = mlx::core::concatenate(std::vector{cosine, cosine}, 3, - device); + stream); auto sine_full = mlx::core::concatenate(std::vector{sine, sine}, 3, - device); + stream); auto first = mlx::core::slice(input, {0, 0, 0, 0}, - {batch, tokens, heads, half}, device); + {batch, tokens, heads, half}, stream); auto second = mlx::core::slice(input, {0, 0, 0, half}, - {batch, tokens, heads, dims}, device); + {batch, tokens, heads, dims}, stream); auto rotated = mlx::core::concatenate( - std::vector{mlx::core::negative(second, device), first}, - 3, device); - return mlx::core::add(mlx::core::multiply(input, cosine_full, device), - mlx::core::multiply(rotated, sine_full, device), device); + std::vector{mlx::core::negative(second, stream), first}, + 3, stream); + return mlx::core::add(mlx::core::multiply(input, cosine_full, stream), + mlx::core::multiply(rotated, sine_full, stream), stream); } bool tensor_offset_attention( const mlx::core::array &query, const mlx::core::array &key, const mlx::core::array &value, const mlx::core::array &k_cache, const mlx::core::array &v_cache, const mlx::core::array &offset, - float scale, int head_dim, float theta, const mlx::core::Device &device, + float scale, int head_dim, float theta, mlx::core::StreamOrDevice stream, mlx::core::array &attention, mlx::core::array &k_updated, mlx::core::array &v_updated, std::string &error) { if (!validate_qkv_cache_attention(query, key, value, k_cache, v_cache, 0, @@ -549,48 +549,48 @@ bool tensor_offset_attention( if (!validate_tensor_offset(offset, capacity, tokens, error)) return false; - auto safe_offset = clamp_offset(offset, capacity - tokens, device); - auto query_rope = rope_with_positions(query, safe_offset, head_dim, theta, device); - auto key_rope = rope_with_positions(key, safe_offset, head_dim, theta, device); - auto query_bn = mlx::core::transpose(query_rope, {0, 2, 1, 3}, device); - auto key_bn = mlx::core::transpose(key_rope, {0, 2, 1, 3}, device); - auto value_bn = mlx::core::transpose(value, {0, 2, 1, 3}, device); - auto start = mlx::core::reshape(safe_offset, {1}, device); - k_updated = mlx::core::slice_update(k_cache, key_bn, start, {2}, device); - v_updated = mlx::core::slice_update(v_cache, value_bn, start, {2}, device); + auto safe_offset = clamp_offset(offset, capacity - tokens, stream); + auto query_rope = rope_with_positions(query, safe_offset, head_dim, theta, stream); + auto key_rope = rope_with_positions(key, safe_offset, head_dim, theta, stream); + auto query_bn = mlx::core::transpose(query_rope, {0, 2, 1, 3}, stream); + auto key_bn = mlx::core::transpose(key_rope, {0, 2, 1, 3}, stream); + auto value_bn = mlx::core::transpose(value, {0, 2, 1, 3}, stream); + auto start = mlx::core::reshape(safe_offset, {1}, stream); + k_updated = mlx::core::slice_update(k_cache, key_bn, start, {2}, stream); + v_updated = mlx::core::slice_update(v_cache, value_bn, start, {2}, stream); auto row = mlx::core::reshape( - mlx::core::add(mlx::core::arange(tokens, mlx::core::int32, device), - safe_offset, device), - {1, 1, tokens, 1}, device); + mlx::core::add(mlx::core::arange(tokens, mlx::core::int32, stream), + safe_offset, stream), + {1, 1, tokens, 1}, stream); auto column = mlx::core::reshape( - mlx::core::arange(capacity, mlx::core::int32, device), - {1, 1, 1, capacity}, device); - auto visible = mlx::core::less_equal(column, row, device); + mlx::core::arange(capacity, mlx::core::int32, stream), + {1, 1, 1, capacity}, stream); + auto visible = mlx::core::less_equal(column, row, stream); auto mask = mlx::core::where( - visible, mlx::core::zeros({}, query.dtype(), device), + visible, mlx::core::zeros({}, query.dtype(), stream), mlx::core::full({}, -std::numeric_limits::infinity(), query.dtype(), - device), - device); + stream), + stream); auto attended = mlx::core::fast::scaled_dot_product_attention( query_bn, k_updated, v_updated, scale, "array", mask, std::nullopt, - device); + stream); attention = mlx::core::reshape( - mlx::core::transpose(attended, {0, 2, 1, 3}, device), - {batch, tokens, query_heads * width}, device); + mlx::core::transpose(attended, {0, 2, 1, 3}, stream), + {batch, tokens, query_heads * width}, stream); (void)kv_heads; return true; } mlx::core::array sdpa(const mlx::core::array &q_rope, const mlx::core::array &k_valid, const mlx::core::array &v_valid, float scale, int T_new, - int valid_len, const mlx::core::Device &device) { + int valid_len, mlx::core::StreamOrDevice stream) { return (T_new == 1) ? mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, device) + q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, stream) : mlx::core::fast::scaled_dot_product_attention( q_rope, k_valid, v_valid, scale, "array", - build_prefill_mask(q_rope, T_new, valid_len, device), std::nullopt, device); + build_prefill_mask(q_rope, T_new, valid_len, stream), std::nullopt, stream); } // ── Generalized (dense-or-quantized) per-layer compute ─────────────────── @@ -599,7 +599,7 @@ mlx::core::array sdpa(const mlx::core::array &q_rope, const mlx::core::array &k_ mlx::core::array layer_core_generalized(const mlx::core::array &hidden, const LayerParamsQ &layer, KVCache &kv, int offset, float scale, int head_dim, float theta, float eps, - const mlx::core::Device &device, + mlx::core::StreamOrDevice stream, mlx::core::array *k_out, mlx::core::array *v_out) { int B = hidden.shape(0); int T_new = hidden.shape(1); @@ -609,61 +609,61 @@ mlx::core::array layer_core_generalized(const mlx::core::array &hidden, int attn_width = N_q * D; int valid_len = offset + T_new; - auto xn = mlx::core::fast::rms_norm(hidden, *layer.norm1, eps, device); - auto q_flat = apply_linear(xn, layer.q_proj, device); - auto k_flat = apply_linear(xn, layer.k_proj, device); - auto v_flat = apply_linear(xn, layer.v_proj, device); + auto xn = mlx::core::fast::rms_norm(hidden, *layer.norm1, eps, stream); + auto q_flat = apply_linear(xn, layer.q_proj, stream); + auto k_flat = apply_linear(xn, layer.k_proj, stream); + auto v_flat = apply_linear(xn, layer.v_proj, stream); - auto q = mlx::core::reshape(q_flat, {B, T_new, N_q, D}, device); - auto k = mlx::core::reshape(k_flat, {B, T_new, N_kv, D}, device); - auto v = mlx::core::reshape(v_flat, {B, T_new, N_kv, D}, device); + auto q = mlx::core::reshape(q_flat, {B, T_new, N_q, D}, stream); + auto k = mlx::core::reshape(k_flat, {B, T_new, N_kv, D}, stream); + auto v = mlx::core::reshape(v_flat, {B, T_new, N_kv, D}, stream); - q = mlx::core::fast::rms_norm(q, *layer.q_norm, eps, device); - k = mlx::core::fast::rms_norm(k, *layer.k_norm, eps, device); + q = mlx::core::fast::rms_norm(q, *layer.q_norm, eps, stream); + k = mlx::core::fast::rms_norm(k, *layer.k_norm, eps, stream); - auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, device); - auto k_bn = mlx::core::transpose(k, {0, 2, 1, 3}, device); - auto v_bn = mlx::core::transpose(v, {0, 2, 1, 3}, device); + auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, stream); + auto k_bn = mlx::core::transpose(k, {0, 2, 1, 3}, stream); + auto v_bn = mlx::core::transpose(v, {0, 2, 1, 3}, stream); - auto q_rope = mlx::core::fast::rope(q_bn, D, false, theta, 1.0f, offset, std::nullopt, device); - auto k_rope = mlx::core::fast::rope(k_bn, D, false, theta, 1.0f, offset, std::nullopt, device); + auto q_rope = mlx::core::fast::rope(q_bn, D, false, theta, 1.0f, offset, std::nullopt, stream); + auto k_rope = mlx::core::fast::rope(k_bn, D, false, theta, 1.0f, offset, std::nullopt, stream); auto k_cache_owned = std::move(*kv.k); auto v_cache_owned = std::move(*kv.v); auto k_upd = mlx::core::slice_update(k_cache_owned, k_rope, to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto v_upd = mlx::core::slice_update(v_cache_owned, v_bn, to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto k_valid = mlx::core::slice(k_upd, to_shape({0, 0, 0, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto v_valid = mlx::core::slice(v_upd, to_shape({0, 0, 0, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); - auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, scale, T_new, valid_len, device); - auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, device); - auto attn_out = mlx::core::reshape(attn_out_bthd, {B, T_new, attn_width}, device); - auto attn_projected = apply_linear(attn_out, layer.o_proj, device); - auto attn_hidden = mlx::core::add(hidden, attn_projected, device); + auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, scale, T_new, valid_len, stream); + auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, stream); + auto attn_out = mlx::core::reshape(attn_out_bthd, {B, T_new, attn_width}, stream); + auto attn_projected = apply_linear(attn_out, layer.o_proj, stream); + auto attn_hidden = mlx::core::add(hidden, attn_projected, stream); - auto xn2 = mlx::core::fast::rms_norm(attn_hidden, *layer.norm2, eps, device); - auto gate = apply_linear(xn2, layer.gate_proj, device); - auto up = apply_linear(xn2, layer.up_proj, device); - auto mlp = mlx::core::multiply(mlx::core::multiply(gate, mlx::core::sigmoid(gate, device), device), - up, device); - auto mlp_out = apply_linear(mlp, layer.down_proj, device); + auto xn2 = mlx::core::fast::rms_norm(attn_hidden, *layer.norm2, eps, stream); + auto gate = apply_linear(xn2, layer.gate_proj, stream); + auto up = apply_linear(xn2, layer.up_proj, stream); + auto mlp = mlx::core::multiply(mlx::core::multiply(gate, mlx::core::sigmoid(gate, stream), stream), + up, stream); + auto mlp_out = apply_linear(mlp, layer.down_proj, stream); if (k_out != nullptr) *k_out = k_upd; if (v_out != nullptr) *v_out = v_upd; - return mlx::core::add(attn_hidden, mlp_out, device); + return mlx::core::add(attn_hidden, mlp_out, stream); } // ── Dense per-layer compute ─────────────────────────────────────────────── mlx::core::array layer_dense_impl(const mlx::core::array &hidden, const LayerParams &layer, KVCache &kv, int offset, float scale, int head_dim, - float theta, float eps, const mlx::core::Device &device, + float theta, float eps, mlx::core::StreamOrDevice stream, mlx::core::array *k_out, mlx::core::array *v_out) { int B = hidden.shape(0); int T_new = hidden.shape(1); @@ -673,55 +673,55 @@ mlx::core::array layer_dense_impl(const mlx::core::array &hidden, const LayerPar int attn_width = N_q * D; int valid_len = offset + T_new; - auto xn = mlx::core::fast::rms_norm(hidden, *layer.norm1, eps, device); - auto q_flat = linear_in_out(xn, *layer.q_proj, device); - auto k_flat = linear_in_out(xn, *layer.k_proj, device); - auto v_flat = linear_in_out(xn, *layer.v_proj, device); + auto xn = mlx::core::fast::rms_norm(hidden, *layer.norm1, eps, stream); + auto q_flat = linear_in_out(xn, *layer.q_proj, stream); + auto k_flat = linear_in_out(xn, *layer.k_proj, stream); + auto v_flat = linear_in_out(xn, *layer.v_proj, stream); - auto q = mlx::core::reshape(q_flat, {B, T_new, N_q, D}, device); - auto k = mlx::core::reshape(k_flat, {B, T_new, N_kv, D}, device); - auto v = mlx::core::reshape(v_flat, {B, T_new, N_kv, D}, device); + auto q = mlx::core::reshape(q_flat, {B, T_new, N_q, D}, stream); + auto k = mlx::core::reshape(k_flat, {B, T_new, N_kv, D}, stream); + auto v = mlx::core::reshape(v_flat, {B, T_new, N_kv, D}, stream); - q = mlx::core::fast::rms_norm(q, *layer.q_norm, eps, device); - k = mlx::core::fast::rms_norm(k, *layer.k_norm, eps, device); + q = mlx::core::fast::rms_norm(q, *layer.q_norm, eps, stream); + k = mlx::core::fast::rms_norm(k, *layer.k_norm, eps, stream); - auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, device); - auto k_bn = mlx::core::transpose(k, {0, 2, 1, 3}, device); - auto v_bn = mlx::core::transpose(v, {0, 2, 1, 3}, device); + auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, stream); + auto k_bn = mlx::core::transpose(k, {0, 2, 1, 3}, stream); + auto v_bn = mlx::core::transpose(v, {0, 2, 1, 3}, stream); - auto q_rope = mlx::core::fast::rope(q_bn, D, false, theta, 1.0f, offset, std::nullopt, device); - auto k_rope = mlx::core::fast::rope(k_bn, D, false, theta, 1.0f, offset, std::nullopt, device); + auto q_rope = mlx::core::fast::rope(q_bn, D, false, theta, 1.0f, offset, std::nullopt, stream); + auto k_rope = mlx::core::fast::rope(k_bn, D, false, theta, 1.0f, offset, std::nullopt, stream); auto k_cache_owned = std::move(*kv.k); auto v_cache_owned = std::move(*kv.v); auto k_upd = mlx::core::slice_update(k_cache_owned, k_rope, to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto v_upd = mlx::core::slice_update(v_cache_owned, v_bn, to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto k_valid = mlx::core::slice(k_upd, to_shape({0, 0, 0, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto v_valid = mlx::core::slice(v_upd, to_shape({0, 0, 0, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); - auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, scale, T_new, valid_len, device); - auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, device); - auto attn_out = mlx::core::reshape(attn_out_bthd, {B, T_new, attn_width}, device); - auto attn_projected = linear_in_out(attn_out, *layer.o_proj, device); - auto attn_hidden = mlx::core::add(hidden, attn_projected, device); + auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, scale, T_new, valid_len, stream); + auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, stream); + auto attn_out = mlx::core::reshape(attn_out_bthd, {B, T_new, attn_width}, stream); + auto attn_projected = linear_in_out(attn_out, *layer.o_proj, stream); + auto attn_hidden = mlx::core::add(hidden, attn_projected, stream); - auto xn2 = mlx::core::fast::rms_norm(attn_hidden, *layer.norm2, eps, device); - auto gate = linear_in_out(xn2, *layer.gate_proj, device); - auto up = linear_in_out(xn2, *layer.up_proj, device); - auto mlp = mlx::core::multiply(mlx::core::multiply(gate, mlx::core::sigmoid(gate, device), device), - up, device); - auto mlp_out = linear_in_out(mlp, *layer.down_proj, device); + auto xn2 = mlx::core::fast::rms_norm(attn_hidden, *layer.norm2, eps, stream); + auto gate = linear_in_out(xn2, *layer.gate_proj, stream); + auto up = linear_in_out(xn2, *layer.up_proj, stream); + auto mlp = mlx::core::multiply(mlx::core::multiply(gate, mlx::core::sigmoid(gate, stream), stream), + up, stream); + auto mlp_out = linear_in_out(mlp, *layer.down_proj, stream); if (k_out != nullptr) *k_out = k_upd; if (v_out != nullptr) *v_out = v_upd; - return mlx::core::add(attn_hidden, mlp_out, device); + return mlx::core::add(attn_hidden, mlp_out, stream); } // ── VTable entrypoints ──────────────────────────────────────────────────── @@ -729,7 +729,7 @@ mlx::core::array layer_dense_impl(const mlx::core::array &hidden, const LayerPar bool v_kv_cache_attention(const mlx::core::array &q, const mlx::core::array &new_k, const mlx::core::array &new_v, mlx::core::array &k_cache, mlx::core::array &v_cache, int offset, double scale, int head_dim, - double theta, const mlx::core::Device &device, mlx::core::array &out, + double theta, mlx::core::StreamOrDevice stream, mlx::core::array &out, mlx::core::array &k_upd, mlx::core::array &v_upd, std::string &error) { try { @@ -745,31 +745,31 @@ bool v_kv_cache_attention(const mlx::core::array &q, const mlx::core::array &new int N_kv = new_k.shape(2); int valid_len = offset + T_new; - auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, device); - auto k_bn = mlx::core::transpose(new_k, {0, 2, 1, 3}, device); - auto v_bn = mlx::core::transpose(new_v, {0, 2, 1, 3}, device); + auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, stream); + auto k_bn = mlx::core::transpose(new_k, {0, 2, 1, 3}, stream); + auto v_bn = mlx::core::transpose(new_v, {0, 2, 1, 3}, stream); auto q_rope = - mlx::core::fast::rope(q_bn, head_dim, false, (float)theta, 1.0f, offset, std::nullopt, device); + mlx::core::fast::rope(q_bn, head_dim, false, (float)theta, 1.0f, offset, std::nullopt, stream); auto k_rope = - mlx::core::fast::rope(k_bn, head_dim, false, (float)theta, 1.0f, offset, std::nullopt, device); + mlx::core::fast::rope(k_bn, head_dim, false, (float)theta, 1.0f, offset, std::nullopt, stream); auto k_cache_owned = std::move(k_cache); auto v_cache_owned = std::move(v_cache); k_upd = mlx::core::slice_update(k_cache_owned, k_rope, to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); v_upd = mlx::core::slice_update(v_cache_owned, v_bn, to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto k_valid = mlx::core::slice(k_upd, to_shape({0, 0, 0, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto v_valid = mlx::core::slice(v_upd, to_shape({0, 0, 0, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); - auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, (float)scale, T_new, valid_len, device); - auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, device); - out = mlx::core::reshape(attn_out_bthd, {B, T_new, N_q * D}, device); + auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, (float)scale, T_new, valid_len, stream); + auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, stream); + out = mlx::core::reshape(attn_out_bthd, {B, T_new, N_q * D}, stream); return true; } catch (const std::exception &e) { error = e.what(); @@ -782,7 +782,7 @@ bool v_kv_cache_attention(const mlx::core::array &q, const mlx::core::array &new bool v_mlp(const mlx::core::array &hidden, const mlx::core::array &norm, const mlx::core::array &gate_proj, const mlx::core::array &up_proj, - const mlx::core::array &down_proj, double eps, const mlx::core::Device &device, + const mlx::core::array &down_proj, double eps, mlx::core::StreamOrDevice stream, mlx::core::array &out, std::string &error) { try { if (!check_rank3_positive(hidden, "hidden", error)) { @@ -801,13 +801,13 @@ bool v_mlp(const mlx::core::array &hidden, const mlx::core::array &norm, return false; } - auto xn = mlx::core::fast::rms_norm(hidden, norm, (float)eps, device); - auto gate = linear_in_out(xn, gate_proj, device); - auto up = linear_in_out(xn, up_proj, device); - auto mlp = mlx::core::multiply(mlx::core::multiply(gate, mlx::core::sigmoid(gate, device), device), - up, device); - auto proj = linear_in_out(mlp, down_proj, device); - out = mlx::core::add(hidden, proj, device); + auto xn = mlx::core::fast::rms_norm(hidden, norm, (float)eps, stream); + auto gate = linear_in_out(xn, gate_proj, stream); + auto up = linear_in_out(xn, up_proj, stream); + auto mlp = mlx::core::multiply(mlx::core::multiply(gate, mlx::core::sigmoid(gate, stream), stream), + up, stream); + auto proj = linear_in_out(mlp, down_proj, stream); + out = mlx::core::add(hidden, proj, stream); return true; } catch (const std::exception &e) { error = e.what(); @@ -819,7 +819,7 @@ bool v_mlp(const mlx::core::array &hidden, const mlx::core::array &norm, } bool v_attention_residual(const mlx::core::array &hidden, const mlx::core::array &attn_out, - const mlx::core::array &o_proj, const mlx::core::Device &device, + const mlx::core::array &o_proj, mlx::core::StreamOrDevice stream, mlx::core::array &out, std::string &error) { try { if (!check_rank3_positive(hidden, "hidden", error) || @@ -836,8 +836,8 @@ bool v_attention_residual(const mlx::core::array &hidden, const mlx::core::array !check_dim(o_proj, 1, H, "o_proj", "output width", error)) { return false; } - auto projected = linear_in_out(attn_out, o_proj, device); - out = mlx::core::add(hidden, projected, device); + auto projected = linear_in_out(attn_out, o_proj, stream); + out = mlx::core::add(hidden, projected, stream); return true; } catch (const std::exception &e) { error = e.what(); @@ -854,7 +854,7 @@ bool v_attention_block(const mlx::core::array &hidden, const mlx::core::array &n const mlx::core::array &q_norm, const mlx::core::array &k_norm, mlx::core::array &k_cache, mlx::core::array &v_cache, int offset, double scale, int head_dim, double theta, double eps, - const mlx::core::Device &device, mlx::core::array &out, + mlx::core::StreamOrDevice stream, mlx::core::array &out, mlx::core::array &k_upd, mlx::core::array &v_upd, std::string &error) { try { if (!check_rank3_positive(hidden, "hidden", error) || @@ -903,45 +903,45 @@ bool v_attention_block(const mlx::core::array &hidden, const mlx::core::array &n } int valid_len = offset + T_new; - auto xn = mlx::core::fast::rms_norm(hidden, norm, (float)eps, device); - auto q_flat = linear_in_out(xn, q_proj, device); - auto k_flat = linear_in_out(xn, k_proj, device); - auto v_flat = linear_in_out(xn, v_proj, device); + auto xn = mlx::core::fast::rms_norm(hidden, norm, (float)eps, stream); + auto q_flat = linear_in_out(xn, q_proj, stream); + auto k_flat = linear_in_out(xn, k_proj, stream); + auto v_flat = linear_in_out(xn, v_proj, stream); - auto q = mlx::core::reshape(q_flat, {B, T_new, N_q, D}, device); - auto k = mlx::core::reshape(k_flat, {B, T_new, N_kv, D}, device); - auto v = mlx::core::reshape(v_flat, {B, T_new, N_kv, D}, device); + auto q = mlx::core::reshape(q_flat, {B, T_new, N_q, D}, stream); + auto k = mlx::core::reshape(k_flat, {B, T_new, N_kv, D}, stream); + auto v = mlx::core::reshape(v_flat, {B, T_new, N_kv, D}, stream); - q = mlx::core::fast::rms_norm(q, q_norm, (float)eps, device); - k = mlx::core::fast::rms_norm(k, k_norm, (float)eps, device); + q = mlx::core::fast::rms_norm(q, q_norm, (float)eps, stream); + k = mlx::core::fast::rms_norm(k, k_norm, (float)eps, stream); - auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, device); - auto k_bn = mlx::core::transpose(k, {0, 2, 1, 3}, device); - auto v_bn = mlx::core::transpose(v, {0, 2, 1, 3}, device); + auto q_bn = mlx::core::transpose(q, {0, 2, 1, 3}, stream); + auto k_bn = mlx::core::transpose(k, {0, 2, 1, 3}, stream); + auto v_bn = mlx::core::transpose(v, {0, 2, 1, 3}, stream); auto q_rope = - mlx::core::fast::rope(q_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, device); + mlx::core::fast::rope(q_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, stream); auto k_rope = - mlx::core::fast::rope(k_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, device); + mlx::core::fast::rope(k_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, stream); auto k_cache_owned = std::move(k_cache); auto v_cache_owned = std::move(v_cache); k_upd = mlx::core::slice_update(k_cache_owned, k_rope, to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); v_upd = mlx::core::slice_update(v_cache_owned, v_bn, to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto k_valid = mlx::core::slice(k_upd, to_shape({0, 0, 0, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); auto v_valid = mlx::core::slice(v_upd, to_shape({0, 0, 0, 0}), - to_shape({B, N_kv, valid_len, D}), device); + to_shape({B, N_kv, valid_len, D}), stream); - auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, (float)scale, T_new, valid_len, device); - auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, device); - auto attn_out = mlx::core::reshape(attn_out_bthd, {B, T_new, attn_width}, device); - auto projected = linear_in_out(attn_out, o_proj, device); - out = mlx::core::add(hidden, projected, device); + auto attn_out_bn = sdpa(q_rope, k_valid, v_valid, (float)scale, T_new, valid_len, stream); + auto attn_out_bthd = mlx::core::transpose(attn_out_bn, {0, 2, 1, 3}, stream); + auto attn_out = mlx::core::reshape(attn_out_bthd, {B, T_new, attn_width}, stream); + auto projected = linear_in_out(attn_out, o_proj, stream); + out = mlx::core::add(hidden, projected, stream); return true; } catch (const std::exception &e) { error = e.what(); @@ -954,14 +954,14 @@ bool v_attention_block(const mlx::core::array &hidden, const mlx::core::array &n bool v_layer_dense(const mlx::core::array &hidden, const LayerParams &layer, KVCache &kv, int offset, double scale, int head_dim, double theta, double eps, - const mlx::core::Device &device, mlx::core::array &out, + mlx::core::StreamOrDevice stream, mlx::core::array &out, mlx::core::array &k_upd, mlx::core::array &v_upd, std::string &error) { try { if (!validate_dense_layer(hidden, layer, kv, offset, head_dim, error)) { return false; } out = layer_dense_impl(hidden, layer, kv, offset, (float)scale, head_dim, (float)theta, - (float)eps, device, &k_upd, &v_upd); + (float)eps, stream, &k_upd, &v_upd); return true; } catch (const std::exception &e) { error = e.what(); @@ -974,14 +974,14 @@ bool v_layer_dense(const mlx::core::array &hidden, const LayerParams &layer, KVC bool v_layer_quantized(const mlx::core::array &hidden, const LayerParamsQ &layer, KVCache &kv, int offset, double scale, int head_dim, double theta, double eps, - const mlx::core::Device &device, mlx::core::array &out, + mlx::core::StreamOrDevice stream, mlx::core::array &out, mlx::core::array &k_upd, mlx::core::array &v_upd, std::string &error) { try { if (!validate_generalized_layer(hidden, layer, kv, offset, head_dim, error)) { return false; } out = layer_core_generalized(hidden, layer, kv, offset, (float)scale, head_dim, (float)theta, - (float)eps, device, &k_upd, &v_upd); + (float)eps, stream, &k_upd, &v_upd); return true; } catch (const std::exception &e) { error = e.what(); @@ -993,7 +993,7 @@ bool v_layer_quantized(const mlx::core::array &hidden, const LayerParamsQ &layer } bool v_final_greedy(const mlx::core::array &hidden, const mlx::core::array &norm, - const mlx::core::array &lm_head, double eps, const mlx::core::Device &device, + const mlx::core::array &lm_head, double eps, mlx::core::StreamOrDevice stream, mlx::core::array &out, std::string &error) { try { if (!check_rank3_positive(hidden, "hidden", error)) { @@ -1008,15 +1008,15 @@ bool v_final_greedy(const mlx::core::array &hidden, const mlx::core::array &norm return false; } - auto last = (T == 1) ? mlx::core::reshape(hidden, {B, H}, device) + auto last = (T == 1) ? mlx::core::reshape(hidden, {B, H}, stream) : mlx::core::reshape(mlx::core::slice(hidden, to_shape({0, T - 1, 0}), - to_shape({B, T, H}), device), - {B, H}, device); + to_shape({B, T, H}), stream), + {B, H}, stream); - auto normed = mlx::core::fast::rms_norm(last, norm, (float)eps, device); + auto normed = mlx::core::fast::rms_norm(last, norm, (float)eps, stream); auto logits = - mlx::core::tensordot(normed, lm_head, std::vector{1}, std::vector{1}, device); - out = mlx::core::argmax(logits, 1, false, device); + mlx::core::tensordot(normed, lm_head, std::vector{1}, std::vector{1}, stream); + out = mlx::core::argmax(logits, 1, false, stream); return true; } catch (const std::exception &e) { error = e.what(); @@ -1031,7 +1031,7 @@ bool v_forward_greedy_from_hidden(const mlx::core::array &hidden, std::vector &kv, const mlx::core::array &norm, const mlx::core::array &lm_head, int offset, double scale, int head_dim, double theta, double eps, bool return_token_id, - const mlx::core::Device &device, mlx::core::array &token_out, + mlx::core::StreamOrDevice stream, mlx::core::array &token_out, int64_t &token_id_out, std::vector &k_out, std::vector &v_out, std::string &error) { try { @@ -1063,7 +1063,7 @@ bool v_forward_greedy_from_hidden(const mlx::core::array &hidden, std::vector &token_out, std::vector &k_out, std::vector &v_out, std::string &error) { @@ -1169,9 +1169,9 @@ bool v_forward_greedy_ids_chunk(const mlx::core::array &input_ids, int B = current_ids.shape(0); int T = current_ids.shape(1); - auto ids = mlx::core::reshape(current_ids, {B * T}, device); - auto current = mlx::core::reshape(mlx::core::take(embed_tokens, ids, 0, device), - {B, T, embed_tokens.shape(1)}, device); + auto ids = mlx::core::reshape(current_ids, {B * T}, stream); + auto current = mlx::core::reshape(mlx::core::take(embed_tokens, ids, 0, stream), + {B, T, embed_tokens.shape(1)}, stream); next_k_cache.clear(); next_v_cache.clear(); @@ -1190,7 +1190,7 @@ bool v_forward_greedy_ids_chunk(const mlx::core::array &input_ids, mlx::core::array k_new = *kv.k; mlx::core::array v_new = *kv.v; current = layer_dense_impl(current, layers[layer_idx], kv, current_offset, (float)scale, - head_dim, (float)theta, (float)eps, device, &k_new, &v_new); + head_dim, (float)theta, (float)eps, stream, &k_new, &v_new); next_k_cache.push_back(k_new); next_v_cache.push_back(v_new); @@ -1202,14 +1202,14 @@ bool v_forward_greedy_ids_chunk(const mlx::core::array &input_ids, auto last = (T_out == 1) - ? mlx::core::reshape(current, {B_out, H_out}, device) + ? mlx::core::reshape(current, {B_out, H_out}, stream) : mlx::core::reshape(mlx::core::slice(current, to_shape({0, T_out - 1, 0}), - to_shape({B_out, T_out, H_out}), device), - {B_out, H_out}, device); + to_shape({B_out, T_out, H_out}), stream), + {B_out, H_out}, stream); - auto normed = mlx::core::fast::rms_norm(last, norm, (float)eps, device); - auto logits = linear_out_in(normed, lm_head, device); - auto token = mlx::core::argmax(logits, 1, false, device); + auto normed = mlx::core::fast::rms_norm(last, norm, (float)eps, stream); + auto logits = linear_out_in(normed, lm_head, stream); + auto token = mlx::core::argmax(logits, 1, false, stream); if (submit_each_step) { std::vector eval_arrays; @@ -1223,7 +1223,7 @@ bool v_forward_greedy_ids_chunk(const mlx::core::array &input_ids, } token_arrays.push_back(token); - current_ids = mlx::core::reshape(token, {B_out, 1}, device); + current_ids = mlx::core::reshape(token, {B_out, 1}, stream); k_cache.swap(next_k_cache); v_cache.swap(next_v_cache); current_offset += 1; @@ -1246,7 +1246,7 @@ bool v_forward_greedy_ids_chunk_quantized( const mlx::core::array &input_ids, const mlx::core::array &embed_tokens, std::vector &layers, std::vector &initial_kv, const mlx::core::array &norm, const LinearWeight &lm_head, int offset, int count, - double scale, int head_dim, double theta, double eps, const mlx::core::Device &device, + double scale, int head_dim, double theta, double eps, mlx::core::StreamOrDevice stream, std::vector &token_out, std::vector &k_out, std::vector &v_out, std::string &error) { try { @@ -1296,9 +1296,9 @@ bool v_forward_greedy_ids_chunk_quantized( int B = current_ids.shape(0); int T = current_ids.shape(1); - auto ids = mlx::core::reshape(current_ids, {B * T}, device); - auto current = mlx::core::reshape(mlx::core::take(embed_tokens, ids, 0, device), - {B, T, embed_tokens.shape(1)}, device); + auto ids = mlx::core::reshape(current_ids, {B * T}, stream); + auto current = mlx::core::reshape(mlx::core::take(embed_tokens, ids, 0, stream), + {B, T, embed_tokens.shape(1)}, stream); next_k_cache.clear(); next_v_cache.clear(); @@ -1318,7 +1318,7 @@ bool v_forward_greedy_ids_chunk_quantized( mlx::core::array v_new = *kv.v; current = layer_core_generalized(current, layers[layer_idx], kv, current_offset, (float)scale, head_dim, (float)theta, (float)eps, - device, &k_new, &v_new); + stream, &k_new, &v_new); next_k_cache.push_back(k_new); next_v_cache.push_back(v_new); @@ -1330,17 +1330,17 @@ bool v_forward_greedy_ids_chunk_quantized( auto last = (T_out == 1) - ? mlx::core::reshape(current, {B_out, H_out}, device) + ? mlx::core::reshape(current, {B_out, H_out}, stream) : mlx::core::reshape(mlx::core::slice(current, to_shape({0, T_out - 1, 0}), - to_shape({B_out, T_out, H_out}), device), - {B_out, H_out}, device); + to_shape({B_out, T_out, H_out}), stream), + {B_out, H_out}, stream); - auto normed = mlx::core::fast::rms_norm(last, norm, (float)eps, device); - auto logits = apply_linear(normed, lm_head, device); - auto token = mlx::core::argmax(logits, 1, false, device); + auto normed = mlx::core::fast::rms_norm(last, norm, (float)eps, stream); + auto logits = apply_linear(normed, lm_head, stream); + auto token = mlx::core::argmax(logits, 1, false, stream); token_arrays.push_back(token); - current_ids = mlx::core::reshape(token, {B_out, 1}, device); + current_ids = mlx::core::reshape(token, {B_out, 1}, stream); k_cache.swap(next_k_cache); v_cache.swap(next_v_cache); current_offset += 1; @@ -1585,7 +1585,7 @@ plugin_mlp(const emlx::plugin::call_t &call, if (!v_mlp(call.operands[0], call.operands[1], call.operands[2], call.operands[3], call.operands[4], f64_from_bits(call.attrs[0]), - call.device, output, error)) + call.stream, output, error)) return error; outputs.push_back(std::move(output)); return std::nullopt; @@ -1612,7 +1612,7 @@ plugin_kv_cache_attention(const emlx::plugin::call_t &call, if (!v_kv_cache_attention( call.operands[0], call.operands[1], call.operands[2], k_cache, v_cache, offset, f64_from_bits(call.attrs[1]), head_dim, - f64_from_bits(call.attrs[3]), call.device, output, + f64_from_bits(call.attrs[3]), call.stream, output, k_updated, v_updated, error)) return error; outputs.push_back(std::move(output)); @@ -1639,7 +1639,7 @@ plugin_kv_cache_attention_tensor(const emlx::plugin::call_t &call, call.operands[0], call.operands[1], call.operands[2], call.operands[3], call.operands[4], call.operands[5], f64_from_bits(call.attrs[0]), head_dim, - f64_from_bits(call.attrs[2]), call.device, attention, + f64_from_bits(call.attrs[2]), call.stream, attention, k_updated, v_updated, error)) return error; outputs.push_back(std::move(attention)); @@ -1658,7 +1658,7 @@ plugin_attention_residual(const emlx::plugin::call_t &call, } auto output = call.operands[0]; if (!v_attention_residual(call.operands[0], call.operands[1], - call.operands[2], call.device, + call.operands[2], call.stream, output, error)) return error; outputs.push_back(std::move(output)); @@ -1689,7 +1689,7 @@ plugin_attention_block(const emlx::plugin::call_t &call, call.operands[6], call.operands[7], k_cache, v_cache, offset, f64_from_bits(call.attrs[1]), head_dim, f64_from_bits(call.attrs[3]), - f64_from_bits(call.attrs[4]), call.device, output, + f64_from_bits(call.attrs[4]), call.stream, output, k_updated, v_updated, error)) return error; outputs.push_back(std::move(output)); @@ -1724,7 +1724,7 @@ plugin_layer_dense(const emlx::plugin::call_t &call, if (!v_layer_dense( operands[0], layer, cache, offset, f64_from_bits(call.attrs[1]), head_dim, f64_from_bits(call.attrs[3]), - f64_from_bits(call.attrs[4]), call.device, output, + f64_from_bits(call.attrs[4]), call.stream, output, k_updated, v_updated, error)) return error; outputs.push_back(std::move(output)); @@ -1782,7 +1782,7 @@ plugin_layer_generalized(const emlx::plugin::call_t &call, call.operands[0], layer, cache, offset, f64_from_bits(call.attrs[2]), head_dim, f64_from_bits(call.attrs[4]), - f64_from_bits(call.attrs[5]), call.device, output, + f64_from_bits(call.attrs[5]), call.stream, output, k_updated, v_updated, error)) return error; outputs.push_back(std::move(output)); @@ -1802,7 +1802,7 @@ plugin_final_greedy(const emlx::plugin::call_t &call, auto output = call.operands[0]; if (!v_final_greedy(call.operands[0], call.operands[1], call.operands[2], f64_from_bits(call.attrs[0]), - call.device, output, error)) + call.stream, output, error)) return error; outputs.push_back(std::move(output)); return std::nullopt; @@ -1898,7 +1898,7 @@ plugin_forward_dense(const emlx::plugin::call_t &call, operands[0], layers, caches, operands[tail], operands[tail + 1], offset, f64_from_bits(call.attrs[2]), head_dim, f64_from_bits(call.attrs[4]), f64_from_bits(call.attrs[5]), - false, call.device, token, ignored_token_id, keys, values, + false, call.stream, token, ignored_token_id, keys, values, error)) return error; outputs.reserve(expected_outputs); @@ -1942,11 +1942,11 @@ plugin_chunk_dense(const emlx::plugin::call_t &call, operands[tail + 1], offset, count, f64_from_bits(call.attrs[3]), head_dim, f64_from_bits(call.attrs[5]), f64_from_bits(call.attrs[6]), call.attrs[7] == 1, - call.device, tokens, keys, values, error)) + call.stream, tokens, keys, values, error)) return error; outputs.reserve(expected_outputs); outputs.push_back(mlx::core::reshape( - mlx::core::stack(tokens, 0, call.device), {count}, call.device)); + mlx::core::stack(tokens, 0, call.stream), {count}, call.stream)); for (size_t index = 0; index < keys.size(); ++index) { outputs.push_back(std::move(keys[index])); outputs.push_back(std::move(values[index])); @@ -2027,12 +2027,12 @@ plugin_chunk_generalized(const emlx::plugin::call_t &call, call.operands[0], call.operands[1], layers, caches, norm, lm_head, offset, count, f64_from_bits(call.attrs[4]), head_dim, f64_from_bits(call.attrs[6]), - f64_from_bits(call.attrs[7]), call.device, tokens, + f64_from_bits(call.attrs[7]), call.stream, tokens, keys, values, error)) return error; outputs.reserve(expected_outputs); outputs.push_back(mlx::core::reshape( - mlx::core::stack(tokens, 0, call.device), {count}, call.device)); + mlx::core::stack(tokens, 0, call.stream), {count}, call.stream)); for (size_t index = 0; index < keys.size(); ++index) { outputs.push_back(std::move(keys[index])); outputs.push_back(std::move(values[index])); From 9a8c420bdef9055dbe5146c15073f04807405574 Mon Sep 17 00:00:00 2001 From: Paulo Valente <16843419+polvalente@users.noreply.github.com> Date: Tue, 21 Jul 2026 22:07:17 -0300 Subject: [PATCH 2/4] fix env --- .github/workflows/emlx.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/emlx.yml b/.github/workflows/emlx.yml index ac7203b..a182ad8 100644 --- a/.github/workflows/emlx.yml +++ b/.github/workflows/emlx.yml @@ -96,6 +96,10 @@ jobs: - { 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 From 774b7bdcec2a1803d8dce87b7a4a207565c04457 Mon Sep 17 00:00:00 2001 From: Paulo Valente <16843419+polvalente@users.noreply.github.com> Date: Tue, 21 Jul 2026 22:50:51 -0300 Subject: [PATCH 3/4] chore: remove flaky test --- emlx/test/emlx/native/expr_test.exs | 44 ----------------------------- 1 file changed, 44 deletions(-) diff --git a/emlx/test/emlx/native/expr_test.exs b/emlx/test/emlx/native/expr_test.exs index 488c14c..bfab11a 100644 --- a/emlx/test/emlx/native/expr_test.exs +++ b/emlx/test/emlx/native/expr_test.exs @@ -3439,44 +3439,6 @@ defmodule EMLX.Native.ExprTest do end end - test "decode-shaped block: fused path improves over primitive replay" do - # A small attention+norm decode step: RMSNorm → causal SDPA → RMSNorm. - scale = 0.125 - - fused = fn q, k, v, w -> - a = EMLX.Fast.scaled_dot_product_attention_causal(q, k, v, scale) - flat = Nx.reshape(a, {1, 16}) - EMLX.Fast.rms_norm(flat, w, 1.0e-5) - end - - primitive = fn q, k, v, w -> - scores = Nx.dot(q, [3], [0, 1], k, [3], [0, 1]) |> Nx.multiply(scale) - a = Nx.dot(normalize_rows(Nx.exp(scores)), [3], [0, 1], v, [2], [0, 1]) - flat = Nx.reshape(a, {1, 16}) - rms = Nx.sqrt(Nx.add(Nx.mean(Nx.pow(flat, 2), axes: [-1], keep_axes: true), 1.0e-5)) - Nx.divide(flat, rms) |> Nx.multiply(w) - end - - q = Nx.iota({1, 2, 1, 8}, type: :f32) |> Nx.divide(100) |> gpu_t() - k = Nx.iota({1, 2, 4, 8}, type: :f32) |> Nx.divide(90) |> gpu_t() - v = Nx.iota({1, 2, 4, 8}, type: :f32) |> Nx.divide(80) |> gpu_t() - w = Nx.broadcast(Nx.tensor(1.0, type: :f32), {16}) |> gpu_t() - - fused_c = Nx.Defn.jit(fused, compiler: EMLX, device: :gpu) - prim_c = Nx.Defn.jit(primitive, compiler: EMLX, device: :gpu) - - # Correctness: same result within fused-kernel tolerance. - assert_all_close(fused_c.(q, k, v, w), prim_c.(q, k, v, w), tol: 1.0e-2) - - # Warm both compiled graphs, then time the replay-only hot path. - for _ <- 1..5, do: fused_c.(q, k, v, w) |> Nx.backend_transfer() - for _ <- 1..5, do: prim_c.(q, k, v, w) |> Nx.backend_transfer() - - fused_us = bench_us(200, fn -> fused_c.(q, k, v, w) |> Nx.backend_transfer() end) - prim_us = bench_us(200, fn -> prim_c.(q, k, v, w) |> Nx.backend_transfer() end) - - assert fused_us <= prim_us * 1.1 - end end describe "prefill RoPE (Metal)" do @@ -4486,10 +4448,4 @@ defmodule EMLX.Native.ExprTest do |> Enum.each(fn {av, bv} -> assert_in_delta(av, bv, tol) end) end - defp bench_us(n, fun) do - t0 = System.monotonic_time(:microsecond) - Enum.each(1..n, fn _ -> fun.() end) - t1 = System.monotonic_time(:microsecond) - (t1 - t0) / n - end end From 5791dc1294ecfe7c7d03e6125014885eb6673f47 Mon Sep 17 00:00:00 2001 From: Paulo Valente <16843419+polvalente@users.noreply.github.com> Date: Tue, 21 Jul 2026 23:07:46 -0300 Subject: [PATCH 4/4] format --- emlx/test/emlx/native/expr_test.exs | 2 -- 1 file changed, 2 deletions(-) diff --git a/emlx/test/emlx/native/expr_test.exs b/emlx/test/emlx/native/expr_test.exs index bfab11a..ccdc27f 100644 --- a/emlx/test/emlx/native/expr_test.exs +++ b/emlx/test/emlx/native/expr_test.exs @@ -3438,7 +3438,6 @@ defmodule EMLX.Native.ExprTest do assert_all_close(eager[[.., .., head..head, ..]], alone, tol: 1.0e-3) end end - end describe "prefill RoPE (Metal)" do @@ -4447,5 +4446,4 @@ defmodule EMLX.Native.ExprTest do Enum.zip(Nx.to_flat_list(a), Nx.to_flat_list(b)) |> Enum.each(fn {av, bv} -> assert_in_delta(av, bv, tol) end) end - end