diff --git a/.github/workflows/emlx.yml b/.github/workflows/emlx.yml index 8fd1392..f98e766 100644 --- a/.github/workflows/emlx.yml +++ b/.github/workflows/emlx.yml @@ -44,7 +44,15 @@ jobs: echo "${ELIXIR_PATH}" >> $GITHUB_PATH - name: Install system dependencies - run: sudo apt-get install -y libopenblas0 + run: | + sudo apt-get install -y libopenblas0 g++-10 + # libmlx (cocoa-xu/mlx-build) bakes its CPU JIT preamble with + # gcc-10 headers; MLX invokes the bare `g++` at runtime for JIT + # kernel compiles, so a newer distro-default g++ (e.g. g++-13 on + # ubuntu-latest) causes builtin-type redeclaration errors + # (ml-explore/mlx#3430). Pin `g++` to match until upstream ships + # a release with ml-explore/mlx#3463. + sudo ln -sf "$(command -v g++-10)" /usr/local/bin/g++ - name: Compile and check warnings working-directory: emlx diff --git a/.gitignore b/.gitignore index 14cee17..95a1ede 100644 --- a/.gitignore +++ b/.gitignore @@ -8,7 +8,7 @@ /deps/ # Where third-party dependencies like ExDoc output generated docs. -/doc/ +/**/docs/ # Ignore .fetch files in case you like to edit your project deps locally. /.fetch diff --git a/emlx/Makefile b/emlx/Makefile index 2e6e7b9..ff43d78 100644 --- a/emlx/Makefile +++ b/emlx/Makefile @@ -10,6 +10,14 @@ PRIV_DIR = $(MIX_APP_PATH)/priv BUILD_DIR = $(EMLX_CACHE_DIR)/emlx-$(EMLX_VERSION)-mlx-$(MLX_VERSION)$(MLX_VARIANT)/objs EMLX_SO = $(PRIV_DIR)/libemlx.so EMLX_LIB_DIR = $(PRIV_DIR)/mlx/lib +# Headers are copied alongside the lib (unlike the lib copy, these are not +# needed by EMLX itself at runtime) so out-of-tree native plugins loaded via +# `EMLX.NIF.load_plugin/2` (e.g. emlx_axon's qwen3 compute plugin) can build +# against the exact same libmlx this EMLX was compiled against, without +# duplicating EMLX's libmlx version/cache-dir resolution logic — see +# emlx_axon's Makefile/mix.exs, which locate these via +# `Application.app_dir(:emlx, "priv/mlx/...")`. +EMLX_INCLUDE_DIR = $(PRIV_DIR)/mlx/include # Only used if MLX_BUILD is true MLX_SRC_ARCHIVE = $(EMLX_CACHE_DIR)/mlx/$(MLX_VERSION)$(MLX_VARIANT).tar.gz @@ -20,8 +28,8 @@ MLX_SO ?= $(MLX_LIB_DIR)/libmlx.dylib $(info LIBMLX_ENABLE_DEBUG=$(LIBMLX_ENABLE_DEBUG)) # Build flags -CFLAGS = -fPIC -I$(call esc,$(ERTS_INCLUDE_DIR)) -I$(call esc,$(MLX_INCLUDE_DIR)) -Wall \ - -std=c++20 +CFLAGS = -fPIC -I$(call esc,$(ERTS_INCLUDE_DIR)) -I$(call esc,$(MLX_INCLUDE_DIR)) \ + -I$(call esc,$(FINE_INCLUDE_DIR)) -Wall -std=c++20 ifeq ($(LIBMLX_ENABLE_DEBUG),true) CFLAGS += -g CMAKE_BUILD_TYPE = Debug @@ -47,8 +55,8 @@ endif MAKE_JOBS ?= $(MAKE_DEFAULT_JOBS) # Source files -SOURCES = c_src/emlx_nif.cpp c_src/emlx_fast.cpp c_src/emlx_fast/qwen3.cpp -HEADERS = c_src/nx_nif_utils.hpp c_src/emlx_worker.hpp c_src/emlx_async.hpp c_src/emlx_nif_shared.hpp c_src/emlx_fast/qwen3.hpp +SOURCES = c_src/emlx_nif.cpp c_src/emlx_fast.cpp c_src/emlx_fast/qwen3.cpp c_src/emlx_compiler.cpp c_src/emlx_plugin_registry.cpp +HEADERS = c_src/nx_nif_utils.hpp c_src/emlx_worker.hpp c_src/emlx_async.hpp c_src/emlx_nif_shared.hpp c_src/emlx_plugin_registry.hpp c_src/emlx_fast/qwen3.hpp c_src/emlx_fast/qwen3_plugin_abi.hpp c_src/emlx_compiler.hpp c_src/emlx_runtime_call_bridge.hpp OBJECTS = $(patsubst c_src/%.cpp,$(call esc,$(BUILD_DIR))/%.o,$(SOURCES)) # Main targets @@ -100,6 +108,9 @@ $(call esc,$(EMLX_SO)): $(call esc,$(PRIV_DIR)) $(call esc,$(MLX_SO)) $(OBJECTS) @ echo "Copying MLX library to $(EMLX_LIB_DIR)" @ mkdir -p "$(EMLX_LIB_DIR)" @ cp -a "$(MLX_LIB_DIR)"/* "$(EMLX_LIB_DIR)" + @ echo "Copying MLX headers to $(EMLX_INCLUDE_DIR)" + @ mkdir -p "$(EMLX_INCLUDE_DIR)" + @ cp -a "$(MLX_INCLUDE_DIR)"/* "$(EMLX_INCLUDE_DIR)" $(CXX) $(OBJECTS) -o "$(EMLX_SO)" $(LDFLAGS) clean: diff --git a/emlx/README.md b/emlx/README.md index d89c6cc..053ea8c 100644 --- a/emlx/README.md +++ b/emlx/README.md @@ -60,6 +60,13 @@ Nx.Defn.default_options(compiler: EMLX) config :nx, :default_defn_options, compiler: EMLX ``` +### Compile-time debug flags + +Development-only assertion flags (`:enable_bounds_check`, `:detect_non_finites`, +`:compiler_debug`) are read at compile time via `Application.compile_env/3`. +See the [EMLX moduledoc](https://hexdocs.pm/emlx/EMLX.html#module-compile-time-debug-flags) +and `config/dev.exs` for how to enable them and what each flag does. + ### MLX binaries EMLX relies on the [MLX](https://github.com/ml-explore/mlx) library to function, and currently EMLX will download precompiled builds from [mlx-build](https://github.com/cocoa-xu/mlx-build). diff --git a/emlx/bench/svd_bench.exs b/emlx/bench/svd_bench.exs new file mode 100644 index 0000000..633fa83 --- /dev/null +++ b/emlx/bench/svd_bench.exs @@ -0,0 +1,255 @@ +# bench/svd_bench.exs +# +# Benchmarks `Nx.LinAlg.svd/1` across two MLX-backed Nx stacks — EMLX and Emily +# (ausimian/emily) — and, within each, across: +# +# - device: :cpu vs :gpu (GPU skipped automatically if Metal is unavailable) +# - compiler: :defn_evaluator — `compiler: Nx.Defn.Evaluator`, dispatching op-by-op +# onto whichever backend the input already lives on +# (EMLX.Backend or Emily.Backend) +# :defn_compiled — `compiler: EMLX`, which lowers the whole graph to +# a single fused native call +# :defn_native — `compiler: Emily.Compiler, native: true`, Emily's +# equivalent single-NIF fused replay +# +# Every lane is built *once* per `{n, device}` via `Nx.Defn.compile/3` and that +# one closure is replayed across all Benchee-measured iterations — the +# realistic "compile once, run many" pattern (see `EMLX`'s moduledoc), not a +# fresh retrace (and, pre-`native_lowerable_block?/2`, dispatch-key rehash) on +# every call, which is what repeated `Nx.Defn.jit_apply/3` would measure. +# +# This isolates two independent costs: library (EMLX vs Emily) and dispatch overhead +# (per-op NIF round-trips vs one fused native call), on both CPU and GPU/Metal. +# +# This script is standalone — it does not need to run inside the `emlx` Mix project. +# `Mix.install/2` pulls in EMLX from this checkout (via a `path:` dependency) and +# Emily from Hex (Apple Silicon only), plus Benchee for the actual measurement. +# +# Run with (from anywhere): +# elixir bench/svd_bench.exs +# +# Tunables (env vars): +# EMLX_SVD_SIZES=128,256,512,1024 (default: 128,256,512,1024) +# EMLX_SVD_TIME=2 (Benchee measurement time per lane, in seconds) +# EMLX_SVD_WARMUP_TIME=1 (Benchee warmup time per lane, in seconds; also +# primes the defn_compiled/defn_native JIT caches) + +apple_silicon? = + match?({:unix, :darwin}, :os.type()) and + :erlang.system_info(:system_architecture) + |> to_string() + |> String.starts_with?("aarch64") + +# Emily's precompiled NIFs are Apple Silicon macOS only (see its README's +# Prerequisites section), so it's only added to the install list there — on any +# other host this script benchmarks EMLX alone. +deps = + [ + {:emlx, path: Path.expand("..", __DIR__)}, + # Pin the same `nx` source as `emlx`'s own mix.exs (git main), overriding + # Emily's Hex-sourced `~> 0.12` requirement, which main satisfies. + {:nx, github: "elixir-nx/nx", branch: "main", sparse: "nx", override: true}, + {:benchee, "~> 1.3"} + ] ++ if(apple_silicon?, do: [{:emily, "~> 0.7"}], else: []) + +Mix.install(deps) + +# MLX's CPU backend JIT-compiles fused kernels via `popen`/`pclose`, which fails with +# `ECHILD` under the BEAM's default SIGCHLD=SIG_IGN disposition — see +# `EMLX.Application`'s moduledoc. `elixir` scripts own this trade-off just like our +# test suite does. +:os.set_signal(:sigchld, :default) + +defmodule Bench.SVD do + # The benchmarked computation is always exactly `Nx.LinAlg.svd/1` + # (`full_matrices?: true`) — every lane only varies the `Nx.Defn.compile/3` + # `opts` (`:compiler`, `:device`, `:native`). Each lane's runner is built + # *once* via `compile/3` (see `build/2`) and replayed across every + # Benchee-measured iteration, mirroring the "compile once, run many" pattern + # `EMLX`'s moduledoc recommends — not a fresh retrace per iteration, which + # is the worst case `Nx.Defn.jit_apply/3` would otherwise measure. + # + # The compiled closure is stashed in `:persistent_term` (keyed by `key`) + # rather than returned directly, because Benchee runs each job in its own + # process: closing over the closure itself would make the *spawn* copy it + # into that process, and `Nx.Defn.compile/3`'s closure keeps a reference to + # the full (un-lowered) `Nx.Block` computation graph — including any + # `default_expr` fallback, e.g. `Nx.LinAlg.svd/1`'s ~100-iteration + # Jacobi-rotation algorithm, a DAG whose nodes heavily share + # sub-expressions across iterations. Erlang's term-copy path doesn't + # preserve that sharing, so flattening the DAG into a tree to copy it + # blows up combinatorially and the spawn effectively never finishes. + # `:persistent_term` values are read by reference (no copy), so fetching + # the closure from inside the already-spawned job process sidesteps this + # entirely. + def build(key, template, opts) do + compiled = Nx.Defn.compile(&Nx.LinAlg.svd(&1, full_matrices?: true), [template], opts) + :persistent_term.put(key, compiled) + key + end + + # The `Nx.sum`/`Nx.to_number` calls aren't part of the compiled closure; + # they just force materialization of all three outputs so a lazily-elided, + # fused-but-unused output can't skew the timing. + def run(key, a) do + {u, s, vt} = :persistent_term.get(key).(a) + Nx.to_number(Nx.sum(u)) + Nx.to_number(Nx.sum(s)) + Nx.to_number(Nx.sum(vt)) + end +end + +sizes = + case System.get_env("EMLX_SVD_SIZES") do + nil -> [128, 256, 512, 1024] + csv -> csv |> String.split(",") |> Enum.map(&String.to_integer(String.trim(&1))) + end + +parse_seconds = fn env, default -> {f, ""} = Float.parse(System.get_env(env, default)); f end + +bench_time = parse_seconds.("EMLX_SVD_TIME", "2") +warmup_time = parse_seconds.("EMLX_SVD_WARMUP_TIME", "1") + +emily_available? = apple_silicon? and Code.ensure_loaded?(Emily.Backend) + +devices = + case EMLX.NIF.command_queue_new(:gpu) do + {:ok, _} -> + [:cpu, :gpu] + + {:error, reason} -> + IO.puts("==> GPU unavailable (#{inspect(reason)}), benchmarking CPU only.\n") + [:cpu] + end + +engines = [:emlx] ++ if(emily_available?, do: [:emily], else: []) + +unless emily_available? do + IO.puts("==> Emily unavailable on this host (Apple Silicon macOS only), skipping.\n") +end + +# Deterministic input, generated once on the host and transferred per-{engine,device} +# below — keeps the matrix identical across every combination. +random_matrix = fn n -> + key = Nx.Random.key(42) + {t, _key} = Nx.Random.uniform(key, shape: {n, n}, type: :f32) + t +end + +backend_for = fn + :emlx, device -> {EMLX.Backend, device: device} + :emily, device -> {Emily.Backend, device: device} +end + +compiled_opts_for = fn + :emlx, device -> [compiler: EMLX, device: device] + :emily, device -> [compiler: Emily.Compiler, native: true, device: device] +end + +# `emlx defn_evaluator` on `:gpu` has no native `full_matrices?: true` kernel +# (see `EMLX.Backend.block/4`'s `:gpu` clause), so `Nx.Defn.Evaluator` runs the +# `default_expr` fallback (a ~100-iteration Jacobi-rotation algorithm) fully +# eagerly, one NIF round-trip per op. That's *correct* but tens of seconds +# slow even at `n=128` and grows with `n`, so it's skipped by default — a +# single stuck-looking lane otherwise dominates the whole suite's wall time. +# Set `EMLX_SVD_INCLUDE_SLOW_GPU_EVALUATOR=1` to include it anyway. +include_slow_gpu_evaluator? = System.get_env("EMLX_SVD_INCLUDE_SLOW_GPU_EVALUATOR") == "1" + +# Each lane builds its `Nx.Defn.compile/3` closure once (for `n`'s template) +# and returns a runner that replays it — the benchmarked function itself +# always lives in `Bench.SVD`, never as an inline anonymous function. +# (`defn_evaluator` has no comparable per-call compile cost to amortize, but +# building it the same way keeps lane-construction code uniform.) +lanes_for = fn engine, device, n -> + native_key = {:svd_bench, engine, device, n, :native} + template = Nx.template({n, n}, {:f, 32}) + Bench.SVD.build(native_key, template, compiled_opts_for.(engine, device)) + + native_lane = [ + {(if engine == :emlx, do: :defn_compiled, else: :defn_native), + &Bench.SVD.run(native_key, &1)} + ] + + if engine == :emlx and device == :gpu and not include_slow_gpu_evaluator? do + native_lane + else + evaluator_key = {:svd_bench, engine, device, n, :defn_evaluator} + Bench.SVD.build(evaluator_key, template, compiler: Nx.Defn.Evaluator) + [{:defn_evaluator, &Bench.SVD.run(evaluator_key, &1)} | native_lane] + end +end + +IO.puts("==> SVD benchmark (EMLX vs Emily)") +IO.puts(" sizes: #{Enum.join(sizes, ", ")}") +IO.puts(" time: #{bench_time}s per lane (warmup: #{warmup_time}s)") +IO.puts(" devices: #{Enum.join(devices, ", ")}") +IO.puts(" engines: #{Enum.join(engines, ", ")}") + +unless include_slow_gpu_evaluator? do + IO.puts( + " (skipping emlx defn_evaluator on :gpu — eager Jacobi fallback, tens of " <> + "seconds+ per call; set EMLX_SVD_INCLUDE_SLOW_GPU_EVALUATOR=1 to include it)" + ) +end + +IO.puts("") + +suites = + for n <- sizes, device <- devices, into: %{} do + IO.puts("--- n=#{n} device=#{device} " <> String.duplicate("-", 40)) + + jobs = + for engine <- engines, + {mode, runner} <- lanes_for.(engine, device, n), + reduce: %{} do + jobs -> + a = Nx.backend_transfer(random_matrix.(n), backend_for.(engine, device)) + label = "#{engine} #{mode}" + Map.put(jobs, label, fn -> runner.(a) end) + end + + suite = + Benchee.run(jobs, + time: bench_time, + warmup: warmup_time, + memory_time: 0, + print: [benchmarking: false, fast_warning: false], + formatters: [Benchee.Formatters.Console] + ) + + {{n, device}, suite} + end + +# Cross-device speedup for the fused/native lane of each engine — Benchee only +# compares jobs within a single run, so CPU vs GPU (two separate runs) is +# summarized here from each run's own statistics. +median_ms = fn suite, job_name -> + case suite && Enum.find(suite.scenarios, &(&1.job_name == job_name)) do + nil -> nil + scenario -> scenario.run_time_data.statistics.median / 1_000_000 + end +end + +if :cpu in devices and :gpu in devices do + IO.puts("\n=== CPU vs GPU speedup (fused/native lane, median ms) ===\n") + + Enum.each(engines, fn engine -> + mode = if engine == :emlx, do: "defn_compiled", else: "defn_native" + job_name = "#{engine} #{mode}" + + Enum.each(sizes, fn n -> + cpu_ms = median_ms.(suites[{n, :cpu}], job_name) + gpu_ms = median_ms.(suites[{n, :gpu}], job_name) + + case {cpu_ms, gpu_ms} do + {c, g} when is_number(c) and is_number(g) and g > 0 -> + speedup = Float.round(c / g, 2) + + IO.puts( + " #{engine} n=#{n}: cpu=#{Float.round(c, 3)} ms, gpu=#{Float.round(g, 3)} ms, speedup=#{speedup}x" + ) + + _ -> + :ok + end + end) + end) +end diff --git a/emlx/c_src/emlx_async.hpp b/emlx/c_src/emlx_async.hpp index 7e7ec16..0447adc 100644 --- a/emlx/c_src/emlx_async.hpp +++ b/emlx/c_src/emlx_async.hpp @@ -60,11 +60,11 @@ #pragma once +#include "emlx_runtime_call_bridge.hpp" #include "emlx_worker.hpp" #include "erl_nif.h" #include "nx_nif_utils.hpp" -#include #include #include @@ -144,6 +144,27 @@ ERL_NIF_TERM async_dispatch(ErlNifEnv *env, int argc, op_argv = std::move(op_argv)]() mutable { ERL_NIF_TERM payload; try { + // Make this job's caller pid available to emlx::native's + // EMLXRuntimeCall primitive for the duration of the call — see + // emlx_runtime_call_bridge.hpp. Jobs *can* nest on a worker's one + // dedicated OS thread: a blocked runtime_call pumps this same + // worker's queue while waiting (Worker::pump_until, + // emlx_worker.hpp), so a job posted by that call's own real + // callback (e.g. EMLX.Quantization reentering EMLX on this same + // worker) may run to completion *before* this outer job does. + // Restore the *previous* pid on exit (not unconditionally + // nullptr) so unwinding back out of a nested job doesn't clobber + // the outer job's still-in-flight caller pid. + struct CallerPidGuard { + ErlNifPid pid; + ErlNifPid *previous; + explicit CallerPidGuard(ErlNifPid p) : pid(p) { + previous = emlx::g_current_caller_pid; + emlx::g_current_caller_pid = &pid; + } + ~CallerPidGuard() { emlx::g_current_caller_pid = previous; } + } caller_pid_guard(caller_pid); + payload = SyncOp(msg_env, static_cast(op_argv.size()), op_argv.data()); } catch (...) { diff --git a/emlx/c_src/emlx_compiler.cpp b/emlx/c_src/emlx_compiler.cpp new file mode 100644 index 0000000..46995fc --- /dev/null +++ b/emlx/c_src/emlx_compiler.cpp @@ -0,0 +1,1938 @@ +// emlx_compiler.cpp — implements emlx::native compile/eval NIF logic. +// +// compile_program bakes the program into a capturing lambda and wraps it with +// mlx::core::compile(), so MLX traces the computation graph on first eval and +// replays the cached compiled graph on subsequent calls. eval_program is a +// thin caller: compiled_fn(inputs) → eval → wrap outputs. +// +// Op dispatch uses a string→function registry instead of an integer opcode enum. +// Adding a new op: register it in `op_registry` below. No enum, no wire +// integers, no lockstep parity table to maintain. + +#include "emlx_compiler.hpp" +#include "emlx_runtime_call_bridge.hpp" +#include "mlx/allocator.h" +#include "mlx/compile_impl.h" +#include "mlx/primitives.h" +#include +#include +#include +#include +#include +#include +#include +#include + +namespace emlx { +namespace native { + +// ── Op registry ─────────────────────────────────────────────────────────────── +// +// Each entry maps an op name string (matching the atom used in the Elixir IR) +// to a C++ function: (resolved_operands, integer_attrs) → result array. +// `operands` are already-resolved mlx::core::arrays; `attrs` are the integer +// attribute channel passed verbatim from the IR. +// +// This is the single source of truth for op semantics on the compiler path. +// No explicit device: ops run on the default stream of the current worker thread. + +using OpFn = std::function< + mlx::core::array(const std::vector &ops, + const std::vector &attrs)>; + +// Float opts (eps/scale/base) ride the int64 attr channel as their IEEE-754 +// double bits (see EMLX.Native.Expr.f64_bits/1). Reverse the reinterpret here. +static inline float attr_to_float(int64_t bits) { + double d; + std::memcpy(&d, &bits, sizeof(double)); + return static_cast(d); +} + +// ── Prefill-RoPE helper ────────────────────────────────────────────────── +// +// mlx::fast::rope's offset argument is either a scalar or one starting +// position per batch example (sequential from there) — it cannot express +// arbitrary per-token positions (e.g. left-padded prefill). For that case we +// compose the same split-half rotation from precomputed per-token `angles` +// directly with primitive ops, matching mlx::fast::rope's non-traditional +// (traditional=false) formula bit-for-bit: +// out[..half] = a[..half]*cos(angles) - a[half..dims]*sin(angles) +// out[half..] = a[half..dims]*cos(angles) + a[..half]*sin(angles) +// (equivalently: rotate_half = concat(-a[half..dims], a[..half]); out = +// a[..dims]*cos_full + rotate_half*sin_full — the form used here, matching +// emlx_fast.cpp's fast_rope_positions eager NIF body). +// `a` is {B, T, H, D}; `angles` is {B, T, half} (already position*inv_freq*scale). +// No new Metal kernel: this is the existing fast_rope_positions composition, +// exposed as a compiled-graph opcode so prefill RoPE stays in one NIF replay. +static mlx::core::array rope_rotate_from_angles(const mlx::core::array &a, + const mlx::core::array &angles, + int dims) { + using namespace mlx::core; + + int B = a.shape(0); + int T = a.shape(1); + int H = a.shape(2); + int D = a.shape(3); + int half = dims / 2; + + auto cos_bt1h = astype(reshape(cos(angles), to_shape({B, T, 1, half})), a.dtype()); + auto sin_bt1h = astype(reshape(sin(angles), to_shape({B, T, 1, half})), a.dtype()); + + auto cos_full = concatenate(std::vector{cos_bt1h, cos_bt1h}, 3); + auto sin_full = concatenate(std::vector{sin_bt1h, sin_bt1h}, 3); + + auto x1 = slice(a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, half})); + auto x2 = slice(a, to_shape({0, 0, 0, half}), to_shape({B, T, H, dims})); + auto rotated = concatenate(std::vector{negative(x2), x1}, 3); + + auto a_head = slice(a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, dims})); + auto rope_head = add(multiply(a_head, cos_full), multiply(rotated, sin_full)); + + if (dims == D) { + return rope_head; + } + auto tail = slice(a, to_shape({0, 0, 0, dims}), to_shape({B, T, H, D})); + return concatenate(std::vector{rope_head, tail}, 3); +} + +// ── Window-op helpers ───────────────────────────────────────────────────── +// +// These mirror the Elixir backend's sliding-window algorithm, which uses +// as_strided to build a view then reduces over the window dims. +// No explicit device needed: MLX ops inherit stream from input arrays. + +// Build a sliding window view. padded has shape [...]; window and strides +// have length n=padded.ndim(). Result shape: [o0,...,on-1, w0,...,wn-1]. +static mlx::core::array compiler_sliding_window_view(const mlx::core::array &padded, + const std::vector &window, + const std::vector &strides) { + int n = padded.ndim(); + auto ps = padded.shape(); + + // Doubled element strides: output dims and window dims share the same strides. + auto orig_strides = padded.strides(); + std::vector view_strides(orig_strides.begin(), orig_strides.end()); + for (auto s : orig_strides) + view_strides.push_back(static_cast(s)); + + std::vector view_shape; + for (int i = 0; i < n; ++i) + view_shape.push_back(ps[i] - window[i] + 1); + for (int w : window) + view_shape.push_back(w); + + auto strided = mlx::core::as_strided(padded, to_shape(view_shape), + to_strides(view_strides), 0); + + std::vector starts(2 * n, 0); + std::vector stops = view_shape; + std::vector sl_strides = strides; + for (int i = 0; i < n; ++i) + sl_strides.push_back(1); + + return mlx::core::slice(strided, to_shape(starts), to_shape(stops), to_shape(sl_strides)); +} + +// Decode attrs and run a window reduction (sum/product/max/min). +// op_code: 0=sum, 1=product, 2=max, 3=min. +// attrs = [n_dims, op_int, lo0, hi0, …, s0, …, w0, …, wd0, …] +static mlx::core::array window_reduce_impl(const mlx::core::array &t, + const std::vector &attrs, + int op_code) { + int n = static_cast(attrs[0]); + // attrs[1] is op_int — matches op_code parameter; keep in sync. + + std::vector low_pad(n), high_pad(n), strides(n), window(n), wd(n); + int off = 2; + for (int i = 0; i < n; ++i) { + low_pad[i] = static_cast(attrs[off++]); + high_pad[i] = static_cast(attrs[off++]); + } + for (int i = 0; i < n; ++i) + strides[i] = static_cast(attrs[off++]); + for (int i = 0; i < n; ++i) + window[i] = static_cast(attrs[off++]); + for (int i = 0; i < n; ++i) + wd[i] = static_cast(attrs[off++]); + + // Apply window dilations by interior-padding the window mask. + // The mask is a bool_ array of shape `window` with interior zeros inserted. + // We achieve this via interior_padding: expanded window dims are + // expanded[i] = window[i] + (window[i]-1)*(wd[i]-1). + std::vector expanded_window(n); + for (int i = 0; i < n; ++i) + expanded_window[i] = window[i] + (window[i] - 1) * (wd[i] - 1); + + // Build a 1-filled window mask (bools), dilated via pad + interior zeros. + auto one_bool = mlx::core::ones(to_shape(window), mlx::core::bool_); + auto zero_bool = mlx::core::zeros({}, mlx::core::bool_); + auto window_mask = one_bool; + for (int i = 0; i < n; ++i) { + if (wd[i] > 1) { + // Interior-pad axis i by (wd[i]-1). + std::vector ax = {i}; + std::vector lo_i(n, 0), hi_i(n, 0); + auto cur_shape = window_mask.shape(); + int interior = wd[i] - 1; + // Reconstruct the shape after interior padding. + // Use as_strided trick: this is complex — use a simpler approach: + // create full expanded shape, fill with zeros, scatter ones. + // Alternatively: concatenate [1, 0, 0, ...] along axis repeatedly. + // Simplest: use pad with interior param — but mlx::core::pad doesn't + // support interior padding. Instead use reshape+broadcast trick: + // reshape from [w] to [w, 1], broadcast to [w, wd], reshape to [w*wd], + // then slice [0, wd-1, 2*wd-1, …] — i.e. stride-wd indices. + // We'll do this per axis sequentially. + int w_i = (int)cur_shape[i]; + // New size after dilation on axis i. + int new_size = w_i + (w_i - 1) * interior; + // Build dilated axis: start with zeros, then set positions 0, wd, 2*wd, ... + auto arange_full = mlx::core::arange(0, new_size, 1, mlx::core::int32); + // Positions occupied by real values: 0, wd, 2*wd, ... + auto divisible = mlx::core::equal( + mlx::core::remainder(arange_full, + mlx::core::full({}, wd[i], mlx::core::int32)), + mlx::core::zeros({}, mlx::core::int32)); + // Build dilated from current using take at strided positions. + // Take the [0, wd, 2*wd, ...] positions from the window_mask axis i. + auto src_idx = mlx::core::astype( + mlx::core::divide(arange_full, mlx::core::full({}, wd[i], mlx::core::int32)), + mlx::core::uint32); + // For interior positions, take from mask (value 1); replace interior zeros. + // It's simpler to: take all, then where(divisible, val, 0). + auto taken = mlx::core::take(window_mask, src_idx, i); + // Mask out non-original positions. + // Broadcast divisible to match taken shape. + std::vector div_shape(window_mask.ndim(), 1); + div_shape[i] = new_size; + auto div_nd = mlx::core::reshape(divisible, to_shape(div_shape)); + auto div_bc = mlx::core::broadcast_to(div_nd, taken.shape()); + window_mask = mlx::core::where(div_bc, taken, mlx::core::zeros({}, mlx::core::bool_)); + } + } + // expanded_window is now the shape of window_mask. + + // Pad the input tensor. + std::vector all_axes(n); + std::iota(all_axes.begin(), all_axes.end(), 0); + + // Build pad_value for this op. + auto make_pad_val = [&]() -> mlx::core::array { + switch (op_code) { + case 0: // sum: 0 + return mlx::core::full({}, 0, t.dtype()); + case 1: // product: 1 + return mlx::core::full({}, 1, t.dtype()); + case 2: // max: lowest representable value + if (mlx::core::issubdtype(t.dtype(), mlx::core::floating)) { + return mlx::core::full({}, -std::numeric_limits::infinity(), t.dtype()); + } else { + return mlx::core::full({}, std::numeric_limits::min(), t.dtype()); + } + case 3: // min: highest representable value + if (mlx::core::issubdtype(t.dtype(), mlx::core::floating)) { + return mlx::core::full({}, std::numeric_limits::infinity(), t.dtype()); + } else { + return mlx::core::full({}, std::numeric_limits::max(), t.dtype()); + } + default: + throw std::runtime_error("window_reduce_impl: unknown op_code"); + } + }; + auto pad_val = make_pad_val(); + + auto padded = mlx::core::pad(t, all_axes, to_shape(low_pad), to_shape(high_pad), pad_val, + "constant"); + + // Sliding window view: shape [o0,...,on-1, w0,...,wn-1]. + auto view = compiler_sliding_window_view(padded, expanded_window, strides); + + // Broadcast window_mask (shape expanded_window) to match view: all batch dims are 1. + std::vector mask_bc_shape(n, 1); + for (int w : expanded_window) + mask_bc_shape.push_back(w); + auto mask_reshaped = mlx::core::reshape(window_mask, to_shape(mask_bc_shape)); + auto mask_bc = mlx::core::broadcast_to(mask_reshaped, view.shape()); + + // Apply mask: replace masked-out positions with pad_val. + auto masked_view = mlx::core::where(mask_bc, view, mlx::core::broadcast_to( + mlx::core::reshape(pad_val, to_shape(std::vector{})), view.shape())); + + // Reduce over the last n dims (the window axes). + std::vector window_axes; + for (int i = n; i < 2 * n; ++i) + window_axes.push_back(i); + + switch (op_code) { + case 0: + return mlx::core::sum(masked_view, window_axes, false); + case 1: + return mlx::core::prod(masked_view, window_axes, false); + case 2: + return mlx::core::max(masked_view, window_axes, false); + case 3: + return mlx::core::min(masked_view, window_axes, false); + default: + throw std::runtime_error("window_reduce_impl: unknown op_code"); + } +} + +// Window scatter (max or min): replicate window_scatter_impl from emlx_nif.cpp. +// attrs = [n_dims, lo0, hi0, …, s0, …, w0, …] +// ops = [tensor_t, source, init_value] +static mlx::core::array window_scatter_impl_compiler(const mlx::core::array &tensor_t, + const mlx::core::array &tensor_source, + const mlx::core::array &tensor_init_value, + const std::vector &attrs, + bool scatter_max) { + int n = static_cast(attrs[0]); + std::vector low_pad(n), high_pad(n), strides(n), window(n); + int off = 1; + for (int i = 0; i < n; ++i) { + low_pad[i] = static_cast(attrs[off++]); + high_pad[i] = static_cast(attrs[off++]); + } + for (int i = 0; i < n; ++i) + strides[i] = static_cast(attrs[off++]); + for (int i = 0; i < n; ++i) + window[i] = static_cast(attrs[off++]); + + auto init_casted = mlx::core::astype(tensor_init_value, tensor_t.dtype()); + std::vector all_axes(n); + std::iota(all_axes.begin(), all_axes.end(), 0); + auto padded = mlx::core::pad(tensor_t, all_axes, to_shape(low_pad), to_shape(high_pad), + init_casted, "constant"); + + auto padded_shape = padded.shape(); + std::vector padded_shape_vec(padded_shape.begin(), padded_shape.end()); + + auto window_view = compiler_sliding_window_view(padded, window, strides); + + std::vector out_shape(window_view.shape().begin(), window_view.shape().begin() + n); + + int K = 1; + for (int w : window) + K *= w; + + std::vector flat_shape = out_shape; + flat_shape.push_back(K); + auto windows_flat = mlx::core::reshape(window_view, to_shape(flat_shape)); + + auto arg_idx = [&]() -> mlx::core::array { + if (scatter_max) { + return mlx::core::argmax(windows_flat, n, false); + } + auto m = mlx::core::min(windows_flat, std::vector{n}, true); + auto mask = + mlx::core::astype(mlx::core::equal(windows_flat, m), mlx::core::uint32); + auto arange_k = mlx::core::astype(mlx::core::arange(0, K, 1), mlx::core::uint32); + std::vector arange_shape(n + 1, 1); + arange_shape[n] = K; + auto arange_nd = mlx::core::reshape(arange_k, to_shape(arange_shape)); + auto weighted = mlx::core::multiply(mask, arange_nd); + return mlx::core::argmax(weighted, n, false); + }(); + + std::vector arg_exp_shape = out_shape; + arg_exp_shape.push_back(1); + auto arg_idx_exp = mlx::core::reshape(arg_idx, to_shape(arg_exp_shape)); + + std::vector abs_indices; + for (int a = 0; a < n; ++a) { + auto arange_a = mlx::core::astype(mlx::core::arange(0, (int)padded_shape[a], 1), + mlx::core::int32); + std::vector iota_shape(n, 1); + iota_shape[a] = (int)padded_shape[a]; + auto iota_nd = mlx::core::reshape(arange_a, to_shape(iota_shape)); + auto iota_bc = mlx::core::broadcast_to(iota_nd, to_shape(padded_shape_vec)); + auto iota_view = compiler_sliding_window_view(iota_bc, window, strides); + auto iota_flat = mlx::core::reshape(iota_view, to_shape(flat_shape)); + auto abs_a = mlx::core::take_along_axis(iota_flat, arg_idx_exp, n); + abs_indices.push_back(mlx::core::reshape(abs_a, to_shape(out_shape))); + } + + auto source_shape_2n = std::vector(tensor_source.shape().begin(), + tensor_source.shape().end()); + for (int i = 0; i < n; ++i) + source_shape_2n.push_back(1); + auto updates = mlx::core::reshape(tensor_source, to_shape(source_shape_2n)); + + auto buffer = mlx::core::broadcast_to( + mlx::core::reshape(init_casted, to_shape(std::vector{})), + to_shape(padded_shape_vec)); + + std::vector scatter_axes(n); + std::iota(scatter_axes.begin(), scatter_axes.end(), 0); + auto scattered = mlx::core::scatter_add(buffer, abs_indices, updates, scatter_axes); + + auto orig_shape = tensor_t.shape(); + std::vector slice_starts = low_pad; + std::vector slice_stops(n); + for (int i = 0; i < n; ++i) + slice_stops[i] = low_pad[i] + (int)orig_shape[i]; + std::vector slice_ones(n, 1); + return mlx::core::slice(scattered, to_shape(slice_starts), to_shape(slice_stops), + to_shape(slice_ones)); +} + +// MLX linalg primitives are CPU-only. Pin them to the CPU device so they +// compose inside a compiled graph regardless of the graph's default (worker) +// stream — validated to work for both :cpu and :gpu default devices. +static const mlx::core::Device k_linalg_cpu(mlx::core::Device::DeviceType::cpu, 0); + +static const std::unordered_map op_registry = { + // ── cast ────────────────────────────────────────────────────────────── + {"astype", + [](const auto &ops, const auto &attrs) { + return mlx::core::astype(ops[0], attrs[0].as_dtype()); + }}, + + // ── unary ops ───────────────────────────────────────────────────────── + {"abs", [](const auto &ops, const auto &) { return mlx::core::abs(ops[0]); }}, + {"ceil", [](const auto &ops, const auto &) { return mlx::core::ceil(ops[0]); }}, + {"floor", [](const auto &ops, const auto &) { return mlx::core::floor(ops[0]); }}, + {"negate", [](const auto &ops, const auto &) { return mlx::core::negative(ops[0]); }}, + {"round", [](const auto &ops, const auto &) { return mlx::core::round(ops[0]); }}, + {"sign", [](const auto &ops, const auto &) { return mlx::core::sign(ops[0]); }}, + {"real", [](const auto &ops, const auto &) { return mlx::core::real(ops[0]); }}, + {"imag", [](const auto &ops, const auto &) { return mlx::core::imag(ops[0]); }}, + {"is_nan", [](const auto &ops, const auto &) { return mlx::core::isnan(ops[0]); }}, + {"is_infinity", [](const auto &ops, const auto &) { return mlx::core::isinf(ops[0]); }}, + {"bitwise_not", + [](const auto &ops, const auto &) { + return mlx::core::bitwise_invert(ops[0]); + }}, + {"conjugate", + [](const auto &ops, const auto &) { return mlx::core::conjugate(ops[0]); }}, + {"logical_not", + [](const auto &ops, const auto &) { return mlx::core::logical_not(ops[0]); }}, + {"sigmoid", [](const auto &ops, const auto &) { return mlx::core::sigmoid(ops[0]); }}, + {"asin", [](const auto &ops, const auto &) { return mlx::core::arcsin(ops[0]); }}, + {"asinh", [](const auto &ops, const auto &) { return mlx::core::arcsinh(ops[0]); }}, + {"acos", [](const auto &ops, const auto &) { return mlx::core::arccos(ops[0]); }}, + {"acosh", [](const auto &ops, const auto &) { return mlx::core::arccosh(ops[0]); }}, + {"atan", [](const auto &ops, const auto &) { return mlx::core::arctan(ops[0]); }}, + {"atanh", [](const auto &ops, const auto &) { return mlx::core::arctanh(ops[0]); }}, + {"cos", [](const auto &ops, const auto &) { return mlx::core::cos(ops[0]); }}, + {"cosh", [](const auto &ops, const auto &) { return mlx::core::cosh(ops[0]); }}, + {"erf", [](const auto &ops, const auto &) { return mlx::core::erf(ops[0]); }}, + {"erf_inv", [](const auto &ops, const auto &) { return mlx::core::erfinv(ops[0]); }}, + {"exp", [](const auto &ops, const auto &) { return mlx::core::exp(ops[0]); }}, + {"expm1", [](const auto &ops, const auto &) { return mlx::core::expm1(ops[0]); }}, + {"log", [](const auto &ops, const auto &) { return mlx::core::log(ops[0]); }}, + {"log1p", [](const auto &ops, const auto &) { return mlx::core::log1p(ops[0]); }}, + {"rsqrt", [](const auto &ops, const auto &) { return mlx::core::rsqrt(ops[0]); }}, + {"sin", [](const auto &ops, const auto &) { return mlx::core::sin(ops[0]); }}, + {"sinh", [](const auto &ops, const auto &) { return mlx::core::sinh(ops[0]); }}, + {"sqrt", [](const auto &ops, const auto &) { return mlx::core::sqrt(ops[0]); }}, + {"tan", [](const auto &ops, const auto &) { return mlx::core::tan(ops[0]); }}, + {"tanh", [](const auto &ops, const auto &) { return mlx::core::tanh(ops[0]); }}, + + // cbrt = x^(1/3); returns NaN for negative real inputs (matches EMLX.Backend) + {"cbrt", + [](const auto &ops, const auto &) { + return mlx::core::power( + ops[0], mlx::core::full({}, static_cast(1.0 / 3.0), mlx::core::float32)); + }}, + + // erfc = 1 - erf(x) + {"erfc", + [](const auto &ops, const auto &) { + auto e = mlx::core::erf(ops[0]); + return mlx::core::subtract(mlx::core::ones({}, e.dtype()), e); + }}, + + // ── binary ops ──────────────────────────────────────────────────────── + {"add", [](const auto &ops, const auto &) { return mlx::core::add(ops[0], ops[1]); }}, + {"subtract", + [](const auto &ops, const auto &) { return mlx::core::subtract(ops[0], ops[1]); }}, + {"multiply", + [](const auto &ops, const auto &) { return mlx::core::multiply(ops[0], ops[1]); }}, + {"divide", + [](const auto &ops, const auto &) { return mlx::core::divide(ops[0], ops[1]); }}, + {"pow", + [](const auto &ops, const auto &) { return mlx::core::power(ops[0], ops[1]); }}, + {"atan2", + [](const auto &ops, const auto &) { return mlx::core::arctan2(ops[0], ops[1]); }}, + {"min", + [](const auto &ops, const auto &) { return mlx::core::minimum(ops[0], ops[1]); }}, + {"max", + [](const auto &ops, const auto &) { return mlx::core::maximum(ops[0], ops[1]); }}, + {"quotient", + [](const auto &ops, const auto &) { return mlx::core::floor_divide(ops[0], ops[1]); }}, + {"left_shift", + [](const auto &ops, const auto &) { return mlx::core::left_shift(ops[0], ops[1]); }}, + {"right_shift", + [](const auto &ops, const auto &) { return mlx::core::right_shift(ops[0], ops[1]); }}, + {"bitwise_and", + [](const auto &ops, const auto &) { return mlx::core::bitwise_and(ops[0], ops[1]); }}, + {"bitwise_or", + [](const auto &ops, const auto &) { return mlx::core::bitwise_or(ops[0], ops[1]); }}, + {"bitwise_xor", + [](const auto &ops, const auto &) { return mlx::core::bitwise_xor(ops[0], ops[1]); }}, + + // remainder: MLX uses floor-division semantics; fix up to truncation (same sign as dividend) + // to match EMLX.Backend.remainder which applies this same correction. + {"remainder", + [](const auto &ops, const auto &) { + auto rem = mlx::core::remainder(ops[0], ops[1]); + auto zero = mlx::core::full({}, 0, ops[0].dtype()); + auto neg_dividend = mlx::core::less(ops[0], zero); + auto adjusted = mlx::core::subtract(rem, ops[1]); + return mlx::core::where(neg_dividend, adjusted, rem); + }}, + + // compare — result is MLX bool_; Elixir lowerer emits astype(bool_, uint8) after + {"equal", + [](const auto &ops, const auto &) { return mlx::core::equal(ops[0], ops[1]); }}, + {"not_equal", + [](const auto &ops, const auto &) { return mlx::core::not_equal(ops[0], ops[1]); }}, + {"greater", + [](const auto &ops, const auto &) { return mlx::core::greater(ops[0], ops[1]); }}, + {"less", + [](const auto &ops, const auto &) { return mlx::core::less(ops[0], ops[1]); }}, + {"greater_equal", + [](const auto &ops, const auto &) { return mlx::core::greater_equal(ops[0], ops[1]); }}, + {"less_equal", + [](const auto &ops, const auto &) { return mlx::core::less_equal(ops[0], ops[1]); }}, + + // logical binary + {"logical_and", + [](const auto &ops, const auto &) { return mlx::core::logical_and(ops[0], ops[1]); }}, + {"logical_or", + [](const auto &ops, const auto &) { return mlx::core::logical_or(ops[0], ops[1]); }}, + // logical_xor = (a || b) && !(a && b) + {"logical_xor", + [](const auto &ops, const auto &) { + auto t1 = mlx::core::logical_or(ops[0], ops[1]); + auto t2 = mlx::core::logical_not(mlx::core::logical_and(ops[0], ops[1])); + return mlx::core::logical_and(t1, t2); + }}, + + // ── shape / movement ops ────────────────────────────────────────────────── + // + // iattrs encoding (must stay in sync with EMLX.Native.Expr moduledoc): + // reshape: attrs = [d0, d1, …] — target shape dims (flat) + // squeeze: attrs = [a0, a1, …] — axes to remove (non-negative) + // transpose: attrs = [p0, p1, …] — permutation (non-negative) + // bitcast: attrs = [dtype_int] — target dtype + // broadcast: attrs = [n, d0..dn-1, m, a0..am-1] — shape then axes (length-delimited) + // pad: attrs = [n_dims, lo0,hi0,int0, …] — n_dims triples per dim + // reverse: attrs = [a0, a1, …] — axes to flip (non-negative) + // concatenate: attrs = [axis]; ops = all input tensors + // stack: attrs = [axis]; ops = all input tensors + + {"reshape", + [](const auto &ops, const auto &attrs) { + std::vector shape; + shape.reserve(attrs.size()); + for (auto d : attrs) + shape.push_back(static_cast(d)); + return mlx::core::reshape(ops[0], to_shape(shape)); + }}, + + {"squeeze", + [](const auto &ops, const auto &attrs) { + if (attrs.empty()) + return ops[0]; + std::vector axes; + axes.reserve(attrs.size()); + for (auto a : attrs) + axes.push_back(static_cast(a)); + return mlx::core::squeeze(ops[0], axes); + }}, + + {"transpose", + [](const auto &ops, const auto &attrs) { + std::vector axes; + axes.reserve(attrs.size()); + for (auto a : attrs) + axes.push_back(static_cast(a)); + return mlx::core::transpose(ops[0], axes); + }}, + + {"bitcast", + [](const auto &ops, const auto &attrs) { + return mlx::core::view(ops[0], attrs[0].as_dtype()); + }}, + + // broadcast: reshape input to place source dims at the axis positions, then broadcast_to. + // Mirrors EMLX.Backend.broadcast / maybe_reshape logic exactly. + {"broadcast", + [](const auto &ops, const auto &attrs) { + int n_shape = static_cast(attrs[0]); + std::vector target_shape; + target_shape.reserve(n_shape); + for (int i = 0; i < n_shape; i++) + target_shape.push_back(static_cast(attrs[1 + i])); + + int n_axes = static_cast(attrs[1 + n_shape]); + std::vector axes; + axes.reserve(n_axes); + for (int i = 0; i < n_axes; i++) + axes.push_back(static_cast(attrs[1 + n_shape + 1 + i])); + + auto tensor = ops[0]; + auto in_shape = tensor.shape(); + + // Build broadcast_shape: 1s everywhere, input dims at axis positions. + std::vector broadcast_shape(n_shape, 1); + for (int i = 0; i < n_axes; i++) { + if (!in_shape.empty()) + broadcast_shape[axes[i]] = static_cast(in_shape[i]); + } + auto reshaped = mlx::core::reshape(tensor, to_shape(broadcast_shape)); + return mlx::core::broadcast_to(reshaped, to_shape(target_shape)); + }}, + + // pad: non-negative lo/hi, interior always 0 (Elixir raises otherwise). + {"pad", + [](const auto &ops, const auto &attrs) { + int n_dims = static_cast(attrs[0]); + std::vector axes, low_pads, high_pads; + axes.reserve(n_dims); + low_pads.reserve(n_dims); + high_pads.reserve(n_dims); + for (int i = 0; i < n_dims; i++) { + axes.push_back(i); + low_pads.push_back(static_cast(attrs[1 + i * 3 + 0])); + high_pads.push_back(static_cast(attrs[1 + i * 3 + 1])); + // attrs[1 + i*3 + 2] = interior, always 0 (validated in Elixir lowerer) + } + return mlx::core::pad(ops[0], axes, to_shape(low_pads), to_shape(high_pads), ops[1], + "constant"); + }}, + + // reverse: implemented via slice with negative strides, matching EMLX.Backend.reverse_mlx. + {"reverse", + [](const auto &ops, const auto &attrs) { + auto tensor = ops[0]; + auto shape = tensor.shape(); + int rank = static_cast(shape.size()); + + std::vector starts(rank), stops(rank), strides(rank); + for (int i = 0; i < rank; i++) { + starts[i] = 0; + stops[i] = static_cast(shape[i]); + strides[i] = 1; + } + for (auto a : attrs) { + int ax = static_cast(a); + int d = static_cast(shape[ax]); + starts[ax] = d - 1; + stops[ax] = -(d + 1); + strides[ax] = -1; + } + return mlx::core::slice(tensor, to_shape(starts), to_shape(stops), to_shape(strides)); + }}, + + {"concatenate", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + return mlx::core::concatenate(ops, axis); + }}, + + {"stack", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + return mlx::core::stack(ops, axis); + }}, + + // ── reductions ──────────────────────────────────────────────────────────── + // + // iattrs = [keep_axes_int, a0, a1, …] + // keep_axes_int: 0 = false, 1 = true. Axis list always explicit (Elixir resolves nil). + // sum, product, reduce_max, reduce_min: emit astype in Elixir lowerer for type changes. + // all, any: Elixir always emits astype since MLX returns bool_. + + {"sum", + [](const auto &ops, const auto &attrs) { + bool keepdims = static_cast(attrs[0]); + std::vector axes; + for (size_t i = 1; i < attrs.size(); i++) + axes.push_back(static_cast(attrs[i])); + return mlx::core::sum(ops[0], axes, keepdims); + }}, + + {"product", + [](const auto &ops, const auto &attrs) { + bool keepdims = static_cast(attrs[0]); + std::vector axes; + for (size_t i = 1; i < attrs.size(); i++) + axes.push_back(static_cast(attrs[i])); + return mlx::core::prod(ops[0], axes, keepdims); + }}, + + {"all", + [](const auto &ops, const auto &attrs) { + bool keepdims = static_cast(attrs[0]); + std::vector axes; + for (size_t i = 1; i < attrs.size(); i++) + axes.push_back(static_cast(attrs[i])); + return mlx::core::all(ops[0], axes, keepdims); + }}, + + {"any", + [](const auto &ops, const auto &attrs) { + bool keepdims = static_cast(attrs[0]); + std::vector axes; + for (size_t i = 1; i < attrs.size(); i++) + axes.push_back(static_cast(attrs[i])); + return mlx::core::any(ops[0], axes, keepdims); + }}, + + {"reduce_max", + [](const auto &ops, const auto &attrs) { + bool keepdims = static_cast(attrs[0]); + std::vector axes; + for (size_t i = 1; i < attrs.size(); i++) + axes.push_back(static_cast(attrs[i])); + return mlx::core::max(ops[0], axes, keepdims); + }}, + + {"reduce_min", + [](const auto &ops, const auto &attrs) { + bool keepdims = static_cast(attrs[0]); + std::vector axes; + for (size_t i = 1; i < attrs.size(); i++) + axes.push_back(static_cast(attrs[i])); + return mlx::core::min(ops[0], axes, keepdims); + }}, + + // argmax / argmin — iattrs = [axis, keep_axis_int]; axis = -1 means global. + {"argmax", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + bool keepdims = static_cast(attrs[1]); + if (axis < 0) + return mlx::core::argmax(ops[0], keepdims); + return mlx::core::argmax(ops[0], axis, keepdims); + }}, + + {"argmin", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + bool keepdims = static_cast(attrs[1]); + if (axis < 0) + return mlx::core::argmin(ops[0], keepdims); + return mlx::core::argmin(ops[0], axis, keepdims); + }}, + + // ── dot ─────────────────────────────────────────────────────────────────── + // + // iattrs = [n_ca, ca…, n_cb, cb…, n_ba, ba…, n_bb, bb…] + // Elixir lowerer casts both operands to computation_type before the dot op. + // Non-batched (n_ba=0, n_bb=0): mlx::core::tensordot(left, right, ca, cb). + // Batched: rebuild einsum spec from shapes + 4 axis lists, call einsum. + + {"dot", + [](const auto &ops, const auto &attrs) { + // Parse 4 length-delimited axis lists. + size_t off = 0; + auto parse_axis_list = [&]() -> std::vector { + int n = static_cast(attrs[off++]); + std::vector v(n); + for (int i = 0; i < n; i++) + v[i] = static_cast(attrs[off++]); + return v; + }; + auto ca = parse_axis_list(); + auto cb = parse_axis_list(); + auto ba = parse_axis_list(); + auto bb = parse_axis_list(); + + if (ba.empty() && bb.empty()) { + return mlx::core::tensordot(ops[0], ops[1], ca, cb); + } + + // Batched: build einsum spec by assigning letter labels. + auto left_shape = ops[0].shape(); + auto right_shape = ops[1].shape(); + int n_left = static_cast(left_shape.size()); + int n_right = static_cast(right_shape.size()); + + // Assign a label to every left and right axis. + const std::string alphabet = "abcdefghijklmnopqrstuvwxyz"; + std::vector ll(n_left), rl(n_right); + for (int i = 0; i < n_left; i++) + ll[i] = alphabet[i]; + for (int i = 0; i < n_right; i++) + rl[i] = alphabet[n_left + i]; + + // Share labels for batch and contraction axes. + auto share = [&](const std::vector &la, const std::vector &ra) { + for (size_t i = 0; i < la.size(); i++) + rl[ra[i]] = ll[la[i]]; + }; + share(ba, bb); + share(ca, cb); + + // Output: batch axes, then free left axes, then free right axes. + auto contains = [](const std::vector &v, int x) { + return std::find(v.begin(), v.end(), x) != v.end(); + }; + std::string output; + for (int b : ba) + output += ll[b]; + for (int i = 0; i < n_left; i++) + if (!contains(ca, i) && !contains(ba, i)) + output += ll[i]; + for (int i = 0; i < n_right; i++) + if (!contains(cb, i) && !contains(bb, i)) + output += rl[i]; + + std::string spec = std::string(ll.begin(), ll.end()) + "," + + std::string(rl.begin(), rl.end()) + "->" + output; + return mlx::core::einsum(spec, {ops[0], ops[1]}); + }}, + + // ── quantized_matmul ────────────────────────────────────────────────────── + // + // iattrs = [group_size, bits, transpose_int, mode_int, has_bias_int] + // Operands: [activation, weight, scales, biases?] — biases present only + // when has_bias_int is 1 (mirrors EMLX.Backend.quantized_dot/4: absent + // for microscaled modes, since mx::fp_quantize doesn't emit biases). + // Emitted only by a call-time-specialized program (see + // EMLX.Native.Expr's "Quantized dot specialization" moduledoc section); + // `weight` is the untouched physical (packed) parameter ref. + {"quantized_matmul", + [](const auto &ops, const auto &attrs) { + int group_size = static_cast(attrs[0]); + int bits = static_cast(attrs[1]); + bool transpose = attrs[2] != 0; + std::string mode = attrs[3].as_mode(); + bool has_bias = attrs[4] != 0; + std::optional biases_opt = + has_bias ? std::make_optional(ops[3]) : std::nullopt; + return mlx::core::quantized_matmul(ops[0], ops[1], ops[2], biases_opt, + transpose, group_size, bits, mode); + }}, + + // ── indexing / selection ───────────────────────────────────────────────── + // + // select: operands = [pred, on_true, on_false]; no attrs. + {"select", + [](const auto &ops, const auto &) { + return mlx::core::where(ops[0], ops[1], ops[2]); + }}, + + // clip: operands = [tensor, min, max]; no attrs. + {"clip", + [](const auto &ops, const auto &) { + return mlx::core::clip(ops[0], ops[1], ops[2]); + }}, + + // slice: attrs = [n_dims, dyn_mask, d0..dn-1, l0..ln-1, str0..strn-1, sv0..svn-1] + // Operands: [tensor, dyn_start_0, dyn_start_1, …] — dynamic starts in axis order. + // + // Strategy: one static mlx::core::slice for all static dims (full span for dynamic dims), + // then mlx::core::take per dynamic dim to apply the dynamic start + strided selection. + {"slice", + [](const auto &ops, const auto &attrs) { + int n_dims = static_cast(attrs[0]); + int64_t dyn_mask = attrs[1]; + + std::vector input_shape(n_dims), lengths(n_dims), strides_v(n_dims), + sv(n_dims); + for (int i = 0; i < n_dims; i++) + input_shape[i] = static_cast(attrs[2 + i]); + for (int i = 0; i < n_dims; i++) + lengths[i] = static_cast(attrs[2 + n_dims + i]); + for (int i = 0; i < n_dims; i++) + strides_v[i] = static_cast(attrs[2 + 2 * n_dims + i]); + for (int i = 0; i < n_dims; i++) + sv[i] = static_cast(attrs[2 + 3 * n_dims + i]); + + // Build the static slice: use clamped static values for static dims, + // full extent for dynamic dims (handled below via take). + std::vector starts(n_dims), stops(n_dims), slice_strides(n_dims); + for (int i = 0; i < n_dims; i++) { + slice_strides[i] = ((dyn_mask >> i) & 1) ? 1 : strides_v[i]; + if (!((dyn_mask >> i) & 1)) { + starts[i] = std::max(0, std::min(sv[i], input_shape[i] - lengths[i])); + stops[i] = starts[i] + lengths[i]; + } else { + starts[i] = 0; + stops[i] = input_shape[i]; + } + } + auto result = + mlx::core::slice(ops[0], to_shape(starts), to_shape(stops), to_shape(slice_strides)); + + // For each dynamic dim, apply take with arange*stride + clamped_start. + int dyn_op_idx = 1; + for (int i = 0; i < n_dims; i++) { + if (!((dyn_mask >> i) & 1)) + continue; + int len_i = lengths[i]; + int str_i = strides_v[i]; + // clamped_start = clip(start_arr, 0, dim_i - len_i) + auto s = mlx::core::astype(ops[dyn_op_idx++], mlx::core::int32); + auto zero = mlx::core::zeros({}, mlx::core::int32); + auto max_s = mlx::core::full({}, input_shape[i] - len_i, mlx::core::int32); + auto clamped = mlx::core::minimum(mlx::core::maximum(s, zero), max_s); + // idx = arange(len_i) * str_i + clamped + auto base = mlx::core::arange(0, len_i, 1, mlx::core::int32); + if (str_i != 1) + base = mlx::core::multiply(base, mlx::core::full({}, str_i, mlx::core::int32)); + auto idx = mlx::core::add(base, clamped); + result = mlx::core::take(result, idx, i); + } + return result; + }}, + + // put_slice: attrs = [n_dims, dyn_mask, d0..dn-1, l0..ln-1, sv0..svn-1] + // Operands: [input, slice, dyn_start_0, …] + // + // Builds a clamped int32 start array and calls the dynamic slice_update overload. + {"put_slice", + [](const auto &ops, const auto &attrs) { + int n_dims = static_cast(attrs[0]); + int64_t dyn_mask = attrs[1]; + + std::vector input_shape(n_dims), lengths(n_dims), sv(n_dims); + for (int i = 0; i < n_dims; i++) + input_shape[i] = static_cast(attrs[2 + i]); + for (int i = 0; i < n_dims; i++) + lengths[i] = static_cast(attrs[2 + n_dims + i]); + for (int i = 0; i < n_dims; i++) + sv[i] = static_cast(attrs[2 + 2 * n_dims + i]); + + // Build per-dim clamped start components (shape [1]) then concatenate. + std::vector start_components; + start_components.reserve(n_dims); + int dyn_op_idx = 2; // ops[0]=input, ops[1]=slice, ops[2..]=dynamic starts + for (int i = 0; i < n_dims; i++) { + int max_start_i = input_shape[i] - lengths[i]; + if ((dyn_mask >> i) & 1) { + auto s = mlx::core::astype(ops[dyn_op_idx++], mlx::core::int32); + s = mlx::core::reshape(s, {1}); + auto max_s = mlx::core::full({1}, max_start_i, mlx::core::int32); + auto zero = mlx::core::zeros({1}, mlx::core::int32); + start_components.push_back(mlx::core::minimum(mlx::core::maximum(s, zero), max_s)); + } else { + int clamped_i = std::max(0, std::min(sv[i], max_start_i)); + start_components.push_back(mlx::core::full({1}, clamped_i, mlx::core::int32)); + } + } + auto start_arr = mlx::core::concatenate(start_components, 0); + + // All axes: [0, 1, ..., n_dims-1] + std::vector all_axes(n_dims); + std::iota(all_axes.begin(), all_axes.end(), 0); + return mlx::core::slice_update(ops[0], ops[1], start_arr, all_axes); + }}, + + // gather: attrs = [n_gather_axes, a0…, n_tensor_dims, ss0…, n_out_dims, od0…] + // Operands: [tensor, indices] — indices shape = [batch…, n_gather_axes] + {"gather", + [](const auto &ops, const auto &attrs) { + int n_gather_axes = static_cast(attrs[0]); + std::vector axes(n_gather_axes); + for (int i = 0; i < n_gather_axes; i++) + axes[i] = static_cast(attrs[1 + i]); + + int n_tensor_dims = static_cast(attrs[1 + n_gather_axes]); + std::vector slice_sizes(n_tensor_dims); + for (int i = 0; i < n_tensor_dims; i++) + slice_sizes[i] = static_cast(attrs[2 + n_gather_axes + i]); + + int n_out_dims = static_cast(attrs[2 + n_gather_axes + n_tensor_dims]); + std::vector out_shape_v(n_out_dims); + for (int i = 0; i < n_out_dims; i++) + out_shape_v[i] = static_cast(attrs[3 + n_gather_axes + n_tensor_dims + i]); + + // Split indices along its last axis: one array per gather axis. + auto indices = ops[1]; + auto idx_shape = indices.shape(); + int last = static_cast(idx_shape.size()) - 1; + + std::vector indices_list; + indices_list.reserve(n_gather_axes); + for (int i = 0; i < n_gather_axes; i++) { + std::vector s_starts(idx_shape.size(), 0), s_stops(idx_shape.size()); + std::vector s_strides(idx_shape.size(), 1); + for (size_t j = 0; j < idx_shape.size(); j++) + s_stops[j] = static_cast(idx_shape[j]); + s_starts[last] = i; + s_stops[last] = i + 1; + auto sl = mlx::core::slice(indices, to_shape(s_starts), to_shape(s_stops), + to_shape(s_strides)); + indices_list.push_back(mlx::core::squeeze(sl, {last})); + } + + auto result = mlx::core::gather(ops[0], indices_list, axes, to_shape(slice_sizes)); + return mlx::core::reshape(result, to_shape(out_shape_v)); + }}, + + // take: operands = [tensor, indices]; attrs = [axis]. + {"take", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + return mlx::core::take(ops[0], ops[1], axis); + }}, + + // take_along_axis: operands = [tensor, indices]; attrs = [axis]. + {"take_along_axis", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + return mlx::core::take_along_axis(ops[0], ops[1], axis); + }}, + + // indexed_add: attrs = [n_axes, a0…, n_updates_shape, us0…] + // Operands: [target, indices, updates_pre_reshape] + // Mirrors EMLX.Backend.indexed_op(:scatter_add, …) + {"indexed_add", + [](const auto &ops, const auto &attrs) { + int n_axes = static_cast(attrs[0]); + std::vector axes(n_axes); + for (int i = 0; i < n_axes; i++) + axes[i] = static_cast(attrs[1 + i]); + int n_upd = static_cast(attrs[1 + n_axes]); + std::vector upd_shape(n_upd); + for (int i = 0; i < n_upd; i++) + upd_shape[i] = static_cast(attrs[2 + n_axes + i]); + + auto indices = ops[1]; + auto idx_shape = indices.shape(); + int last = static_cast(idx_shape.size()) - 1; + + std::vector indices_list; + indices_list.reserve(n_axes); + for (int i = 0; i < n_axes; i++) { + std::vector s_starts(idx_shape.size(), 0), s_stops(idx_shape.size()); + std::vector s_strides(idx_shape.size(), 1); + for (size_t j = 0; j < idx_shape.size(); j++) + s_stops[j] = static_cast(idx_shape[j]); + s_starts[last] = i; + s_stops[last] = i + 1; + auto sl = mlx::core::slice(indices, to_shape(s_starts), to_shape(s_stops), + to_shape(s_strides)); + indices_list.push_back(mlx::core::squeeze(sl, {last})); + } + + auto updates_reshaped = mlx::core::reshape(ops[2], to_shape(upd_shape)); + return mlx::core::scatter_add(ops[0], indices_list, updates_reshaped, axes); + }}, + + // indexed_put: same structure as indexed_add but uses scatter (replace semantics). + {"indexed_put", + [](const auto &ops, const auto &attrs) { + int n_axes = static_cast(attrs[0]); + std::vector axes(n_axes); + for (int i = 0; i < n_axes; i++) + axes[i] = static_cast(attrs[1 + i]); + int n_upd = static_cast(attrs[1 + n_axes]); + std::vector upd_shape(n_upd); + for (int i = 0; i < n_upd; i++) + upd_shape[i] = static_cast(attrs[2 + n_axes + i]); + + auto indices = ops[1]; + auto idx_shape = indices.shape(); + int last = static_cast(idx_shape.size()) - 1; + + std::vector indices_list; + indices_list.reserve(n_axes); + for (int i = 0; i < n_axes; i++) { + std::vector s_starts(idx_shape.size(), 0), s_stops(idx_shape.size()); + std::vector s_strides(idx_shape.size(), 1); + for (size_t j = 0; j < idx_shape.size(); j++) + s_stops[j] = static_cast(idx_shape[j]); + s_starts[last] = i; + s_stops[last] = i + 1; + auto sl = mlx::core::slice(indices, to_shape(s_starts), to_shape(s_stops), + to_shape(s_strides)); + indices_list.push_back(mlx::core::squeeze(sl, {last})); + } + + auto updates_reshaped = mlx::core::reshape(ops[2], to_shape(upd_shape)); + return mlx::core::scatter(ops[0], indices_list, updates_reshaped, axes); + }}, + + // ── sort / argsort ───────────────────────────────────────────────────────── + // + // iattrs = [axis, asc_int] (asc_int: 1=ascending, 0=descending) + // + // Replicates EMLX.Backend.sort/argsort NaN-aware algorithm: + // 1. argsort(t, axis) (negate t first if descending) + // 2. sorted_values = take_along_axis(t, sort_indices, axis) + // 3. is_nan = isnan(sorted_values) + // 4. partition_indices = argsort(is_nan, axis) (ascending) + // OR argsort(!is_nan, axis) (descending) + // (cast bool→uint8 to avoid Metal metallib gap) + // 5. return take_along_axis(sorted_values, partition, axis) + // argsort: same but step 5 applies partition to sort_indices. + + {"sort", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + bool asc = (attrs[1] != 0); + + // Step 1: get initial argsort indices. + auto t = ops[0]; + auto sort_input = asc ? t : mlx::core::negative(t); + auto sort_idx = mlx::core::argsort(sort_input, axis); + + // Step 2: gather sorted values. + auto sorted_vals = mlx::core::take_along_axis(t, sort_idx, axis); + + // Step 3: NaN mask. + auto is_nan = mlx::core::isnan(sorted_vals); + + // Step 4: partition indices (move NaNs to correct end). + auto is_nan_u8 = mlx::core::astype(is_nan, mlx::core::uint8); + mlx::core::array partition_input = + asc ? is_nan_u8 + : mlx::core::astype(mlx::core::logical_not(is_nan), mlx::core::uint8); + auto partition = mlx::core::argsort(partition_input, axis); + + // Step 5: reorder sorted_values by partition. + return mlx::core::take_along_axis(sorted_vals, partition, axis); + }}, + + {"argsort", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + bool asc = (attrs[1] != 0); + + // Step 1: get initial argsort indices. + auto t = ops[0]; + auto sort_input = asc ? t : mlx::core::negative(t); + auto sort_idx = mlx::core::argsort(sort_input, axis); + + // Step 2: gather sorted values to check NaN. + auto sorted_vals = mlx::core::take_along_axis(t, sort_idx, axis); + + // Step 3: NaN mask. + auto is_nan = mlx::core::isnan(sorted_vals); + + // Step 4: partition indices. + auto is_nan_u8 = mlx::core::astype(is_nan, mlx::core::uint8); + mlx::core::array partition_input = + asc ? is_nan_u8 + : mlx::core::astype(mlx::core::logical_not(is_nan), mlx::core::uint8); + auto partition = mlx::core::argsort(partition_input, axis); + + // Step 5: reorder sort_idx by partition. + return mlx::core::take_along_axis(sort_idx, partition, axis); + }}, + + // ── window reductions ───────────────────────────────────────────────────── + // + // Shared helper: build a sliding-window view of a padded tensor. + // Returns shape [o0,...,on-1, w0,...,wn-1] where oi = (ps[i]-w[i])/s[i]+1. + // Uses as_strided + slice, matching EMLX.Backend.sliding_window_view. + // + // iattrs = [n_dims, op_int, lo0, hi0, …, s0, …, w0, …, wd0, …] + // op_int: 0=sum, 1=product, 2=max, 3=min + // lo/hi: n_dims pairs of padding (2*n_dims values starting at offset 2) + // strides: n_dims values + // window: n_dims values + // window_dilations: n_dims values + + {"window_sum", + [](const auto &ops, const auto &attrs) { + return window_reduce_impl(ops[0], attrs, 0); + }}, + {"window_product", + [](const auto &ops, const auto &attrs) { + return window_reduce_impl(ops[0], attrs, 1); + }}, + {"window_max", + [](const auto &ops, const auto &attrs) { + return window_reduce_impl(ops[0], attrs, 2); + }}, + {"window_min", + [](const auto &ops, const auto &attrs) { + return window_reduce_impl(ops[0], attrs, 3); + }}, + + // ── window_scatter_max / min ─────────────────────────────────────────────── + // + // iattrs = [n_dims, lo0, hi0, …, s0, …, w0, …] + // Operands: [tensor_t, source, init_value] + // Replicates window_scatter_impl from emlx_nif.cpp. + + {"window_scatter_max", + [](const auto &ops, const auto &attrs) { + return window_scatter_impl_compiler(ops[0], ops[1], ops[2], attrs, true); + }}, + {"window_scatter_min", + [](const auto &ops, const auto &attrs) { + return window_scatter_impl_compiler(ops[0], ops[1], ops[2], attrs, false); + }}, + + // ── cumulative reductions ───────────────────────────────────────────────── + // + // iattrs = [axis, reverse_int] + // inclusive is always 1 (matches Nx semantics). + + {"cumulative_sum", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + bool rev = (attrs[1] != 0); + return mlx::core::cumsum(ops[0], axis, rev, true); + }}, + {"cumulative_product", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + bool rev = (attrs[1] != 0); + return mlx::core::cumprod(ops[0], axis, rev, true); + }}, + {"cumulative_min", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + bool rev = (attrs[1] != 0); + return mlx::core::cummin(ops[0], axis, rev, true); + }}, + {"cumulative_max", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + bool rev = (attrs[1] != 0); + return mlx::core::cummax(ops[0], axis, rev, true); + }}, + + // ── fft / ifft ──────────────────────────────────────────────────────────── + // + // iattrs = [axis, n] where n is the FFT length. + + {"fft", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + int n = static_cast(attrs[1]); + return mlx::core::fft::fft(ops[0], n, axis, mlx::core::fft::FFTNorm::Backward); + }}, + {"ifft", + [](const auto &ops, const auto &attrs) { + int axis = static_cast(attrs[0]); + int n = static_cast(attrs[1]); + return mlx::core::fft::ifft(ops[0], n, axis, mlx::core::fft::FFTNorm::Backward); + }}, + + // ── fft2 / ifft2 ────────────────────────────────────────────────────────── + // + // iattrs = [ax0, ax1, n0, n1] + + {"fft2", + [](const auto &ops, const auto &attrs) { + std::vector axes = {static_cast(attrs[0]), static_cast(attrs[1])}; + std::vector ns = {static_cast(attrs[2]), static_cast(attrs[3])}; + return mlx::core::fft::fft2(ops[0], to_shape(ns), axes, + mlx::core::fft::FFTNorm::Backward); + }}, + {"ifft2", + [](const auto &ops, const auto &attrs) { + std::vector axes = {static_cast(attrs[0]), static_cast(attrs[1])}; + std::vector ns = {static_cast(attrs[2]), static_cast(attrs[3])}; + return mlx::core::fft::ifft2(ops[0], to_shape(ns), axes, + mlx::core::fft::FFTNorm::Backward); + }}, + + // ── creation ops ───────────────────────────────────────────────────────── + // + // iota: attrs = [dtype_int, n_dims, axis_int, d0..dn-1] + // No operands. axis_int = -1 means flat enumeration (no axis). + {"iota", + [](const auto & /*ops*/, const auto &attrs) { + auto dtype = attrs[0].as_dtype(); + int n = static_cast(attrs[1]); + int axis = static_cast(attrs[2]); + + std::vector shape(n); + for (int i = 0; i < n; i++) + shape[i] = static_cast(attrs[3 + i]); + + if (n == 0) { + // scalar iota: always 0 + return mlx::core::astype(mlx::core::full({}, 0, mlx::core::int32), dtype); + } + + if (axis == -1) { + // flat enumeration: arange(0, product(shape)) reshaped + int total = 1; + for (int d : shape) + total *= d; + auto flat = mlx::core::arange(0, total, 1, mlx::core::int32); + auto reshaped = mlx::core::reshape(flat, to_shape(shape)); + return mlx::core::astype(reshaped, dtype); + } else { + // axis-specific iota: arange(0, shape[axis]), broadcast to full shape + int dim = shape[axis]; + auto linear = mlx::core::arange(0, dim, 1, mlx::core::int32); + std::vector rs(n, 1); + rs[axis] = dim; + auto r = mlx::core::reshape(linear, to_shape(rs)); + auto bc = mlx::core::broadcast_to(r, to_shape(shape)); + return mlx::core::astype(bc, dtype); + } + }}, + + // eye: attrs = [dtype_int, m, n]. No operands. + {"eye", + [](const auto & /*ops*/, const auto &attrs) { + auto dtype = attrs[0].as_dtype(); + int m = static_cast(attrs[1]); + int n = static_cast(attrs[2]); + return mlx::core::eye(m, n, 0, dtype); + }}, + + // ── conv_general ───────────────────────────────────────────────────────── + {"conv_general", + [](const auto &ops, const auto &attrs) { + int n_dims = static_cast(attrs[0]); + int off = 1; + + std::vector strides(n_dims), padding_lo(n_dims), padding_hi(n_dims), + kernel_dilation(n_dims), input_dilation(n_dims); + + for (int i = 0; i < n_dims; i++) + strides[i] = static_cast(attrs[off++]); + for (int i = 0; i < n_dims; i++) { + padding_lo[i] = static_cast(attrs[off++]); + padding_hi[i] = static_cast(attrs[off++]); + } + for (int i = 0; i < n_dims; i++) + kernel_dilation[i] = static_cast(attrs[off++]); + for (int i = 0; i < n_dims; i++) + input_dilation[i] = static_cast(attrs[off++]); + int fgs = static_cast(attrs[off]); + + return mlx::core::conv_general(ops[0], ops[1], strides, padding_lo, padding_hi, + kernel_dilation, input_dilation, fgs); + }}, + + // ── linalg (single-output) ──────────────────────────────────────────── + // + // MLX linalg primitives run on the CPU stream only; we pin them to + // k_linalg_cpu so they compose inside the compiled graph on any default + // device. Multi-output factorizations (qr/eigh/svd/lu) live in + // multi_op_registry below. + + // cholesky: operands = [a]; no attrs. Lower-triangular factor (upper=false). + {"cholesky", + [](const auto &ops, const auto &) { + return mlx::core::contiguous( + mlx::core::linalg::cholesky(ops[0], /*upper=*/false, k_linalg_cpu), + false, k_linalg_cpu); + }}, + + // solve: operands = [a, b]; no attrs. Solves A x = b. + {"solve", + [](const auto &ops, const auto &) { + return mlx::core::contiguous( + mlx::core::linalg::solve(ops[0], ops[1], k_linalg_cpu), false, k_linalg_cpu); + }}, + + // solve_triangular: operands = [a, b]; attrs = [upper_int]. + {"solve_triangular", + [](const auto &ops, const auto &attrs) { + bool upper = (attrs[0] != 0); + return mlx::core::contiguous( + mlx::core::linalg::solve_triangular(ops[0], ops[1], upper, k_linalg_cpu), + false, k_linalg_cpu); + }}, + + // ── EMLX.Fast fused kernels (mlx::core::fast) ────────────────────────── + // + // Recognized from `EMLX.Fast.*`'s `:__EMLX__`-tagged `Nx.Defn.Expr.metadata` + // nodes (see EMLX.Native.Expr's `:metadata` expand_node clause and + // EMLX.Fast's moduledoc). Run on the worker's default stream — these are + // Metal kernels, so the defn must be compiled/replayed on a GPU worker. + // Float opts (eps/scale/base) arrive as IEEE-754 bits via attr_to_float. + + // rms_norm: operands = [x, weight]; attrs = [eps_bits]. + {"fast_rms_norm", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::rms_norm(ops[0], ops[1], attr_to_float(attrs[0])); + }}, + + // layer_norm (with bias): operands = [x, weight, bias]; attrs = [eps_bits]. + {"fast_layer_norm", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::layer_norm(ops[0], ops[1], ops[2], attr_to_float(attrs[0])); + }}, + + // layer_norm (weight-only): operands = [x, weight]; attrs = [eps_bits]. + {"fast_layer_norm_no_bias", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::layer_norm(ops[0], ops[1], std::nullopt, + attr_to_float(attrs[0])); + }}, + + // swiglu: operands = [gate, up]; no attrs. silu(gate) * up. + {"fast_swiglu", + [](const auto &ops, const auto &) { + return mlx::core::multiply( + mlx::core::multiply(ops[0], mlx::core::sigmoid(ops[0])), ops[1]); + }}, + + // sdpa (no mask): operands = [q, k, v]; attrs = [scale_bits]. + {"fast_sdpa", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::scaled_dot_product_attention( + ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "", std::nullopt, + std::nullopt); + }}, + + // sdpa (no mask, + sinks): operands = [q, k, v, sinks]; attrs = [scale_bits]. + {"fast_sdpa_sinks", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::scaled_dot_product_attention( + ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "", std::nullopt, + ops[3]); + }}, + + // sdpa (array mask): operands = [q, k, v, mask]; attrs = [scale_bits]. + {"fast_sdpa_masked", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::scaled_dot_product_attention( + ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "array", ops[3], + std::nullopt); + }}, + + // sdpa (array mask, + sinks): operands = [q, k, v, mask, sinks]; + // attrs = [scale_bits]. + {"fast_sdpa_masked_sinks", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::scaled_dot_product_attention( + ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "array", ops[3], + ops[4]); + }}, + + // sdpa (causal): operands = [q, k, v]; attrs = [scale_bits]. + {"fast_sdpa_causal", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::scaled_dot_product_attention( + ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "causal", + std::nullopt, std::nullopt); + }}, + + // sdpa (causal, + sinks): operands = [q, k, v, sinks]; attrs = [scale_bits]. + {"fast_sdpa_causal_sinks", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::scaled_dot_product_attention( + ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "causal", + std::nullopt, ops[3]); + }}, + + // sdpa (causal + key_mask): operands = [q, k, v, key_mask]; + // attrs = [scale_bits, kv_offset]. Unlike the eager NIF, the compiled + // graph cannot branch on a runtime `all(key_mask)` (that forces an eval), + // so we always build the combined causal+key_mask additive mask in-graph. + {"fast_sdpa_causal_key_masked", + [](const auto &ops, const auto &attrs) { + const auto &q = ops[0]; + const auto &k = ops[1]; + const auto &v = ops[2]; + const auto &key_mask = ops[3]; + float scale = attr_to_float(attrs[0]); + int kv_offset = static_cast(attrs[1]); + + auto km = mlx::core::reshape( + key_mask, {key_mask.shape(0), 1, 1, key_mask.shape(1)}); + int T_q = q.shape(2); + int T_kv = k.shape(2); + + auto row = mlx::core::reshape(mlx::core::arange(T_q, mlx::core::int32), + {1, 1, T_q, 1}); + auto col = mlx::core::reshape(mlx::core::arange(T_kv, mlx::core::int32), + {1, 1, 1, T_kv}); + auto causal_bool = mlx::core::less_equal( + col, mlx::core::add(row, mlx::core::array(kv_offset, mlx::core::int32))); + auto keep = mlx::core::logical_and(km, causal_bool); + + auto mask_dtype = q.dtype(); + auto zero_val = mlx::core::zeros({}, mask_dtype); + auto neginf_val = mlx::core::full( + {}, -std::numeric_limits::infinity(), mask_dtype); + auto additive = mlx::core::where(keep, zero_val, neginf_val); + + return mlx::core::fast::scaled_dot_product_attention( + q, k, v, scale, "array", additive, std::nullopt); + }}, + + // sdpa (causal + key_mask, + sinks): operands = [q, k, v, key_mask, + // sinks]; attrs = [scale_bits, kv_offset]. Same in-graph combined mask as + // "fast_sdpa_causal_key_masked" above, plus the sinks operand. + {"fast_sdpa_causal_key_masked_sinks", + [](const auto &ops, const auto &attrs) { + const auto &q = ops[0]; + const auto &k = ops[1]; + const auto &v = ops[2]; + const auto &key_mask = ops[3]; + const auto &sinks = ops[4]; + float scale = attr_to_float(attrs[0]); + int kv_offset = static_cast(attrs[1]); + + auto km = mlx::core::reshape( + key_mask, {key_mask.shape(0), 1, 1, key_mask.shape(1)}); + int T_q = q.shape(2); + int T_kv = k.shape(2); + + auto row = mlx::core::reshape(mlx::core::arange(T_q, mlx::core::int32), + {1, 1, T_q, 1}); + auto col = mlx::core::reshape(mlx::core::arange(T_kv, mlx::core::int32), + {1, 1, 1, T_kv}); + auto causal_bool = mlx::core::less_equal( + col, mlx::core::add(row, mlx::core::array(kv_offset, mlx::core::int32))); + auto keep = mlx::core::logical_and(km, causal_bool); + + auto mask_dtype = q.dtype(); + auto zero_val = mlx::core::zeros({}, mask_dtype); + auto neginf_val = mlx::core::full( + {}, -std::numeric_limits::infinity(), mask_dtype); + auto additive = mlx::core::where(keep, zero_val, neginf_val); + + return mlx::core::fast::scaled_dot_product_attention( + q, k, v, scale, "array", additive, sinks); + }}, + + // rope (scalar offset): operands = [a]; + // attrs = [dims, traditional, base_bits, scale_bits, offset]. + {"fast_rope", + [](const auto &ops, const auto &attrs) { + return mlx::core::fast::rope( + ops[0], static_cast(attrs[0]), attrs[1] != 0, + attr_to_float(attrs[2]), attr_to_float(attrs[3]), + static_cast(attrs[4]), std::nullopt); + }}, + + // rope (per-batch offset array): operands = [a, position_ids]; + // attrs = [dims, traditional, base_bits, scale_bits]. Offsets are + // position_ids[:, 0] (matches EMLX.Fast.rope_with_positions_fast_callback). + {"fast_rope_ids", + [](const auto &ops, const auto &attrs) { + const auto &pos = ops[1]; + int B = pos.shape(0); + auto offsets = mlx::core::reshape( + mlx::core::slice(pos, to_shape({0, 0}), to_shape({B, 1}), + to_shape({1, 1})), + to_shape({B})); + return mlx::core::fast::rope( + ops[0], static_cast(attrs[0]), attrs[1] != 0, + attr_to_float(attrs[2]), attr_to_float(attrs[3]), offsets, + std::nullopt); + }}, + + // rope (precomputed freqs): operands = [a, position_ids, freqs]; + // attrs = [dims, traditional, scale_bits]. base = nullopt (freqs supplied). + {"fast_rope_with_freqs", + [](const auto &ops, const auto &attrs) { + const auto &pos = ops[1]; + int B = pos.shape(0); + auto offsets = mlx::core::reshape( + mlx::core::slice(pos, to_shape({0, 0}), to_shape({B, 1}), + to_shape({1, 1})), + to_shape({B})); + return mlx::core::fast::rope( + ops[0], static_cast(attrs[0]), attrs[1] != 0, std::nullopt, + attr_to_float(attrs[2]), offsets, ops[2]); + }}, + + // rope (arbitrary per-token position_ids, base-derived inv_freq): + // operands = [a, position_ids]; attrs = [dims, traditional, base_bits, + // scale_bits]. Prefill path (T>1) / high-base decode — mirrors + // emlx_fast.cpp's fast_rope_positions eager NIF. + {"fast_rope_positions", + [](const auto &ops, const auto &attrs) { + int dims = static_cast(attrs[0]); + bool traditional = attrs[1] != 0; + if (traditional) { + throw std::invalid_argument( + "fast_rope_positions: traditional=true not supported"); + } + float base = attr_to_float(attrs[2]); + float scale = attr_to_float(attrs[3]); + const auto &a = ops[0]; + const auto &pos = ops[1]; + int half = dims / 2; + + std::vector inv_freq_host(half); + for (int i = 0; i < half; ++i) { + float expo = (2.0f * static_cast(i)) / static_cast(dims); + inv_freq_host[i] = 1.0f / std::pow(base, expo); + } + auto inv_freq = + mlx::core::array(inv_freq_host.begin(), {half}, mlx::core::float32); + + int B = pos.shape(0); + int T = pos.shape(1); + auto pos_f = mlx::core::astype(pos, mlx::core::float32); + auto pos_bt1 = mlx::core::reshape(pos_f, to_shape({B, T, 1})); + auto inv_11h = mlx::core::reshape(inv_freq, to_shape({1, 1, half})); + auto angles = mlx::core::multiply( + mlx::core::multiply(pos_bt1, inv_11h), + mlx::core::array(scale, mlx::core::float32)); + + return rope_rotate_from_angles(a, angles, dims); + }}, + + // rope (arbitrary per-token position_ids, precomputed freqs): + // operands = [a, position_ids, freqs]; attrs = [dims, traditional, + // scale_bits]. Prefill path (T>1) for RoPE-scaling strategies that bake a + // freqs tensor (e.g. :llama3) — mlx::fast::rope's freqs overload takes + // `inv_freqs = reciprocal(freqs)` (see mlx/fast.cpp default_inv_freqs vs + // the inputs.size()==3 branch), so we replicate that reciprocal here + // rather than using `freqs` directly, to stay bit-for-bit with the eager + // EMLX.Fast.rope_with_freqs_callback path this replaces. + {"fast_rope_with_freqs_positions", + [](const auto &ops, const auto &attrs) { + int dims = static_cast(attrs[0]); + bool traditional = attrs[1] != 0; + if (traditional) { + throw std::invalid_argument( + "fast_rope_with_freqs_positions: traditional=true not supported"); + } + float scale = attr_to_float(attrs[2]); + const auto &a = ops[0]; + const auto &pos = ops[1]; + const auto &freqs = ops[2]; + int half = dims / 2; + + int B = pos.shape(0); + int T = pos.shape(1); + auto pos_f = mlx::core::astype(pos, mlx::core::float32); + auto pos_bt1 = mlx::core::reshape(pos_f, to_shape({B, T, 1})); + auto inv_freq = mlx::core::reciprocal(mlx::core::astype(freqs, mlx::core::float32)); + auto inv_11h = mlx::core::reshape(inv_freq, to_shape({1, 1, half})); + auto angles = mlx::core::multiply( + mlx::core::multiply(pos_bt1, inv_11h), + mlx::core::array(scale, mlx::core::float32)); + + return rope_rotate_from_angles(a, angles, dims); + }}, +}; + +// ── Multi-output op registry ────────────────────────────────────────────────── +// +// Ops that produce more than one result array (linalg factorizations). Their +// outputs are appended to the flat `results` accumulator in the returned order; +// the Elixir lowerer assigns consecutive result indices to the matching output +// refs (see to_native/1's list-result handling). Pinned to the CPU device. +using MultiOpFn = std::function< + std::vector(const std::vector &ops, + const std::vector &attrs)>; + +// Force each linalg output to contiguous layout (on the CPU stream). MLX's +// factorizations can return strided views; if such a view is a program output +// (or otherwise materialized directly), MLX tries to JIT a strided fused CPU +// kernel that can fail to compile. A plain contiguous Copy avoids that. +static std::vector +contiguous_all(std::vector arrs) { + for (auto &a : arrs) + a = mlx::core::contiguous(a, false, k_linalg_cpu); + return arrs; +} + +// A genuine `mlx::core::Primitive` for `Nx.runtime_call/4` — see +// EMLX.Native.Expr's moduledoc "Runtime calls" section. Unlike every other +// entry in this file (a plain function called once, during the single +// interpreter-lambda trace `mlx::core::detail::compile()` performs — see +// this file's header comment), a Primitive's `eval_cpu`/`eval_gpu` genuinely +// re-executes on every replay of the cached compiled tape, which is what +// lets the real Elixir callback fire once per `eval_program` call instead of +// once ever. `eval()` blocks the calling worker OS thread inside +// `invoke_runtime_call` (emlx_runtime_call_bridge.hpp) until the Elixir side +// replies via `EMLX.NIF.resolve_runtime_call/3`. +class EMLXRuntimeCall : public mlx::core::Primitive { +public: + EMLXRuntimeCall(mlx::core::Stream stream, int64_t callback_index) + : mlx::core::Primitive(stream), callback_index_(callback_index) {} + + void eval_cpu(const std::vector &inputs, + std::vector &outputs) override { + eval(inputs, outputs); + } + void eval_gpu(const std::vector &inputs, + std::vector &outputs) override { + eval(inputs, outputs); + } + + // Side-effecting (fires an Elixir callback with observable effects) — + // never dedup/CSE. This is also `Primitive::is_equivalent`'s own default, + // kept explicit here for clarity rather than relied upon. + bool is_equivalent(const mlx::core::Primitive &) const override { return false; } + + const char *name() const override { return "EMLXRuntimeCall"; } + +private: + void eval(const std::vector &inputs, + std::vector &outputs) { + emlx::native::invoke_runtime_call(callback_index_, inputs, outputs); + } + + int64_t callback_index_; +}; + +static const std::unordered_map multi_op_registry = { + // runtime_call: operands = flattened callback-argument leaves, in order. + // attrs = [callback_index, n_outputs, dtype0, n_dims0, d0.., dtype1, + // n_dims1, d1.., ...] — see EMLX.Native.Expr's moduledoc iattrs table. + // outputs = one array per declared output, produced by the + // EMLXRuntimeCall primitive above (never eagerly computed here). + {"runtime_call", + [](const auto &ops, const auto &attrs) -> std::vector { + size_t off = 0; + int64_t callback_index = attrs[off++]; + int64_t n_outputs = attrs[off++]; + + std::vector shapes; + std::vector dtypes; + shapes.reserve(static_cast(n_outputs)); + dtypes.reserve(static_cast(n_outputs)); + + for (int64_t i = 0; i < n_outputs; i++) { + dtypes.push_back(attrs[off++].as_dtype()); + int64_t n_dims = attrs[off++]; + std::vector dims(static_cast(n_dims)); + for (int64_t d = 0; d < n_dims; d++) { + dims[static_cast(d)] = static_cast(attrs[off++]); + } + shapes.push_back(to_shape(dims)); + } + + // Pinned to the CPU stream (like the linalg factorizations above), + // regardless of the compiled graph's default device. eval_gpu is + // never actually reached this way: the primitive does no Metal work + // of its own (it only blocks the worker thread on the Elixir + // round-trip and memcpy's the reply into the output buffer), and + // running it under mlx::core::gpu::eval's Metal command-buffer + // bookkeeping segfaults — see workdir/native-compiler/32a for + // details. MLX handles the cross-stream data dependencies (GPU + // operand arrays feeding a CPU-pinned primitive, and vice versa) + // the same way it does for the linalg ops. + auto primitive = std::make_shared( + mlx::core::default_stream(k_linalg_cpu), callback_index); + return mlx::core::array::make_arrays(shapes, dtypes, primitive, ops); + }}, + + // qr (reduced mode): operands = [a]; outputs = [q, r]. + {"qr", + [](const auto &ops, const auto &) -> std::vector { + auto [q, r] = mlx::core::linalg::qr(ops[0], k_linalg_cpu); + return contiguous_all({q, r}); + }}, + + // eigh (lower triangle): operands = [a]; outputs = [eigenvalues, eigenvectors]. + {"eigh", + [](const auto &ops, const auto &) -> std::vector { + auto [w, v] = mlx::core::linalg::eigh(ops[0], "L", k_linalg_cpu); + return contiguous_all({w, v}); + }}, + + // svd (full matrices): operands = [a]; outputs = [u, s, vt]. + {"svd", + [](const auto &ops, const auto &) -> std::vector { + return contiguous_all(mlx::core::linalg::svd(ops[0], /*compute_uv=*/true, k_linalg_cpu)); + }}, + + // lu: operands = [a]; outputs = [pivots, l, u] (pivots is a uint32 index + // vector; the lowerer rebuilds the permutation matrix via eye + take). + {"lu", + [](const auto &ops, const auto &) -> std::vector { + return contiguous_all(mlx::core::linalg::lu(ops[0], k_linalg_cpu)); + }}, +}; + +// ── Global compile-cache mutex ──────────────────────────────────────────────── +// +// mlx::core::detail::compile and compile_erase both mutate MLX's process-wide +// compile cache. compile_program runs on the worker thread; ~Expr runs on +// whichever BEAM scheduler/GC thread drops the last resource reference. These +// two paths can race, corrupting the cache and causing stale graph replay (e.g. +// a static-indexed put_slice graph replayed for a dynamic-indexed program). +// +// All three cache-touching calls (compile, erase, and the first compiled_fn +// invocation that inserts the traced graph) are serialised through this mutex. +static std::mutex s_mlx_compile_mutex; + +// ── Expr destructor ─────────────────────────────────────────────────────────── +// +// Evicts the per-Expr entry from MLX's global compile cache so stale compiled +// graphs don't accumulate. Called by default_dtor when the BEAM resource +// is GC'd. + +Expr::~Expr() { + if (compile_id != 0) { + std::lock_guard lk(s_mlx_compile_mutex); + mlx::core::detail::compile_erase(compile_id); + } +} + +// ── NIF implementations ─────────────────────────────────────────────────────── + +// compile_program — decodes the wire Program (see EMLX.Native.Program / +// EMLX.Native.Expr.to_native/1, decoded directly by fine::Decoder in +// emlx_compiler.hpp), builds a capturing interpreter lambda backed by the op +// registry, wraps it with mlx::core::compile(), and stores the result as an +// opaque Expr BEAM resource. +fine::Term compile_program_impl(ErlNifEnv *env, Program program) { + // Validate all op names against the registry up front so that any unknown + // op surfaces here rather than inside the lambda at (first) eval time. + bool has_runtime_call = false; + for (const auto &instr : program.instructions) { + const std::string &name = instr.op.to_string(); + if (op_registry.find(name) == op_registry.end() && + multi_op_registry.find(name) == multi_op_registry.end()) + throw std::runtime_error("emlx::native: unknown op \"" + name + "\""); + if (name == "runtime_call") + has_runtime_call = true; + } + + // Build constant arrays on the current (worker) thread using its default stream. + std::vector constants; + constants.reserve(program.constants.size()); + for (const auto &[value, dtype] : program.constants) { + constants.push_back(mlx::core::full({}, value, dtype)); + } + + int num_inputs_val = program.num_inputs; + + // Build the interpreter lambda capturing all program data, then pass it + // through mlx::core::compile(). MLX traces the lambda on the first + // eval_program call (building a compiled computation graph) and replays + // the cached graph on every subsequent call — no repeated graph construction. + emlx::function fn = + [captures = std::move(program.captures), + constants = std::move(constants), + instructions = std::move(program.instructions), + output_refs = std::move(program.outputs)]( + const std::vector &inputs) + -> std::vector { + std::vector results; + results.reserve(instructions.size()); + + auto resolve = [&](const Ref &ref) -> mlx::core::array { + switch (ref.kind) { + case RefKind::Input: + return inputs.at(static_cast(ref.index)); + case RefKind::Capture: + return captures.at(static_cast(ref.index)); + case RefKind::Const: + return constants.at(static_cast(ref.index)); + case RefKind::Result: + return results.at(static_cast(ref.index)); + } + throw std::runtime_error("emlx::native: invalid ref kind"); + }; + + for (const auto &instr : instructions) { + std::vector op_inputs; + op_inputs.reserve(instr.operands.size()); + for (const auto &ref : instr.operands) { + op_inputs.push_back(resolve(ref)); + } + + const std::string &name = instr.op.to_string(); + auto multi_it = multi_op_registry.find(name); + if (multi_it != multi_op_registry.end()) { + // Multi-output op: append each result in order to the flat accumulator. + auto outs = multi_it->second(op_inputs, instr.attrs); + for (auto &o : outs) + results.push_back(o); + } else { + results.push_back(op_registry.at(name)(op_inputs, instr.attrs)); + } + } + + std::vector outputs; + outputs.reserve(output_refs.size()); + for (const auto &ref : output_refs) { + outputs.push_back(resolve(ref)); + } + return outputs; + }; + + // Allocate the program resource. + auto *ptr = static_cast( + enif_alloc_resource(resource_object::type, sizeof(Expr))); + if (!ptr) + throw std::runtime_error("Failed to allocate Expr resource"); + + // Assign a unique ID so MLX's global compile cache has a distinct entry per + // Expr resource. All our lambdas share the same C++ type (same capture + // types), so the public mlx::core::compile() would map them all to the same + // cache key — causing stale graph reuse across different compiled programs. + static std::atomic next_id{1}; + std::uintptr_t unique_id = next_id.fetch_add(1, std::memory_order_relaxed); + + new (ptr) Expr(); + ptr->num_inputs = num_inputs_val; + ptr->compile_id = unique_id; + ptr->has_runtime_call = has_runtime_call; + { + std::lock_guard lk(s_mlx_compile_mutex); + ptr->compiled_fn = mlx::core::detail::compile(std::move(fn), unique_id); + } + + ERL_NIF_TERM ret = enif_make_resource(env, ptr); + enif_release_resource(ptr); + return fine::Term(ret); +} +FINE_ASYNC_NIF(compile_program) + +// eval_program — calls the MLX-compiled function against runtime inputs. +// MLX traces on the first call and replays the cached graph on subsequent calls. +// Returns lazy output array refs — materialization is deferred to the caller +// (to_binary / Nx.to_number), matching the Evaluator's deferred-eval pattern. +// +// argv[0] : program_ref (emlx::native::Expr resource) +// argv[1] : input_refs (list of MLX array resource refs — runtime inputs) +ERL_NIF_TERM eval_program(ErlNifEnv *env, int argc, + const ERL_NIF_TERM argv[]) { + try { + Expr *prog; + if (!enif_get_resource(env, argv[0], resource_object::type, + reinterpret_cast(&prog))) + return nx::nif::error(env, "Invalid Expr resource"); + + LIST_PARAM(1, std::vector, inputs); + + + // Force-evaluate all inputs on the worker thread before passing them to the + // compiled function. MLX tensors created via Elixir (e.g. Nx.tensor/2) may + // arrive as unevaluated lazy nodes associated with a different computation. + // If they are consumed inside compiled_fn while still lazy, MLX may + // propagate stale or garbage data from a previously-evaluated graph that + // happened to share the same underlying buffer. Forcing eval here ensures + // all inputs are materialized scalars/arrays on the worker's stream before + // the compiled graph is built or replayed. + mlx::core::eval(inputs); + + std::vector outputs; + { + std::lock_guard lk(s_mlx_compile_mutex); + outputs = prog->compiled_fn(inputs); + } + + // A program containing an inlined `:runtime_call` node must be forced to + // materialize now, while the caller pid this NIF call was dispatched + // with (emlx::g_current_caller_pid, set by emlx::async_dispatch — see + // emlx_async.hpp) is still in scope: EMLXRuntimeCall::eval_cpu/eval_gpu + // reads it to know which BEAM process to send the round-trip request + // to. Deliberately outside `s_mlx_compile_mutex`'s scope above — that + // lock is process-wide (shared with every other worker's + // compile_program/eval_program calls), and a runtime_call's blocking + // wait can run arbitrary, unbounded Elixir code, including calls that + // themselves need that same lock on another worker thread. Every other + // program keeps today's lazy/deferred return (no eval() call here). + if (prog->has_runtime_call) { + mlx::core::eval(outputs); + mlx::core::synchronize(); + } + + size_t n = outputs.size(); + std::vector terms; + terms.reserve(n); + for (size_t i = 0; i < n; i++) { + terms.push_back(create_tensor_resource(env, outputs[i])); + } + ERL_NIF_TERM list = + enif_make_list_from_array(env, terms.data(), static_cast(n)); + return nx::nif::ok(env, list); + } + CATCH() +} + +} // namespace native +} // namespace emlx diff --git a/emlx/c_src/emlx_compiler.hpp b/emlx/c_src/emlx_compiler.hpp new file mode 100644 index 0000000..4d667d3 --- /dev/null +++ b/emlx/c_src/emlx_compiler.hpp @@ -0,0 +1,200 @@ +#pragma once + +// emlx_compiler.hpp — emlx::native namespace: Op enum, Expr program struct, +// and NIF implementation declarations for compile_program / eval_program / +// native_expr_opcode_table. +// +// The three *_impl functions are implemented in emlx_compiler.cpp and called +// from thin NIF wrappers in emlx_nif.cpp. + +#include "emlx_nif_shared.hpp" +#include + +namespace emlx { +namespace native { + +// ── Program wire types ──────────────────────────────────────────────────── +// +// Elixir counterparts: EMLX.Native.Instruction / EMLX.Native.Program +// (lib/emlx/native/{instruction,program}.ex), produced by EMLX.Native.Expr.to_native/1. Decoded +// by the fine::Decoder specializations below instead of the old to_native/1 +// format's positional NIF args + bit-packed-int refs. + +// A reference to an already-produced value ({:input, i} / {:capture, i} / +// {:const, i} / {:result, i} on the Elixir side), resolved against the +// runtime inputs / closed-over captures / closed-over constants / the flat +// per-eval results accumulator, respectively. +enum class RefKind { Input, Capture, Const, Result }; + +struct Ref { + RefKind kind; + int64_t index; +}; + +// One instruction attribute. Most are plain integers (shapes, axes, flags, +// f64_bits-encoded floats); a handful are MLX dtypes or quantized_matmul +// mode strings, sent as atoms so no int<->meaning lookup table needs to be +// kept in sync with Elixir (see string2dtype/dtype_map in +// emlx_nif_shared.hpp). The implicit int64_t conversion below keeps every +// existing `attrs[i]`-as-int64_t use site in emlx_compiler.cpp's op registry +// compiling unchanged. +class Attr { +public: + Attr(int64_t v) : value_(v) {} + Attr(fine::Atom a) : value_(std::move(a)) {} + + operator int64_t() const { return std::get(value_); } + + mlx::core::Dtype as_dtype() const { + return string2dtype(std::get(value_).to_string()); + } + + std::string as_mode() const { return std::get(value_).to_string(); } + +private: + std::variant value_; +}; + +struct Instruction { + fine::Atom op; + std::vector operands; + std::vector attrs; +}; + +struct Program { + int num_inputs; + std::vector captures; + std::vector> constants; + std::vector instructions; + std::vector outputs; +}; + +// Compiled representation of an EMLX.Native.Expr program. +// Stored as an opaque BEAM resource; one instance per compiled defn cache entry. +// compile_program bakes the program into a capturing lambda, wraps it with +// mlx::core::detail::compile() (using a unique per-Expr ID) so MLX traces and +// caches the graph. eval_program just calls compiled_fn(inputs). +// The destructor evicts the per-ID entry from MLX's global compile cache. +struct Expr { + int num_inputs = 0; + std::uintptr_t compile_id = 0; // unique key for mlx::core::detail compile cache + // True when this program contains an inlined `:runtime_call` node — see + // eval_program: such a program is force-evaluated (mlx::core::eval on the + // outputs) before returning, instead of the usual deferred/lazy return, so + // EMLXRuntimeCall::eval_cpu/eval_gpu fires while this NIF call's caller + // pid (emlx::g_current_caller_pid) is still in scope. + bool has_runtime_call = false; + emlx::function compiled_fn; + + ~Expr(); +}; + +// compile_program is defined via FINE_ASYNC_NIF(compile_program) in +// emlx_compiler.cpp (declares `compile_program`/`compile_program_async` +// here); eval_program is a plain hand-written NIF, called from a thin +// wrapper in emlx_nif.cpp. +ERL_NIF_TERM compile_program(ErlNifEnv *env, int argc, + const ERL_NIF_TERM argv[]); +ERL_NIF_TERM compile_program_async(ErlNifEnv *env, int argc, + const ERL_NIF_TERM argv[]); +ERL_NIF_TERM eval_program(ErlNifEnv *env, int argc, + const ERL_NIF_TERM argv[]); + +} // namespace native +} // namespace emlx + +namespace fine { + +template <> struct Decoder { + static emlx::native::Ref decode(ErlNifEnv *env, const ERL_NIF_TERM &term) { + auto [kind_atom, index] = + fine::decode>(env, term); + if (kind_atom == "input") + return {emlx::native::RefKind::Input, index}; + if (kind_atom == "capture") + return {emlx::native::RefKind::Capture, index}; + if (kind_atom == "const") + return {emlx::native::RefKind::Const, index}; + if (kind_atom == "result") + return {emlx::native::RefKind::Result, index}; + throw std::invalid_argument("decode failed, unknown ref kind: " + + kind_atom.to_string()); + } +}; + +template <> struct Decoder { + static emlx::native::Attr decode(ErlNifEnv *env, const ERL_NIF_TERM &term) { + // Try int64 first (the overwhelmingly common case: shapes, axes, flags, + // f64_bits-encoded floats); fall back to atom (dtypes, quant modes). + ErlNifSInt64 v; + if (enif_get_int64(env, term, &v)) { + return emlx::native::Attr(static_cast(v)); + } + return emlx::native::Attr(fine::decode(env, term)); + } +}; + +// Custom (rather than the generic T::module/T::fields struct-decode +// mechanism) so field lookup failures produce a clear, specific error +// message without needing constexpr member-pointer/Atom-pointer tables. +template <> struct Decoder { + static emlx::native::Instruction decode(ErlNifEnv *env, + const ERL_NIF_TERM &term) { + static const fine::Atom op_atom("op"); + static const fine::Atom operands_atom("operands"); + static const fine::Atom attrs_atom("attrs"); + + ERL_NIF_TERM op_term, operands_term, attrs_term; + if (!enif_get_map_value(env, term, fine::encode(env, op_atom), &op_term) || + !enif_get_map_value(env, term, fine::encode(env, operands_atom), + &operands_term) || + !enif_get_map_value(env, term, fine::encode(env, attrs_atom), + &attrs_term)) { + throw std::invalid_argument( + "decode failed, expected an EMLX.Native.Instruction struct, " + "got: " + + format_term(env, term)); + } + + return emlx::native::Instruction{ + fine::decode(env, op_term), + fine::decode>(env, operands_term), + fine::decode>(env, attrs_term)}; + } +}; + +template <> struct Decoder { + static emlx::native::Program decode(ErlNifEnv *env, const ERL_NIF_TERM &term) { + static const fine::Atom num_inputs_atom("num_inputs"); + static const fine::Atom captures_atom("captures"); + static const fine::Atom constants_atom("constants"); + static const fine::Atom instructions_atom("instructions"); + static const fine::Atom outputs_atom("outputs"); + + auto get_field = [&](const fine::Atom &key) -> ERL_NIF_TERM { + ERL_NIF_TERM value; + if (!enif_get_map_value(env, term, fine::encode(env, key), &value)) { + throw std::invalid_argument( + "decode failed, expected an EMLX.Native.Program struct with " + "field " + + key.to_string() + ", got: " + format_term(env, term)); + } + return value; + }; + + emlx::native::Program program; + program.num_inputs = fine::decode(env, get_field(num_inputs_atom)); + program.captures = fine::decode>( + env, get_field(captures_atom)); + program.constants = + fine::decode>>( + env, get_field(constants_atom)); + program.instructions = fine::decode>( + env, get_field(instructions_atom)); + program.outputs = + fine::decode>(env, get_field(outputs_atom)); + return program; + } +}; + +} // namespace fine diff --git a/emlx/c_src/emlx_fast.cpp b/emlx/c_src/emlx_fast.cpp index 4531448..705133b 100644 --- a/emlx/c_src/emlx_fast.cpp +++ b/emlx/c_src/emlx_fast.cpp @@ -2,129 +2,112 @@ // mlx::fast ops — single fused Metal shaders // ============================================================================ +// +// Uses `fine`'s typed decode/encode (see the "`fine` bridging" section of +// emlx_nif_shared.hpp) instead of the hand-rolled PARAM/TENSOR_PARAM/TENSOR/ +// CATCH macros used elsewhere. Each NIF is a typed `NAME##_impl` function +// plus a one-line `FINE_ASYNC_NIF(NAME)` — the registered NIF name, arity, +// and the ASYNC_NIF/nif_funcs[] wiring in emlx_nif.cpp are all unchanged. // fast_rms_norm — fused RMS normalisation // MLX: mlx::fast::rms_norm(x, weight, eps, stream) → array, same shape as x -// weight is optional in the C++ API; TENSOR_PARAM gives array* which -// implicitly converts to optional. -NIF(fast_rms_norm) { - TENSOR_PARAM(0, x); - TENSOR_PARAM(1, weight); - PARAM(2, double, eps); - DEVICE_PARAM(3, device); - - TENSOR(fast::rms_norm(*x, *weight, (float)eps, device)); +// weight is optional in the C++ API; a plain TensorArg implicitly +// converts to optional via array's own converting constructor. +mlx::core::array fast_rms_norm_impl(ErlNifEnv *env, TensorArg x, + TensorArg weight, double eps, + mlx::core::Device device) { + return fast::rms_norm(*x, *weight, (float)eps, device); } -ASYNC_NIF(fast_rms_norm) +FINE_ASYNC_NIF(fast_rms_norm) // fast_rope — fused rotary position embedding (scalar offset arity) // MLX: mlx::fast::rope(x, dims, traditional, base, scale, offset, freqs, stream) // base is optional; we always supply it. // traditional=false → split-half (Qwen3 convention). -NIF(fast_rope) { - TENSOR_PARAM(0, a); - PARAM(1, int, dims); - PARAM(2, bool, traditional); - PARAM(3, double, base); - PARAM(4, double, scale); - PARAM(5, int, offset); - DEVICE_PARAM(6, device); - - TENSOR(fast::rope(*a, dims, traditional, (float)base, (float)scale, - offset, std::nullopt, device)); +mlx::core::array fast_rope_impl(ErlNifEnv *env, TensorArg a, int dims, + bool traditional, double base, double scale, + int offset, mlx::core::Device device) { + return fast::rope(*a, dims, traditional, (float)base, (float)scale, offset, + std::nullopt, device); } -ASYNC_NIF(fast_rope) +FINE_ASYNC_NIF(fast_rope) // fast_sdpa — flash-attention SDPA, no mask // MLX: mlx::fast::scaled_dot_product_attention(q, k, v, scale, mask_mode, mask_arr, sinks, stream) // GQA-native: k/v may have fewer heads than q. -NIF(fast_sdpa) { - TENSOR_PARAM(0, q); - TENSOR_PARAM(1, k); - TENSOR_PARAM(2, v); - PARAM(3, double, scale); - DEVICE_PARAM(4, device); - - TENSOR(fast::scaled_dot_product_attention( - *q, *k, *v, (float)scale, "", std::nullopt, std::nullopt, device)); +// sinks (optional, {N_q} or {N_q, ...} learned per-head attention-sink logits) +// is `nil` from Elixir when absent — see the std::optional decoder. +mlx::core::array fast_sdpa_impl(ErlNifEnv *env, TensorArg q, TensorArg k, + TensorArg v, double scale, + std::optional sinks, + mlx::core::Device device) { + std::optional sinks_opt = + sinks ? std::make_optional(**sinks) : std::nullopt; + + return fast::scaled_dot_product_attention( + *q, *k, *v, (float)scale, "", std::nullopt, sinks_opt, device); } -ASYNC_NIF(fast_sdpa) +FINE_ASYNC_NIF(fast_sdpa) // fast_sdpa_masked — flash-attention SDPA with additive or boolean mask // mask_mode="array" tells MLX to treat mask_arr as an additive bias/bool mask. -NIF(fast_sdpa_masked) { - TENSOR_PARAM(0, q); - TENSOR_PARAM(1, k); - TENSOR_PARAM(2, v); - PARAM(3, double, scale); - TENSOR_PARAM(4, mask); - DEVICE_PARAM(5, device); - - TENSOR(fast::scaled_dot_product_attention( - *q, *k, *v, (float)scale, "array", *mask, std::nullopt, device)); +mlx::core::array fast_sdpa_masked_impl(ErlNifEnv *env, TensorArg q, + TensorArg k, TensorArg v, double scale, + TensorArg mask, + std::optional sinks, + mlx::core::Device device) { + std::optional sinks_opt = + sinks ? std::make_optional(**sinks) : std::nullopt; + + return fast::scaled_dot_product_attention( + *q, *k, *v, (float)scale, "array", *mask, sinks_opt, device); } -ASYNC_NIF(fast_sdpa_masked) +FINE_ASYNC_NIF(fast_sdpa_masked) // fast_layer_norm — fused layer normalisation // MLX: mlx::fast::layer_norm(x, weight?, bias?, eps, stream) // weight and bias are optional; we always provide both. -NIF(fast_layer_norm) { - TENSOR_PARAM(0, x); - TENSOR_PARAM(1, weight); - TENSOR_PARAM(2, bias); - PARAM(3, double, eps); - DEVICE_PARAM(4, device); - - TENSOR(fast::layer_norm(*x, *weight, *bias, (float)eps, device)); +mlx::core::array fast_layer_norm_impl(ErlNifEnv *env, TensorArg x, + TensorArg weight, TensorArg bias, + double eps, mlx::core::Device device) { + return fast::layer_norm(*x, *weight, *bias, (float)eps, device); } -ASYNC_NIF(fast_layer_norm) +FINE_ASYNC_NIF(fast_layer_norm) // fast_layer_norm_no_bias — fused layer norm without bias (weight-only variant) -NIF(fast_layer_norm_no_bias) { - TENSOR_PARAM(0, x); - TENSOR_PARAM(1, weight); - PARAM(2, double, eps); - DEVICE_PARAM(3, device); - - TENSOR(fast::layer_norm(*x, *weight, std::nullopt, (float)eps, device)); +mlx::core::array fast_layer_norm_no_bias_impl(ErlNifEnv *env, TensorArg x, + TensorArg weight, double eps, + mlx::core::Device device) { + return fast::layer_norm(*x, *weight, std::nullopt, (float)eps, device); } -ASYNC_NIF(fast_layer_norm_no_bias) +FINE_ASYNC_NIF(fast_layer_norm_no_bias) // fast_rope_ids — fused RoPE with per-batch offset array (position_ids) // Calls the array-offset overload of mlx::fast::rope. // offset must be shape {B} — one starting position per batch example. // Assumes positions are sequential within each example: [offset[b], offset[b]+1, ..., offset[b]+T-1]. -NIF(fast_rope_ids) { - TENSOR_PARAM(0, a); - PARAM(1, int, dims); - PARAM(2, bool, traditional); - PARAM(3, double, base); - PARAM(4, double, scale); - TENSOR_PARAM(5, offset); - DEVICE_PARAM(6, device); - - TENSOR(fast::rope(*a, dims, traditional, (float)base, (float)scale, - *offset, std::nullopt, device)); +mlx::core::array fast_rope_ids_impl(ErlNifEnv *env, TensorArg a, int dims, + bool traditional, double base, + double scale, TensorArg offset, + mlx::core::Device device) { + return fast::rope(*a, dims, traditional, (float)base, (float)scale, *offset, + std::nullopt, device); } -ASYNC_NIF(fast_rope_ids) +FINE_ASYNC_NIF(fast_rope_ids) // fast_rope_with_freqs — fused RoPE with precomputed inv-frequency vector // Calls the freqs overload of mlx::fast::rope (base=nullopt, freqs supplied). // offset must be shape {B} — one starting position per batch example. // freqs must be shape {dims/2} — precomputed inverse frequencies. -NIF(fast_rope_with_freqs) { - TENSOR_PARAM(0, a); - PARAM(1, int, dims); - PARAM(2, bool, traditional); - PARAM(3, double, scale); - TENSOR_PARAM(4, offset); - TENSOR_PARAM(5, freqs); - DEVICE_PARAM(6, device); - - TENSOR(fast::rope(*a, dims, traditional, std::nullopt, (float)scale, - *offset, *freqs, device)); +mlx::core::array fast_rope_with_freqs_impl(ErlNifEnv *env, TensorArg a, + int dims, bool traditional, + double scale, TensorArg offset, + TensorArg freqs, + mlx::core::Device device) { + return fast::rope(*a, dims, traditional, std::nullopt, (float)scale, + *offset, *freqs, device); } -ASYNC_NIF(fast_rope_with_freqs) +FINE_ASYNC_NIF(fast_rope_with_freqs) // fast_rope_positions — RoPE for arbitrary per-token position_ids. // @@ -141,99 +124,100 @@ ASYNC_NIF(fast_rope_with_freqs) // Notes: // - `traditional=true` is currently unsupported in this fallback NIF. // - Intended for high-base / padded paths where fast_rope_ids is not correct. -NIF(fast_rope_positions) { - TENSOR_PARAM(0, a); - PARAM(1, int, dims); - PARAM(2, bool, traditional); - PARAM(3, double, base); - PARAM(4, double, scale); - TENSOR_PARAM(5, position_ids); - DEVICE_PARAM(6, device); - - try { - if (a->ndim() != 4) { - return nx::nif::error(env, "fast_rope_positions expects a rank-4 tensor a {B,T,H,D}"); - } - if (position_ids->ndim() != 2) { - return nx::nif::error(env, "fast_rope_positions expects rank-2 position_ids {B,T}"); - } +mlx::core::array fast_rope_positions_impl(ErlNifEnv *env, TensorArg a, + int dims, bool traditional, + double base, double scale, + TensorArg position_ids, + mlx::core::Device device) { + if (a->ndim() != 4) { + throw std::invalid_argument( + "fast_rope_positions expects a rank-4 tensor a {B,T,H,D}"); + } + if (position_ids->ndim() != 2) { + throw std::invalid_argument( + "fast_rope_positions expects rank-2 position_ids {B,T}"); + } - int B = a->shape(0); - int T = a->shape(1); - int H = a->shape(2); - int D = a->shape(3); + int B = a->shape(0); + int T = a->shape(1); + int H = a->shape(2); + int D = a->shape(3); - if (position_ids->shape(0) != B || position_ids->shape(1) != T) { - return nx::nif::error(env, "fast_rope_positions: position_ids shape must match {B,T} from a"); - } - if (dims <= 0 || dims > D || (dims % 2) != 0) { - return nx::nif::error(env, "fast_rope_positions: dims must be even and <= last dimension"); - } - if (traditional) { - return nx::nif::error(env, "fast_rope_positions: traditional=true not supported"); - } + if (position_ids->shape(0) != B || position_ids->shape(1) != T) { + throw std::invalid_argument( + "fast_rope_positions: position_ids shape must match {B,T} from a"); + } + if (dims <= 0 || dims > D || (dims % 2) != 0) { + throw std::invalid_argument( + "fast_rope_positions: dims must be even and <= last dimension"); + } + if (traditional) { + throw std::invalid_argument( + "fast_rope_positions: traditional=true not supported"); + } - int half = dims / 2; + int half = dims / 2; - std::vector inv_freq_host(half); - float base_f = static_cast(base); - for (int i = 0; i < half; ++i) { - float expo = (2.0f * static_cast(i)) / static_cast(dims); - inv_freq_host[i] = 1.0f / std::pow(base_f, expo); - } + std::vector inv_freq_host(half); + float base_f = static_cast(base); + for (int i = 0; i < half; ++i) { + float expo = (2.0f * static_cast(i)) / static_cast(dims); + inv_freq_host[i] = 1.0f / std::pow(base_f, expo); + } - auto inv_freq = array(inv_freq_host.begin(), {half}, float32); + auto inv_freq = array(inv_freq_host.begin(), {half}, float32); - auto pos_f = astype(*position_ids, float32, device); - auto pos_bt1 = reshape(pos_f, {B, T, 1}, device); - auto inv_11h = reshape(inv_freq, {1, 1, half}, device); - auto scale_arr = array(static_cast(scale), float32); - auto angles = multiply(multiply(pos_bt1, inv_11h, device), scale_arr, device); + auto pos_f = astype(*position_ids, float32, device); + auto pos_bt1 = reshape(pos_f, {B, T, 1}, device); + auto inv_11h = reshape(inv_freq, {1, 1, half}, device); + auto scale_arr = array(static_cast(scale), float32); + auto angles = multiply(multiply(pos_bt1, inv_11h, device), scale_arr, device); - auto cos_bt1h = astype(reshape(cos(angles, device), {B, T, 1, half}, device), a->dtype(), device); - auto sin_bt1h = astype(reshape(sin(angles, device), {B, T, 1, half}, device), a->dtype(), device); + auto cos_bt1h = astype(reshape(cos(angles, device), {B, T, 1, half}, device), a->dtype(), device); + auto sin_bt1h = astype(reshape(sin(angles, device), {B, T, 1, half}, device), a->dtype(), device); - auto cos_full = concatenate(std::vector{cos_bt1h, cos_bt1h}, 3, device); - auto sin_full = concatenate(std::vector{sin_bt1h, sin_bt1h}, 3, device); + auto cos_full = concatenate(std::vector{cos_bt1h, cos_bt1h}, 3, device); + auto sin_full = concatenate(std::vector{sin_bt1h, sin_bt1h}, 3, device); - auto x1 = slice(*a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, half}), device); - auto x2 = slice(*a, to_shape({0, 0, 0, half}), to_shape({B, T, H, dims}), device); - auto rotated = concatenate(std::vector{negative(x2, device), x1}, 3, device); + auto x1 = slice(*a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, half}), device); + auto x2 = slice(*a, to_shape({0, 0, 0, half}), to_shape({B, T, H, dims}), device); + auto rotated = concatenate(std::vector{negative(x2, device), x1}, 3, device); - auto a_head = slice(*a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, dims}), device); - auto rope_head = add(multiply(a_head, cos_full, device), multiply(rotated, sin_full, device), device); + auto a_head = slice(*a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, dims}), device); + auto rope_head = add(multiply(a_head, cos_full, device), multiply(rotated, sin_full, device), device); - if (dims == D) { - TENSOR(rope_head); - } else { - auto tail = slice(*a, to_shape({0, 0, 0, dims}), to_shape({B, T, H, D}), device); - TENSOR(concatenate(std::vector{rope_head, tail}, 3, device)); - } + if (dims == D) { + return rope_head; + } else { + auto tail = slice(*a, to_shape({0, 0, 0, dims}), to_shape({B, T, H, D}), device); + return concatenate(std::vector{rope_head, tail}, 3, device); } - CATCH() } -ASYNC_NIF(fast_rope_positions) +FINE_ASYNC_NIF(fast_rope_positions) // fast_sdpa_causal_key_masked — causal SDPA that checks key_mask at C++ level. // key_mask shape: {B, T_kv} — 1 = attend, 0 = padding. // If all values are 1 (no padding), dispatches to fast causal SDPA (no mask alloc). // Otherwise builds a combined causal + key_mask additive float mask and uses // masked SDPA. The all-ones check forces eval of only the small key_mask subgraph. -NIF(fast_sdpa_causal_key_masked) { - TENSOR_PARAM(0, q); // {B, N_q, T_q, D} - TENSOR_PARAM(1, k); // {B, N_kv, T_kv, D} - TENSOR_PARAM(2, v); // {B, N_kv, T_kv, D} - PARAM(3, double, scale); - TENSOR_PARAM(4, key_mask); // {B, T_kv} boolean / int - PARAM(5, int, kv_offset); // caller-controlled: decode → T_kv-1, prefill → 0 - DEVICE_PARAM(6, device); +// sinks (optional) — see fast_sdpa's comment above. +mlx::core::array fast_sdpa_causal_key_masked_impl( + ErlNifEnv *env, TensorArg q, // {B, N_q, T_q, D} + TensorArg k, // {B, N_kv, T_kv, D} + TensorArg v, // {B, N_kv, T_kv, D} + double scale, + TensorArg key_mask, // {B, T_kv} boolean / int + int kv_offset, // caller-controlled: decode → T_kv-1, prefill → 0 + std::optional sinks, mlx::core::Device device) { + std::optional sinks_opt = + sinks ? std::make_optional(**sinks) : std::nullopt; // key_mask values are 0/1 (int or bool); astype→bool_ then all() checks all non-zero. bool trivial = all(astype(*key_mask, bool_)).item(); if (trivial) { - TENSOR(fast::scaled_dot_product_attention( - *q, *k, *v, (float)scale, "causal", std::nullopt, std::nullopt, device)); + return fast::scaled_dot_product_attention( + *q, *k, *v, (float)scale, "causal", std::nullopt, sinks_opt, device); } else { // Expand key_mask {B, T_kv} → {B, 1, 1, T_kv} for broadcasting. auto km = reshape(*key_mask, {key_mask->shape(0), 1, 1, key_mask->shape(1)}); @@ -256,37 +240,36 @@ NIF(fast_sdpa_causal_key_masked) { auto neginf_val = full({}, -std::numeric_limits::infinity(), mask_dtype); auto additive = where(keep, zero_val, neginf_val); - TENSOR(fast::scaled_dot_product_attention( - *q, *k, *v, (float)scale, "array", additive, std::nullopt, device)); + return fast::scaled_dot_product_attention( + *q, *k, *v, (float)scale, "array", additive, sinks_opt, device); } } -ASYNC_NIF(fast_sdpa_causal_key_masked) +FINE_ASYNC_NIF(fast_sdpa_causal_key_masked) // fast_swiglu — fused SwiGLU: silu(gate) * up // SiLU (Sigmoid Linear Unit): silu(x) = x * sigmoid(x). // gate and up must have the same shape; output has the same shape and dtype. -NIF(fast_swiglu) { - TENSOR_PARAM(0, gate); - TENSOR_PARAM(1, up); - DEVICE_PARAM(2, device); - +mlx::core::array fast_swiglu_impl(ErlNifEnv *env, TensorArg gate, + TensorArg up, mlx::core::Device device) { // silu(gate) * up where silu(x) = x * sigmoid(x). // MLX's lazy graph evaluation fuses these into a single kernel dispatch. - TENSOR(multiply(multiply(*gate, sigmoid(*gate, device), device), *up, device)); + return multiply(multiply(*gate, sigmoid(*gate, device), device), *up, + device); } -ASYNC_NIF(fast_swiglu) - -NIF(fast_sdpa_causal) { - TENSOR_PARAM(0, q); - TENSOR_PARAM(1, k); - TENSOR_PARAM(2, v); - PARAM(3, double, scale); - DEVICE_PARAM(4, device); - - TENSOR(fast::scaled_dot_product_attention( - *q, *k, *v, (float)scale, "causal", std::nullopt, std::nullopt, device)); +FINE_ASYNC_NIF(fast_swiglu) + +// sinks (optional) — see fast_sdpa's comment above. +mlx::core::array fast_sdpa_causal_impl(ErlNifEnv *env, TensorArg q, + TensorArg k, TensorArg v, double scale, + std::optional sinks, + mlx::core::Device device) { + std::optional sinks_opt = + sinks ? std::make_optional(**sinks) : std::nullopt; + + return fast::scaled_dot_product_attention( + *q, *k, *v, (float)scale, "causal", std::nullopt, sinks_opt, device); } -ASYNC_NIF(fast_sdpa_causal) +FINE_ASYNC_NIF(fast_sdpa_causal) // kv_cache_attention — fused KV cache update + variable-length SDPA in one Metal pass. // @@ -306,95 +289,81 @@ ASYNC_NIF(fast_sdpa_causal) // // All three are eval'd in a single MLX command buffer — slice_update → valid_slice // → SDPA form one fused Metal encoder submission. -NIF(kv_cache_attention) { - TENSOR_PARAM(0, q); - TENSOR_PARAM(1, new_k); - TENSOR_PARAM(2, new_v); - TENSOR_PARAM(3, k_cache); - TENSOR_PARAM(4, v_cache); - PARAM(5, int, offset); - PARAM(6, double, scale); - DEVICE_PARAM(7, device); - - try { - int B = q->shape(0); - int T_q = q->shape(1); - int D = q->shape(3); - int N_kv = new_k->shape(2); - int T_new = new_k->shape(1); - int valid_len = offset + T_new; - - // Donate the cache buffers so MLX can reuse them in-place. - // After std::move, the ENIF resource blocks hold moved-from arrays. - // k_cache_owned / v_cache_owned are the sole shared_ptr owners until - // slice_update copies them into its inputs list (count rises to 2). - // When this function returns, those locals destruct → count drops to 1 - // → SliceUpdate::eval_gpu detects is_donatable() → no new 4 MB buffer. - auto k_cache_owned = std::move(*k_cache); - auto v_cache_owned = std::move(*v_cache); - - // 1. Insert new K/V at cache position `offset`. - // Output has the same shape as k_cache: {B, T_max, N_kv, D}. - auto k_upd = mlx::core::slice_update( - k_cache_owned, *new_k, - to_shape({0, offset, 0, 0}), - to_shape({B, valid_len, N_kv, D}), - device); - auto v_upd = mlx::core::slice_update( - v_cache_owned, *new_v, - to_shape({0, offset, 0, 0}), - to_shape({B, valid_len, N_kv, D}), - device); - - // 2. Slice to valid portion: {B, valid_len, N_kv, D}. - auto k_valid = mlx::core::slice( - k_upd, to_shape({0, 0, 0, 0}), to_shape({B, valid_len, N_kv, D}), device); - auto v_valid = mlx::core::slice( - v_upd, to_shape({0, 0, 0, 0}), to_shape({B, valid_len, N_kv, D}), device); - - // 3. Transpose from Bumblebee {B, T, N, D} to MLX SDPA {B, N, T, D}. - auto q_t = mlx::core::transpose(*q, {0, 2, 1, 3}, device); - auto k_valid_t = mlx::core::transpose(k_valid, {0, 2, 1, 3}, device); - auto v_valid_t = mlx::core::transpose(v_valid, {0, 2, 1, 3}, device); - - // 4. SDPA over the valid slice. - // T_new == 1 (decode): no mask — single query is trivially causal. - // T_new > 1 (prefill): build an additive causal mask in q's dtype to avoid - // any float32 promotion that would mismatch BF16 Q/K/V tensors. - auto build_prefill_mask = [&]() -> mlx::core::array { - 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); - - // Causal: query position i can attend to key position j iff j <= i + kv_offset. - int kv_offset = valid_len - T_q; - auto row = mlx::core::reshape(mlx::core::arange(T_q, mlx::core::int32, device), {1, 1, T_q, 1}, device); - auto col = mlx::core::reshape(mlx::core::arange(valid_len, mlx::core::int32, device), {1, 1, 1, valid_len}, device); - auto causal_bool = mlx::core::less_equal( - col, mlx::core::add(row, mlx::core::array(kv_offset, mlx::core::int32), device), device); +std::tuple +kv_cache_attention_impl(ErlNifEnv *env, TensorArg q, TensorArg new_k, + TensorArg new_v, TensorArg k_cache, TensorArg v_cache, + int offset, double scale, mlx::core::Device device) { + int B = q->shape(0); + int T_q = q->shape(1); + int D = q->shape(3); + int N_kv = new_k->shape(2); + int T_new = new_k->shape(1); + int valid_len = offset + T_new; + + // Donate the cache buffers so MLX can reuse them in-place. + // After std::move, the ENIF resource blocks hold moved-from arrays. + // k_cache_owned / v_cache_owned are the sole shared_ptr owners until + // slice_update copies them into its inputs list (count rises to 2). + // When this function returns, those locals destruct → count drops to 1 + // → SliceUpdate::eval_gpu detects is_donatable() → no new 4 MB buffer. + auto k_cache_owned = std::move(*k_cache); + auto v_cache_owned = std::move(*v_cache); + + // 1. Insert new K/V at cache position `offset`. + // Output has the same shape as k_cache: {B, T_max, N_kv, D}. + auto k_upd = mlx::core::slice_update( + k_cache_owned, *new_k, + to_shape({0, offset, 0, 0}), + to_shape({B, valid_len, N_kv, D}), + device); + auto v_upd = mlx::core::slice_update( + v_cache_owned, *new_v, + to_shape({0, offset, 0, 0}), + to_shape({B, valid_len, N_kv, D}), + device); + + // 2. Slice to valid portion: {B, valid_len, N_kv, D}. + auto k_valid = mlx::core::slice( + k_upd, to_shape({0, 0, 0, 0}), to_shape({B, valid_len, N_kv, D}), device); + auto v_valid = mlx::core::slice( + v_upd, to_shape({0, 0, 0, 0}), to_shape({B, valid_len, N_kv, D}), device); + + // 3. Transpose from Bumblebee {B, T, N, D} to MLX SDPA {B, N, T, D}. + auto q_t = mlx::core::transpose(*q, {0, 2, 1, 3}, device); + auto k_valid_t = mlx::core::transpose(k_valid, {0, 2, 1, 3}, device); + auto v_valid_t = mlx::core::transpose(v_valid, {0, 2, 1, 3}, device); + + // 4. SDPA over the valid slice. + // T_new == 1 (decode): no mask — single query is trivially causal. + // T_new > 1 (prefill): build an additive causal mask in q's dtype to avoid + // any float32 promotion that would mismatch BF16 Q/K/V tensors. + auto build_prefill_mask = [&]() -> mlx::core::array { + 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); - return mlx::core::where(causal_bool, zero_val, neginf_val, device); - }; + // Causal: query position i can attend to key position j iff j <= i + kv_offset. + int kv_offset = valid_len - T_q; + auto row = mlx::core::reshape(mlx::core::arange(T_q, mlx::core::int32, device), {1, 1, T_q, 1}, device); + auto col = mlx::core::reshape(mlx::core::arange(valid_len, mlx::core::int32, device), {1, 1, 1, valid_len}, device); + auto causal_bool = mlx::core::less_equal( + col, mlx::core::add(row, mlx::core::array(kv_offset, mlx::core::int32), device), device); - auto attn_t = (T_new == 1) - ? mlx::core::fast::scaled_dot_product_attention( - q_t, k_valid_t, v_valid_t, (float)scale, "", std::nullopt, std::nullopt, device) - : mlx::core::fast::scaled_dot_product_attention( - q_t, k_valid_t, v_valid_t, (float)scale, "array", build_prefill_mask(), std::nullopt, device); + return mlx::core::where(causal_bool, zero_val, neginf_val, device); + }; - // 5. Transpose output back: {B, N_q, T_q, D} → {B, T_q, N_q, D} (Bumblebee format). - auto attn_out = mlx::core::transpose(attn_t, {0, 2, 1, 3}, device); + auto attn_t = (T_new == 1) + ? mlx::core::fast::scaled_dot_product_attention( + q_t, k_valid_t, v_valid_t, (float)scale, "", std::nullopt, std::nullopt, device) + : mlx::core::fast::scaled_dot_product_attention( + q_t, k_valid_t, v_valid_t, (float)scale, "array", build_prefill_mask(), std::nullopt, device); - ERL_NIF_TERM result_tuple[3]; - result_tuple[0] = create_tensor_resource(env, attn_out); - result_tuple[1] = create_tensor_resource(env, k_upd); - result_tuple[2] = create_tensor_resource(env, v_upd); + // 5. Transpose output back: {B, N_q, T_q, D} → {B, T_q, N_q, D} (Bumblebee format). + auto attn_out = mlx::core::transpose(attn_t, {0, 2, 1, 3}, device); - return nx::nif::ok(env, enif_make_tuple3(env, result_tuple[0], result_tuple[1], result_tuple[2])); - } - CATCH() + return {attn_out, k_upd, v_upd}; } -ASYNC_NIF(kv_cache_attention) +FINE_ASYNC_NIF(kv_cache_attention) // kv_cache_attention_masked — same as kv_cache_attention but applies a key_mask // in addition to the causal mask. This is required for padded prefill where the @@ -409,129 +378,119 @@ ASYNC_NIF(kv_cache_attention) // kernels that may not be available for all tensor shapes. // // Returns same 3-tuple as kv_cache_attention: {attn_out, k_upd, v_upd}. -NIF(kv_cache_attention_masked) { - TENSOR_PARAM(0, q); // {B, T_q, N_q, D} Bumblebee format - TENSOR_PARAM(1, new_k); // {B, T_new, N_kv, D} - TENSOR_PARAM(2, new_v); // {B, T_new, N_kv, D} - TENSOR_PARAM(3, k_cache); // {B, T_max, N_kv, D} - TENSOR_PARAM(4, v_cache); // {B, T_max, N_kv, D} - PARAM(5, int, offset); - PARAM(6, double, scale); - TENSOR_PARAM(7, key_mask); // {B, T_kv} — 1=attend, 0=skip - DEVICE_PARAM(8, device); - - try { - int B = q->shape(0); - int T_q = q->shape(1); - int D = q->shape(3); - int N_kv = new_k->shape(2); - int T_new = new_k->shape(1); - int valid_len = offset + T_new; - - // Donate cache buffers (same pattern as kv_cache_attention). - auto k_cache_owned = std::move(*k_cache); - auto v_cache_owned = std::move(*v_cache); - - // 1. Insert new K/V at cache position `offset`. - auto k_upd = mlx::core::slice_update( - k_cache_owned, *new_k, - to_shape({0, offset, 0, 0}), - to_shape({B, valid_len, N_kv, D}), - device); - auto v_upd = mlx::core::slice_update( - v_cache_owned, *new_v, - to_shape({0, offset, 0, 0}), - to_shape({B, valid_len, N_kv, D}), - device); - - // 2. Slice to valid portion. - auto k_valid = mlx::core::slice( - k_upd, to_shape({0, 0, 0, 0}), to_shape({B, valid_len, N_kv, D}), device); - auto v_valid = mlx::core::slice( - v_upd, to_shape({0, 0, 0, 0}), to_shape({B, valid_len, N_kv, D}), device); - - // 3. Transpose from Bumblebee {B, T, N, D} to MLX SDPA {B, N, T, D}. - auto q_t = mlx::core::transpose(*q, {0, 2, 1, 3}, device); - auto k_valid_t = mlx::core::transpose(k_valid, {0, 2, 1, 3}, device); - auto v_valid_t = mlx::core::transpose(v_valid, {0, 2, 1, 3}, device); - - // 4. Compute attention output. - // - // For decode (T_q == 1), the causal constraint is trivially satisfied: a single - // query position always sees all valid keys. We skip the arange/reshape/less_equal - // construction and dispatch to the cheapest SDPA variant: - // - trivial key_mask (all-ones, non-padded batch): pure causal Metal kernel. - // - non-trivial key_mask (padded batch): key_mask-only additive mask. - // The all().item() sync is negligible at {B, valid_len} (1–256 elements). - // - // For prefill (T_q > 1), build the full causal + key_mask combined additive mask. - 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); +std::tuple +kv_cache_attention_masked_impl(ErlNifEnv *env, TensorArg q, // {B, T_q, N_q, D} Bumblebee format + TensorArg new_k, // {B, T_new, N_kv, D} + TensorArg new_v, // {B, T_new, N_kv, D} + TensorArg k_cache, // {B, T_max, N_kv, D} + TensorArg v_cache, // {B, T_max, N_kv, D} + int offset, double scale, + TensorArg key_mask, // {B, T_kv} — 1=attend, 0=skip + mlx::core::Device device) { + int B = q->shape(0); + int T_q = q->shape(1); + int D = q->shape(3); + int N_kv = new_k->shape(2); + int T_new = new_k->shape(1); + int valid_len = offset + T_new; + + // Donate cache buffers (same pattern as kv_cache_attention). + auto k_cache_owned = std::move(*k_cache); + auto v_cache_owned = std::move(*v_cache); + + // 1. Insert new K/V at cache position `offset`. + auto k_upd = mlx::core::slice_update( + k_cache_owned, *new_k, + to_shape({0, offset, 0, 0}), + to_shape({B, valid_len, N_kv, D}), + device); + auto v_upd = mlx::core::slice_update( + v_cache_owned, *new_v, + to_shape({0, offset, 0, 0}), + to_shape({B, valid_len, N_kv, D}), + device); + + // 2. Slice to valid portion. + auto k_valid = mlx::core::slice( + k_upd, to_shape({0, 0, 0, 0}), to_shape({B, valid_len, N_kv, D}), device); + auto v_valid = mlx::core::slice( + v_upd, to_shape({0, 0, 0, 0}), to_shape({B, valid_len, N_kv, D}), device); + + // 3. Transpose from Bumblebee {B, T, N, D} to MLX SDPA {B, N, T, D}. + auto q_t = mlx::core::transpose(*q, {0, 2, 1, 3}, device); + auto k_valid_t = mlx::core::transpose(k_valid, {0, 2, 1, 3}, device); + auto v_valid_t = mlx::core::transpose(v_valid, {0, 2, 1, 3}, device); + + // 4. Compute attention output. + // + // For decode (T_q == 1), the causal constraint is trivially satisfied: a single + // query position always sees all valid keys. We skip the arange/reshape/less_equal + // construction and dispatch to the cheapest SDPA variant: + // - trivial key_mask (all-ones, non-padded batch): pure causal Metal kernel. + // - non-trivial key_mask (padded batch): key_mask-only additive mask. + // The all().item() sync is negligible at {B, valid_len} (1–256 elements). + // + // For prefill (T_q > 1), build the full causal + key_mask combined additive mask. + 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); + + mlx::core::array attn_t = [&]() -> mlx::core::array { + if (T_q == 1) { + // For single-query decode, the causal constraint is trivially satisfied + // (one query always attends to all preceding keys). Build a key_mask-only + // additive mask — 3 GPU ops vs 8 in the original full causal+mask path. + // + // We intentionally skip the all().item() trivial check here: that + // check forces a GPU→CPU sync on every layer call (28×/step), whose + // latency cost exceeds the savings from choosing "causal" mode. For + // non-padded inference the additive mask is all-zeros, which is + // functionally identical to pure causal mode. + auto km = mlx::core::reshape( + mlx::core::astype(*key_mask, mlx::core::bool_, device), + {key_mask->shape(0), 1, 1, key_mask->shape(1)}, + device); + auto additive = mlx::core::where(km, zero_val, neginf_val, device); + return mlx::core::fast::scaled_dot_product_attention( + q_t, k_valid_t, v_valid_t, (float)scale, "array", + additive, std::nullopt, device); + } else { + // Prefill (T_q > 1): full causal + key_mask combined additive mask. + // kv_offset: valid_len - T_q (= 0 for a fresh prefill of length T_q). + int kv_offset = valid_len - T_q; + auto row = mlx::core::reshape( + mlx::core::arange(T_q, mlx::core::int32, device), {1, 1, T_q, 1}, device); + auto col = mlx::core::reshape( + mlx::core::arange(valid_len, mlx::core::int32, device), {1, 1, 1, valid_len}, device); + auto causal_bool = mlx::core::less_equal( + col, mlx::core::add(row, mlx::core::array(kv_offset, mlx::core::int32), device), device); - mlx::core::array attn_t = [&]() -> mlx::core::array { - if (T_q == 1) { - // For single-query decode, the causal constraint is trivially satisfied - // (one query always attends to all preceding keys). Build a key_mask-only - // additive mask — 3 GPU ops vs 8 in the original full causal+mask path. - // - // We intentionally skip the all().item() trivial check here: that - // check forces a GPU→CPU sync on every layer call (28×/step), whose - // latency cost exceeds the savings from choosing "causal" mode. For - // non-padded inference the additive mask is all-zeros, which is - // functionally identical to pure causal mode. - auto km = mlx::core::reshape( - mlx::core::astype(*key_mask, mlx::core::bool_, device), - {key_mask->shape(0), 1, 1, key_mask->shape(1)}, - device); - auto additive = mlx::core::where(km, zero_val, neginf_val, device); - return mlx::core::fast::scaled_dot_product_attention( - q_t, k_valid_t, v_valid_t, (float)scale, "array", - additive, std::nullopt, device); - } else { - // Prefill (T_q > 1): full causal + key_mask combined additive mask. - // kv_offset: valid_len - T_q (= 0 for a fresh prefill of length T_q). - int kv_offset = valid_len - T_q; - auto row = mlx::core::reshape( - mlx::core::arange(T_q, mlx::core::int32, device), {1, 1, T_q, 1}, device); - auto col = mlx::core::reshape( - mlx::core::arange(valid_len, mlx::core::int32, device), {1, 1, 1, valid_len}, device); - auto causal_bool = mlx::core::less_equal( - col, mlx::core::add(row, mlx::core::array(kv_offset, mlx::core::int32), device), device); - - auto km = mlx::core::reshape( - mlx::core::astype(*key_mask, mlx::core::bool_, device), - {key_mask->shape(0), 1, 1, key_mask->shape(1)}, - device); - auto keep = mlx::core::logical_and(km, causal_bool, device); - auto additive = mlx::core::where(keep, zero_val, neginf_val, device); - - return mlx::core::fast::scaled_dot_product_attention( - q_t, k_valid_t, v_valid_t, (float)scale, "array", - additive, std::nullopt, device); - } - }(); - - // Replace NaN with 0 for all-masked rows (softmax(-inf,...,-inf) = NaN in - // Flash-Attention when seq_len >= Metal tile size, but semantically = 0). - auto attn_safe = mlx::core::where( - mlx::core::isnan(attn_t), - mlx::core::zeros_like(attn_t, device), - attn_t, device); - - // 6. Transpose output back: {B, N_q, T_q, D} → {B, T_q, N_q, D}. - auto attn_out = mlx::core::transpose(attn_safe, {0, 2, 1, 3}, device); - - ERL_NIF_TERM result_tuple[3]; - result_tuple[0] = create_tensor_resource(env, attn_out); - result_tuple[1] = create_tensor_resource(env, k_upd); - result_tuple[2] = create_tensor_resource(env, v_upd); - - return nx::nif::ok(env, enif_make_tuple3(env, result_tuple[0], result_tuple[1], result_tuple[2])); - } - CATCH() + auto km = mlx::core::reshape( + mlx::core::astype(*key_mask, mlx::core::bool_, device), + {key_mask->shape(0), 1, 1, key_mask->shape(1)}, + device); + auto keep = mlx::core::logical_and(km, causal_bool, device); + auto additive = mlx::core::where(keep, zero_val, neginf_val, device); + + return mlx::core::fast::scaled_dot_product_attention( + q_t, k_valid_t, v_valid_t, (float)scale, "array", + additive, std::nullopt, device); + } + }(); + + // Replace NaN with 0 for all-masked rows (softmax(-inf,...,-inf) = NaN in + // Flash-Attention when seq_len >= Metal tile size, but semantically = 0). + auto attn_safe = mlx::core::where( + mlx::core::isnan(attn_t), + mlx::core::zeros_like(attn_t, device), + attn_t, device); + + // 6. Transpose output back: {B, N_q, T_q, D} → {B, T_q, N_q, D}. + auto attn_out = mlx::core::transpose(attn_safe, {0, 2, 1, 3}, device); + + return {attn_out, k_upd, v_upd}; } -ASYNC_NIF(kv_cache_attention_masked) +FINE_ASYNC_NIF(kv_cache_attention_masked) // kv_cache_sdpa_update — fused donation-optimised KV cache update + SDPA // for the native NIF loop (BNHD layout: {B, N, T, D}). @@ -548,7 +507,7 @@ ASYNC_NIF(kv_cache_attention_masked) // detects is_donatable() → the existing 4 MB Metal buffer is reused in-place, // no new allocation needed. // -// Inputs (argv indices inside the sync NIF body, after worker is stripped): +// Inputs: // q — {B, N_q, T_q, D} post-RoPE query (native BNHD) // new_k — {B, N_kv, T_new, D} post-RoPE key // new_v — {B, N_kv, T_new, D} value @@ -562,82 +521,68 @@ ASYNC_NIF(kv_cache_attention_masked) // attn_out — {B, N_q, T_q, D} (native BNHD, caller transposes/reshapes) // k_upd — {B, N_kv, T_max, D} same Metal buffer as k_cache, updated // v_upd — {B, N_kv, T_max, D} same Metal buffer as v_cache, updated -NIF(kv_cache_sdpa_update) { - TENSOR_PARAM(0, q); - TENSOR_PARAM(1, new_k); - TENSOR_PARAM(2, new_v); - TENSOR_PARAM(3, k_cache); - TENSOR_PARAM(4, v_cache); - PARAM(5, int, offset); - PARAM(6, double, scale); - DEVICE_PARAM(7, device); - - try { - int B = q->shape(0); - int N_kv = new_k->shape(1); - int T_new = new_k->shape(2); - int D = q->shape(3); - int valid_len = offset + T_new; - - // Move-extract: after this the ENIF resource blocks hold moved-from arrays. - // k_cache_owned / v_cache_owned are the sole owners (use_count == 1) until - // slice_update copies them into its inputs (use_count rises to 2). - // On function return both locals destruct → use_count drops to 1 again - // → is_donatable() is true at eval time → no new 4 MB Metal buffer. - auto k_cache_owned = std::move(*k_cache); - auto v_cache_owned = std::move(*v_cache); - - // 1. Insert new_k / new_v at cache position `offset` along axis 2 (T axis). - // Layout: {B, N_kv, T_max, D} — offset along dimension 2. - auto k_upd = mlx::core::slice_update( - k_cache_owned, *new_k, - to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), - device); - auto v_upd = mlx::core::slice_update( - v_cache_owned, *new_v, - to_shape({0, 0, offset, 0}), - to_shape({B, N_kv, valid_len, D}), - device); - - // 2. Slice valid prefix: k_upd[0:B, 0:N_kv, 0:valid_len, 0:D]. - 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); - - // 3. SDPA: q / k_valid / v_valid are already in {B, N, T, D} format. - // Decode (T_new == 1): no mask — single query is trivially causal. - // Prefill (T_new > 1): additive causal mask in q's dtype. - auto build_prefill_mask = [&]() -> mlx::core::array { - 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); - 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 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); - }; - - auto attn_out = (T_new == 1) - ? mlx::core::fast::scaled_dot_product_attention( - *q, k_valid, v_valid, (float)scale, "", std::nullopt, std::nullopt, device) - : mlx::core::fast::scaled_dot_product_attention( - *q, k_valid, v_valid, (float)scale, "array", build_prefill_mask(), std::nullopt, device); - // k_cache_owned and v_cache_owned destruct here → use_count 2 → 1. - - ERL_NIF_TERM result_tuple[3]; - result_tuple[0] = create_tensor_resource(env, attn_out); - result_tuple[1] = create_tensor_resource(env, k_upd); - result_tuple[2] = create_tensor_resource(env, v_upd); - - return nx::nif::ok(env, enif_make_tuple3(env, result_tuple[0], result_tuple[1], result_tuple[2])); - } - CATCH() +std::tuple +kv_cache_sdpa_update_impl(ErlNifEnv *env, TensorArg q, TensorArg new_k, + TensorArg new_v, TensorArg k_cache, + TensorArg v_cache, int offset, double scale, + mlx::core::Device device) { + int B = q->shape(0); + int N_kv = new_k->shape(1); + int T_new = new_k->shape(2); + int D = q->shape(3); + int valid_len = offset + T_new; + + // Move-extract: after this the ENIF resource blocks hold moved-from arrays. + // k_cache_owned / v_cache_owned are the sole owners (use_count == 1) until + // slice_update copies them into its inputs (use_count rises to 2). + // On function return both locals destruct → use_count drops to 1 again + // → is_donatable() is true at eval time → no new 4 MB Metal buffer. + auto k_cache_owned = std::move(*k_cache); + auto v_cache_owned = std::move(*v_cache); + + // 1. Insert new_k / new_v at cache position `offset` along axis 2 (T axis). + // Layout: {B, N_kv, T_max, D} — offset along dimension 2. + auto k_upd = mlx::core::slice_update( + k_cache_owned, *new_k, + to_shape({0, 0, offset, 0}), + to_shape({B, N_kv, valid_len, D}), + device); + auto v_upd = mlx::core::slice_update( + v_cache_owned, *new_v, + to_shape({0, 0, offset, 0}), + to_shape({B, N_kv, valid_len, D}), + device); + + // 2. Slice valid prefix: k_upd[0:B, 0:N_kv, 0:valid_len, 0:D]. + 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); + + // 3. SDPA: q / k_valid / v_valid are already in {B, N, T, D} format. + // Decode (T_new == 1): no mask — single query is trivially causal. + // Prefill (T_new > 1): additive causal mask in q's dtype. + auto build_prefill_mask = [&]() -> mlx::core::array { + 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); + 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 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); + }; + + auto attn_out = (T_new == 1) + ? mlx::core::fast::scaled_dot_product_attention( + *q, k_valid, v_valid, (float)scale, "", std::nullopt, std::nullopt, device) + : mlx::core::fast::scaled_dot_product_attention( + *q, k_valid, v_valid, (float)scale, "array", build_prefill_mask(), std::nullopt, device); + // k_cache_owned and v_cache_owned destruct here → use_count 2 → 1. + + return {attn_out, k_upd, v_upd}; } -ASYNC_NIF(kv_cache_sdpa_update) - +FINE_ASYNC_NIF(kv_cache_sdpa_update) diff --git a/emlx/c_src/emlx_fast/qwen3.cpp b/emlx/c_src/emlx_fast/qwen3.cpp index f91f3a3..c72ea56 100644 --- a/emlx/c_src/emlx_fast/qwen3.cpp +++ b/emlx/c_src/emlx_fast/qwen3.cpp @@ -1,245 +1,273 @@ #include "qwen3.hpp" #include "../emlx_nif_shared.hpp" +#include "../emlx_plugin_registry.hpp" +#include "qwen3_plugin_abi.hpp" #include -// Qwen3 model accelerators used by emlx_axon. These live in a separate -// translation unit so the model native code has a clear boundary from the -// generic EMLX fast operations and can be extracted later. - -static bool qwen3_check_rank( - const mlx::core::array &tensor, - int expected, - const char *name, - std::string &error) { - if (tensor.ndim() != expected) { - std::ostringstream msg; - msg << name << " expects rank " << expected << ", got rank " << tensor.ndim(); - error = msg.str(); - return false; +// Qwen3 model accelerators used by emlx_axon. This file is the *host* side +// of the qwen3 NIF/plugin split: it owns everything that touches Erlang +// terms (decoding args, wrapping `mlx::core::array` results back into +// tensor resources) and calls through to the standalone "qwen3" plugin — a +// dynamically loaded shared library with no Erlang dependency at all, +// living in emlx_axon (c_src/qwen3_plugin.cpp there) — for every actual MLX +// computation. See qwen3_plugin_abi.hpp for the ABI. +// +// This split exists so the qwen3 compute can live in emlx_axon as its own +// build artifact without dragging erl_nif/resource-type plumbing along +// with it. The plugin is loaded generically via `EMLX.NIF.load_plugin/2` +// (see emlx_plugin_registry.hpp) under the name `"qwen3"`; this file only +// knows how to *decode* qwen3's specific argument shapes, not how a plugin +// gets loaded. + +// qwen3_plugin — every qwen3_* NIF calls this first to fetch the "qwen3" +// plugin's vtable; it must have been loaded via +// `EMLX.NIF.load_plugin("qwen3", path)` before any of these can run (see +// EMLXAxon.Application, which loads it eagerly at boot). Returns `nullptr` +// (and fills `out_error`) if it hasn't been loaded (yet). +static const emlx_qwen3_plugin::VTable *qwen3_plugin(ErlNifEnv *env, ERL_NIF_TERM *out_error) { + const void *vtable = emlx_get_plugin("qwen3"); + if (vtable != nullptr) { + return reinterpret_cast(vtable); } - - return true; + *out_error = + nx::nif::error(env, "qwen3 plugin not loaded — call EMLX.NIF.load_plugin(\"qwen3\", path) first"); + return nullptr; } -static bool qwen3_check_positive(int value, const char *name, std::string &error) { - if (value <= 0) { - std::ostringstream msg; - msg << name << " must be positive"; - error = msg.str(); - return false; - } +// ── Term decoding helpers ───────────────────────────────────────────────── +// These stay host-side: they read directly off `TENSOR_TYPE`-backed +// resources and the refcounting scheme `enif_alloc_resource` lays out for +// them (see `Qwen3TensorHandle` below), so they cannot move into the +// plugin without also moving Erlang's resource machinery there — which is +// exactly the cross-library resource-type problem this split avoids. - return true; -} +class Qwen3TensorHandle { +public: + explicit Qwen3TensorHandle(mlx::core::array *ptr) : ptr_(ptr) { + refcount_ = reinterpret_cast *>(ptr_ + 1); -static bool qwen3_check_non_negative(int value, const char *name, std::string &error) { - if (value < 0) { - std::ostringstream msg; - msg << name << " must be non-negative"; - error = msg.str(); - return false; - } + if (refcount_->load() == 0) { + ptr_ = nullptr; + return; + } - return true; -} + ++(*refcount_); + } -static bool qwen3_check_dim( - const mlx::core::array &tensor, - int axis, - int expected, - const char *name, - const char *dim_name, - std::string &error) { - if (tensor.shape(axis) != expected) { - std::ostringstream msg; - msg << name << " " << dim_name << " must be " << expected - << ", got " << tensor.shape(axis); - error = msg.str(); - return false; + ~Qwen3TensorHandle() { + if (is_valid()) { + if (refcount_->fetch_sub(1) == 0) { + ptr_->~array(); + } + } } - return true; -} + bool is_valid() const { return ptr_ != nullptr; } + mlx::core::array *data() const { return ptr_; } -static bool qwen3_check_rank4_positive( - const mlx::core::array &tensor, - const char *name, - std::string &error) { - if (!qwen3_check_rank(tensor, 4, name, error)) { - return false; - } +private: + mlx::core::array *ptr_; + std::atomic *refcount_; +}; - for (int axis = 0; axis < 4; ++axis) { - if (tensor.shape(axis) <= 0) { - std::ostringstream msg; - msg << name << " dimensions must be positive"; - error = msg.str(); - return false; - } - } +using Qwen3TensorHandles = std::vector>; - return true; -} +static bool qwen3_get_tensor( + ErlNifEnv *env, + ERL_NIF_TERM term, + mlx::core::array **out, + Qwen3TensorHandles &handles, + ERL_NIF_TERM *error) { + mlx::core::array *raw = nullptr; -static bool qwen3_check_rank3_positive( - const mlx::core::array &tensor, - const char *name, - std::string &error) { - if (!qwen3_check_rank(tensor, 3, name, error)) { + if (!enif_get_resource( + env, term, resource_object::type, + reinterpret_cast(&raw))) { return false; } - for (int axis = 0; axis < 3; ++axis) { - if (tensor.shape(axis) <= 0) { - std::ostringstream msg; - msg << name << " dimensions must be positive"; - error = msg.str(); - return false; - } + auto handle = std::make_unique(raw); + if (!handle->is_valid()) { + *error = nx::nif::error(env, "Tensor has been deallocated"); + return false; } + *out = handle->data(); + handles.push_back(std::move(handle)); return true; } -static bool qwen3_check_rank2_positive( - const mlx::core::array &tensor, - const char *name, - std::string &error) { - if (!qwen3_check_rank(tensor, 2, name, error)) { +static bool qwen3_get_tensor_or_device_ref( + ErlNifEnv *env, + ERL_NIF_TERM term, + mlx::core::array **out, + Qwen3TensorHandles &handles, + ERL_NIF_TERM *error) { + if (qwen3_get_tensor(env, term, out, handles, error)) { + return true; + } + if (*error != 0) { return false; } - for (int axis = 0; axis < 2; ++axis) { - if (tensor.shape(axis) <= 0) { - std::ostringstream msg; - msg << name << " dimensions must be positive"; - error = msg.str(); - return false; - } + int arity = 0; + const ERL_NIF_TERM *items = nullptr; + if (!enif_get_tuple(env, term, &arity, &items) || arity != 2) { + return false; } - return true; + return qwen3_get_tensor(env, items[1], out, handles, error); } -static bool qwen3_check_rank1_dim( - const mlx::core::array &tensor, - int expected, - const char *name, - std::string &error) { - if (!qwen3_check_rank(tensor, 1, name, error)) { +// Decodes a linear weight term — either `{:dense, tensor_ref}` or +// `{:quantized, weight_ref, scales_ref, biases_ref_or_nil, group_size, bits, +// mode, transpose}` (mirrors the tuple built by `EMLX.Native.Qwen3.linear_weight_term/1` +// in emlx.ex) — into an `emlx_qwen3_plugin::LinearWeight`. `dense_transpose` +// selects the dense orientation when the term is `{:dense, ref}` (false for +// the usual {H,out} q/k/v/o/gate/up/down projections, true for the {out,H} +// lm_head convention); quantized terms carry their own explicit `transpose` +// flag. +static bool qwen3_get_linear_weight( + ErlNifEnv *env, + ERL_NIF_TERM term, + bool dense_transpose, + emlx_qwen3_plugin::LinearWeight &out, + Qwen3TensorHandles &handles, + ERL_NIF_TERM *error) { + int arity = 0; + const ERL_NIF_TERM *items = nullptr; + if (!enif_get_tuple(env, term, &arity, &items) || arity < 2) { return false; } - return qwen3_check_dim(tensor, 0, expected, name, "size", error); -} - -static bool qwen3_validate_projection_width( - const mlx::core::array &projection, - int input_width, - int head_dim, - const char *name, - std::string &error) { - if (!qwen3_check_rank2_positive(projection, name, error)) { + std::string tag; + if (!nx::nif::get_atom(env, items[0], tag)) { return false; } - if (!qwen3_check_dim(projection, 0, input_width, name, "input width", error)) { - return false; + + if (tag == "dense" && arity == 2) { + out.quantized = false; + out.transpose = dense_transpose; + return qwen3_get_tensor(env, items[1], &out.weight, handles, error); } - if ((projection.shape(1) % head_dim) != 0) { - std::ostringstream msg; - msg << name << " output width must be divisible by head_dim"; - error = msg.str(); - return false; + + if (tag == "quantized" && arity == 8) { + out.quantized = true; + + if (!qwen3_get_tensor(env, items[1], &out.weight, handles, error) || + !qwen3_get_tensor(env, items[2], &out.scales, handles, error)) { + return false; + } + + std::string nil_atom; + bool biases_nil = nx::nif::get_atom(env, items[3], nil_atom) && nil_atom == "nil"; + if (biases_nil) { + out.biases = nullptr; + } else if (!qwen3_get_tensor(env, items[3], &out.biases, handles, error)) { + return false; + } + + return nx::nif::get(env, items[4], &out.group_size) && + nx::nif::get(env, items[5], &out.bits) && + nx::nif::get(env, items[6], out.mode) && + nx::nif::get(env, items[7], &out.transpose); } - return true; + return false; } -static bool qwen3_validate_kv_cache_bn( - const mlx::core::array &k_cache, - const mlx::core::array &v_cache, - int batch, - int num_kv_heads, - int offset, - int token_count, - int head_dim, - std::string &error) { - if (!qwen3_check_rank4_positive(k_cache, "k_cache", error) || - !qwen3_check_rank4_positive(v_cache, "v_cache", error)) { +static bool qwen3_get_layer( + ErlNifEnv *env, + ERL_NIF_TERM term, + emlx_qwen3_plugin::LayerParams &layer, + Qwen3TensorHandles &handles, + ERL_NIF_TERM *error) { + int arity = 0; + const ERL_NIF_TERM *items = nullptr; + if (!enif_get_tuple(env, term, &arity, &items) || arity != 11) { return false; } - if (!qwen3_check_dim(k_cache, 0, batch, "k_cache", "batch", error) || - !qwen3_check_dim(v_cache, 0, batch, "v_cache", "batch", error) || - !qwen3_check_dim(k_cache, 1, num_kv_heads, "k_cache", "heads", error) || - !qwen3_check_dim(v_cache, 1, num_kv_heads, "v_cache", "heads", error) || - !qwen3_check_dim(k_cache, 3, head_dim, "k_cache", "head_dim", error) || - !qwen3_check_dim(v_cache, 3, head_dim, "v_cache", "head_dim", error)) { - return false; - } + return qwen3_get_tensor(env, items[0], &layer.norm1, handles, error) && + qwen3_get_tensor(env, items[1], &layer.norm2, handles, error) && + qwen3_get_tensor(env, items[2], &layer.q_norm, handles, error) && + qwen3_get_tensor(env, items[3], &layer.k_norm, handles, error) && + qwen3_get_tensor(env, items[4], &layer.q_proj, handles, error) && + qwen3_get_tensor(env, items[5], &layer.k_proj, handles, error) && + qwen3_get_tensor(env, items[6], &layer.v_proj, handles, error) && + qwen3_get_tensor(env, items[7], &layer.o_proj, handles, error) && + qwen3_get_tensor(env, items[8], &layer.gate_proj, handles, error) && + qwen3_get_tensor(env, items[9], &layer.up_proj, handles, error) && + qwen3_get_tensor(env, items[10], &layer.down_proj, handles, error); +} - if (v_cache.shape(2) != k_cache.shape(2)) { - error = "k_cache and v_cache capacity must match"; +static bool qwen3_get_kv( + ErlNifEnv *env, + ERL_NIF_TERM term, + emlx_qwen3_plugin::KVCache &kv, + Qwen3TensorHandles &handles, + ERL_NIF_TERM *error) { + int arity = 0; + const ERL_NIF_TERM *items = nullptr; + if (!enif_get_tuple(env, term, &arity, &items) || arity != 2) { return false; } - int64_t required_len = static_cast(offset) + static_cast(token_count); - int capacity = k_cache.shape(2); + return qwen3_get_tensor_or_device_ref(env, items[0], &kv.k, handles, error) && + qwen3_get_tensor_or_device_ref(env, items[1], &kv.v, handles, error); +} - if (required_len > capacity) { - std::ostringstream msg; - msg << "KV cache capacity " << capacity - << " is smaller than required length " << required_len; - error = msg.str(); +static bool qwen3_get_layer_generalized( + ErlNifEnv *env, + ERL_NIF_TERM term, + emlx_qwen3_plugin::LayerParamsQ &layer, + Qwen3TensorHandles &handles, + ERL_NIF_TERM *error) { + int arity = 0; + const ERL_NIF_TERM *items = nullptr; + if (!enif_get_tuple(env, term, &arity, &items) || arity != 11) { return false; } - return true; + return qwen3_get_tensor(env, items[0], &layer.norm1, handles, error) && + qwen3_get_tensor(env, items[1], &layer.norm2, handles, error) && + qwen3_get_tensor(env, items[2], &layer.q_norm, handles, error) && + qwen3_get_tensor(env, items[3], &layer.k_norm, handles, error) && + qwen3_get_linear_weight(env, items[4], false, layer.q_proj, handles, error) && + qwen3_get_linear_weight(env, items[5], false, layer.k_proj, handles, error) && + qwen3_get_linear_weight(env, items[6], false, layer.v_proj, handles, error) && + qwen3_get_linear_weight(env, items[7], false, layer.o_proj, handles, error) && + qwen3_get_linear_weight(env, items[8], false, layer.gate_proj, handles, error) && + qwen3_get_linear_weight(env, items[9], false, layer.up_proj, handles, error) && + qwen3_get_linear_weight(env, items[10], false, layer.down_proj, handles, error); } -static bool qwen3_validate_qkv_cache_attention( - const mlx::core::array &q, - const mlx::core::array &new_k, - const mlx::core::array &new_v, - const mlx::core::array &k_cache, - const mlx::core::array &v_cache, - int offset, - int head_dim, - std::string &error) { - if (!qwen3_check_rank4_positive(q, "q", error) || - !qwen3_check_rank4_positive(new_k, "new_k", error) || - !qwen3_check_rank4_positive(new_v, "new_v", error) || - !qwen3_check_non_negative(offset, "offset", error) || - !qwen3_check_positive(head_dim, "head_dim", error)) { - return false; +static ERL_NIF_TERM qwen3_ref_error_or( + ErlNifEnv *env, ERL_NIF_TERM error, const char *fallback) { + if (error != 0) { + return error; } - int B = q.shape(0); - int T_new = q.shape(1); - int N_q = q.shape(2); - int D = q.shape(3); - int N_kv = new_k.shape(2); + return nx::nif::error(env, fallback); +} - if (D != head_dim) { - error = "q last dimension must match head_dim"; +// Minimal rank/positivity check for `input_ids`/`embed_tokens`, ahead of the +// host-side embedding lookup (`mlx::core::take`) that every +// `qwen3_forward_greedy_*` NIF performs before delegating to the plugin. +// The plugin re-validates everything downstream of the embedded hidden +// state; this only guards the `shape(0)`/`shape(1)` accesses below it. +static bool qwen3_require_rank2_positive(const mlx::core::array &tensor, const char *name, + std::string &error) { + if (tensor.ndim() != 2) { + error = std::string(name) + " expects rank 2, got rank " + std::to_string(tensor.ndim()); return false; } - if ((N_q % N_kv) != 0) { - error = "query heads must be divisible by key/value heads"; + if (tensor.shape(0) <= 0 || tensor.shape(1) <= 0) { + error = std::string(name) + " dimensions must be positive"; return false; } - if (!qwen3_check_dim(new_k, 0, B, "new_k", "batch", error) || - !qwen3_check_dim(new_v, 0, B, "new_v", "batch", error) || - !qwen3_check_dim(new_k, 1, T_new, "new_k", "sequence length", error) || - !qwen3_check_dim(new_v, 1, T_new, "new_v", "sequence length", error) || - !qwen3_check_dim(new_v, 2, N_kv, "new_v", "heads", error) || - !qwen3_check_dim(new_k, 3, D, "new_k", "head_dim", error) || - !qwen3_check_dim(new_v, 3, D, "new_v", "head_dim", error)) { - return false; - } - - return qwen3_validate_kv_cache_bn(k_cache, v_cache, B, N_kv, offset, T_new, D, error); + return true; } // qwen3_kv_cache_attention — Qwen3 fused RoPE + KV update + SDPA. @@ -256,11 +284,14 @@ static bool qwen3_validate_qkv_cache_attention( // theta — float RoPE base // device — atom // -// Returns {attn_out, k_upd, v_upd}: -// attn_out — {B, T_new, N_q * D} ready for projection BTH layout -// k_upd — {B, N_kv, T_max, D} -// v_upd — {B, N_kv, T_max, D} +// Returns {attn_out, k_upd, v_upd}. NIF(qwen3_kv_cache_attention) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + TENSOR_PARAM(0, q); TENSOR_PARAM(1, new_k); TENSOR_PARAM(2, new_v); @@ -273,71 +304,15 @@ NIF(qwen3_kv_cache_attention) { DEVICE_PARAM(9, device); try { + mlx::core::array out(0), k_upd(0), v_upd(0); std::string error; - if (!qwen3_validate_qkv_cache_attention( - *q, *new_k, *new_v, *k_cache, *v_cache, offset, head_dim, error)) { + if (!plugin->kv_cache_attention(*q, *new_k, *new_v, *k_cache, *v_cache, offset, scale, + head_dim, theta, device, out, k_upd, v_upd, error)) { return nx::nif::error(env, error.c_str()); } - int B = q->shape(0); - int T_new = q->shape(1); - int N_q = q->shape(2); - int D = q->shape(3); - 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_rope = mlx::core::fast::rope( - q_bn, head_dim, false, (float)theta, 1.0f, offset, std::nullopt, device); - auto k_rope = mlx::core::fast::rope( - k_bn, head_dim, false, (float)theta, 1.0f, offset, std::nullopt, device); - - auto k_cache_owned = std::move(*k_cache); - auto v_cache_owned = std::move(*v_cache); - - 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); - 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); - - 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 build_prefill_mask = [&]() -> mlx::core::array { - 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); - 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 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); - }; - - auto attn_out_bn = (T_new == 1) - ? mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, (float)scale, "", std::nullopt, std::nullopt, device) - : mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, (float)scale, "array", build_prefill_mask(), std::nullopt, 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, N_q * D}, device); - ERL_NIF_TERM result_tuple[3]; - result_tuple[0] = create_tensor_resource(env, attn_out); + result_tuple[0] = create_tensor_resource(env, out); result_tuple[1] = create_tensor_resource(env, k_upd); result_tuple[2] = create_tensor_resource(env, v_upd); @@ -347,62 +322,14 @@ NIF(qwen3_kv_cache_attention) { } ASYNC_NIF(qwen3_kv_cache_attention) -static mlx::core::array qwen3_linear_in_out( - const mlx::core::array &x, - const mlx::core::array &weight, - const mlx::core::Device &device) { - 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); - } - - return mlx::core::matmul(x, weight, device); -} - -static mlx::core::array qwen3_linear_out_in( - const mlx::core::array &x, - const mlx::core::array &weight, - const mlx::core::Device &device) { - return mlx::core::tensordot( - x, weight, std::vector{static_cast(x.ndim()) - 1}, std::vector{1}, device); -} - -static int64_t qwen3_token_to_int64(mlx::core::array &token) { - mlx::core::eval(token); - auto dtype = token.dtype(); - - if (dtype == mlx::core::uint8) { - return static_cast(token.item()); - } else if (dtype == mlx::core::uint16) { - return static_cast(token.item()); - } else if (dtype == mlx::core::uint32) { - return static_cast(token.item()); - } else if (dtype == mlx::core::uint64) { - return static_cast(token.item()); - } else if (dtype == mlx::core::int8) { - return static_cast(token.item()); - } else if (dtype == mlx::core::int16) { - return static_cast(token.item()); - } else if (dtype == mlx::core::int32) { - return static_cast(token.item()); - } else { - return token.item(); - } -} - // qwen3_mlp — dense Qwen3 MLP block: RMSNorm + gate/up + SwiGLU + down + residual. -// -// Inputs: -// hidden — {B, T, H} -// norm — {H} -// gate_proj — {H, I} -// up_proj — {H, I} -// down_proj — {I, H} -// eps — RMSNorm epsilon -// -// Returns hidden + out {B, T, H}. NIF(qwen3_mlp) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + TENSOR_PARAM(0, hidden); TENSOR_PARAM(1, norm); TENSOR_PARAM(2, gate_proj); @@ -412,35 +339,14 @@ NIF(qwen3_mlp) { DEVICE_PARAM(6, device); try { + mlx::core::array out(0); std::string error; - if (!qwen3_check_rank3_positive(*hidden, "hidden", error)) { - return nx::nif::error(env, error.c_str()); - } - - int H = hidden->shape(2); - if (!qwen3_check_rank1_dim(*norm, H, "norm", error) || - !qwen3_check_rank2_positive(*gate_proj, "gate_proj", error) || - !qwen3_check_rank2_positive(*up_proj, "up_proj", error) || - !qwen3_check_rank2_positive(*down_proj, "down_proj", error) || - !qwen3_check_dim(*gate_proj, 0, H, "gate_proj", "input width", error) || - !qwen3_check_dim(*up_proj, 0, H, "up_proj", "input width", error) || - !qwen3_check_dim(*up_proj, 1, gate_proj->shape(1), "up_proj", "output width", error) || - !qwen3_check_dim(*down_proj, 0, gate_proj->shape(1), "down_proj", "input width", error) || - !qwen3_check_dim(*down_proj, 1, H, "down_proj", "output width", error)) { + if (!plugin->mlp(*hidden, *norm, *gate_proj, *up_proj, *down_proj, eps, device, out, + error)) { return nx::nif::error(env, error.c_str()); } - auto xn = mlx::core::fast::rms_norm(*hidden, *norm, (float)eps, device); - auto gate = qwen3_linear_in_out(xn, *gate_proj, device); - auto up = qwen3_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 out = qwen3_linear_in_out(mlp, *down_proj, device); - auto residual = mlx::core::add(*hidden, out, device); - - TENSOR(residual); + TENSOR(out); } CATCH() } @@ -450,29 +356,14 @@ ASYNC_NIF(qwen3_mlp) // attention input RMSNorm + dense attention block + RMSNorm after attention // + dense MLP + residual add. // -// Inputs: -// hidden — {B, T_new, H} -// norm1 — {H} -// q_proj — {H, N_q * D} -// k_proj — {H, N_kv * D} -// v_proj — {H, N_kv * D} -// o_proj — {N_q * D, H} -// q_norm — {D} -// k_norm — {D} -// k_cache — {B, N_kv, T_max, D} -// v_cache — {B, N_kv, T_max, D} -// norm2 — {H} -// gate_proj — {H, I} -// up_proj — {H, I} -// down_proj — {I, H} -// offset — int -// scale — float -// head_dim — int -// theta — float -// eps — RMSNorm epsilon -// // Returns {hidden_out, k_upd, v_upd}. NIF(qwen3_layer) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + TENSOR_PARAM(0, hidden); TENSOR_PARAM(1, norm1); TENSOR_PARAM(2, q_proj); @@ -495,572 +386,138 @@ NIF(qwen3_layer) { DEVICE_PARAM(19, device); try { - std::string error; - if (!qwen3_check_rank3_positive(*hidden, "hidden", error) || - !qwen3_check_non_negative(offset, "offset", error) || - !qwen3_check_positive(head_dim, "head_dim", error)) { - return nx::nif::error(env, error.c_str()); - } - - int B = hidden->shape(0); - int T_new = hidden->shape(1); - int H = hidden->shape(2); - int D = head_dim; - if (!qwen3_check_rank1_dim(*norm1, H, "norm1", error) || - !qwen3_check_rank1_dim(*norm2, H, "norm2", error) || - !qwen3_validate_projection_width(*q_proj, H, D, "q_proj", error) || - !qwen3_validate_projection_width(*k_proj, H, D, "k_proj", error) || - !qwen3_validate_projection_width(*v_proj, H, D, "v_proj", error) || - !qwen3_check_dim(*v_proj, 1, k_proj->shape(1), "v_proj", "output width", error) || - !qwen3_check_rank1_dim(*q_norm, D, "q_norm", error) || - !qwen3_check_rank1_dim(*k_norm, D, "k_norm", error)) { - return nx::nif::error(env, error.c_str()); - } - - int N_q = q_proj->shape(1) / D; - int N_kv = k_proj->shape(1) / D; - int attn_width = N_q * D; + emlx_qwen3_plugin::LayerParams layer{norm1, norm2, q_norm, k_norm, q_proj, + k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj}; + emlx_qwen3_plugin::KVCache kv{k_cache, v_cache}; - if ((N_q % N_kv) != 0) { - return nx::nif::error(env, "query heads must be divisible by key/value heads"); - } - if (!qwen3_check_rank2_positive(*o_proj, "o_proj", error) || - !qwen3_check_dim(*o_proj, 0, attn_width, "o_proj", "input width", error) || - !qwen3_check_dim(*o_proj, 1, H, "o_proj", "output width", error) || - !qwen3_validate_kv_cache_bn(*k_cache, *v_cache, B, N_kv, offset, T_new, D, error) || - !qwen3_check_rank2_positive(*gate_proj, "gate_proj", error) || - !qwen3_check_rank2_positive(*up_proj, "up_proj", error) || - !qwen3_check_rank2_positive(*down_proj, "down_proj", error) || - !qwen3_check_dim(*gate_proj, 0, H, "gate_proj", "input width", error) || - !qwen3_check_dim(*up_proj, 0, H, "up_proj", "input width", error) || - !qwen3_check_dim(*up_proj, 1, gate_proj->shape(1), "up_proj", "output width", error) || - !qwen3_check_dim(*down_proj, 0, gate_proj->shape(1), "down_proj", "input width", error) || - !qwen3_check_dim(*down_proj, 1, H, "down_proj", "output width", error)) { + mlx::core::array out(0), k_upd(0), v_upd(0); + std::string error; + if (!plugin->layer_dense(*hidden, layer, kv, offset, scale, head_dim, theta, eps, + device, out, k_upd, v_upd, error)) { return nx::nif::error(env, error.c_str()); } - int valid_len = offset + T_new; - - auto xn = mlx::core::fast::rms_norm(*hidden, *norm1, (float)eps, device); - auto q_flat = qwen3_linear_in_out(xn, *q_proj, device); - auto k_flat = qwen3_linear_in_out(xn, *k_proj, device); - auto v_flat = qwen3_linear_in_out(xn, *v_proj, device); - - 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); - - q = mlx::core::fast::rms_norm(q, *q_norm, (float)eps, device); - k = mlx::core::fast::rms_norm(k, *k_norm, (float)eps, device); - - 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_rope = mlx::core::fast::rope( - q_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, device); - auto k_rope = mlx::core::fast::rope( - k_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, device); - - auto k_cache_owned = std::move(*k_cache); - auto v_cache_owned = std::move(*v_cache); - - 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); - 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); - - 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 build_prefill_mask = [&]() -> mlx::core::array { - 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); - 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 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); - }; - - auto attn_out_bn = (T_new == 1) - ? mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, (float)scale, "", std::nullopt, std::nullopt, device) - : mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, (float)scale, "array", build_prefill_mask(), std::nullopt, 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 = qwen3_linear_in_out(attn_out, *o_proj, device); - auto attn_hidden = mlx::core::add(*hidden, attn_projected, device); - - auto xn2 = mlx::core::fast::rms_norm(attn_hidden, *norm2, (float)eps, device); - auto gate = qwen3_linear_in_out(xn2, *gate_proj, device); - auto up = qwen3_linear_in_out(xn2, *up_proj, device); - auto mlp = mlx::core::multiply( - mlx::core::multiply(gate, mlx::core::sigmoid(gate, device), device), - up, - device); - auto mlp_out = qwen3_linear_in_out(mlp, *down_proj, device); - auto out = mlx::core::add(attn_hidden, mlp_out, device); ERL_NIF_TERM result_tuple[3]; - result_tuple[0] = create_tensor_resource(env, out); - result_tuple[1] = create_tensor_resource(env, k_upd); - result_tuple[2] = create_tensor_resource(env, v_upd); - - return nx::nif::ok(env, enif_make_tuple3(env, result_tuple[0], result_tuple[1], result_tuple[2])); - } - CATCH() -} -ASYNC_NIF(qwen3_layer) - -struct Qwen3LayerParams { - mlx::core::array *norm1; - mlx::core::array *norm2; - mlx::core::array *q_norm; - mlx::core::array *k_norm; - mlx::core::array *q_proj; - mlx::core::array *k_proj; - mlx::core::array *v_proj; - mlx::core::array *o_proj; - mlx::core::array *gate_proj; - mlx::core::array *up_proj; - mlx::core::array *down_proj; -}; - -struct Qwen3KVCache { - mlx::core::array *k; - mlx::core::array *v; -}; - -class Qwen3TensorHandle { -public: - explicit Qwen3TensorHandle(mlx::core::array *ptr) : ptr_(ptr) { - refcount_ = reinterpret_cast *>(ptr_ + 1); - - if (refcount_->load() == 0) { - ptr_ = nullptr; - return; - } - - ++(*refcount_); - } - - ~Qwen3TensorHandle() { - if (is_valid()) { - if (refcount_->fetch_sub(1) == 0) { - ptr_->~array(); - } - } - } - - bool is_valid() const { return ptr_ != nullptr; } - mlx::core::array *data() const { return ptr_; } - -private: - mlx::core::array *ptr_; - std::atomic *refcount_; -}; - -using Qwen3TensorHandles = std::vector>; - -static bool qwen3_get_tensor( - ErlNifEnv *env, - ERL_NIF_TERM term, - mlx::core::array **out, - Qwen3TensorHandles &handles, - ERL_NIF_TERM *error) { - mlx::core::array *raw = nullptr; - - if (!enif_get_resource( - env, term, resource_object::type, - reinterpret_cast(&raw))) { - return false; - } - - auto handle = std::make_unique(raw); - if (!handle->is_valid()) { - *error = nx::nif::error(env, "Tensor has been deallocated"); - return false; - } - - *out = handle->data(); - handles.push_back(std::move(handle)); - return true; -} - -static bool qwen3_get_tensor_or_device_ref( - ErlNifEnv *env, - ERL_NIF_TERM term, - mlx::core::array **out, - Qwen3TensorHandles &handles, - ERL_NIF_TERM *error) { - if (qwen3_get_tensor(env, term, out, handles, error)) { - return true; - } - if (*error != 0) { - return false; - } - - int arity = 0; - const ERL_NIF_TERM *items = nullptr; - if (!enif_get_tuple(env, term, &arity, &items) || arity != 2) { - return false; - } - - return qwen3_get_tensor(env, items[1], out, handles, error); -} - -static bool qwen3_get_layer( - ErlNifEnv *env, - ERL_NIF_TERM term, - Qwen3LayerParams &layer, - Qwen3TensorHandles &handles, - ERL_NIF_TERM *error) { - int arity = 0; - const ERL_NIF_TERM *items = nullptr; - if (!enif_get_tuple(env, term, &arity, &items) || arity != 11) { - return false; - } - - return qwen3_get_tensor(env, items[0], &layer.norm1, handles, error) && - qwen3_get_tensor(env, items[1], &layer.norm2, handles, error) && - qwen3_get_tensor(env, items[2], &layer.q_norm, handles, error) && - qwen3_get_tensor(env, items[3], &layer.k_norm, handles, error) && - qwen3_get_tensor(env, items[4], &layer.q_proj, handles, error) && - qwen3_get_tensor(env, items[5], &layer.k_proj, handles, error) && - qwen3_get_tensor(env, items[6], &layer.v_proj, handles, error) && - qwen3_get_tensor(env, items[7], &layer.o_proj, handles, error) && - qwen3_get_tensor(env, items[8], &layer.gate_proj, handles, error) && - qwen3_get_tensor(env, items[9], &layer.up_proj, handles, error) && - qwen3_get_tensor(env, items[10], &layer.down_proj, handles, error); -} - -static bool qwen3_get_kv( - ErlNifEnv *env, - ERL_NIF_TERM term, - Qwen3KVCache &kv, - Qwen3TensorHandles &handles, - ERL_NIF_TERM *error) { - int arity = 0; - const ERL_NIF_TERM *items = nullptr; - if (!enif_get_tuple(env, term, &arity, &items) || arity != 2) { - return false; - } - - return qwen3_get_tensor_or_device_ref(env, items[0], &kv.k, handles, error) && - qwen3_get_tensor_or_device_ref(env, items[1], &kv.v, handles, error); -} + result_tuple[0] = create_tensor_resource(env, out); + result_tuple[1] = create_tensor_resource(env, k_upd); + result_tuple[2] = create_tensor_resource(env, v_upd); -static ERL_NIF_TERM qwen3_ref_error_or( - ErlNifEnv *env, ERL_NIF_TERM error, const char *fallback) { - if (error != 0) { - return error; + return nx::nif::ok(env, enif_make_tuple3(env, result_tuple[0], result_tuple[1], result_tuple[2])); } - - return nx::nif::error(env, fallback); + CATCH() } +ASYNC_NIF(qwen3_layer) -static bool qwen3_validate_dense_layer( - const mlx::core::array &hidden, - const Qwen3LayerParams &layer, - const Qwen3KVCache &kv, - int offset, - int head_dim, - std::string &error) { - if (!qwen3_check_rank3_positive(hidden, "hidden", error) || - !qwen3_check_non_negative(offset, "offset", error) || - !qwen3_check_positive(head_dim, "head_dim", error)) { - return false; +// qwen3_layer_quantized — generalized Qwen3 transformer layer: same fusion as +// `qwen3_layer`, but each of the 7 projections (q/k/v/o/gate/up/down) +// independently accepts a dense-or-quantized weight term (see +// `qwen3_get_linear_weight`). +// +// Returns {hidden_out, k_upd, v_upd}. +NIF(qwen3_layer_quantized) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; } - int B = hidden.shape(0); - int T_new = hidden.shape(1); - int H = hidden.shape(2); - int D = head_dim; - - if (!qwen3_check_rank1_dim(*layer.norm1, H, "norm1", error) || - !qwen3_check_rank1_dim(*layer.norm2, H, "norm2", error) || - !qwen3_validate_projection_width(*layer.q_proj, H, D, "q_proj", error) || - !qwen3_validate_projection_width(*layer.k_proj, H, D, "k_proj", error) || - !qwen3_validate_projection_width(*layer.v_proj, H, D, "v_proj", error) || - !qwen3_check_dim(*layer.v_proj, 1, layer.k_proj->shape(1), "v_proj", "output width", error) || - !qwen3_check_rank1_dim(*layer.q_norm, D, "q_norm", error) || - !qwen3_check_rank1_dim(*layer.k_norm, D, "k_norm", error)) { - return false; - } + TENSOR_PARAM(0, hidden); - int N_q = layer.q_proj->shape(1) / D; - int N_kv = layer.k_proj->shape(1) / D; - int attn_width = N_q * D; + Qwen3TensorHandles handles; + ERL_NIF_TERM ref_error = 0; + emlx_qwen3_plugin::LayerParamsQ layer; + mlx::core::array *k_cache = nullptr; + mlx::core::array *v_cache = nullptr; - if ((N_q % N_kv) != 0) { - error = "query heads must be divisible by key/value heads"; - return false; + if (!qwen3_get_tensor(env, argv[1], &layer.norm1, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized expects norm1 tensor ref"); } - if (!qwen3_check_rank2_positive(*layer.o_proj, "o_proj", error) || - !qwen3_check_dim(*layer.o_proj, 0, attn_width, "o_proj", "input width", error) || - !qwen3_check_dim(*layer.o_proj, 1, H, "o_proj", "output width", error) || - !qwen3_validate_kv_cache_bn(*kv.k, *kv.v, B, N_kv, offset, T_new, D, error) || - !qwen3_check_rank2_positive(*layer.gate_proj, "gate_proj", error) || - !qwen3_check_rank2_positive(*layer.up_proj, "up_proj", error) || - !qwen3_check_rank2_positive(*layer.down_proj, "down_proj", error) || - !qwen3_check_dim(*layer.gate_proj, 0, H, "gate_proj", "input width", error) || - !qwen3_check_dim(*layer.up_proj, 0, H, "up_proj", "input width", error) || - !qwen3_check_dim(*layer.up_proj, 1, layer.gate_proj->shape(1), "up_proj", "output width", error) || - !qwen3_check_dim(*layer.down_proj, 0, layer.gate_proj->shape(1), "down_proj", "input width", error) || - !qwen3_check_dim(*layer.down_proj, 1, H, "down_proj", "output width", error)) { - return false; + if (!qwen3_get_linear_weight(env, argv[2], false, layer.q_proj, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized got invalid q_proj weight term"); } - - return true; -} - -static mlx::core::array qwen3_dense_layer_impl( - const mlx::core::array &hidden, - const Qwen3LayerParams &layer, - Qwen3KVCache &kv, - int offset, - float scale, - int head_dim, - float theta, - float eps, - const mlx::core::Device &device, - ERL_NIF_TERM &k_term, - ERL_NIF_TERM &v_term, - ErlNifEnv *env, - mlx::core::array *k_out = nullptr, - mlx::core::array *v_out = nullptr) { - int B = hidden.shape(0); - int T_new = hidden.shape(1); - int D = head_dim; - int N_q = layer.q_proj->shape(1) / D; - int N_kv = layer.k_proj->shape(1) / D; - 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 = qwen3_linear_in_out(xn, *layer.q_proj, device); - auto k_flat = qwen3_linear_in_out(xn, *layer.k_proj, device); - auto v_flat = qwen3_linear_in_out(xn, *layer.v_proj, device); - - 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); - - q = mlx::core::fast::rms_norm(q, *layer.q_norm, eps, device); - k = mlx::core::fast::rms_norm(k, *layer.k_norm, eps, device); - - 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_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 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); - 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); - - 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 build_prefill_mask = [&]() -> mlx::core::array { - 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); - 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 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); - }; - - auto attn_out_bn = (T_new == 1) - ? mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, device) - : mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, scale, "array", build_prefill_mask(), std::nullopt, 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 = qwen3_linear_in_out(attn_out, *layer.o_proj, device); - auto attn_hidden = mlx::core::add(hidden, attn_projected, device); - - auto xn2 = mlx::core::fast::rms_norm(attn_hidden, *layer.norm2, eps, device); - auto gate = qwen3_linear_in_out(xn2, *layer.gate_proj, device); - auto up = qwen3_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 = qwen3_linear_in_out(mlp, *layer.down_proj, device); - - if (k_out != nullptr) { - *k_out = k_upd; + if (!qwen3_get_linear_weight(env, argv[3], false, layer.k_proj, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized got invalid k_proj weight term"); } - if (v_out != nullptr) { - *v_out = v_upd; + if (!qwen3_get_linear_weight(env, argv[4], false, layer.v_proj, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized got invalid v_proj weight term"); } - - if (env != nullptr) { - k_term = create_tensor_resource(env, k_upd); - v_term = create_tensor_resource(env, v_upd); + if (!qwen3_get_linear_weight(env, argv[5], false, layer.o_proj, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized got invalid o_proj weight term"); } - - return mlx::core::add(attn_hidden, mlp_out, device); -} - -static ERL_NIF_TERM qwen3_forward_greedy_from_hidden( - ErlNifEnv *env, - const mlx::core::array &hidden, - ERL_NIF_TERM layers_arg, - ERL_NIF_TERM kv_cache_arg, - mlx::core::array *norm, - 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) { - unsigned int layer_count = 0; - unsigned int kv_count = 0; - - if (!enif_get_list_length(env, layers_arg, &layer_count)) { - return nx::nif::error(env, "Qwen3 greedy forward expects layers to be a list"); + if (!qwen3_get_tensor(env, argv[6], &layer.q_norm, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized expects q_norm tensor ref"); } - - if (!enif_get_list_length(env, kv_cache_arg, &kv_count)) { - return nx::nif::error(env, "Qwen3 greedy forward expects kv_cache to be a list"); + if (!qwen3_get_tensor(env, argv[7], &layer.k_norm, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized expects k_norm tensor ref"); } - - if (layer_count != kv_count) { - return nx::nif::error(env, "Qwen3 greedy forward layers and kv_cache length mismatch"); + if (!qwen3_get_tensor(env, argv[8], &k_cache, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized expects k_cache tensor ref"); } - std::string input_error; - if (!qwen3_check_rank3_positive(hidden, "hidden", input_error) || - !qwen3_check_non_negative(offset, "offset", input_error) || - !qwen3_check_positive(head_dim, "head_dim", input_error)) { - return nx::nif::error(env, input_error.c_str()); + if (!qwen3_get_tensor(env, argv[9], &v_cache, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized expects v_cache tensor ref"); + } + if (!qwen3_get_tensor(env, argv[10], &layer.norm2, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized expects norm2 tensor ref"); + } + if (!qwen3_get_linear_weight(env, argv[11], false, layer.gate_proj, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized got invalid gate_proj weight term"); + } + if (!qwen3_get_linear_weight(env, argv[12], false, layer.up_proj, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized got invalid up_proj weight term"); + } + if (!qwen3_get_linear_weight(env, argv[13], false, layer.down_proj, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "qwen3_layer_quantized got invalid down_proj weight term"); } - auto current = hidden; - std::vector kv_terms; - std::vector eval_arrays; - kv_terms.reserve(layer_count); - eval_arrays.reserve((layer_count * 2) + 1); - - ERL_NIF_TERM layer_head, layer_tail = layers_arg; - ERL_NIF_TERM kv_head, kv_tail = kv_cache_arg; - Qwen3TensorHandles handles; - ERL_NIF_TERM ref_error = 0; + PARAM(14, int, offset); + PARAM(15, double, scale); + PARAM(16, int, head_dim); + PARAM(17, double, theta); + PARAM(18, double, eps); + DEVICE_PARAM(19, device); - while (enif_get_list_cell(env, layer_tail, &layer_head, &layer_tail) && - enif_get_list_cell(env, kv_tail, &kv_head, &kv_tail)) { - Qwen3LayerParams layer; - Qwen3KVCache kv; + try { + emlx_qwen3_plugin::KVCache kv{k_cache, v_cache}; - if (!qwen3_get_layer(env, layer_head, layer, handles, &ref_error)) { - return qwen3_ref_error_or( - env, ref_error, "Qwen3 greedy forward got invalid layer tuple"); - } - if (!qwen3_get_kv(env, kv_head, kv, handles, &ref_error)) { - return qwen3_ref_error_or( - env, ref_error, "Qwen3 greedy forward got invalid kv_cache tuple"); - } + mlx::core::array out(0), k_upd(0), v_upd(0); std::string error; - if (!qwen3_validate_dense_layer(current, layer, kv, offset, head_dim, error)) { + if (!plugin->layer_quantized(*hidden, layer, kv, offset, scale, head_dim, theta, eps, + device, out, k_upd, v_upd, error)) { return nx::nif::error(env, error.c_str()); } - auto k_new = *kv.k; - auto v_new = *kv.v; - ERL_NIF_TERM unused_k_term; - ERL_NIF_TERM unused_v_term; - - current = qwen3_dense_layer_impl( - current, layer, kv, offset, static_cast(scale), head_dim, - static_cast(theta), static_cast(eps), device, - unused_k_term, unused_v_term, nullptr, &k_new, &v_new); + ERL_NIF_TERM result_tuple[3]; + result_tuple[0] = create_tensor_resource(env, out); + result_tuple[1] = create_tensor_resource(env, k_upd); + result_tuple[2] = create_tensor_resource(env, v_upd); - eval_arrays.push_back(k_new); - eval_arrays.push_back(v_new); - kv_terms.push_back( - enif_make_tuple2( - env, - create_tensor_resource(env, k_new), - create_tensor_resource(env, v_new))); + return nx::nif::ok(env, enif_make_tuple3(env, result_tuple[0], result_tuple[1], result_tuple[2])); } + CATCH() +} +ASYNC_NIF(qwen3_layer_quantized) - int B = current.shape(0); - int T = current.shape(1); - int H = current.shape(2); - if (return_token_id && B != 1) { - return nx::nif::error(env, "token_id return paths require batch size 1"); - } - std::string final_error; - if (!qwen3_check_rank1_dim(*norm, H, "norm", final_error) || - !qwen3_check_rank2_positive(*lm_head, "lm_head", final_error) || - !qwen3_check_dim(*lm_head, 1, H, "lm_head", "hidden width", final_error)) { - return nx::nif::error(env, final_error.c_str()); +static ERL_NIF_TERM qwen3_wrap_kv_terms( + ErlNifEnv *env, const std::vector &k, const std::vector &v) { + std::vector kv_terms; + kv_terms.reserve(k.size()); + for (size_t i = 0; i < k.size(); ++i) { + kv_terms.push_back( + enif_make_tuple2(env, create_tensor_resource(env, k[i]), create_tensor_resource(env, v[i]))); } - - auto last = (T == 1) - ? mlx::core::reshape(current, {B, H}, device) - : mlx::core::reshape( - mlx::core::slice(current, to_shape({0, T - 1, 0}), to_shape({B, T, H}), device), - {B, H}, - device); - - auto normed = mlx::core::fast::rms_norm(last, *norm, static_cast(eps), device); - auto logits = qwen3_linear_out_in(normed, *lm_head, device); - auto token = mlx::core::argmax(logits, 1, false, device); - - eval_arrays.push_back(token); - mlx::core::async_eval(eval_arrays); - - ERL_NIF_TERM token_term = - return_token_id - ? nx::nif::make(env, qwen3_token_to_int64(token)) - : create_tensor_resource(env, token); - ERL_NIF_TERM kv_list = enif_make_list_from_array(env, kv_terms.data(), kv_terms.size()); - - return nx::nif::ok(env, enif_make_tuple2(env, token_term, kv_list)); + return enif_make_list_from_array(env, kv_terms.data(), kv_terms.size()); } // qwen3_forward_greedy_ids — embedding lookup + dense forward through all layers + // final greedy token. Returns {token_ids, kv_cache} where token_ids has // shape {B}. NIF(qwen3_forward_greedy_ids) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + TENSOR_PARAM(0, input_ids); TENSOR_PARAM(1, embed_tokens); PARAM(6, int, offset); @@ -1085,26 +542,62 @@ NIF(qwen3_forward_greedy_ids) { } std::string error; - if (!qwen3_check_rank2_positive(*input_ids, "input_ids", error) || - !qwen3_check_rank2_positive(*embed_tokens, "embed_tokens", error) || - !qwen3_check_non_negative(offset, "offset", error) || - !qwen3_check_positive(head_dim, "head_dim", error)) { + if (!qwen3_require_rank2_positive(*input_ids, "input_ids", error) || + !qwen3_require_rank2_positive(*embed_tokens, "embed_tokens", error)) { return nx::nif::error(env, error.c_str()); } + unsigned int layer_count = 0; + unsigned int kv_count = 0; + if (!enif_get_list_length(env, argv[2], &layer_count)) { + return nx::nif::error(env, "Qwen3 greedy forward expects layers to be a list"); + } + if (!enif_get_list_length(env, argv[3], &kv_count)) { + return nx::nif::error(env, "Qwen3 greedy forward expects kv_cache to be a list"); + } + if (layer_count != kv_count) { + return nx::nif::error(env, "Qwen3 greedy forward layers and kv_cache length mismatch"); + } + + std::vector layers; + layers.reserve(layer_count); + ERL_NIF_TERM layer_head, layer_tail = argv[2]; + while (enif_get_list_cell(env, layer_tail, &layer_head, &layer_tail)) { + emlx_qwen3_plugin::LayerParams layer; + if (!qwen3_get_layer(env, layer_head, layer, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "Qwen3 greedy forward got invalid layer tuple"); + } + layers.push_back(layer); + } + + std::vector kvs; + kvs.reserve(kv_count); + ERL_NIF_TERM kv_head, kv_tail = argv[3]; + while (enif_get_list_cell(env, kv_tail, &kv_head, &kv_tail)) { + emlx_qwen3_plugin::KVCache kv; + if (!qwen3_get_kv(env, kv_head, kv, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "Qwen3 greedy forward got invalid kv_cache tuple"); + } + kvs.push_back(kv); + } + int B = input_ids->shape(0); int T = input_ids->shape(1); - auto ids = mlx::core::reshape(*input_ids, {B * T}, device); - auto embedded = 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, device), {B, T, embed_tokens->shape(1)}, device); + + mlx::core::array token_out(0); + int64_t token_id_out = 0; + std::vector k_out, v_out; + if (!plugin->forward_greedy_from_hidden( + embedded, layers, kvs, *norm, *lm_head, offset, scale, head_dim, theta, eps, false, + device, token_out, token_id_out, k_out, v_out, error)) { + return nx::nif::error(env, error.c_str()); + } - return qwen3_forward_greedy_from_hidden( - env, embedded, argv[2], argv[3], norm, lm_head, offset, scale, head_dim, - theta, eps, false, device); + ERL_NIF_TERM token_term = create_tensor_resource(env, token_out); + return nx::nif::ok(env, enif_make_tuple2(env, token_term, qwen3_wrap_kv_terms(env, k_out, v_out))); } CATCH() } @@ -1112,9 +605,14 @@ ASYNC_NIF(qwen3_forward_greedy_ids) // qwen3_forward_greedy_ids_chunk — repeatedly decode greedy tokens from a // single token id tensor without returning to Elixir between decode steps. -// Returns {token_id_refs, kv_cache}, where token_id_refs is a list of token -// tensors in generation order and kv_cache is the final updated raw cache. +// Returns {token_id_refs, kv_cache}. NIF(qwen3_forward_greedy_ids_chunk) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + TENSOR_PARAM(0, input_ids); TENSOR_PARAM(1, embed_tokens); PARAM(6, int, offset); @@ -1126,10 +624,6 @@ NIF(qwen3_forward_greedy_ids_chunk) { DEVICE_PARAM(12, device); try { - if (count <= 0) { - return nx::nif::error(env, "qwen3_forward_greedy_ids_chunk expects positive count"); - } - Qwen3TensorHandles handles; ERL_NIF_TERM ref_error = 0; mlx::core::array *norm = nullptr; @@ -1145,7 +639,6 @@ NIF(qwen3_forward_greedy_ids_chunk) { unsigned int layer_count = 0; unsigned int kv_count = 0; - if (!enif_get_list_length(env, argv[2], &layer_count)) { return nx::nif::error(env, "qwen3_forward_greedy_ids_chunk expects layers to be a list"); } @@ -1156,13 +649,11 @@ NIF(qwen3_forward_greedy_ids_chunk) { return nx::nif::error(env, "qwen3_forward_greedy_ids_chunk layers and kv_cache length mismatch"); } - std::vector layers; + std::vector layers; layers.reserve(layer_count); - - ERL_NIF_TERM layer_head; - ERL_NIF_TERM layer_tail = argv[2]; + ERL_NIF_TERM layer_head, layer_tail = argv[2]; while (enif_get_list_cell(env, layer_tail, &layer_head, &layer_tail)) { - Qwen3LayerParams layer; + emlx_qwen3_plugin::LayerParams layer; if (!qwen3_get_layer(env, layer_head, layer, handles, &ref_error)) { return qwen3_ref_error_or( env, ref_error, "qwen3_forward_greedy_ids_chunk got invalid layer tuple"); @@ -1170,13 +661,11 @@ NIF(qwen3_forward_greedy_ids_chunk) { layers.push_back(layer); } - std::vector initial_kv; + std::vector initial_kv; initial_kv.reserve(kv_count); - - ERL_NIF_TERM kv_head; - ERL_NIF_TERM kv_tail = argv[3]; + ERL_NIF_TERM kv_head, kv_tail = argv[3]; while (enif_get_list_cell(env, kv_tail, &kv_head, &kv_tail)) { - Qwen3KVCache kv; + emlx_qwen3_plugin::KVCache kv; if (!qwen3_get_kv(env, kv_head, kv, handles, &ref_error)) { return qwen3_ref_error_or( env, ref_error, "qwen3_forward_greedy_ids_chunk got invalid kv_cache tuple"); @@ -1185,154 +674,35 @@ NIF(qwen3_forward_greedy_ids_chunk) { } std::string error; - if (!qwen3_check_rank2_positive(*input_ids, "input_ids", error) || - !qwen3_check_rank2_positive(*embed_tokens, "embed_tokens", error) || - !qwen3_check_non_negative(offset, "offset", error) || - !qwen3_check_positive(head_dim, "head_dim", error) || - !qwen3_check_rank1_dim(*norm, embed_tokens->shape(1), "norm", error) || - !qwen3_check_rank2_positive(*lm_head, "lm_head", error) || - !qwen3_check_dim(*lm_head, 1, embed_tokens->shape(1), "lm_head", "hidden width", error)) { + std::vector token_out, k_out, v_out; + if (!plugin->forward_greedy_ids_chunk( + *input_ids, *embed_tokens, layers, initial_kv, *norm, *lm_head, offset, count, scale, + head_dim, theta, eps, device, token_out, k_out, v_out, error)) { return nx::nif::error(env, error.c_str()); } - if (input_ids->shape(0) != 1) { - return nx::nif::error(env, "qwen3_forward_greedy_ids_chunk requires batch size 1"); - } - if (input_ids->shape(1) != 1) { - return nx::nif::error(env, "qwen3_forward_greedy_ids_chunk requires sequence length 1"); - } - std::vector k_cache; - std::vector v_cache; - std::vector next_k_cache; - std::vector next_v_cache; std::vector token_terms; - std::vector token_arrays; - k_cache.reserve(layer_count); - v_cache.reserve(layer_count); - next_k_cache.reserve(layer_count); - next_v_cache.reserve(layer_count); - token_terms.reserve(count); - token_arrays.reserve(count); - - auto current_ids = *input_ids; - int current_offset = offset; - - 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 current = mlx::core::reshape( - mlx::core::take(*embed_tokens, ids, 0, device), - {B, T, embed_tokens->shape(1)}, - device); - - next_k_cache.clear(); - next_v_cache.clear(); - - for (unsigned int layer_idx = 0; layer_idx < layer_count; ++layer_idx) { - Qwen3KVCache kv = - (step == 0) - ? initial_kv[layer_idx] - : Qwen3KVCache{&k_cache[layer_idx], &v_cache[layer_idx]}; - - std::string layer_error; - if (!qwen3_validate_dense_layer( - current, layers[layer_idx], kv, current_offset, head_dim, layer_error)) { - return nx::nif::error(env, layer_error.c_str()); - } - - ERL_NIF_TERM unused_k_term; - ERL_NIF_TERM unused_v_term; - auto k_new = *kv.k; - auto v_new = *kv.v; - - current = qwen3_dense_layer_impl( - current, - layers[layer_idx], - kv, - current_offset, - static_cast(scale), - head_dim, - static_cast(theta), - static_cast(eps), - device, - unused_k_term, - unused_v_term, - nullptr, - &k_new, - &v_new); - - next_k_cache.push_back(k_new); - next_v_cache.push_back(v_new); - } - - int B_out = current.shape(0); - int T_out = current.shape(1); - int H_out = current.shape(2); - - auto last = (T_out == 1) - ? mlx::core::reshape(current, {B_out, H_out}, device) - : 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); - - auto normed = mlx::core::fast::rms_norm(last, *norm, static_cast(eps), device); - auto logits = qwen3_linear_out_in(normed, *lm_head, device); - auto token = mlx::core::argmax(logits, 1, false, device); - - token_arrays.push_back(token); + token_terms.reserve(token_out.size()); + for (const auto &token : token_out) { token_terms.push_back(create_tensor_resource(env, token)); - current_ids = mlx::core::reshape(token, {B_out, 1}, device); - k_cache.swap(next_k_cache); - v_cache.swap(next_v_cache); - current_offset += 1; - } - - // Queue token materialisation and final cache ownership before Elixir - // starts copying/decoding tokens or returns the cache to the next chunk. - std::vector eval_arrays; - eval_arrays.reserve(token_arrays.size() + (layer_count * 2)); - eval_arrays.insert(eval_arrays.end(), token_arrays.begin(), token_arrays.end()); - - for (unsigned int layer_idx = 0; layer_idx < layer_count; ++layer_idx) { - eval_arrays.push_back(k_cache[layer_idx]); - eval_arrays.push_back(v_cache[layer_idx]); - } - - mlx::core::async_eval(eval_arrays); - - std::vector kv_terms; - kv_terms.reserve(layer_count); - - for (unsigned int layer_idx = 0; layer_idx < layer_count; ++layer_idx) { - kv_terms.push_back( - enif_make_tuple2( - env, - create_tensor_resource(env, k_cache[layer_idx]), - create_tensor_resource(env, v_cache[layer_idx]))); } - ERL_NIF_TERM token_list = - enif_make_list_from_array(env, token_terms.data(), token_terms.size()); - ERL_NIF_TERM kv_list = - enif_make_list_from_array(env, kv_terms.data(), kv_terms.size()); - - return nx::nif::ok(env, enif_make_tuple2(env, token_list, kv_list)); + ERL_NIF_TERM token_list = enif_make_list_from_array(env, token_terms.data(), token_terms.size()); + return nx::nif::ok(env, enif_make_tuple2(env, token_list, qwen3_wrap_kv_terms(env, k_out, v_out))); } CATCH() } ASYNC_NIF(qwen3_forward_greedy_ids_chunk) // qwen3_forward_greedy_ids_token_id — same as qwen3_forward_greedy_ids, but -// returns the sampled token id as a BEAM integer. This is for decode paths that -// already need each token on the host for streaming/callbacks. +// returns the sampled token id as a BEAM integer. NIF(qwen3_forward_greedy_ids_token_id) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + TENSOR_PARAM(0, input_ids); TENSOR_PARAM(1, embed_tokens); PARAM(6, int, offset); @@ -1356,27 +726,63 @@ NIF(qwen3_forward_greedy_ids_token_id) { env, ref_error, "qwen3_forward_greedy_ids_token_id expects lm_head tensor ref"); } + unsigned int layer_count = 0; + unsigned int kv_count = 0; + if (!enif_get_list_length(env, argv[2], &layer_count)) { + return nx::nif::error(env, "Qwen3 greedy forward expects layers to be a list"); + } + if (!enif_get_list_length(env, argv[3], &kv_count)) { + return nx::nif::error(env, "Qwen3 greedy forward expects kv_cache to be a list"); + } + if (layer_count != kv_count) { + return nx::nif::error(env, "Qwen3 greedy forward layers and kv_cache length mismatch"); + } + + std::vector layers; + layers.reserve(layer_count); + ERL_NIF_TERM layer_head, layer_tail = argv[2]; + while (enif_get_list_cell(env, layer_tail, &layer_head, &layer_tail)) { + emlx_qwen3_plugin::LayerParams layer; + if (!qwen3_get_layer(env, layer_head, layer, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "Qwen3 greedy forward got invalid layer tuple"); + } + layers.push_back(layer); + } + + std::vector kvs; + kvs.reserve(kv_count); + ERL_NIF_TERM kv_head, kv_tail = argv[3]; + while (enif_get_list_cell(env, kv_tail, &kv_head, &kv_tail)) { + emlx_qwen3_plugin::KVCache kv; + if (!qwen3_get_kv(env, kv_head, kv, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "Qwen3 greedy forward got invalid kv_cache tuple"); + } + kvs.push_back(kv); + } + std::string error; - if (!qwen3_check_rank2_positive(*input_ids, "input_ids", error) || - !qwen3_check_rank2_positive(*embed_tokens, "embed_tokens", error) || - !qwen3_check_non_negative(offset, "offset", error) || - !qwen3_check_positive(head_dim, "head_dim", error)) { + if (!qwen3_require_rank2_positive(*input_ids, "input_ids", error) || + !qwen3_require_rank2_positive(*embed_tokens, "embed_tokens", error)) { return nx::nif::error(env, error.c_str()); } int B = input_ids->shape(0); int T = input_ids->shape(1); - auto ids = mlx::core::reshape(*input_ids, {B * T}, device); - auto embedded = 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, device), {B, T, embed_tokens->shape(1)}, device); + + mlx::core::array token_out(0); + int64_t token_id_out = 0; + std::vector k_out, v_out; + if (!plugin->forward_greedy_from_hidden( + embedded, layers, kvs, *norm, *lm_head, offset, scale, head_dim, theta, eps, true, + device, token_out, token_id_out, k_out, v_out, error)) { + return nx::nif::error(env, error.c_str()); + } - return qwen3_forward_greedy_from_hidden( - env, embedded, argv[2], argv[3], norm, lm_head, offset, scale, head_dim, - theta, eps, true, device); + return nx::nif::ok( + env, enif_make_tuple2(env, nx::nif::make(env, token_id_out), qwen3_wrap_kv_terms(env, k_out, v_out))); } CATCH() } @@ -1386,6 +792,12 @@ ASYNC_NIF(qwen3_forward_greedy_ids_token_id) // token as a BEAM integer, avoiding host Nx tensor construction and backend // transfer for the single token greedy decode hot path. NIF(qwen3_forward_greedy_token_id) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + PARAM(0, int64_t, token_id); TENSOR_PARAM(1, embed_tokens); PARAM(6, int, offset); @@ -1408,40 +820,77 @@ NIF(qwen3_forward_greedy_token_id) { return qwen3_ref_error_or( env, ref_error, "qwen3_forward_greedy_token_id expects lm_head tensor ref"); } - std::string error; - if (!qwen3_check_rank2_positive(*embed_tokens, "embed_tokens", error) || - !qwen3_check_non_negative(offset, "offset", error) || - !qwen3_check_positive(head_dim, "head_dim", error)) { - return nx::nif::error(env, error.c_str()); + std::string rank_error; + if (!qwen3_require_rank2_positive(*embed_tokens, "embed_tokens", rank_error)) { + return nx::nif::error(env, rank_error.c_str()); } if (token_id < 0 || token_id >= embed_tokens->shape(0)) { return nx::nif::error(env, "token_id is outside the embedding vocabulary"); } + unsigned int layer_count = 0; + unsigned int kv_count = 0; + if (!enif_get_list_length(env, argv[2], &layer_count)) { + return nx::nif::error(env, "Qwen3 greedy forward expects layers to be a list"); + } + if (!enif_get_list_length(env, argv[3], &kv_count)) { + return nx::nif::error(env, "Qwen3 greedy forward expects kv_cache to be a list"); + } + if (layer_count != kv_count) { + return nx::nif::error(env, "Qwen3 greedy forward layers and kv_cache length mismatch"); + } + + std::vector layers; + layers.reserve(layer_count); + ERL_NIF_TERM layer_head, layer_tail = argv[2]; + while (enif_get_list_cell(env, layer_tail, &layer_head, &layer_tail)) { + emlx_qwen3_plugin::LayerParams layer; + if (!qwen3_get_layer(env, layer_head, layer, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "Qwen3 greedy forward got invalid layer tuple"); + } + layers.push_back(layer); + } + + std::vector kvs; + kvs.reserve(kv_count); + ERL_NIF_TERM kv_head, kv_tail = argv[3]; + while (enif_get_list_cell(env, kv_tail, &kv_head, &kv_tail)) { + emlx_qwen3_plugin::KVCache kv; + if (!qwen3_get_kv(env, kv_head, kv, handles, &ref_error)) { + return qwen3_ref_error_or(env, ref_error, "Qwen3 greedy forward got invalid kv_cache tuple"); + } + kvs.push_back(kv); + } + + std::string error; auto ids = mlx::core::array(token_id, mlx::core::int64); auto embedded = mlx::core::reshape( - mlx::core::take(*embed_tokens, ids, 0, device), - {1, 1, embed_tokens->shape(1)}, - device); + mlx::core::take(*embed_tokens, ids, 0, device), {1, 1, embed_tokens->shape(1)}, device); + + mlx::core::array token_out(0); + int64_t token_id_out = 0; + std::vector k_out, v_out; + if (!plugin->forward_greedy_from_hidden( + embedded, layers, kvs, *norm, *lm_head, offset, scale, head_dim, theta, eps, true, + device, token_out, token_id_out, k_out, v_out, error)) { + return nx::nif::error(env, error.c_str()); + } - return qwen3_forward_greedy_from_hidden( - env, embedded, argv[2], argv[3], norm, lm_head, offset, scale, head_dim, - theta, eps, true, device); + return nx::nif::ok( + env, enif_make_tuple2(env, nx::nif::make(env, token_id_out), qwen3_wrap_kv_terms(env, k_out, v_out))); } CATCH() } ASYNC_NIF(qwen3_forward_greedy_token_id) // qwen3_final_greedy — final RMSNorm + dense lm_head + argmax for greedy decode. -// -// Inputs: -// hidden — {B, T, H} -// norm — {H} -// lm_head — {V, H} -// eps — RMSNorm epsilon -// -// Returns token ids with shape {B}. NIF(qwen3_final_greedy) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + TENSOR_PARAM(0, hidden); TENSOR_PARAM(1, norm); TENSOR_PARAM(2, lm_head); @@ -1449,72 +898,39 @@ NIF(qwen3_final_greedy) { DEVICE_PARAM(4, device); try { + mlx::core::array out(0); std::string error; - if (!qwen3_check_rank3_positive(*hidden, "hidden", error)) { + if (!plugin->final_greedy(*hidden, *norm, *lm_head, eps, device, out, error)) { return nx::nif::error(env, error.c_str()); } - int B = hidden->shape(0); - int T = hidden->shape(1); - int H = hidden->shape(2); - if (!qwen3_check_rank1_dim(*norm, H, "norm", error) || - !qwen3_check_rank2_positive(*lm_head, "lm_head", error) || - !qwen3_check_dim(*lm_head, 1, H, "lm_head", "hidden width", error)) { - return nx::nif::error(env, error.c_str()); - } - - auto last = (T == 1) - ? mlx::core::reshape(*hidden, {B, H}, device) - : mlx::core::reshape( - mlx::core::slice(*hidden, to_shape({0, T - 1, 0}), to_shape({B, T, H}), device), - {B, H}, - device); - - auto normed = mlx::core::fast::rms_norm(last, *norm, (float)eps, device); - auto logits = mlx::core::tensordot(normed, *lm_head, std::vector{1}, std::vector{1}, device); - - TENSOR(mlx::core::argmax(logits, 1, false, device)); + TENSOR(out); } CATCH() } ASYNC_NIF(qwen3_final_greedy) // qwen3_attention_residual — dense attention output projection + residual add. -// -// Inputs: -// hidden — {B, T, H} -// attn_out — {B, T, I} -// o_proj — {I, H} -// -// Returns hidden + attn_out @ o_proj as {B, T, H}. NIF(qwen3_attention_residual) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + TENSOR_PARAM(0, hidden); TENSOR_PARAM(1, attn_out); TENSOR_PARAM(2, o_proj); DEVICE_PARAM(3, device); try { + mlx::core::array out(0); std::string error; - if (!qwen3_check_rank3_positive(*hidden, "hidden", error) || - !qwen3_check_rank3_positive(*attn_out, "attn_out", error)) { - return nx::nif::error(env, error.c_str()); - } - - int B = hidden->shape(0); - int T = hidden->shape(1); - int H = hidden->shape(2); - if (!qwen3_check_dim(*attn_out, 0, B, "attn_out", "batch", error) || - !qwen3_check_dim(*attn_out, 1, T, "attn_out", "sequence length", error) || - !qwen3_check_rank2_positive(*o_proj, "o_proj", error) || - !qwen3_check_dim(*o_proj, 0, attn_out->shape(2), "o_proj", "input width", error) || - !qwen3_check_dim(*o_proj, 1, H, "o_proj", "output width", error)) { + if (!plugin->attention_residual(*hidden, *attn_out, *o_proj, device, out, error)) { return nx::nif::error(env, error.c_str()); } - auto projected = qwen3_linear_in_out(*attn_out, *o_proj, device); - auto residual = mlx::core::add(*hidden, projected, device); - - TENSOR(residual); + TENSOR(out); } CATCH() } @@ -1524,25 +940,14 @@ ASYNC_NIF(qwen3_attention_residual) // input RMSNorm + Q/K/V projections + Q/K RMSNorm + RoPE + KV update + SDPA // + output projection + residual add. // -// Inputs: -// hidden — {B, T_new, H} residual hidden state before attention -// norm — {H} input RMSNorm weight -// q_proj — {H, N_q * D} -// k_proj — {H, N_kv * D} -// v_proj — {H, N_kv * D} -// o_proj — {N_q * D, H} -// q_norm — {D} -// k_norm — {D} -// k_cache — {B, N_kv, T_max, D} -// v_cache — {B, N_kv, T_max, D} -// offset — int -// scale — float -// head_dim — int -// theta — float -// eps — RMSNorm epsilon -// // Returns {hidden_out, k_upd, v_upd}. NIF(qwen3_attention_block) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + TENSOR_PARAM(0, hidden); TENSOR_PARAM(1, norm); TENSOR_PARAM(2, q_proj); @@ -1561,122 +966,120 @@ NIF(qwen3_attention_block) { DEVICE_PARAM(15, device); try { + mlx::core::array out(0), k_upd(0), v_upd(0); std::string error; - if (!qwen3_check_rank3_positive(*hidden, "hidden", error) || - !qwen3_check_non_negative(offset, "offset", error) || - !qwen3_check_positive(head_dim, "head_dim", error)) { + if (!plugin->attention_block(*hidden, *norm, *q_proj, *k_proj, *v_proj, *o_proj, + *q_norm, *k_norm, *k_cache, *v_cache, offset, scale, + head_dim, theta, eps, device, out, k_upd, v_upd, + error)) { return nx::nif::error(env, error.c_str()); } - int B = hidden->shape(0); - int T_new = hidden->shape(1); - int H = hidden->shape(2); - int D = head_dim; - if (!qwen3_check_rank1_dim(*norm, H, "norm", error) || - !qwen3_check_rank1_dim(*q_norm, D, "q_norm", error) || - !qwen3_check_rank1_dim(*k_norm, D, "k_norm", error) || - !qwen3_check_rank2_positive(*q_proj, "q_proj", error) || - !qwen3_check_rank2_positive(*k_proj, "k_proj", error) || - !qwen3_check_rank2_positive(*v_proj, "v_proj", error) || - !qwen3_check_dim(*q_proj, 0, H, "q_proj", "input width", error) || - !qwen3_check_dim(*k_proj, 0, H, "k_proj", "input width", error) || - !qwen3_check_dim(*v_proj, 0, H, "v_proj", "input width", error)) { - return nx::nif::error(env, error.c_str()); + ERL_NIF_TERM result_tuple[3]; + result_tuple[0] = create_tensor_resource(env, out); + result_tuple[1] = create_tensor_resource(env, k_upd); + result_tuple[2] = create_tensor_resource(env, v_upd); + + return nx::nif::ok(env, enif_make_tuple3(env, result_tuple[0], result_tuple[1], result_tuple[2])); + } + CATCH() +} +ASYNC_NIF(qwen3_attention_block) + +// qwen3_forward_greedy_ids_chunk_quantized — generalized variant of +// `qwen3_forward_greedy_ids_chunk`: every layer's 7 projections and the +// final `lm_head` each independently accept a dense-or-quantized weight +// term (see `qwen3_get_linear_weight`). +// +// Returns {token_id_refs, kv_cache}. +NIF(qwen3_forward_greedy_ids_chunk_quantized) { + ERL_NIF_TERM plugin_error; + const emlx_qwen3_plugin::VTable *plugin = qwen3_plugin(env, &plugin_error); + if (plugin == nullptr) { + return plugin_error; + } + + TENSOR_PARAM(0, input_ids); + TENSOR_PARAM(1, embed_tokens); + PARAM(6, int, offset); + PARAM(7, int, count); + PARAM(8, double, scale); + PARAM(9, int, head_dim); + PARAM(10, double, theta); + PARAM(11, double, eps); + DEVICE_PARAM(12, device); + + try { + Qwen3TensorHandles handles; + ERL_NIF_TERM ref_error = 0; + mlx::core::array *norm = nullptr; + emlx_qwen3_plugin::LinearWeight lm_head; + if (!qwen3_get_tensor(env, argv[4], &norm, handles, &ref_error)) { + return qwen3_ref_error_or( + env, ref_error, "qwen3_forward_greedy_ids_chunk_quantized expects norm tensor ref"); + } + if (!qwen3_get_linear_weight(env, argv[5], true, lm_head, handles, &ref_error)) { + return qwen3_ref_error_or( + env, ref_error, + "qwen3_forward_greedy_ids_chunk_quantized got invalid lm_head weight term"); } - if ((q_proj->shape(1) % D) != 0 || (k_proj->shape(1) % D) != 0) { - return nx::nif::error(env, "projection output widths must be divisible by head_dim"); + unsigned int layer_count = 0; + unsigned int kv_count = 0; + if (!enif_get_list_length(env, argv[2], &layer_count)) { + return nx::nif::error( + env, "qwen3_forward_greedy_ids_chunk_quantized expects layers to be a list"); + } + if (!enif_get_list_length(env, argv[3], &kv_count)) { + return nx::nif::error( + env, "qwen3_forward_greedy_ids_chunk_quantized expects kv_cache to be a list"); } - if (v_proj->shape(1) != k_proj->shape(1)) { - return nx::nif::error(env, "v_proj output width must match k_proj output width"); + if (layer_count != kv_count) { + return nx::nif::error( + env, "qwen3_forward_greedy_ids_chunk_quantized layers and kv_cache length mismatch"); } - int N_q = q_proj->shape(1) / D; - int N_kv = k_proj->shape(1) / D; - int attn_width = N_q * D; + std::vector layers; + layers.reserve(layer_count); + ERL_NIF_TERM layer_head, layer_tail = argv[2]; + while (enif_get_list_cell(env, layer_tail, &layer_head, &layer_tail)) { + emlx_qwen3_plugin::LayerParamsQ layer; + if (!qwen3_get_layer_generalized(env, layer_head, layer, handles, &ref_error)) { + return qwen3_ref_error_or( + env, ref_error, "qwen3_forward_greedy_ids_chunk_quantized got invalid layer tuple"); + } + layers.push_back(layer); + } - if ((N_q % N_kv) != 0) { - return nx::nif::error(env, "query heads must be divisible by key/value heads"); + std::vector initial_kv; + initial_kv.reserve(kv_count); + ERL_NIF_TERM kv_head, kv_tail = argv[3]; + while (enif_get_list_cell(env, kv_tail, &kv_head, &kv_tail)) { + emlx_qwen3_plugin::KVCache kv; + if (!qwen3_get_kv(env, kv_head, kv, handles, &ref_error)) { + return qwen3_ref_error_or( + env, ref_error, "qwen3_forward_greedy_ids_chunk_quantized got invalid kv_cache tuple"); + } + initial_kv.push_back(kv); } - if (!qwen3_check_rank2_positive(*o_proj, "o_proj", error) || - !qwen3_check_dim(*o_proj, 0, attn_width, "o_proj", "input width", error) || - !qwen3_check_dim(*o_proj, 1, H, "o_proj", "output width", error) || - !qwen3_validate_kv_cache_bn(*k_cache, *v_cache, B, N_kv, offset, T_new, D, error)) { + + std::string error; + std::vector token_out, k_out, v_out; + if (!plugin->forward_greedy_ids_chunk_quantized( + *input_ids, *embed_tokens, layers, initial_kv, *norm, lm_head, offset, count, scale, + head_dim, theta, eps, device, token_out, k_out, v_out, error)) { return nx::nif::error(env, error.c_str()); } - int valid_len = offset + T_new; - - auto xn = mlx::core::fast::rms_norm(*hidden, *norm, (float)eps, device); - auto q_flat = qwen3_linear_in_out(xn, *q_proj, device); - auto k_flat = qwen3_linear_in_out(xn, *k_proj, device); - auto v_flat = qwen3_linear_in_out(xn, *v_proj, device); - - 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); - - q = mlx::core::fast::rms_norm(q, *q_norm, (float)eps, device); - k = mlx::core::fast::rms_norm(k, *k_norm, (float)eps, device); - - 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_rope = mlx::core::fast::rope( - q_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, device); - auto k_rope = mlx::core::fast::rope( - k_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, device); - - auto k_cache_owned = std::move(*k_cache); - auto v_cache_owned = std::move(*v_cache); - - 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); - 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); - - 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 build_prefill_mask = [&]() -> mlx::core::array { - 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); - 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 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); - }; - - auto attn_out_bn = (T_new == 1) - ? mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, (float)scale, "", std::nullopt, std::nullopt, device) - : mlx::core::fast::scaled_dot_product_attention( - q_rope, k_valid, v_valid, (float)scale, "array", build_prefill_mask(), std::nullopt, 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 = qwen3_linear_in_out(attn_out, *o_proj, device); - auto out = mlx::core::add(*hidden, projected, device); - ERL_NIF_TERM result_tuple[3]; - result_tuple[0] = create_tensor_resource(env, out); - result_tuple[1] = create_tensor_resource(env, k_upd); - result_tuple[2] = create_tensor_resource(env, v_upd); + std::vector token_terms; + token_terms.reserve(token_out.size()); + for (const auto &token : token_out) { + token_terms.push_back(create_tensor_resource(env, token)); + } - return nx::nif::ok(env, enif_make_tuple3(env, result_tuple[0], result_tuple[1], result_tuple[2])); + ERL_NIF_TERM token_list = enif_make_list_from_array(env, token_terms.data(), token_terms.size()); + return nx::nif::ok(env, enif_make_tuple2(env, token_list, qwen3_wrap_kv_terms(env, k_out, v_out))); } CATCH() } -ASYNC_NIF(qwen3_attention_block) +ASYNC_NIF(qwen3_forward_greedy_ids_chunk_quantized) diff --git a/emlx/c_src/emlx_fast/qwen3.hpp b/emlx/c_src/emlx_fast/qwen3.hpp index 5dca1e8..56aa4ba 100644 --- a/emlx/c_src/emlx_fast/qwen3.hpp +++ b/emlx/c_src/emlx_fast/qwen3.hpp @@ -8,8 +8,10 @@ ERL_NIF_TERM qwen3_kv_cache_attention_async(ErlNifEnv *, int, const ERL_NIF_TERM []); ERL_NIF_TERM qwen3_mlp_async(ErlNifEnv *, int, const ERL_NIF_TERM []); ERL_NIF_TERM qwen3_layer_async(ErlNifEnv *, int, const ERL_NIF_TERM []); +ERL_NIF_TERM qwen3_layer_quantized_async(ErlNifEnv *, int, const ERL_NIF_TERM []); ERL_NIF_TERM qwen3_forward_greedy_ids_async(ErlNifEnv *, int, const ERL_NIF_TERM []); ERL_NIF_TERM qwen3_forward_greedy_ids_chunk_async(ErlNifEnv *, int, const ERL_NIF_TERM []); +ERL_NIF_TERM qwen3_forward_greedy_ids_chunk_quantized_async(ErlNifEnv *, int, const ERL_NIF_TERM []); ERL_NIF_TERM qwen3_forward_greedy_ids_token_id_async(ErlNifEnv *, int, const ERL_NIF_TERM []); ERL_NIF_TERM qwen3_forward_greedy_token_id_async(ErlNifEnv *, int, const ERL_NIF_TERM []); ERL_NIF_TERM qwen3_final_greedy_async(ErlNifEnv *, int, const ERL_NIF_TERM []); diff --git a/emlx/c_src/emlx_fast/qwen3_plugin_abi.hpp b/emlx/c_src/emlx_fast/qwen3_plugin_abi.hpp new file mode 100644 index 0000000..61a8d4b --- /dev/null +++ b/emlx/c_src/emlx_fast/qwen3_plugin_abi.hpp @@ -0,0 +1,179 @@ +#pragma once + +// ABI shared between the host NIF library (emlx's c_src/emlx_fast/qwen3.cpp, +// statically linked into libemlx.so) and the standalone qwen3 compute +// plugin (emlx_axon's c_src/qwen3_plugin.cpp, built as its own shared +// library and loaded at runtime via `EMLX.NIF.load_plugin("qwen3", path)`). +// +// This header lives in emlx (not emlx_axon) because it is the contract the +// *host* decode/encode side depends on — emlx_axon's plugin build adds +// emlx's `c_src/emlx_fast` to its include path to pick it up (see +// emlx_axon/Makefile). +// +// This header intentionally has NO dependency on erl_nif.h: every type here +// is plain C++/MLX so the plugin can be built and `dlopen`'d without ever +// linking against the Erlang runtime. The host decodes Erlang terms into +// these plain structs (see qwen3.cpp's `qwen3_get_*` helpers) before calling +// through `VTable`, and re-encodes the `mlx::core::array` results back into +// tensor resources afterwards. This split works only because the boundary +// carries raw `mlx::core::array`/struct values, never Erlang resource +// terms — a truly independent NIF module can't share tensor resource types +// with the rest of EMLX (Erlang ties resource types to the specific +// module/`.so` that opened them). + +#include "mlx/mlx.h" + +#include +#include +#include + +namespace emlx_qwen3_plugin { + +// A dense-or-quantized projection weight — mirrors `EMLX.Native.Qwen3`'s +// `{:dense, ref}` / `{:quantized, ...}` linear weight term after decoding. +struct LinearWeight { + bool quantized = false; + mlx::core::array *weight = nullptr; + mlx::core::array *scales = nullptr; + mlx::core::array *biases = nullptr; + int group_size = 0; + int bits = 0; + std::string mode; + bool transpose = false; +}; + +// Dense-only per-layer projection set (qwen3_layer / qwen3_forward_greedy_*). +struct LayerParams { + mlx::core::array *norm1; + mlx::core::array *norm2; + mlx::core::array *q_norm; + mlx::core::array *k_norm; + mlx::core::array *q_proj; + mlx::core::array *k_proj; + mlx::core::array *v_proj; + mlx::core::array *o_proj; + mlx::core::array *gate_proj; + mlx::core::array *up_proj; + mlx::core::array *down_proj; +}; + +// Generalized (dense-or-quantized) per-layer projection set +// (qwen3_layer_quantized / qwen3_forward_greedy_ids_chunk_quantized). +struct LayerParamsQ { + mlx::core::array *norm1; + mlx::core::array *norm2; + mlx::core::array *q_norm; + mlx::core::array *k_norm; + LinearWeight q_proj; + LinearWeight k_proj; + LinearWeight v_proj; + LinearWeight o_proj; + LinearWeight gate_proj; + LinearWeight up_proj; + LinearWeight down_proj; +}; + +struct KVCache { + mlx::core::array *k; + mlx::core::array *v; +}; + +// Every entrypoint validates its own inputs and returns `false` + fills +// `error` on failure instead of throwing — C++ exceptions are not relied +// on to cross the `dlopen` boundary. +struct VTable { + bool (*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, + mlx::core::array &k_upd, mlx::core::array &v_upd, std::string &error); + + bool (*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, mlx::core::array &out, + std::string &error); + + bool (*attention_residual)( + const mlx::core::array &hidden, const mlx::core::array &attn_out, + const mlx::core::array &o_proj, const mlx::core::Device &device, + mlx::core::array &out, std::string &error); + + bool (*attention_block)( + const mlx::core::array &hidden, const mlx::core::array &norm, + const mlx::core::array &q_proj, const mlx::core::array &k_proj, + const mlx::core::array &v_proj, const mlx::core::array &o_proj, + 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::array &k_upd, mlx::core::array &v_upd, std::string &error); + + bool (*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::array &k_upd, mlx::core::array &v_upd, + std::string &error); + + bool (*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::array &k_upd, + mlx::core::array &v_upd, std::string &error); + + bool (*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, mlx::core::array &out, + std::string &error); + + // Dense greedy forward from an already-embedded hidden state through + // every layer, then final norm + lm_head + argmax. `layers`/`kv` are + // parallel per-layer vectors; `kv[i]` is updated in place with the new + // K/V cache and also returned via `k_out`/`v_out` for the caller to wrap + // as tensor resources. Exactly one of `token_out`/`token_id_out` is + // meaningful, selected by `return_token_id`. + bool (*forward_greedy_from_hidden)( + const mlx::core::array &hidden, std::vector &layers, + 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, + int64_t &token_id_out, std::vector &k_out, + std::vector &v_out, std::string &error); + + // Dense chunked decode: `count` greedy steps starting from `input_ids` + // ({1,1}), threading the KV cache across steps without returning to + // Elixir in between. + bool (*forward_greedy_ids_chunk)( + 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 mlx::core::array &lm_head, + int offset, int count, double scale, int head_dim, double theta, + double eps, const mlx::core::Device &device, + std::vector &token_out, + std::vector &k_out, + std::vector &v_out, std::string &error); + + // Generalized (dense-or-quantized) chunked decode. + bool (*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, + std::vector &token_out, + std::vector &k_out, + std::vector &v_out, std::string &error); +}; + +} // namespace emlx_qwen3_plugin + +// Sole exported entrypoint — `EMLX.NIF.load_plugin("qwen3", path)` +// `dlopen`s the plugin and `dlsym`s the generic `emlx_plugin_vtable` symbol +// (see emlx/c_src/emlx_plugin_registry.hpp) to obtain this vtable. +extern "C" const emlx_qwen3_plugin::VTable *emlx_plugin_vtable(); diff --git a/emlx/c_src/emlx_nif.cpp b/emlx/c_src/emlx_nif.cpp index ab47462..f6afcfa 100644 --- a/emlx/c_src/emlx_nif.cpp +++ b/emlx/c_src/emlx_nif.cpp @@ -1,5 +1,7 @@ +#include "emlx_compiler.hpp" #include "emlx_nif_shared.hpp" #include "emlx_fast/qwen3.hpp" +#include "emlx_plugin_registry.hpp" #include #include @@ -10,15 +12,6 @@ #include #include -std::map dtypes = { - {"bool", mlx::core::bool_}, {"uint8", mlx::core::uint8}, - {"uint16", mlx::core::uint16}, {"uint32", mlx::core::uint32}, - {"uint64", mlx::core::uint64}, {"int8", mlx::core::int8}, - {"int16", mlx::core::int16}, {"int32", mlx::core::int32}, - {"int64", mlx::core::int64}, {"float16", mlx::core::float16}, - {"float32", mlx::core::float32}, {"bfloat16", mlx::core::bfloat16}, - {"complex64", mlx::core::complex64}}; - std::map dtype_sizes = { {"bool", mlx::core::bool_.size()}, {"uint8", mlx::core::uint8.size()}, @@ -34,23 +27,6 @@ std::map dtype_sizes = { {"bfloat16", mlx::core::bfloat16.size()}, {"complex64", mlx::core::complex64.size()}}; -inline mlx::core::Dtype string2dtype(const std::string &atom) { - auto it = dtypes.find(atom); - if (it != dtypes.end()) { - return it->second; - } - throw std::runtime_error("Unknown dtype: " + atom); -} - -inline const std::string *dtype2string(const mlx::core::Dtype dtype) { - for (const auto &pair : dtypes) { - if (pair.second == dtype) { - return &pair.first; - } - } - return nullptr; -} - ERL_NIF_TERM create_tensor_resource(ErlNifEnv *env, mlx::core::array tensor) { @@ -75,29 +51,6 @@ create_tensor_resource(ErlNifEnv *env, mlx::core::array tensor) { return ret; } -ERL_NIF_TERM create_function_resource(ErlNifEnv *env, emlx::function function) { - ERL_NIF_TERM ret; - std::atomic *refcount; - auto function_ptr = (emlx::function *)enif_alloc_resource( - resource_object::type, - sizeof(std::function(const std::vector &)>) + - sizeof(std::atomic) + sizeof(std::atomic_flag)); - - if (function_ptr == NULL) { - return enif_make_badarg(env); - } - - new (function_ptr) emlx::function(function); - refcount = new (function_ptr + 1) std::atomic(1); - new (refcount + 1) std::atomic_flag(); - - ret = enif_make_resource(env, function_ptr); - enif_release_resource(function_ptr); - - return ret; -} - - NIF(deallocate) { TensorP t(env, argv[0]); if (t.deallocate()) { @@ -137,14 +90,6 @@ NIF(ones) { TENSOR(mlx::core::ones(to_shape(shape), type, device)); } -NIF(zeros) { - SHAPE_PARAM(0, shape); - TYPE_PARAM(1, type); - DEVICE_PARAM(2, device); - - TENSOR(mlx::core::zeros(to_shape(shape), type, device)); -} - NIF(reshape) { TENSOR_PARAM(0, t); SHAPE_PARAM(1, shape); @@ -332,18 +277,19 @@ NIF(tensordot) { } NIF(einsum) { - TENSOR_PARAM(0, a); - TENSOR_PARAM(1, b); + // Variadic operand count (2+), e.g. for a 3-operand contraction like + // "ij,jk,kl->il" — see stack/concatenate above for the same + // list-of-tensor-resources decode pattern. + LIST_PARAM(0, std::vector, arrays); std::string spec_string; - if (!nx::nif::get(env, argv[2], spec_string)) { + if (!nx::nif::get(env, argv[1], spec_string)) { return nx::nif::error(env, "Unable to get spec_string param."); } - DEVICE_PARAM(3, device); + DEVICE_PARAM(2, device); - TENSOR(mlx::core::einsum(spec_string, std::vector({*a, *b}), - device)); + TENSOR(mlx::core::einsum(spec_string, arrays, device)); } NIF(tri_inv) { @@ -500,6 +446,21 @@ NIF(argsort) { NIF(eval) { TENSOR_PARAM(0, t); mlx::core::eval(*t); + mlx::core::synchronize(); + return nx::nif::ok(env); +} + +// Evaluates several tensors in one round trip instead of one `eval` NIF +// call per ref. `eval_program` itself returns lazy refs (see its comment); +// this batches their materialization. Currently unused from Elixir (kept +// as a general-purpose multi-ref eval primitive) — the earlier in-graph +// `:host_callback` round-trip opcode that motivated it was removed in +// favor of graph-splitting on bare `Nx.runtime_call` (see `EMLX.__compile__/3` +// and `EMLX.Defn.Tree`'s `:__EMLX__` metadata handling in `expr.ex`). +NIF(eval_many) { + LIST_PARAM(0, std::vector, inputs); + mlx::core::eval(inputs); + mlx::core::synchronize(); return nx::nif::ok(env); } @@ -1037,10 +998,6 @@ static int open_resources(ErlNifEnv *env) { return -1; } - if (!open_resource(env, mod, "CompiledFunction")) { - return -1; - } - // emlx::Worker — backs EMLX.CommandQueue and the application default // worker. Default destructor (~Worker) signals stop, drains pending // jobs, and joins the OS thread. @@ -1048,6 +1005,17 @@ static int open_resources(ErlNifEnv *env) { return -1; } + // emlx::native::Expr — opaque compiled program resource for the defn compiler. + if (!open_resource(env, mod, "NativeProgram")) { + return -1; + } + + // emlx::native::PendingRuntimeCall — opaque handle for one in-flight + // Nx.runtime_call/4 round trip (see emlx_runtime_call_bridge.hpp). + if (!open_resource(env, mod, "PendingRuntimeCall")) { + return -1; + } + return 0; } @@ -1381,52 +1349,76 @@ NIF(as_strided) { // quantized_matmul - Multiplies x with a quantized weight matrix w // This is the key operation for efficient 4-bit inference // MLX API: quantized_matmul(x, w, scales, biases, transpose, group_size, bits, stream) +// mode: "affine" (default, real biases) or a microscaled variant +// ("mxfp4"/"mxfp8"/"nvfp4" — no biases; mx::fp_quantize returns only +// (wq, scales)). biases is `nil` from Elixir for microscaled modes. NIF(quantized_matmul) { TENSOR_PARAM(0, x); // Input tensor [batch, seq, hidden] TENSOR_PARAM(1, w); // Quantized weights [out/8, in] (uint32 packed) - TENSOR_PARAM(2, scales); // Scales [out/group_size, in] (bfloat16) - TENSOR_PARAM(3, biases); // Biases [out/group_size, in] (bfloat16) + TENSOR_PARAM(2, scales); // Scales [out/group_size, in] (bfloat16, or u8 for microscaled) + OPTIONAL_TENSOR_PARAM(3, biases); // Biases (bfloat16); nil for microscaled modes PARAM(4, bool, transpose); PARAM(5, int, group_size); PARAM(6, int, bits); - DEVICE_PARAM(7, device); + std::string mode; + if (!nx::nif::get(env, argv[7], mode)) { + return nx::nif::error(env, "Unable to get mode param."); + } + DEVICE_PARAM(8, device); + + std::optional biases_opt = + biases ? std::make_optional(*biases) : std::nullopt; TENSOR(mlx::core::quantized_matmul( - *x, *w, *scales, *biases, transpose, group_size, bits, "affine", device)); + *x, *w, *scales, biases_opt, transpose, group_size, bits, mode, device)); } // dequantize - Converts quantized weights back to float // Useful for debugging and verification -// MLX API: dequantize(w, scales, biases, group_size, bits, stream) +// MLX API: dequantize(w, scales, biases, group_size, bits, mode, stream) NIF(dequantize) { TENSOR_PARAM(0, w); // Quantized weights (uint32 packed) - TENSOR_PARAM(1, scales); // Scales (bfloat16) - TENSOR_PARAM(2, biases); // Biases (bfloat16) + TENSOR_PARAM(1, scales); // Scales (bfloat16, or u8 for microscaled) + OPTIONAL_TENSOR_PARAM(2, biases); // Biases (bfloat16); nil for microscaled modes PARAM(3, int, group_size); PARAM(4, int, bits); - DEVICE_PARAM(5, device); + std::string mode; + if (!nx::nif::get(env, argv[5], mode)) { + return nx::nif::error(env, "Unable to get mode param."); + } + DEVICE_PARAM(6, device); + + std::optional biases_opt = + biases ? std::make_optional(*biases) : std::nullopt; - TENSOR(mlx::core::dequantize(*w, *scales, *biases, group_size, bits, "affine", std::nullopt, std::nullopt, device)); + TENSOR(mlx::core::dequantize(*w, *scales, biases_opt, group_size, bits, mode, std::nullopt, std::nullopt, device)); } // quantize - Quantizes a float tensor to packed format -// Returns tuple of {weights, scales, biases} -// MLX API: quantize(w, group_size, bits, stream) -> tuple +// Returns a 3-tuple {weights, scales, biases}; biases is the atom `nil` for +// microscaled modes ("mxfp4"/"mxfp8"/"nvfp4"), which don't produce a biases +// array (mx::quantize's `result` vector has 2 elements instead of 3 there). +// MLX API: quantize(w, group_size, bits, mode, stream) -> vector NIF(quantize) { TENSOR_PARAM(0, w); // Float weights to quantize PARAM(1, int, group_size); PARAM(2, int, bits); - DEVICE_PARAM(3, device); + std::string mode; + if (!nx::nif::get(env, argv[3], mode)) { + return nx::nif::error(env, "Unable to get mode param."); + } + DEVICE_PARAM(4, device); try { - auto result = mlx::core::quantize(*w, group_size, bits, "affine", std::nullopt, device); + auto result = mlx::core::quantize(*w, group_size, bits, mode, std::nullopt, device); - ERL_NIF_TERM result_tuple[3]; - result_tuple[0] = create_tensor_resource(env, result[0]); - result_tuple[1] = create_tensor_resource(env, result[1]); - result_tuple[2] = create_tensor_resource(env, result[2]); + ERL_NIF_TERM weights_term = create_tensor_resource(env, result[0]); + ERL_NIF_TERM scales_term = create_tensor_resource(env, result[1]); + ERL_NIF_TERM biases_term = result.size() > 2 + ? create_tensor_resource(env, result[2]) + : enif_make_atom(env, "nil"); - return nx::nif::ok(env, enif_make_tuple3(env, result_tuple[0], result_tuple[1], result_tuple[2])); + return nx::nif::ok(env, enif_make_tuple3(env, weights_term, scales_term, biases_term)); } CATCH() } @@ -1682,6 +1674,7 @@ NIF(window_scatter_min) { // limits, evaluated buffer pointers). ASYNC_NIF(eval) +ASYNC_NIF(eval_many) ASYNC_NIF(to_blob) ASYNC_NIF(tensor_to_shm) ASYNC_NIF(item) @@ -1813,8 +1806,19 @@ ASYNC_NIF(tensordot) ASYNC_NIF(window_scatter_max) ASYNC_NIF(window_scatter_min) +// ── Native compiler NIFs (logic lives in emlx_compiler.cpp) ────────────────── +// +// compile_program is defined directly in emlx_compiler.cpp via +// FINE_ASYNC_NIF(compile_program) (see emlx_compiler.hpp for the +// compile_program/compile_program_async declarations); referenced fully +// qualified in nif_funcs[] below, same as resolve_runtime_call. + +NIF(eval_program) { return emlx::native::eval_program(env, argc, argv); } +ASYNC_NIF(eval_program) + static ErlNifFunc nif_funcs[] = { - {"eval", 2, eval_async}, + {"eval", 2, eval_async, ERL_NIF_DIRTY_JOB_IO_BOUND}, + {"eval_many", 2, eval_many_async, ERL_NIF_DIRTY_JOB_IO_BOUND}, {"to_device", 3, to_device_async}, {"to_blob", 2, to_blob_async}, {"to_blob", 3, to_blob_async}, @@ -1944,7 +1948,7 @@ static ErlNifFunc nif_funcs[] = { {"linalg_solve", 4, linalg_solve_async}, {"linalg_solve_triangular", 5, linalg_solve_triangular_async}, {"conv_general", 10, conv_general_async}, - {"einsum", 5, einsum_async}, + {"einsum", 4, einsum_async}, {"tensordot", 6, tensordot_async}, {"window_scatter_max", 9, window_scatter_max_async}, @@ -1969,35 +1973,53 @@ static ErlNifFunc nif_funcs[] = { {"command_queue_new", 1, command_queue_new}, {"command_queue_synchronize", 1, command_queue_synchronize}, // Quantization operations (async — must run on a worker thread) - {"quantized_matmul", 9, quantized_matmul_async}, - {"dequantize", 7, dequantize_async}, - {"quantize", 5, quantize_async}, + {"quantized_matmul", 10, quantized_matmul_async}, + {"dequantize", 8, dequantize_async}, + {"quantize", 6, quantize_async}, // mlx::fast ops (worker arity includes queue ref as argv[0]) {"fast_rms_norm", 5, fast_rms_norm_async}, {"fast_rope", 8, fast_rope_async}, - {"fast_sdpa", 6, fast_sdpa_async}, - {"fast_sdpa_masked", 7, fast_sdpa_masked_async}, + {"fast_sdpa", 7, fast_sdpa_async}, + {"fast_sdpa_masked", 8, fast_sdpa_masked_async}, {"fast_rope_ids", 8, fast_rope_ids_async}, {"fast_rope_with_freqs", 8, fast_rope_with_freqs_async}, {"fast_rope_positions", 8, fast_rope_positions_async}, - {"fast_sdpa_causal_key_masked", 8, fast_sdpa_causal_key_masked_async}, - {"fast_sdpa_causal", 6, fast_sdpa_causal_async}, + {"fast_sdpa_causal_key_masked", 9, fast_sdpa_causal_key_masked_async}, + {"fast_sdpa_causal", 7, fast_sdpa_causal_async}, {"fast_layer_norm", 6, fast_layer_norm_async}, {"fast_layer_norm_no_bias", 5, fast_layer_norm_no_bias_async}, {"fast_swiglu", 4, fast_swiglu_async}, {"kv_cache_attention", 9, kv_cache_attention_async}, {"kv_cache_attention_masked", 10, kv_cache_attention_masked_async}, {"kv_cache_sdpa_update", 9, kv_cache_sdpa_update_async}, + + // ── Native compiler NIFs. + {"compile_program", 2, emlx::native::compile_program_async, + ERL_NIF_DIRTY_JOB_CPU_BOUND}, + {"eval_program", 3, eval_program_async, ERL_NIF_DIRTY_JOB_IO_BOUND}, + // resolve_runtime_call is NOT worker-routed: it only decodes a reply and + // memcpy's it into pre-registered buffers/notifies a condvar — no MLX + // graph work, so it can run directly on the calling BEAM scheduler. + {"resolve_runtime_call", 3, emlx::native::resolve_runtime_call}, + + // ── Qwen3 model accelerators (emlx_fast/qwen3.cpp). {"qwen3_kv_cache_attention", 11, qwen3_kv_cache_attention_async}, {"qwen3_mlp", 8, qwen3_mlp_async}, {"qwen3_layer", 21, qwen3_layer_async}, + {"qwen3_layer_quantized", 21, qwen3_layer_quantized_async}, {"qwen3_forward_greedy_ids", 13, qwen3_forward_greedy_ids_async}, {"qwen3_forward_greedy_ids_chunk", 14, qwen3_forward_greedy_ids_chunk_async}, + {"qwen3_forward_greedy_ids_chunk_quantized", 14, + qwen3_forward_greedy_ids_chunk_quantized_async}, {"qwen3_forward_greedy_ids_token_id", 13, qwen3_forward_greedy_ids_token_id_async}, {"qwen3_forward_greedy_token_id", 13, qwen3_forward_greedy_token_id_async}, {"qwen3_final_greedy", 6, qwen3_final_greedy_async}, {"qwen3_attention_residual", 5, qwen3_attention_residual_async}, - {"qwen3_attention_block", 17, qwen3_attention_block_async}}; + {"qwen3_attention_block", 17, qwen3_attention_block_async}, + // load_plugin `dlopen`s a named, standalone native plugin (see + // emlx_plugin_registry.hpp); not worker-routed since it does no MLX + // graph work. + {"load_plugin", 2, load_plugin}}; ERL_NIF_INIT(Elixir.EMLX.NIF, nif_funcs, load, NULL, upgrade, NULL) diff --git a/emlx/c_src/emlx_nif_shared.hpp b/emlx/c_src/emlx_nif_shared.hpp index b028790..83e67d5 100644 --- a/emlx/c_src/emlx_nif_shared.hpp +++ b/emlx/c_src/emlx_nif_shared.hpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -65,6 +66,28 @@ class TensorP { } } + // Movable (needed so a TensorP can be constructed inside a fine::Decoder + // and returned by value into the async-dispatched NIF's argument list — + // see the TensorArg bridge below), but never copyable: a bitwise copy + // would double-decrement the atomic refcount on destruction. + TensorP(TensorP &&other) noexcept + : ptr(other.ptr), refcount(other.refcount), deleted(other.deleted), + err(other.err) { + other.ptr = nullptr; + } + TensorP &operator=(TensorP &&other) noexcept { + if (this != &other) { + ptr = other.ptr; + refcount = other.refcount; + deleted = other.deleted; + err = other.err; + other.ptr = nullptr; + } + return *this; + } + TensorP(const TensorP &) = delete; + TensorP &operator=(const TensorP &) = delete; + ~TensorP() { if (is_valid()) { // decrease reference count @@ -99,42 +122,268 @@ class TensorP { ERL_NIF_TERM err; }; -#define CATCH() \ - catch (const std::exception &e) { \ - std::ostringstream msg; \ - msg << e.what() << " in NIF." << __func__ << "/" << argc; \ - return nx::nif::error(env, msg.str().c_str()); \ - } \ - catch (...) { \ - return nx::nif::error(env, "Unknown error occurred"); \ +#define CATCH() \ + catch (const std::exception &e) { \ + std::ostringstream msg; \ + msg << e.what() << " in NIF." << __func__ << "/" << argc; \ + return nx::nif::error(env, msg.str().c_str()); \ + } \ + catch (...) { \ + return nx::nif::error(env, "Unknown error occurred"); \ } -#define TENSOR(A) \ - try { \ - return nx::nif::ok(env, create_tensor_resource(env, A)); \ - } \ +#define TENSOR(A) \ + try { \ + return nx::nif::ok(env, create_tensor_resource(env, A)); \ + } \ CATCH() -#define NIF(NAME) \ +#define NIF(NAME) \ ERL_NIF_TERM NAME(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) // One-line async wrapper: declare `NIF(OP) { ... }` then `ASYNC_NIF(OP)`. // Register in `nif_funcs[]` at `original_arity + 1` (command queue is argv[0]). // Example: {"add", 4, add_async} // was {"add", 3, add} -#define ASYNC_NIF(OP) \ - ERL_NIF_TERM OP##_async(ErlNifEnv *env, int argc, \ - const ERL_NIF_TERM argv[]) { \ - return emlx::async_dispatch(env, argc, argv); \ +#define ASYNC_NIF(OP) \ + ERL_NIF_TERM OP##_async(ErlNifEnv *env, int argc, \ + const ERL_NIF_TERM argv[]) { \ + return emlx::async_dispatch(env, argc, argv); \ } -#define TENSOR_PARAM(ARGN, VAR) \ - TensorP VAR##_tp(env, argv[ARGN]); \ - mlx::core::array *VAR; \ - if (!VAR##_tp.is_valid()) { \ - return VAR##_tp.error(); \ - } else { \ - VAR = VAR##_tp.data(); \ +#define TENSOR_PARAM(ARGN, VAR) \ + TensorP VAR##_tp(env, argv[ARGN]); \ + mlx::core::array *VAR; \ + if (!VAR##_tp.is_valid()) { \ + return VAR##_tp.error(); \ + } else { \ + VAR = VAR##_tp.data(); \ } -// Forward declaration — defined in emlx_nif.cpp, used in emlx_fast.cpp. +// Optional tensor argument: the Elixir caller passes the atom `nil` when the +// tensor is absent (e.g. biases for a microscaled quantization mode, or +// sinks for a plain SDPA call); any other term is decoded as a tensor +// resource. VAR is `nullptr` when absent. +#define OPTIONAL_TENSOR_PARAM(ARGN, VAR) \ + std::optional VAR##_tp; \ + mlx::core::array *VAR = nullptr; \ + { \ + std::string VAR##_nil_check; \ + bool VAR##_is_nil = nx::nif::get_atom(env, argv[ARGN], VAR##_nil_check) &&\ + VAR##_nil_check == "nil"; \ + if (!VAR##_is_nil) { \ + VAR##_tp.emplace(env, argv[ARGN]); \ + if (!VAR##_tp->is_valid()) { \ + return VAR##_tp->error(); \ + } \ + VAR = VAR##_tp->data(); \ + } \ + } + +// Forward declaration — defined in emlx_nif.cpp, used in emlx_fast.cpp and +// emlx_compiler.cpp. ERL_NIF_TERM create_tensor_resource(ErlNifEnv *env, mlx::core::array tensor); + +// Dtype name ↔ mlx::core::Dtype mapping — shared across emlx_nif.cpp and +// emlx_compiler.cpp. +inline const std::map &dtype_map() { + static const std::map table = { + {"bool", mlx::core::bool_}, {"uint8", mlx::core::uint8}, + {"uint16", mlx::core::uint16}, {"uint32", mlx::core::uint32}, + {"uint64", mlx::core::uint64}, {"int8", mlx::core::int8}, + {"int16", mlx::core::int16}, {"int32", mlx::core::int32}, + {"int64", mlx::core::int64}, {"float16", mlx::core::float16}, + {"float32", mlx::core::float32}, {"bfloat16", mlx::core::bfloat16}, + {"complex64", mlx::core::complex64}}; + return table; +} + +inline mlx::core::Dtype string2dtype(const std::string &atom) { + const auto &table = dtype_map(); + auto it = table.find(atom); + if (it != table.end()) + return it->second; + throw std::runtime_error("Unknown dtype: " + atom); +} + +inline const std::string *dtype2string(const mlx::core::Dtype dtype) { + for (const auto &pair : dtype_map()) { + if (pair.second == dtype) + return &pair.first; + } + return nullptr; +} + +// ─── `fine` bridging ───────────────────────────────────────────────────── +// +// `fine::nif()`/`FINE_NIF` (the library's own typed-dispatch entry points) +// translate a thrown C++ exception into a *raised* Elixir exception via +// `enif_raise_exception`. That only makes sense as the return value of a +// live NIF call serviced directly by the BEAM scheduler. EMLX's ASYNC_NIF / +// `emlx::async_dispatch` convention instead runs the sync NIF body on a +// worker thread and ships its `{:ok, _}` / `{:error, _}` tagged-tuple +// result back to the caller as an ordinary term via `enif_send` (see +// emlx_async.hpp) — `enif_raise_exception`'s special return sentinel is +// meaningless in that context (there is no live NIF call for the BEAM to +// attach the exception to), so `fine::nif()` cannot be used as-is here. +// +// Bridge: reuse `fine`'s `Decoder...`/`Encoder` machinery +// (typed arg decode, exception-safe dispatch) via our own dispatcher that +// funnels failures through the *existing* `nx::nif::error(env, msg)` +// tuple convention instead of `enif_raise_exception`. This keeps every +// existing async/registration/error-surfacing mechanism (and the +// `EMLX.NIFError` Elixir-side contract) unchanged; `fine` supplies argument +// marshalling only. +// +// Tensor resources are bridged the same way: fine's `Decoder`/`Encoder` +// are specialized below for `TensorArg`/`mlx::core::array` against the +// *existing* `TensorP` class and `create_tensor_resource` — not +// `fine::ResourcePtr`. `fine::ResourcePtr` only wraps ERTS's own +// `enif_keep_resource`/`enif_release_resource` refcounting; `TensorP` adds +// a second, independent atomic refcount + "deleted" flag on top, used by +// the explicit `deallocate` NIF (emlx_nif.cpp) to free GPU memory ahead of +// BEAM GC. That's a real semantic layer `fine::ResourcePtr` does not +// provide — `TensorP` cannot be a mechanical `fine::ResourcePtr` swap and +// must stay as a bridged custom type. + +#include + +// Decoded tensor argument: owns a TensorP (RAII refcount bump/decrement for +// the duration of the call, mirroring TENSOR_PARAM's local today) plus the +// raw `array*` for ergonomic `*x` / `x->...` use in NIF bodies. +struct TensorArg { + TensorP tp; + mlx::core::array *ptr; + + mlx::core::array &operator*() const { return *ptr; } + mlx::core::array *operator->() const { return ptr; } +}; + +namespace fine { + +template <> struct Decoder { + static TensorArg decode(ErlNifEnv *env, const ERL_NIF_TERM &term) { + TensorP tp(env, term); + if (!tp.is_valid()) { + throw std::invalid_argument("Unable to get tensor param in NIF"); + } + auto *ptr = tp.data(); + return TensorArg{std::move(tp), ptr}; + } +}; + +// `std::optional` mirrors OPTIONAL_TENSOR_PARAM: the Elixir +// caller passes the atom `nil` when the tensor is absent. +template <> struct Decoder> { + static std::optional decode(ErlNifEnv *env, + const ERL_NIF_TERM &term) { + std::string atom_val; + if (nx::nif::get_atom(env, term, atom_val) && atom_val == "nil") { + return std::nullopt; + } + return Decoder::decode(env, term); + } +}; + +// Device atom (`:cpu` / `:gpu`) -> mlx::core::Device, matching DEVICE_PARAM. +template <> struct Decoder { + static mlx::core::Device decode(ErlNifEnv *env, const ERL_NIF_TERM &term) { + std::string atom_val; + if (!nx::nif::get_atom(env, term, atom_val)) { + throw std::invalid_argument("Unable to get device atom in NIF"); + } + return string2device(atom_val); + } +}; + +// Dtype atom (`:float32`, `:bool`, …) -> mlx::core::Dtype. Lets NIF argument +// types (e.g. emlx_compiler.hpp's `Program::constants`) decode dtypes +// directly instead of going through an int<->dtype lookup table shared with +// Elixir (see EMLX.Native.to_mlx_type/1 on the Elixir side). +template <> struct Decoder { + static mlx::core::Dtype decode(ErlNifEnv *env, const ERL_NIF_TERM &term) { + std::string atom_val; + if (!nx::nif::get_atom(env, term, atom_val)) { + throw std::invalid_argument("Unable to get dtype atom in NIF"); + } + return string2dtype(atom_val); + } +}; + +// Plain mlx::core::array argument (as opposed to TensorArg, which also keeps +// the TensorP RAII wrapper alive for the duration of the call) — a straight +// copy of the resource's array value, matching the pre-`fine` LIST_PARAM(..., +// std::vector, ...) semantics used e.g. for compile_program's +// captures. +template <> struct Decoder { + static mlx::core::array decode(ErlNifEnv *env, const ERL_NIF_TERM &term) { + return *Decoder::decode(env, term); + } +}; + +// `int` — fine only ships int64_t/uint64_t decoders; match the existing +// nx::nif::get(..., int*)/enif_get_int semantics exactly (used for e.g. +// RoPE `dims`/`offset` params). +template <> struct Decoder { + static int decode(ErlNifEnv *env, const ERL_NIF_TERM &term) { + int value; + if (!enif_get_int(env, term, &value)) { + throw std::invalid_argument("Unable to get int param in NIF"); + } + return value; + } +}; + +// Tensor return value -> the existing create_tensor_resource (unchanged +// resource type / refcount scheme, not fine::ResourcePtr). +template <> struct Encoder { + static ERL_NIF_TERM encode(ErlNifEnv *env, const mlx::core::array &value) { + return create_tensor_resource(env, value); + } +}; + +} // namespace fine + +namespace emlx_fine { + +// Mirrors fine::__private__::nif_impl, but on failure returns EMLX's own +// `{:error, msg}` tuple (nx::nif::error) instead of raising — see the +// rationale comment above. +template +ERL_NIF_TERM nif_impl(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[], + Return (*fun)(ErlNifEnv *, Args...), const char *name, + std::index_sequence) { + try { + auto result = fun(env, fine::decode(env, argv[Indices])...); + return nx::nif::ok(env, fine::encode(env, result)); + } catch (const std::exception &e) { + std::ostringstream msg; + msg << e.what() << " in NIF." << name << "/" << argc; + return nx::nif::error(env, msg.str().c_str()); + } catch (...) { + return nx::nif::error(env, "Unknown error occurred"); + } +} + +template +ERL_NIF_TERM nif(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[], + Return (*fun)(ErlNifEnv *, Args...), const char *name) { + if (static_cast(sizeof...(Args)) != argc) { + return nx::nif::error(env, "wrong number of arguments"); + } + return nif_impl(env, argc, argv, fun, name, + std::make_index_sequence()); +} + +} // namespace emlx_fine + +// Declares `NAME##_impl(ErlNifEnv*, ...)` (typed body, written +// by the caller right above this macro) as a `fine`-dispatched sync NIF +// named `NAME`, then wires it into the existing ASYNC_NIF/nif_funcs[] +// machinery exactly as `NIF(NAME) { ... } ASYNC_NIF(NAME)` did — the +// generated `NAME`/`NAME_async` symbols and registered arity are +// unchanged, so no emlx_nif.cpp registration-table edits are needed. +#define FINE_ASYNC_NIF(NAME) \ + ERL_NIF_TERM NAME(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { \ + return emlx_fine::nif(env, argc, argv, NAME##_impl, #NAME); \ + } \ + ASYNC_NIF(NAME) diff --git a/emlx/c_src/emlx_plugin_registry.cpp b/emlx/c_src/emlx_plugin_registry.cpp new file mode 100644 index 0000000..90d4298 --- /dev/null +++ b/emlx/c_src/emlx_plugin_registry.cpp @@ -0,0 +1,60 @@ +#include "emlx_plugin_registry.hpp" +#include "nx_nif_utils.hpp" + +#include +#include +#include + +namespace { + +std::unordered_map g_plugin_handles; +std::unordered_map g_plugin_vtables; + +} // namespace + +const void *emlx_get_plugin(const std::string &name) { + auto it = g_plugin_vtables.find(name); + return it == g_plugin_vtables.end() ? nullptr : it->second; +} + +ERL_NIF_TERM load_plugin(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { + std::string name; + std::string path; + if (!nx::nif::get(env, argv[0], name)) { + return nx::nif::error(env, "load_plugin expects a name (as the 1st argument)"); + } + if (!nx::nif::get(env, argv[1], path)) { + return nx::nif::error(env, "load_plugin expects a path string (as the 2nd argument)"); + } + + auto existing_handle = g_plugin_handles.find(name); + if (existing_handle != g_plugin_handles.end()) { + dlclose(existing_handle->second); + g_plugin_handles.erase(existing_handle); + g_plugin_vtables.erase(name); + } + + void *handle = dlopen(path.c_str(), RTLD_NOW | RTLD_LOCAL); + if (handle == nullptr) { + std::ostringstream msg; + msg << "Failed to load plugin \"" << name << "\" at " << path << ": " << dlerror(); + return nx::nif::error(env, msg.str().c_str()); + } + + using VTableFn = void *(*)(); + auto get_vtable = reinterpret_cast(dlsym(handle, "emlx_plugin_vtable")); + if (get_vtable == nullptr) { + dlclose(handle); + return nx::nif::error(env, "plugin is missing the emlx_plugin_vtable symbol"); + } + + void *vtable = get_vtable(); + if (vtable == nullptr) { + dlclose(handle); + return nx::nif::error(env, "plugin returned a null vtable"); + } + + g_plugin_handles[name] = handle; + g_plugin_vtables[name] = vtable; + return nx::nif::ok(env); +} diff --git a/emlx/c_src/emlx_plugin_registry.hpp b/emlx/c_src/emlx_plugin_registry.hpp new file mode 100644 index 0000000..c17a449 --- /dev/null +++ b/emlx/c_src/emlx_plugin_registry.hpp @@ -0,0 +1,26 @@ +#pragma once + +// Generic, name-keyed native plugin loader — `dlopen`s a standalone shared +// library (no erl_nif dependency) and caches the `void*` it returns from a +// well-known exported symbol, keyed by an arbitrary caller-chosen name. +// +// This has no knowledge of any specific plugin's ABI: callers (e.g. +// `emlx_fast/qwen3.cpp`) look their plugin up by name via `emlx_get_plugin` +// and `reinterpret_cast` the result to whatever vtable struct they expect — +// that contract lives entirely in the caller (see +// `emlx_fast/qwen3_plugin_abi.hpp`), not here. + +#include "erl_nif.h" + +#include + +// load_plugin(name, path) — `dlopen`s `path` and `dlsym`s the fixed +// `emlx_plugin_vtable` entry point, storing the resulting pointer under +// `name` for later retrieval via `emlx_get_plugin`. Reloading the same +// `name` replaces the previous entry (and `dlclose`s its handle). Not +// worker-routed: `dlopen`/`dlsym` do no MLX graph work. +ERL_NIF_TERM load_plugin(ErlNifEnv *, int, const ERL_NIF_TERM[]); + +// Returns the cached vtable pointer for `name`, or `nullptr` if no plugin +// has been successfully loaded under that name yet. +const void *emlx_get_plugin(const std::string &name); diff --git a/emlx/c_src/emlx_runtime_call_bridge.hpp b/emlx/c_src/emlx_runtime_call_bridge.hpp new file mode 100644 index 0000000..6826320 --- /dev/null +++ b/emlx/c_src/emlx_runtime_call_bridge.hpp @@ -0,0 +1,262 @@ +#pragma once + +// emlx_runtime_call_bridge.hpp — blocking bridge between the worker OS +// thread executing a compiled program's `EMLXRuntimeCall` primitive +// (emlx_compiler.cpp) and the BEAM process that actually runs the real +// Elixir callback for `Nx.runtime_call/4` (see EMLX.await_worker/2's +// `:emlx_runtime_call` receive clause). +// +// Architecture mirrors EXLA's `nx/exla/c_src/exla/custom_calls/ +// runtime_callback_bridge.{h,cc}`: a `PendingRuntimeCall` resource carries a +// mutex/condvar pair plus pre-allocated output destinations. After +// `enif_send`ing a request, the worker thread does NOT simply park on the +// condvar — it calls `Worker::pump_until` (emlx_worker.hpp) to keep +// draining and running its own job queue while waiting. This matters +// because the real Elixir callback that will eventually call +// `resolve_runtime_call/3` may itself need to run further work on this +// exact worker (EMLX.Quantization's callbacks, for one, reenter EMLX on +// the same device/worker they were called from) — parking this, the +// worker's one dedicated OS thread, on a bare wait would deadlock against +// that reentrant job sitting behind it in the very queue it can no longer +// service. See emlx_worker.hpp's `pump_until` doc for the full argument. +// +// `g_current_caller_pid` is set by `emlx::async_dispatch` +// (emlx_async.hpp) for the duration of one worker-thread job — including +// any `mlx::core::eval(...)` call the job body makes — so it is in scope +// whenever `EMLXRuntimeCall::eval_cpu`/`eval_gpu` actually fires. One +// `EMLX.CommandQueue` worker owns exactly one dedicated OS thread (see +// emlx_async.hpp's header comment), so a plain `thread_local` is enough to +// disambiguate multiple workers evaluating concurrently. + +#include "emlx_worker.hpp" +#include "erl_nif.h" +#include "mlx/allocator.h" +#include "mlx/array.h" +#include "nx_nif_utils.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace emlx { + +// Set by `async_dispatch` around the call to `SyncOp` (e.g. +// `eval_program`) — see emlx_async.hpp. +inline thread_local ErlNifPid *g_current_caller_pid = nullptr; + +namespace native { + +// Per-call state shared between the worker thread blocked inside +// `EMLXRuntimeCall::eval()` and the BEAM process that delivers the reply via +// `resolve_runtime_call/3`. Exposed as an ERTS resource so it can ride as an +// opaque handle inside the `{:emlx_runtime_call, pending, callback_index, +// args}` message instead of an integer id. +struct PendingRuntimeCall { + std::mutex mu; + std::condition_variable cv; + bool done = false; + bool ok = false; + std::string error; + // (destination pointer, byte size) per output, in order — already + // allocated by `invoke_runtime_call` before the message is sent, per + // `mlx::core::Primitive`'s "the evaluation function is responsible for + // allocating space for the array" contract. + std::vector> outputs; +}; + +// Blocks the calling (worker) thread until the Elixir side invokes +// `EMLX.NIF.resolve_runtime_call/3` for this call. Called from +// `EMLXRuntimeCall::eval()` (emlx_compiler.cpp). Throws `std::runtime_error` +// on any failure (no caller pid in scope, send failure, or a callback that +// errored) — the caller (`eval_cpu`/`eval_gpu`, reached through +// `eval_program`'s `CATCH()` macro) turns that into `{:error, _}`. +inline void invoke_runtime_call(int64_t callback_index, + const std::vector &inputs, + std::vector &outputs) { + if (g_current_caller_pid == nullptr) { + throw std::runtime_error( + "emlx::native: runtime_call primitive fired with no caller pid in " + "scope (a compiled program containing a runtime_call must be " + "force-evaluated inside eval_program's async job body)"); + } + ErlNifPid caller_pid = *g_current_caller_pid; + + auto *pending = static_cast(enif_alloc_resource( + resource_object::type, sizeof(PendingRuntimeCall))); + if (pending == nullptr) { + throw std::runtime_error( + "emlx::native: failed to allocate runtime_call pending resource"); + } + new (pending) PendingRuntimeCall(); + + // Pre-allocate every output buffer now so the destination pointers are + // stable before we hand the request to the Elixir side. + pending->outputs.reserve(outputs.size()); + for (auto &out : outputs) { + out.set_data(mlx::core::allocator::malloc(out.nbytes())); + pending->outputs.emplace_back(out.data(), out.nbytes()); + } + + ErlNifEnv *msg_env = enif_alloc_env(); + if (msg_env == nullptr) { + enif_release_resource(pending); + throw std::runtime_error("emlx::native: failed to allocate msg env for runtime_call"); + } + + std::vector arg_terms; + arg_terms.reserve(inputs.size()); + for (const auto &in : inputs) { + size_t nbytes = in.nbytes(); + ErlNifBinary bin; + if (!enif_alloc_binary(nbytes, &bin)) { + enif_release_resource(pending); + enif_free_env(msg_env); + throw std::runtime_error( + "emlx::native: failed to allocate binary for runtime_call arg"); + } + if (nbytes > 0) { + std::memcpy(bin.data, in.data(), nbytes); + } + arg_terms.push_back(enif_make_binary(msg_env, &bin)); + } + + ERL_NIF_TERM pending_term = enif_make_resource(msg_env, pending); + // enif_make_resource/enif_make_copy below keep their own references; drop + // ours now — the resource stays alive via the message term (and later, + // whatever term resolve_runtime_call's caller decodes it into) until its + // own destructor runs. + enif_release_resource(pending); + + ERL_NIF_TERM args_list = enif_make_list_from_array( + msg_env, arg_terms.data(), static_cast(arg_terms.size())); + + ERL_NIF_TERM msg = + enif_make_tuple4(msg_env, enif_make_atom(msg_env, "emlx_runtime_call"), + pending_term, enif_make_int64(msg_env, callback_index), + args_list); + + ErlNifPid send_pid = caller_pid; + int sent = enif_send(NULL, &send_pid, msg_env, msg); + enif_free_env(msg_env); + + if (!sent) { + throw std::runtime_error( + "emlx::native: failed to send runtime_call request (target process not alive?)"); + } + + auto is_done = [pending] { + std::lock_guard lock(pending->mu); + return pending->done; + }; + + // g_current_worker is set for the whole lifetime of this OS thread (see + // emlx_worker.hpp's thread_main) — always non-null here, since this + // function only ever runs from inside a job posted to some worker (via + // EMLXRuntimeCall::eval_cpu/eval_gpu, reached through eval_program). + // The plain-condvar fallback exists only in case that invariant is ever + // violated (e.g. a future caller that isn't worker-dispatched) rather + // than silently deadlocking with no queue to pump. + if (Worker *worker = g_current_worker) { + worker->pump_until(is_done); + } else { + std::unique_lock lock(pending->mu); + pending->cv.wait(lock, [pending] { return pending->done; }); + } + + std::lock_guard lock(pending->mu); + if (!pending->ok) { + throw std::runtime_error("emlx::native: runtime_call callback failed: " + pending->error); + } +} + +// resolve_runtime_call/3 NIF — called from the BEAM process that ran the +// real Elixir callback (EMLX.await_worker/2's `:emlx_runtime_call` receive +// clause) to deliver the reply and wake the blocked worker thread. +// +// argv[0] : pending resource ref +// argv[1] : status atom (:ok | :error) +// argv[2] : on :ok, a list of binaries (one per output, in declared order, +// each exactly `out.nbytes()` long); on :error, a binary error +// message. +inline ERL_NIF_TERM resolve_runtime_call(ErlNifEnv *env, int argc, + const ERL_NIF_TERM argv[]) { + (void)argc; + PendingRuntimeCall *pending; + if (!enif_get_resource(env, argv[0], resource_object::type, + reinterpret_cast(&pending))) { + return nx::nif::error(env, "Invalid runtime_call pending resource"); + } + + char status_buf[8]; + if (enif_get_atom(env, argv[1], status_buf, sizeof(status_buf), ERL_NIF_LATIN1) <= 0) { + return nx::nif::error(env, "Invalid runtime_call status atom"); + } + bool is_ok = std::string(status_buf) == "ok"; + + bool decode_ok = true; + std::string message; + + { + std::lock_guard lock(pending->mu); + + if (is_ok) { + unsigned n; + if (!enif_get_list_length(env, argv[2], &n)) { + decode_ok = false; + message = "runtime_call reply must be a list of binaries"; + } else if (static_cast(n) != pending->outputs.size()) { + decode_ok = false; + message = "runtime_call reply has " + std::to_string(n) + + " outputs, expected " + std::to_string(pending->outputs.size()); + } else { + ERL_NIF_TERM head, tail, list = argv[2]; + size_t i = 0; + while (decode_ok && enif_get_list_cell(env, list, &head, &tail)) { + ErlNifBinary bin; + if (!enif_inspect_binary(env, head, &bin)) { + decode_ok = false; + message = "runtime_call reply element " + std::to_string(i) + " is not a binary"; + break; + } + auto &[dst, size] = pending->outputs[i]; + if (bin.size != size) { + decode_ok = false; + message = "runtime_call reply output " + std::to_string(i) + + " has unexpected byte size (" + std::to_string(bin.size) + + ", expected " + std::to_string(size) + ")"; + break; + } + if (size > 0) { + std::memcpy(dst, bin.data, size); + } + list = tail; + i++; + } + } + } else { + ErlNifBinary bin; + if (!enif_inspect_binary(env, argv[2], &bin)) { + decode_ok = false; + message = "runtime_call error reply is not a binary"; + } else { + message = std::string(reinterpret_cast(bin.data), bin.size); + } + } + + pending->ok = is_ok && decode_ok; + pending->error = message; + pending->done = true; + } + + pending->cv.notify_one(); + return nx::nif::ok(env); +} + +} // namespace native +} // namespace emlx diff --git a/emlx/c_src/emlx_worker.hpp b/emlx/c_src/emlx_worker.hpp index fabbd25..5b4b714 100644 --- a/emlx/c_src/emlx_worker.hpp +++ b/emlx/c_src/emlx_worker.hpp @@ -22,6 +22,7 @@ #include "mlx/stream.h" #include +#include #include #include #include @@ -32,6 +33,16 @@ namespace emlx { +class Worker; + +// Points at the `Worker` whose dedicated OS thread is currently running, +// for the entire lifetime of that thread (set once in `thread_main`, not +// per-job — unlike `emlx::g_current_caller_pid`, which changes per job). +// Lets code running deep inside a job (e.g. `invoke_runtime_call` in +// emlx_runtime_call_bridge.hpp) find its own worker to pump its queue +// while blocked, without threading a `Worker*` through every call site. +inline thread_local Worker *g_current_worker = nullptr; + class Worker { public: using Job = std::function; @@ -105,6 +116,53 @@ class Worker { mlx::core::Device device() const { return device_; } mlx::core::Stream stream() const { return stream_; } + // Called *from inside* a currently-executing job, on this worker's own + // OS thread — e.g. `emlx::native::invoke_runtime_call` + // (emlx_runtime_call_bridge.hpp), blocked inside `EMLXRuntimeCall:: + // eval_cpu`/`eval_gpu` waiting for the Elixir side's reply. Rather than + // parking this thread on a bare wait (which would starve the rest of + // this worker's queue — including any job the pending reply's own + // real callback needs, e.g. `EMLX.Quantization`'s callbacks reentering + // EMLX on this very worker/device to run the actual op), keep draining + // and running this worker's own queue until `done()` returns true. + // Safe to do because we're still on the one OS thread that owns this + // worker's MLX stream thread-local state (see the class comment above + // and emlx_async.hpp's header comment) — no cross-thread MLX dispatch + // occurs, we're just running the *next* queued job a little earlier + // than `thread_main`'s own loop would have. Jobs run this way nest: + // if one of them itself blocks on another runtime_call, it recurses + // into another `pump_until` a level deeper, which is why + // `CallerPidGuard` (emlx_async.hpp) restores the *previous* caller pid + // on exit rather than unconditionally clearing it. + // + // `done` may be polled many times and must be cheap and never block. + void pump_until(const std::function &done) { + while (!done()) { + Job job; + { + std::unique_lock lock(mutex_); + // Bounded wait: re-checks `done()` even if no job ever arrives + // (e.g. the reply is delivered without going through this + // worker's queue at all). + cv_.wait_for(lock, std::chrono::milliseconds(5), [this] { + return !queue_.empty() || stop_.load(std::memory_order_acquire); + }); + if (queue_.empty()) { + continue; + } + job = std::move(queue_.front()); + queue_.pop_front(); + } + + // Same contract as thread_main's loop below: a job must never let + // an exception escape. + try { + job(); + } catch (...) { + } + } + } + private: void thread_main(std::promise stream_promise) { // Allocate the stream on THIS thread (MLX 0.31.2 thread-locality @@ -123,6 +181,8 @@ class Worker { return; } + g_current_worker = this; + while (true) { Job job; { @@ -150,6 +210,8 @@ class Worker { } } + g_current_worker = nullptr; + // Release any per-thread MLX resources we hold (the StreamThread // for stream_ in the global scheduler will be torn down when the // last reference to its index drops). diff --git a/emlx/c_src/nx_nif_utils.hpp b/emlx/c_src/nx_nif_utils.hpp index 76f6f85..a2ca647 100644 --- a/emlx/c_src/nx_nif_utils.hpp +++ b/emlx/c_src/nx_nif_utils.hpp @@ -4,7 +4,6 @@ #include "mlx/mlx.h" inline ErlNifResourceType *TENSOR_TYPE; -inline ErlNifResourceType *FUNCTION_TYPE; namespace emlx { typedef std::function( @@ -113,7 +112,6 @@ int open_resource(ErlNifEnv *env, const char *mod, const char *name, } ERL_NIF_TERM create_tensor_resource(ErlNifEnv *env, mlx::core::array tensor); -ERL_NIF_TERM create_function_resource(ErlNifEnv *env, emlx::function function); namespace nx { namespace nif { @@ -284,13 +282,6 @@ inline int get(ErlNifEnv *env, ERL_NIF_TERM term, bool *var) { return 1; } -// function - -inline int get(ErlNifEnv *env, ERL_NIF_TERM term, emlx::function *&var) { - return enif_get_resource(env, term, resource_object::type, - reinterpret_cast(&var)); -} - // Containers template diff --git a/emlx/config/config.exs b/emlx/config/config.exs index 03e55cb..73afa1e 100644 --- a/emlx/config/config.exs +++ b/emlx/config/config.exs @@ -2,4 +2,12 @@ import Config if config_env() == :test do config :emlx, :add_backend_on_inspect, false + + # Opt-in: recompile with both debug-assertion flags on so + # debug_flags_functional_test.exs can exercise their actual raise behavior. + # compile_env is baked in at compile time, so this can't be toggled at + # runtime — see that file's moduledoc for the invocation. + if System.get_env("EMLX_DEBUG_FLAGS") == "1" do + config :emlx, detect_non_finites: true, enable_bounds_check: true, compiler_debug: true + end end diff --git a/emlx/config/dev.exs b/emlx/config/dev.exs index c8e23f7..bfa75dd 100644 --- a/emlx/config/dev.exs +++ b/emlx/config/dev.exs @@ -14,7 +14,15 @@ import Config # indexed_add, and indexed_put before the NIF call. # config :emlx, enable_bounds_check: true -# Raises ArgumentError if dot (matmul / einsum) produces NaN or Inf. -# When EMLX.Fast is implemented (task 05), rms_norm, layer_norm, and -# scaled_dot_product_attention will also be checked. +# Raises ArgumentError if dot (matmul / einsum), conv, or EMLX.Fast's +# rms_norm/layer_norm/scaled_dot_product_attention kernels produce NaN or Inf. # config :emlx, detect_non_finites: true + +# Raises ArgumentError on internal lowering/to_native invariant violations in +# EMLX.Native.Expr (ref id collisions, forward/self-referencing instructions, +# double-bound result refs) that would otherwise silently miscompile instead +# of failing loudly. Unlike the two flags above, this one is cheap (no extra +# eval syncs) — it's off by default only because these bugs should never +# happen in a working compiler and the checks add wasted work on the hot +# compile-cache-miss path. +# config :emlx, compiler_debug: true diff --git a/emlx/lib/emlx.ex b/emlx/lib/emlx.ex index 8448e1f..fdb5e07 100644 --- a/emlx/lib/emlx.ex +++ b/emlx/lib/emlx.ex @@ -134,6 +134,103 @@ defmodule EMLX.Macro do end defmodule EMLX do + @moduledoc """ + Low-level MLX NIF wrappers and the native `Nx.Defn.Compiler` for + `EMLX.Backend` tensors. + + Most users don't call this module directly — set `EMLX.Backend` as your + `Nx` backend (see the [README](readme.html)) and, optionally, `EMLX` as + your `Nx.Defn` compiler for the JIT-compiled native path: + + Nx.Defn.default_options(compiler: EMLX) + + ## Compile once, replay many times + + `Nx.Defn.jit/2` and `Nx.Defn.jit_apply/3` **retrace** the given function + from scratch on every call. For most ops that's cheap relative to the + actual computation, but for a hot loop (e.g. a decode step, or any call + site invoked with the same input shapes/types repeatedly) it means paying + full retrace + dispatch-key computation cost on every single call — for + some ops (e.g. `Nx.LinAlg.svd/1`, which traces a large, normally-unused + fallback algorithm as part of every trace) this cost can dominate over the + actual native computation. + + If you know you'll call the same function repeatedly with the same input + shapes/types, prefer `Nx.Defn.compile/3`, which traces and lowers **once** + and returns a plain closure that replays the already-compiled program on + every subsequent call — no retrace, no re-lowering, no dispatch-key + recomputation: + + svd = Nx.Defn.compile(&Nx.LinAlg.svd/1, [Nx.template({128, 128}, {:f, 32})], compiler: EMLX) + # Hold onto `svd` (e.g. in a GenServer, or a module attribute at startup) + # and call it repeatedly — each call below is pure replay: + {u, s, vt} = svd.(a) + + This is exactly the strategy `Nx.Serving` and Bumblebee already rely on for + other compilers, and it works identically here — no `EMLX`-specific code + needed. EMLX additionally keeps a persistent, structural (shape/op-based, + not object-identity) dispatch-key cache across calls (see `dispatch_key/3` + in the source), which is what makes `Nx.Defn.Graph.run/3`'s per-call + re-tracing and structurally-identical-but-distinct call sites (e.g. many + copies of the same layer in a model) cheap too — but a caller-held + `Nx.Defn.compile/3` closure is always cheaper still, since it skips + retracing entirely. + + ## Compile-time debug flags + + Several development-only checks are gated by `Application.compile_env/3` + at compile time — when a flag is `false`, the check is erased entirely + (no BEAM opcodes, zero runtime cost). After changing a flag, run + `mix compile --force`; values are baked in at compile time, not read at + runtime. See `config/dev.exs` for commented examples. + + config :emlx, enable_bounds_check: true + config :emlx, detect_non_finites: true + config :emlx, compiler_debug: true + + * `:enable_bounds_check` — raises on out-of-bounds indices in gather, + take, take_along_axis, indexed_add, and indexed_put. + * `:detect_non_finites` — raises when dot, conv, or `EMLX.Fast` fused + kernels produce NaN or Inf. Forces extra `EMLX.eval` syncs and breaks + MLX lazy-graph fusion; never enable in production. + * `:compiler_debug` — raises on internal `EMLX.Native.Expr` lowering / + `to_native` invariant violations that would otherwise silently miscompile. + Cheap (no extra eval syncs); off by default. + + WARNING: `:enable_bounds_check` and `:detect_non_finites` break MLX + lazy-graph fusion. On non-unified-memory targets (Linux GPU), + `:enable_bounds_check` also incurs an extra GPU→CPU copy per indexed op. + + ## CPU JIT compilation and SIGCHLD + + On the CPU backend, MLX JIT-compiles fused kernels the first time it sees + a new graph shape by shelling out to `popen("g++ ...")` and reading the + result back via `pclose()`. The BEAM sets `SIGCHLD` to `SIG_IGN` by + default (so it can run as PID 1 in a container without leaking zombies); + under Linux/POSIX semantics that makes the kernel auto-reap any child the + instant it exits, so by the time MLX's `pclose()` calls `waitpid()` there + is nothing left to collect — it fails with `ECHILD` + (`** (EMLX.NIFError) ... pclose() failed.`), independent of whether the + compile itself would have succeeded. + + EMLX does not change this VM-wide setting itself — doing so from a + dependency would silently change zombie-reaping behavior for the whole + host application. If you hit this error (typically on Linux CPU backend, + the first time a given fused-op shape runs), restore the default + disposition yourself, as early as possible in your own application + (e.g. in your own `Application.start/2`, before `:emlx` or `:nx` start + doing real work): + + :os.set_signal(:sigchld, :default) + + This is the same fix TensorFlow's Erlang bindings needed for the + identical reason — see + https://erlang.org/pipermail/erlang-questions/2020-November/100109.html. + Only skip this if your VM runs as PID 1 in a container and can't + tolerate the (small, `g++`-subprocess-shaped) risk of zombie processes. + See `EMLX.Application` for more detail. + """ + use EMLX.Macro @profile_eval Application.compile_env(:emlx, :profile_eval, false) @@ -221,7 +318,7 @@ defmodule EMLX do deftensor isclose(tensorA, tensorB, rtol, atol, equal_nan) deftensor tensordot(tensorA, tensorB, axesA, axesB) - deftensor einsum(tensorA, tensorB, spec_string) + deftensor einsum(tensors, spec_string) deftensor transpose(tensor, axes) deftensor pad(tensor, axes, low_pad_size, high_pad_size, tensor_pad_value) deftensor sort(tensor, axis) @@ -342,31 +439,33 @@ defmodule EMLX do Performs quantized matrix multiplication. This is the key operation for efficient 4-bit inference. It multiplies `x` with - quantized weights `w` (packed as uint32), using scales and biases for - dequantization during the computation. + quantized weights `w` (packed as uint32), using scales and (for `"affine"`) + biases for dequantization during the computation. ## Parameters - `x` - Input tensor (e.g., {batch, seq, hidden}) - `w` - Quantized weights as uint32 (8 int4 values packed per uint32) - - `scales` - Per-group scale factors (bfloat16) - - `biases` - Per-group zero points (bfloat16) + - `scales` - Per-group scale factors (bfloat16, or u8 for microscaled modes) + - `biases` - Per-group zero points (bfloat16), or `nil` for microscaled modes - `transpose` - Whether to transpose weights (default: true) - `group_size` - Number of weights per scale/bias group (default: 64) - `bits` - Quantization bits (default: 4) + - `mode` - `"affine"` (default), `"mxfp4"`, `"mxfp8"`, or `"nvfp4"` """ - @mlx_function {:quantized_matmul, 9} + @mlx_function {:quantized_matmul, 10} def quantized_matmul( {dev_x, ref_x} = _tensor_x, {dev_w, ref_w} = _tensor_w, {dev_s, ref_s} = _tensor_scales, - {dev_b, ref_b} = _tensor_biases, + biases, transpose \\ true, group_size \\ 64, - bits \\ 4 + bits \\ 4, + mode \\ "affine" ) - when is_tensor(dev_x, ref_x) and is_tensor(dev_w, ref_w) and - is_tensor(dev_s, ref_s) and is_tensor(dev_b, ref_b) do - device = merge_device(merge_device(dev_x, dev_w), merge_device(dev_s, dev_b)) + when is_tensor(dev_x, ref_x) and is_tensor(dev_w, ref_w) and is_tensor(dev_s, ref_s) do + {ref_b, biases_device} = unwrap_optional_tensor(biases) + device = merge_device(merge_device(dev_x, dev_w), merge_device(dev_s, biases_device)) {worker, effective_device} = resolve_worker(device) job_ref = @@ -379,6 +478,7 @@ defmodule EMLX do transpose, group_size, bits, + mode, effective_device ) |> unwrap!() @@ -394,25 +494,28 @@ defmodule EMLX do ## Parameters - `w` - Quantized weights as uint32 (packed int4 values) - - `scales` - Per-group scale factors - - `biases` - Per-group zero points + - `scales` - Per-group scale factors (or u8 for microscaled modes) + - `biases` - Per-group zero points, or `nil` for microscaled modes - `group_size` - Number of weights per group (default: 64) - `bits` - Quantization bits (default: 4) + - `mode` - `"affine"` (default), `"mxfp4"`, `"mxfp8"`, or `"nvfp4"` """ - @mlx_function {:dequantize, 7} + @mlx_function {:dequantize, 8} def dequantize( {dev_w, ref_w} = _tensor_w, {dev_s, ref_s} = _tensor_scales, - {dev_b, ref_b} = _tensor_biases, + biases, group_size, - bits + bits, + mode \\ "affine" ) - when is_tensor(dev_w, ref_w) and is_tensor(dev_s, ref_s) and is_tensor(dev_b, ref_b) do - device = merge_device(dev_w, merge_device(dev_s, dev_b)) + when is_tensor(dev_w, ref_w) and is_tensor(dev_s, ref_s) do + {ref_b, biases_device} = unwrap_optional_tensor(biases) + device = merge_device(dev_w, merge_device(dev_s, biases_device)) {worker, effective_device} = resolve_worker(device) job_ref = - EMLX.NIF.dequantize(worker, ref_w, ref_s, ref_b, group_size, bits, effective_device) + EMLX.NIF.dequantize(worker, ref_w, ref_s, ref_b, group_size, bits, mode, effective_device) |> unwrap!() await_worker(job_ref) |> wrap_tensor(effective_device) @@ -423,29 +526,39 @@ defmodule EMLX do Returns a tuple of `{quantized_weights, scales, biases}` where: - `quantized_weights` - Packed uint32 tensor (8 int4 values per uint32) - - `scales` - Per-group scale factors - - `biases` - Per-group zero points + - `scales` - Per-group scale factors (or u8 for microscaled modes) + - `biases` - Per-group zero points, or `nil` for microscaled modes + (`"mxfp4"`/`"mxfp8"`/`"nvfp4"` — `mx::fp_quantize` doesn't emit biases) ## Parameters - `w` - Float tensor to quantize - `group_size` - Number of weights per group (default: 64) - `bits` - Quantization bits (default: 4) + - `mode` - `"affine"` (default), `"mxfp4"`, `"mxfp8"`, or `"nvfp4"` """ - @mlx_function {:quantize, 5} - def quantize({dev_w, ref_w}, group_size, bits) + @mlx_function {:quantize, 6} + def quantize({dev_w, ref_w}, group_size, bits, mode \\ "affine") when is_tensor(dev_w, ref_w) do device = dev_w {worker, effective_device} = resolve_worker(device) {weights_ref, scales_ref, biases_ref} = - EMLX.NIF.quantize(worker, ref_w, group_size, bits, effective_device) + EMLX.NIF.quantize(worker, ref_w, group_size, bits, mode, effective_device) |> unwrap!() |> await_worker() {{effective_device, weights_ref}, {effective_device, scales_ref}, - {effective_device, biases_ref}} + wrap_optional_tensor(biases_ref, effective_device)} end + # `nil` (microscaled modes have no biases) passes through as the atom `nil` + # to the NIF layer; a real tensor unwraps to its raw ref for device merging. + defp unwrap_optional_tensor(nil), do: {nil, nil} + defp unwrap_optional_tensor({device, ref}), do: {ref, device} + + defp wrap_optional_tensor(nil, _device), do: nil + defp wrap_optional_tensor(ref, device), do: {device, ref} + # ── mlx::fast ops ─────────────────────────────────────────────────────────── @doc """ @@ -514,17 +627,19 @@ defmodule EMLX do - `k` — `{B, N_kv, T_kv, D}` - `v` — `{B, N_kv, T_kv, D}` - `scale` — scalar (typically `1 / sqrt(D)`) + - `sinks` — optional learned per-head attention-sink logits tensor, or `nil` Prefer `EMLX.Fast.scaled_dot_product_attention/4` inside `defn`. """ - @mlx_function {:fast_sdpa, 6} - def fast_sdpa({dev_q, ref_q}, {dev_k, ref_k}, {dev_v, ref_v}, scale) + @mlx_function {:fast_sdpa, 7} + def fast_sdpa({dev_q, ref_q}, {dev_k, ref_k}, {dev_v, ref_v}, scale, sinks \\ nil) when is_tensor(dev_q, ref_q) and is_tensor(dev_k, ref_k) and is_tensor(dev_v, ref_v) do - device = merge_device(dev_q, merge_device(dev_k, dev_v)) + {ref_s, sinks_device} = unwrap_optional_tensor(sinks) + device = merge_device(dev_q, merge_device(dev_k, merge_device(dev_v, sinks_device))) {worker, effective_device} = resolve_worker(device) job_ref = - EMLX.NIF.fast_sdpa(worker, ref_q, ref_k, ref_v, scale * 1.0, effective_device) + EMLX.NIF.fast_sdpa(worker, ref_q, ref_k, ref_v, scale * 1.0, ref_s, effective_device) |> unwrap!() await_worker(job_ref) |> wrap_tensor(effective_device) @@ -536,23 +651,42 @@ defmodule EMLX do `mask` must be broadcast-compatible with `{B, N_q, T_q, T_kv}`. Boolean `false` entries are treated as `-∞`. + `sinks` — optional learned per-head attention-sink logits tensor, or `nil`. + Prefer `EMLX.Fast.scaled_dot_product_attention/5` inside `defn`. """ - @mlx_function {:fast_sdpa_masked, 7} + @mlx_function {:fast_sdpa_masked, 8} def fast_sdpa_masked( {dev_q, ref_q}, {dev_k, ref_k}, {dev_v, ref_v}, {dev_m, ref_m}, - scale + scale, + sinks \\ nil ) when is_tensor(dev_q, ref_q) and is_tensor(dev_k, ref_k) and is_tensor(dev_v, ref_v) and is_tensor(dev_m, ref_m) do - device = merge_device(dev_q, merge_device(dev_k, merge_device(dev_v, dev_m))) + {ref_s, sinks_device} = unwrap_optional_tensor(sinks) + + device = + merge_device( + dev_q, + merge_device(dev_k, merge_device(dev_v, merge_device(dev_m, sinks_device))) + ) + {worker, effective_device} = resolve_worker(device) job_ref = - EMLX.NIF.fast_sdpa_masked(worker, ref_q, ref_k, ref_v, scale * 1.0, ref_m, effective_device) + EMLX.NIF.fast_sdpa_masked( + worker, + ref_q, + ref_k, + ref_v, + scale * 1.0, + ref_m, + ref_s, + effective_device + ) |> unwrap!() await_worker(job_ref) |> wrap_tensor(effective_device) @@ -729,16 +863,19 @@ defmodule EMLX do - `v` — `{B, N_kv, T_kv, D}` - `scale` — scalar (typically `1 / sqrt(D)`) + `sinks` — optional learned per-head attention-sink logits tensor, or `nil`. + Prefer `EMLX.Fast.scaled_dot_product_attention_causal/4` inside `defn`. """ - @mlx_function {:fast_sdpa_causal, 6} - def fast_sdpa_causal({dev_q, ref_q}, {dev_k, ref_k}, {dev_v, ref_v}, scale) + @mlx_function {:fast_sdpa_causal, 7} + def fast_sdpa_causal({dev_q, ref_q}, {dev_k, ref_k}, {dev_v, ref_v}, scale, sinks \\ nil) when is_tensor(dev_q, ref_q) and is_tensor(dev_k, ref_k) and is_tensor(dev_v, ref_v) do - device = merge_device(dev_q, merge_device(dev_k, dev_v)) + {ref_s, sinks_device} = unwrap_optional_tensor(sinks) + device = merge_device(dev_q, merge_device(dev_k, merge_device(dev_v, sinks_device))) {worker, effective_device} = resolve_worker(device) job_ref = - EMLX.NIF.fast_sdpa_causal(worker, ref_q, ref_k, ref_v, scale * 1.0, effective_device) + EMLX.NIF.fast_sdpa_causal(worker, ref_q, ref_k, ref_v, scale * 1.0, ref_s, effective_device) |> unwrap!() await_worker(job_ref) |> wrap_tensor(effective_device) @@ -757,21 +894,30 @@ defmodule EMLX do - `v` — `{B, N_kv, T_kv, D}` - `scale` — scalar (typically `1 / sqrt(D)`) - `key_mask` — `{B, T_kv}` boolean/int tensor (1 = attend, 0 = padding) + - `sinks` — optional learned per-head attention-sink logits tensor, or `nil` Prefer `EMLX.Fast.scaled_dot_product_attention_causal_key_masked/5` inside `defn`. """ - @mlx_function {:fast_sdpa_causal_key_masked, 8} + @mlx_function {:fast_sdpa_causal_key_masked, 9} def fast_sdpa_causal_key_masked( {dev_q, ref_q}, {dev_k, ref_k}, {dev_v, ref_v}, scale, {dev_m, ref_m}, - kv_offset + kv_offset, + sinks \\ nil ) when is_tensor(dev_q, ref_q) and is_tensor(dev_k, ref_k) and is_tensor(dev_v, ref_v) and is_tensor(dev_m, ref_m) and is_integer(kv_offset) do - device = merge_device(dev_q, merge_device(dev_k, merge_device(dev_v, dev_m))) + {ref_s, sinks_device} = unwrap_optional_tensor(sinks) + + device = + merge_device( + dev_q, + merge_device(dev_k, merge_device(dev_v, merge_device(dev_m, sinks_device))) + ) + {worker, effective_device} = resolve_worker(device) job_ref = @@ -783,6 +929,7 @@ defmodule EMLX do scale * 1.0, ref_m, kv_offset, + ref_s, effective_device ) |> unwrap!() @@ -953,578 +1100,41 @@ defmodule EMLX do {{effective_device, attn_ref}, {effective_device, k_upd_ref}, {effective_device, v_upd_ref}} end - @doc """ - Qwen3 fused RoPE + KV cache update + SDPA for the native decode path. - - Accepts `query`, `new_key`, and `new_value` in projection layout - `{B, T, N, D}`. The NIF transposes Q/K/V, applies Qwen3 RoPE to Q/K, updates - the owned KV cache, runs SDPA, and returns flattened attention output - `{B, T, N * D}` ready for projection, plus updated cache refs. - """ - @mlx_function {:qwen3_kv_cache_attention, 11} - def qwen3_kv_cache_attention( - {dev_q, ref_q}, - {_dev_k, ref_k}, - {_dev_v, ref_v}, - {_dev_kc, ref_kc}, - {_dev_vc, ref_vc}, - offset, - scale, - head_dim, - theta - ) - when is_tensor(dev_q, ref_q) and is_integer(offset) and is_float(scale) and - is_integer(head_dim) and is_number(theta) do - device = dev_q - {worker, effective_device} = resolve_worker(device) - - {attn_ref, k_upd_ref, v_upd_ref} = - EMLX.NIF.qwen3_kv_cache_attention( - worker, - ref_q, - ref_k, - ref_v, - ref_kc, - ref_vc, - offset, - scale, - head_dim, - theta * 1.0, - effective_device - ) - |> unwrap!() - |> await_worker() - - {{effective_device, attn_ref}, {effective_device, k_upd_ref}, {effective_device, v_upd_ref}} - end - - @doc """ - Qwen3 dense MLP helper. - - Accepts hidden states `{B, T, H}`, RMSNorm weight after attention `{H}`, - dense gate/up projections `{H, I}`, dense down projection `{I, H}`, and RMSNorm - epsilon. Returns `hidden + mlp_output` as `{B, T, H}`. - """ - @mlx_function {:qwen3_mlp, 8} - def qwen3_mlp( - {dev_h, ref_h}, - {_dev_norm, ref_norm}, - {_dev_gate, ref_gate}, - {_dev_up, ref_up}, - {_dev_down, ref_down}, - eps - ) - when is_tensor(dev_h, ref_h) and is_float(eps) do - device = dev_h - {worker, effective_device} = resolve_worker(device) - - out_ref = - EMLX.NIF.qwen3_mlp( - worker, - ref_h, - ref_norm, - ref_gate, - ref_up, - ref_down, - eps, - effective_device - ) - |> unwrap!() - |> await_worker() - - {effective_device, out_ref} - end - - @doc """ - Qwen3 dense attention output projection plus residual add. - - Accepts residual hidden state `{B, T, H}`, flattened attention output - `{B, T, I}`, and dense output projection `{I, H}`. Returns - `hidden + projected_attention`. - """ - @mlx_function {:qwen3_attention_residual, 5} - def qwen3_attention_residual( - {dev_h, ref_h}, - {_dev_attn, ref_attn}, - {_dev_o, ref_o} - ) - when is_tensor(dev_h, ref_h) do - device = dev_h - {worker, effective_device} = resolve_worker(device) - - out_ref = - EMLX.NIF.qwen3_attention_residual( - worker, - ref_h, - ref_attn, - ref_o, - effective_device - ) - |> unwrap!() - |> await_worker() - - {effective_device, out_ref} - end - - @doc """ - Qwen3 dense transformer layer helper. - - Accepts hidden states `{B, T, H}`, input RMSNorm and RMSNorm weights after attention, - dense attention projections, Q/K RMSNorm weights, owned KV cache refs, dense - MLP projections, offset, scale, RoPE parameters, and RMSNorm epsilon. Returns - `{hidden_out, k_cache, v_cache}`. - """ - @mlx_function {:qwen3_layer, 21} - def qwen3_layer( - {dev_h, ref_h}, - {_dev_norm1, ref_norm1}, - {_dev_q, ref_q}, - {_dev_k, ref_k}, - {_dev_v, ref_v}, - {_dev_o, ref_o}, - {_dev_qn, ref_qn}, - {_dev_kn, ref_kn}, - {_dev_kc, ref_kc}, - {_dev_vc, ref_vc}, - {_dev_norm2, ref_norm2}, - {_dev_gate, ref_gate}, - {_dev_up, ref_up}, - {_dev_down, ref_down}, - offset, - scale, - head_dim, - theta, - eps - ) - when is_tensor(dev_h, ref_h) and is_integer(offset) and is_float(scale) and - is_integer(head_dim) and is_number(theta) and is_float(eps) do - device = dev_h - {worker, effective_device} = resolve_worker(device) - - {out_ref, k_upd_ref, v_upd_ref} = - EMLX.NIF.qwen3_layer( - worker, - ref_h, - ref_norm1, - ref_q, - ref_k, - ref_v, - ref_o, - ref_qn, - ref_kn, - ref_kc, - ref_vc, - ref_norm2, - ref_gate, - ref_up, - ref_down, - offset, - scale, - head_dim, - theta * 1.0, - eps, - effective_device - ) - |> unwrap!() - |> await_worker() - - {{effective_device, out_ref}, {effective_device, k_upd_ref}, {effective_device, v_upd_ref}} - end - - @doc """ - Qwen3 embedding lookup plus dense forward through all layers and greedy token. - - This accepts token ids and embedding weights directly, so the dense greedy - path can avoid constructing a separate embedding lookup graph before entering - the native Qwen3 worker call. - """ - @mlx_function {:qwen3_forward_greedy_ids, 13} - def qwen3_forward_greedy_ids( - {dev_ids, ref_ids}, - {_dev_embed, ref_embed}, - layers, - kv_cache, - {_dev_norm, ref_norm}, - {_dev_lm_head, ref_lm_head}, - offset, - scale, - head_dim, - theta, - eps - ) - when is_tensor(dev_ids, ref_ids) and is_list(layers) and is_list(kv_cache) and - is_integer(offset) and is_float(scale) and is_integer(head_dim) and - is_number(theta) and is_float(eps) do - device = dev_ids - {worker, effective_device} = resolve_worker(device) - - layer_refs = Enum.map(layers, &qwen3_layer_refs!/1) - kv_refs = qwen3_kv_refs(kv_cache) - - {token_ref, kv_updated_refs} = - EMLX.NIF.qwen3_forward_greedy_ids( - worker, - ref_ids, - ref_embed, - layer_refs, - kv_refs, - ref_norm, - ref_lm_head, - offset, - scale, - head_dim, - theta * 1.0, - eps, - effective_device - ) - |> unwrap!() - |> await_worker() - - {{effective_device, token_ref}, - Enum.map(kv_updated_refs, fn {k_ref, v_ref} -> - {{effective_device, k_ref}, {effective_device, v_ref}} - end)} - end - - @doc """ - Qwen3 greedy decode chunk helper. - - Starting from a `{1, 1}` token id tensor, runs `count` greedy decode steps - without returning to Elixir between steps. Returns `{token_refs, kv_cache}`, - where `token_refs` is a list of raw EMLX token refs in generation order and - `kv_cache` is the final updated raw cache. - """ - @mlx_function {:qwen3_forward_greedy_ids_chunk, 14} - def qwen3_forward_greedy_ids_chunk( - {dev_ids, ref_ids}, - {_dev_embed, ref_embed}, - layers, - kv_cache, - {_dev_norm, ref_norm}, - {_dev_lm_head, ref_lm_head}, - offset, - count, - scale, - head_dim, - theta, - eps - ) - when is_tensor(dev_ids, ref_ids) and is_list(layers) and is_list(kv_cache) and - is_integer(offset) and is_integer(count) and count > 0 and is_float(scale) and - is_integer(head_dim) and is_number(theta) and is_float(eps) do - device = dev_ids - {worker, effective_device} = resolve_worker(device) - - qwen3_assert_decode_ids!({dev_ids, ref_ids}, "qwen3_forward_greedy_ids_chunk") - - layer_refs = Enum.map(layers, &qwen3_layer_refs!/1) - kv_refs = qwen3_kv_refs(kv_cache) - - {token_refs, kv_updated_refs} = - EMLX.NIF.qwen3_forward_greedy_ids_chunk( - worker, - ref_ids, - ref_embed, - layer_refs, - kv_refs, - ref_norm, - ref_lm_head, - offset, - count, - scale, - head_dim, - theta * 1.0, - eps, - effective_device - ) - |> unwrap!() - |> await_worker() - - tokens = Enum.map(token_refs, &{effective_device, &1}) - - kv_cache = - Enum.map(kv_updated_refs, fn {k_ref, v_ref} -> - {{effective_device, k_ref}, {effective_device, v_ref}} - end) - - {tokens, kv_cache} - end - - @doc """ - Qwen3 embedding lookup plus dense forward through all layers and greedy token. - - This variant returns the selected token id as a BEAM integer while keeping the - updated KV cache as raw EMLX refs. It is intended for streaming decode paths - that need the token on the host anyway. - """ - @mlx_function {:qwen3_forward_greedy_ids_token_id, 13} - def qwen3_forward_greedy_ids_token_id( - {dev_ids, ref_ids}, - {_dev_embed, ref_embed}, - layers, - kv_cache, - {_dev_norm, ref_norm}, - {_dev_lm_head, ref_lm_head}, - offset, - scale, - head_dim, - theta, - eps - ) - when is_tensor(dev_ids, ref_ids) and is_list(layers) and is_list(kv_cache) and - is_integer(offset) and is_float(scale) and is_integer(head_dim) and - is_number(theta) and is_float(eps) do - device = dev_ids - {worker, effective_device} = resolve_worker(device) - - qwen3_assert_batch_size_one!({dev_ids, ref_ids}, "qwen3_forward_greedy_ids_token_id") - - layer_refs = Enum.map(layers, &qwen3_layer_refs!/1) - kv_refs = qwen3_kv_refs(kv_cache) - - {token_id, kv_updated_refs} = - EMLX.NIF.qwen3_forward_greedy_ids_token_id( - worker, - ref_ids, - ref_embed, - layer_refs, - kv_refs, - ref_norm, - ref_lm_head, - offset, - scale, - head_dim, - theta * 1.0, - eps, - effective_device - ) - |> unwrap!() - |> await_worker() - - {token_id, - Enum.map(kv_updated_refs, fn {k_ref, v_ref} -> - {{effective_device, k_ref}, {effective_device, v_ref}} - end)} - end - - @doc """ - Qwen3 dense forward and greedy decode for one token. - - This variant accepts the previous token id as a BEAM integer and returns the - selected token id as a BEAM integer. It is intended for decode loops that - already synchronize once per token for streaming or EOS checks. By default it - runs on the embedding tensor's device; pass `device` explicitly to override. - """ - @mlx_function {:qwen3_forward_greedy_token_id, 13} - def qwen3_forward_greedy_token_id( - token_id, - {dev_embed, ref_embed}, - layers, - kv_cache, - {_dev_norm, ref_norm}, - {_dev_lm_head, ref_lm_head}, - offset, - scale, - head_dim, - theta, - eps, - device \\ nil - ) - when is_integer(token_id) and is_list(layers) and is_list(kv_cache) and - is_integer(offset) and is_float(scale) and is_integer(head_dim) and - is_number(theta) and is_float(eps) do - device = device || dev_embed - {worker, effective_device} = resolve_worker(device) + # Microscaled modes pin an exact {group_size, bits} pair (MLX's + # `fp_quantize` — see `mlx::core::quantize` in `mlx/ops.h`), checked + # directly against the same MLX version this repo vendors. + @microscaled_constraints %{"mxfp4" => {32, 4}, "mxfp8" => {32, 8}, "nvfp4" => {16, 4}} + @valid_quantization_modes ~w(affine mxfp4 mxfp8 nvfp4) - layer_refs = Enum.map(layers, &qwen3_layer_refs!/1) - kv_refs = qwen3_kv_refs(kv_cache) + defp validate_quantization_mode!(mode) when mode in @valid_quantization_modes, do: :ok - {next_token_id, kv_updated_refs} = - EMLX.NIF.qwen3_forward_greedy_token_id( - worker, - token_id, - ref_embed, - layer_refs, - kv_refs, - ref_norm, - ref_lm_head, - offset, - scale, - head_dim, - theta * 1.0, - eps, - effective_device - ) - |> unwrap!() - |> await_worker() - - {next_token_id, - Enum.map(kv_updated_refs, fn {k_ref, v_ref} -> - {{effective_device, k_ref}, {effective_device, v_ref}} - end)} + defp validate_quantization_mode!(mode) do + raise ArgumentError, + "EMLX.quantize/2: :mode must be one of #{inspect(@valid_quantization_modes)}, " <> + "got: #{inspect(mode)}" end - @doc """ - Qwen3 dense final RMSNorm + lm_head + greedy argmax helper. - - Accepts hidden states `{B, T, H}`, final RMSNorm weight `{H}`, dense lm_head - `{V, H}`, and RMSNorm epsilon. Returns token ids as `{B}`. - """ - @mlx_function {:qwen3_final_greedy, 6} - def qwen3_final_greedy( - {dev_h, ref_h}, - {_dev_norm, ref_norm}, - {_dev_lm_head, ref_lm_head}, - eps - ) - when is_tensor(dev_h, ref_h) and is_float(eps) do - device = dev_h - {worker, effective_device} = resolve_worker(device) + defp validate_microscaled_constraints!("affine", _group_size, _bits), do: :ok - out_ref = - EMLX.NIF.qwen3_final_greedy( - worker, - ref_h, - ref_norm, - ref_lm_head, - eps, - effective_device - ) - |> unwrap!() - |> await_worker() + defp validate_microscaled_constraints!(mode, group_size, bits) do + {expected_gs, expected_bits} = Map.fetch!(@microscaled_constraints, mode) - {effective_device, out_ref} - end - - @doc """ - Qwen3 dense attention block helper. - - Accepts residual hidden state before attention `{B, T, H}`, input RMSNorm weight - `{H}`, dense Q/K/V/O projections, Q/K RMSNorm weights, owned KV cache refs, - offset, scale, RoPE parameters, and RMSNorm epsilon. Returns - `{hidden_out, k_cache, v_cache}`. - """ - @mlx_function {:qwen3_attention_block, 17} - def qwen3_attention_block( - {dev_h, ref_h}, - {_dev_norm, ref_norm}, - {_dev_q, ref_q}, - {_dev_k, ref_k}, - {_dev_v, ref_v}, - {_dev_o, ref_o}, - {_dev_qn, ref_qn}, - {_dev_kn, ref_kn}, - {_dev_kc, ref_kc}, - {_dev_vc, ref_vc}, - offset, - scale, - head_dim, - theta, - eps - ) - when is_tensor(dev_h, ref_h) and is_integer(offset) and is_float(scale) and - is_integer(head_dim) and is_number(theta) and is_float(eps) do - device = dev_h - {worker, effective_device} = resolve_worker(device) - - {out_ref, k_upd_ref, v_upd_ref} = - EMLX.NIF.qwen3_attention_block( - worker, - ref_h, - ref_norm, - ref_q, - ref_k, - ref_v, - ref_o, - ref_qn, - ref_kn, - ref_kc, - ref_vc, - offset, - scale, - head_dim, - theta * 1.0, - eps, - effective_device - ) - |> unwrap!() - |> await_worker() - - {{effective_device, out_ref}, {effective_device, k_upd_ref}, {effective_device, v_upd_ref}} - end - - defp qwen3_layer_refs!( - {norm1, norm2, q_norm, k_norm, q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, - down_proj} - ) do - { - tensor_ref!(norm1), - tensor_ref!(norm2), - tensor_ref!(q_norm), - tensor_ref!(k_norm), - tensor_ref!(q_proj), - tensor_ref!(k_proj), - tensor_ref!(v_proj), - tensor_ref!(o_proj), - tensor_ref!(gate_proj), - tensor_ref!(up_proj), - tensor_ref!(down_proj) - } - end - - defp qwen3_kv_refs!({{_device_k, ref_k}, {_device_v, ref_v}}) - when is_reference(ref_k) and is_reference(ref_v), - do: {ref_k, ref_v} - - defp qwen3_kv_refs!({k_cache, v_cache}), do: {tensor_ref!(k_cache), tensor_ref!(v_cache)} - - defp qwen3_kv_refs([{{_device_k, ref_k}, {_device_v, ref_v}} | _rest] = kv_cache) - when is_reference(ref_k) and is_reference(ref_v), - do: kv_cache - - defp qwen3_kv_refs(kv_cache), do: Enum.map(kv_cache, &qwen3_kv_refs!/1) - - defp qwen3_assert_batch_size_one!(tensor, function_name) do - case EMLX.shape(tensor) do - {1, _seq_len} -> - :ok - - {batch_size, _seq_len} -> + cond do + group_size != expected_gs -> raise ArgumentError, - "#{function_name} requires batch size 1, got batch size #{batch_size}" + "EMLX.quantize/2: mode #{inspect(mode)} requires group_size=#{expected_gs}, " <> + "got: #{inspect(group_size)}" - shape -> + bits != expected_bits -> raise ArgumentError, - "#{function_name} expects rank-2 input ids, got shape #{inspect(shape)}" - end - end + "EMLX.quantize/2: mode #{inspect(mode)} requires bits=#{expected_bits}, " <> + "got: #{inspect(bits)}" - defp qwen3_assert_decode_ids!(tensor, function_name) do - case EMLX.shape(tensor) do - {1, 1} -> + true -> :ok - - {1, seq_len} -> - raise ArgumentError, - "#{function_name} requires sequence length 1, got sequence length #{seq_len}" - - {batch_size, _seq_len} -> - raise ArgumentError, - "#{function_name} requires batch size 1, got batch size #{batch_size}" - - shape -> - raise ArgumentError, - "#{function_name} expects rank-2 input ids, got shape #{inspect(shape)}" end end - defp tensor_ref!(%Nx.Tensor{data: %EMLX.Backend{ref: {_device, ref}}}), do: ref - - defp tensor_ref!(tensor) do - {_device, ref} = EMLX.Backend.from_nx(tensor) - ref - end - @doc """ Stable wrapper for causal self attention with an owned KV cache. @@ -1550,15 +1160,6 @@ defmodule EMLX do store and reuse device resident cache buffers without converting them back to `Nx.Tensor` between decode steps. """ - @spec causal_kv_attention( - tensor_ref(), - tensor_ref(), - tensor_ref(), - tensor_ref(), - tensor_ref(), - non_neg_integer(), - keyword() - ) :: {tensor_ref(), tensor_ref(), tensor_ref()} def causal_kv_attention(query, new_key, new_value, key_cache, value_cache, offset, opts) when is_integer(offset) and offset >= 0 and is_list(opts) do scale = Keyword.fetch!(opts, :scale) @@ -1592,13 +1193,22 @@ defmodule EMLX do * `:type` — storage type: `{:s, 2}`, `{:s, 4}` (default), or `{:s, 8}`. * `:group_size` — 32, 64, or 128 (default 64). Must evenly divide the last - dimension of `tensor`. + dimension of `tensor`. Microscaled modes pin this to a specific value + (see `:mode` below). + * `:mode` — `"affine"` (default, real biases), or a microscaled mode — + `"mxfp4"` (group_size 32, bits 4), `"mxfp8"` (group_size 32, bits 8), or + `"nvfp4"` (group_size 16, bits 4). Microscaled modes have no biases + (`mx::fp_quantize` returns only `(wq, scales)`); the returned tensor's + `EMLX.Quantization.Config.biases` is `nil`. """ - @spec quantize(Nx.Tensor.t(), keyword()) :: Nx.Tensor.t() def quantize(%Nx.Tensor{} = tensor, opts) when is_list(opts) do type = Keyword.get(opts, :type, {:s, 4}) {_, bits} = type group_size = Keyword.get(opts, :group_size, 64) + mode = Keyword.get(opts, :mode, "affine") + + validate_quantization_mode!(mode) + validate_microscaled_constraints!(mode, group_size, bits) unless Nx.rank(tensor) == 2 do raise ArgumentError, @@ -1614,20 +1224,21 @@ defmodule EMLX do end device_ref = EMLX.Backend.from_nx(tensor) - {weight_ref, scales_ref, biases_ref} = EMLX.quantize(device_ref, group_size, bits) + {weight_ref, scales_ref, biases_ref} = EMLX.quantize(device_ref, group_size, bits, mode) scales = EMLX.Backend.to_nx(scales_ref) - biases = EMLX.Backend.to_nx(biases_ref) + biases = biases_ref && EMLX.Backend.to_nx(biases_ref) config = %EMLX.Quantization.Config{ scales: scales, biases: biases, group_size: group_size, - bits: bits + bits: bits, + mode: mode } weight_shape = EMLX.shape(weight_ref) - template = Nx.template(Nx.shape(tensor), type) + %Nx.Tensor{} = template = Nx.template(Nx.shape(tensor), type) %Nx.Tensor{ template @@ -1642,21 +1253,24 @@ defmodule EMLX do @doc """ Dequantize a quantized `Nx.Tensor` (created by `EMLX.quantize/2`) to a - dense float tensor by calling `mx::dequantize`. + dense float tensor by calling `mx::dequantize`. Supports every mode + `EMLX.quantize/2` accepts (`"affine"` and the microscaled variants). """ - @spec dequantize(Nx.Tensor.t()) :: Nx.Tensor.t() def dequantize( %Nx.Tensor{ data: %EMLX.Backend{ref: weight_ref, quantization_config: cfg} } = _qw ) when not is_nil(cfg) do + biases_ref = cfg.biases && EMLX.Backend.from_nx(cfg.biases) + EMLX.dequantize( weight_ref, EMLX.Backend.from_nx(cfg.scales), - EMLX.Backend.from_nx(cfg.biases), + biases_ref, cfg.group_size, - cfg.bits + cfg.bits, + cfg.mode ) |> EMLX.Backend.to_nx() end @@ -1667,7 +1281,6 @@ defmodule EMLX do `qw` must be a quantized tensor produced by `EMLX.quantize/2`. Raises `ArgumentError` if both arguments are quantized. """ - @spec quantized_matmul(Nx.Tensor.t(), Nx.Tensor.t()) :: Nx.Tensor.t() def quantized_matmul(%Nx.Tensor{} = activation, %Nx.Tensor{} = qw) do cfg = qw.data.quantization_config @@ -1682,15 +1295,18 @@ defmodule EMLX do "argument; got two quantized tensors. Dequantize one of them first." end + biases_ref = cfg.biases && EMLX.Backend.from_nx(cfg.biases) + result = EMLX.quantized_matmul( EMLX.Backend.from_nx(activation), qw.data.ref, EMLX.Backend.from_nx(cfg.scales), - EMLX.Backend.from_nx(cfg.biases), + biases_ref, true, cfg.group_size, - cfg.bits + cfg.bits, + cfg.mode ) EMLX.Backend.to_nx(result) @@ -1762,9 +1378,10 @@ defmodule EMLX do EMLX.NIF.shm_unlink_handle(name) |> unwrap!() end - defp unwrap!(:ok), do: :ok - defp unwrap!({:ok, result}), do: result - defp unwrap!({:error, error}), do: raise(EMLX.NIFError, List.to_string(error)) + @doc false + def unwrap!(:ok), do: :ok + def unwrap!({:ok, result}), do: result + def unwrap!({:error, error}), do: raise(EMLX.NIFError, List.to_string(error)) # Wraps a worker-thread payload in {device, ref} envelopes. # Already-unwrapped (no leading {:ok, _}) — `await_worker/1` peels that @@ -1820,9 +1437,12 @@ defmodule EMLX do """ def eval({device, ref}) when is_tensor(device, ref) do maybe_profile(EMLX.Profiling.inc_eval()) - {worker, _effective_device} = resolve_worker(device) - job_ref = EMLX.NIF.eval(worker, ref) |> unwrap!() - await_worker(job_ref) + + EMLX.Telemetry.span_eval(fn -> + {worker, _effective_device} = resolve_worker(device) + job_ref = EMLX.NIF.eval(worker, ref) |> unwrap!() + await_worker(job_ref) + end) end @doc """ @@ -1889,17 +1509,112 @@ defmodule EMLX do end end - defp await_worker(job_ref) do + # `runtime_calls` (see `EMLX.Native.Expr`'s moduledoc "Runtime calls" + # section) is only ever non-empty for an `eval_program` job whose compiled + # program contains an inlined `:runtime_call` node — every other call site + # keeps the 1-arg form and can never receive an `:emlx_runtime_call` + # request in the first place. A request is served here, mid-flight, before + # the final `{^job_ref, _}` reply for *this* `eval_program` call arrives: + # the worker thread is blocked inside `EMLXRuntimeCall::eval_cpu`/`eval_gpu` + # (emlx_compiler.cpp) waiting on exactly the reply this loop sends back via + # `EMLX.NIF.resolve_runtime_call/3`. + @doc false + def await_worker(job_ref, runtime_calls \\ [], tensors \\ [], dev \\ nil) do receive do # Worker NIFs (sync bodies) return one of: # nx::nif::ok(env) => :ok # nx::nif::ok(env, term) => {:ok, term} # nx::nif::error(env, msg) => {:error, msg} # The async wrapper forwards the payload as-is in {ref, payload}. - {^job_ref, :ok} -> :ok - {^job_ref, {:ok, result}} -> result - {^job_ref, {:error, reason}} -> raise(EMLX.NIFError, List.to_string(reason)) + {^job_ref, :ok} -> + :ok + + {^job_ref, {:ok, result}} -> + result + + {^job_ref, {:error, reason}} -> + raise(EMLX.NIFError, List.to_string(reason)) + + {:emlx_runtime_call, pending, callback_index, args_binaries} -> + handle_runtime_call(pending, callback_index, args_binaries, runtime_calls, tensors, dev) + await_worker(job_ref, runtime_calls, tensors, dev) + end + end + + # Runs the real Elixir callback for one `:emlx_runtime_call` request and + # replies via `EMLX.NIF.resolve_runtime_call/3`, waking the worker thread + # blocked inside `EMLXRuntimeCall::eval_cpu`/`eval_gpu`. A raising callback + # is caught and reported as an `:error` reply instead of crashing this + # process (which would otherwise leave the worker thread blocked forever). + # + # `tensors` is *this* `eval_program` call's own materialised input list + # (`build_native_eval_fn/7`'s own `tensors`, matching this program's own + # `:parameter` numbering exactly — including for a `:runtime_call` nested + # inside a `while` body, whose own re-entrant compile has its own + # independent parameter numbering over its own materialised inputs). Used + # only to substitute back a quantized argument's original bound tensor — + # see `EMLX.Native.Expr`'s moduledoc "Runtime calls" section and + # `arg_param_positions`'s doc. + # + # The real callback runs inside `EMLX.CommandQueue.with_queue/2` bound to + # `EMLX.Application.runtime_call_worker(dev)` — see that function's doc + # for why: any eager EMLX call the callback itself makes (e.g. + # `EMLX.Quantization.dequantize_callback/2` calling `EMLX.dequantize/1`) + # must never be routed back to the worker that is, right now, blocked + # inside `EMLXRuntimeCall::eval_cpu`/`eval_gpu` waiting for exactly this + # callback to return. + defp handle_runtime_call(pending, callback_index, args_binaries, runtime_calls, tensors, dev) do + %{ + callback: callback, + args_template: args_template, + arg_param_positions: positions, + 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 + } + + reply = + try do + result = + EMLX.CommandQueue.with_queue(callback_queue, fn -> callback.(args_container, opts) end) + + binaries = + [result] + |> Nx.Defn.Composite.flatten_list() + |> Enum.map(&Nx.to_binary/1) + + {:ok, binaries} + rescue + e -> {:error, Exception.format(:error, e, __STACKTRACE__)} + end + + case reply do + {:ok, binaries} -> EMLX.NIF.resolve_runtime_call(pending, :ok, binaries) + {:error, message} -> EMLX.NIF.resolve_runtime_call(pending, :error, message) end + |> unwrap!() end deftensor slice(tensor, starts, stops, strides) @@ -1940,32 +1655,821 @@ defmodule EMLX do # that manage their own queues (equivalent to a manual `with_queue`). @valid_compiler_keys [:device, :max_concurrency, :command_queue] + # Process-lifetime dispatch cache backing `dispatch_key/3` + + # `get_or_compile_program/6` (see their docs) — a compiled program is keyed + # by a *structural* signature of its `Expr` (not object identity), so + # it survives across `Nx.Defn.Graph.run/3`'s per-call re-tracing and is + # shared across structurally-identical call sites (e.g. every one of + # Qwen3's 28 attention layers), not just within one closure's lifetime. + @native_dispatch_cache_table :emlx_native_dispatch_cache + + # A second, process-lifetime cache in front of `dispatch_key/3`'s own + # (expensive — O(nodes), plus per-opaque-scope SHA256 hashing) structural + # walk, keyed by `output_expr`'s own node identity rather than its + # structural signature. This matters for `run_while_loop/3`'s host-driven + # `cond_fn`/`body_fn`: `Nx.Defn.jit/2` retraces `fn _ -> body_expr end` + # once and caches *that* trace by argument template, so every subsequent + # call re-enters `__jit__`/`build_eval_fn` with the *exact same* `Expr` + # (identical ids, not just structurally identical) — walking it again on + # every decode step is pure waste. (This is unlike `Nx.Defn.Graph.run/3`'s + # per-stage re-tracing, which *does* mint fresh ids each call — hence + # `dispatch_key/3` still has to fall back to the structural walk on a miss.) + @dispatch_key_by_id_table :emlx_dispatch_key_by_id + @impl Nx.Defn.Compiler def __jit__(key, vars, fun, args_list, opts) do __compile__(key, vars, fun, opts).(args_list) end @impl Nx.Defn.Compiler - def __compile__(key, vars, fun, opts) do + def __compile__(_key, vars, fun, opts) do Keyword.validate!(opts, @valid_compiler_keys) - {compiler_opts, rest_opts} = split_compiler_opts(opts) - queue = Keyword.get(compiler_opts, :command_queue) - - inner = - Nx.Defn.Evaluator.__compile__( - key, - vars, - fun, - Keyword.put(rest_opts, :compiler, Nx.Defn.Evaluator) - ) + queue = Keyword.get(opts, :command_queue) + device = Keyword.get(opts, :device, default_device()) + + wrap_with_queue(queue, native_compile(vars, fun, device)) + end + + # Attempts to lower `fun.(vars)` to an `EMLX.Native.Expr` program and build a + # compiled program resource, returning the resulting eval closure. Single-mode: + # any op not yet implemented raises `ArgumentError` ("does not yet lower op + # ...") straight through to the caller — there is no whole-`defn` fallback lane. + defp native_compile(vars, fun, device) do + output_expr = fun.(vars) + {worker, effective_device} = resolve_worker(device) + # `vars`' flattened leaf count is the true call arity — passed through so + # `lower/2` can densify its input list even when `output_expr` doesn't + # reference every parameter position (see EMLX.Native.Expr.lower/2 doc). + num_inputs = vars |> Nx.Defn.Composite.flatten_list() |> length() + build_eval_fn(output_expr, worker, effective_device, num_inputs, fun, vars) + end + + # Wraps a compiled eval closure so each invocation routes through `queue`. + # + # A CommandQueue is a *transient* execution scope: it is bound only for the + # duration of the call and torn down immediately after. The compiled program + # therefore runs on the queue's worker thread, producing a lazy result whose + # MLX graph is tied to that thread's stream. We must materialise the outputs + # *before* the binding is restored — otherwise a later read (e.g. `to_blob`) + # resolves to the device-default worker, whose thread does not own the + # queue's stream ("There is no Stream(gpu, N) in current thread"). + defp wrap_with_queue(nil, eval_fn), do: eval_fn + + defp wrap_with_queue(queue, eval_fn) do + fn inputs -> + EMLX.CommandQueue.with_queue(queue, fn -> + result = eval_fn.(inputs) + force_eval_outputs(result) + result + end) + end + end - if queue do - # Capture the queue ref in a closure so each invocation of the compiled - # function routes through the correct CommandQueue. The queue lives as - # long as the Nx.Serving module_state that holds this compiled function. - fn inputs -> EMLX.CommandQueue.with_queue(queue, fn -> inner.(inputs) end) end + # Evals every EMLX-backed output tensor on the currently-bound worker so the + # result is materialised (not a dangling lazy graph) once the queue unbinds. + defp force_eval_outputs(result) do + result + |> List.wrap() + |> Enum.each(fn container -> + Nx.Defn.Composite.traverse(container, fn + %Nx.Tensor{data: %EMLX.Backend{ref: ref}} = tensor -> + eval(ref) + tensor + + leaf -> + leaf + end) + end) + end + + # Routes a traced expression to the right eval-closure builder. `while` is + # the only remaining structural split point (`Nx.Defn.Graph`) — the loop + # runs from Elixir while every straight-line segment still compiles to a + # single-NIF native program. `:runtime_call` no longer splits the graph: + # it lowers in-graph to a real `:runtime_call` opcode backed by a genuine + # `mx::core::Primitive` (see `EMLX.Native.Expr`'s moduledoc "Runtime + # calls" section and `emlx_compiler.cpp`'s `EMLXRuntimeCall`), whose + # callback fires from `await_worker/2`'s `:emlx_runtime_call` receive + # clause while the single `eval_program` NIF call for that stage is still + # in flight. An `EMLX.Fast.*` fused kernel is not a `:runtime_call` node at + # all (it's a plain `:metadata`-tagged expr — see `EMLX.Fast`'s + # moduledoc), so it never splits either; it lowers in-graph to a single + # fused opcode (`EMLX.Native.Expr`'s `:metadata` clause). + # + # * no split point in the parent scope -> one flat native program + # (possibly containing one or more inlined `:runtime_call` nodes). + # * a bare tail `while` (base case) -> host-driven; the condition/body + # run by re-entering this compiler, so a nested split point recurses + # through the same path. + # * a split point with surrounding work -> `Nx.Defn.Graph.split/2` on + # every `while`, replayed by `Nx.Defn.Graph.run/3` with + # `compiler: EMLX`; each stage re-enters this compiler (flat stages + # compile flat, isolated split-point stages hit the base case above). + defp build_eval_fn(output_expr, worker, effective_device, num_inputs, fun, vars) do + cond do + bare_while?(output_expr) -> + build_while_base_eval_fn(output_expr, effective_device) + + not contains_split_point?(output_expr) -> + base_key = dispatch_key(output_expr, num_inputs, effective_device) + + {resource, hooks, runtime_calls} = + get_or_compile_program( + base_key, + %{}, + fn -> output_expr end, + num_inputs, + worker, + effective_device + ) + + build_native_eval_fn( + base_key, + resource, + hooks, + runtime_calls, + output_expr, + num_inputs, + effective_device, + fun, + vars + ) + + true -> + build_split_chain_eval_fn(output_expr, effective_device) + end + end + + # True when the parent scope contains a `while` split point. `post_order/2` + # treats it as an opaque leaf, so this only sees parent-scope split + # points — a nested one inside a sub-scope surfaces when that sub-scope is + # compiled. + defp contains_split_point?(output_expr) do + output_expr + |> EMLX.Defn.Tree.post_order(&EMLX.Native.Expr.scope_dependencies/1) + |> Enum.any?(&split_point?/1) + end + + defp split_point?(%Nx.Tensor{data: %Nx.Defn.Expr{op: :while}}), do: true + defp split_point?(%Nx.Tensor{}), do: false + + # Compiles an `EMLX.Native.Expr` program to a NIF resource on `worker`. + # Captured host tensors are copied onto `device` first: a `defn`-embedded + # constant tensor (e.g. an RNG algorithm constant) is traced with the default + # backend, so it must be moved to EMLX before `to_native` can extract its ref. + defp compile_native_program(worker, device, %EMLX.Native.Expr{} = program) do + program = ensure_emlx_captures(program, device) + wire_program = EMLX.Native.Expr.to_native(program) + + EMLX.NIF.compile_program(worker, wire_program) + |> unwrap!() + |> await_worker() + end + + # Ensures every captured tensor is EMLX-backed on `device` (copies any that + # were traced with another backend), so `to_native` can read a NIF ref per + # capture. + defp ensure_emlx_captures(%EMLX.Native.Expr{captures: captures} = program, device) do + captures = + Enum.map(captures, fn + {ref, %Nx.Tensor{data: %EMLX.Backend{}} = tensor} -> + {ref, tensor} + + {ref, %Nx.Tensor{} = tensor} -> + {ref, Nx.backend_copy(tensor, {EMLX.Backend, device: device})} + end) + + %{program | captures: captures} + end + + # Materialises defn input lazy refs to real bound %Nx.Tensor{} values on + # `dev` (copying any non-EMLX-backed tensor). + 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}) + end + end) + end + + # Derives the "quantization signature" a call's bound inputs need. A + # quantized weight's Nx-visible `.shape`/`.type` deliberately mirror its + # logical dense shape (so eager `Nx.dot` can transparently reroute to + # `EMLX.quantized_matmul` via `EMLX.Backend.dot/7`'s runtime dispatch) — + # invisible to `EMLX.Native.Expr.lower/2` at trace time. Once real tensors + # are bound (here, at call time) their `quantization_config` is visible, + # so a specialized program can be built that lowers + # the specific `:dot` nodes consuming these positions to + # `:quantized_matmul` instead of plain `:dot`. Returns `%{}` when nothing + # is quantized — the common, zero-overhead case. + # + # `allowed_positions` (`EMLX.Native.Expr.quantizable_param_positions/1`, + # computed once from `output_expr`) restricts this to positions actually + # consumed as a `:dot` right operand somewhere in the program: a + # quantized tensor merely *passed through* to something else (e.g. an + # `EMLX.Fast.*` fused kernel's `:__EMLX__` metadata operands, or + # `EMLX.Quantization.dequantize/1`'s own `:runtime_call` operand) is never + # specialized on, so it must not affect the cache key + # either — otherwise every distinct quantized tensor identity would mint + # its own permanently-unreused dispatch-cache entry for an + # otherwise structurally-identical program. + defp quant_signature(tensors, allowed_positions) do + tensors + |> Enum.with_index() + |> Enum.reduce(%{}, fn {tensor, pos}, sig -> + cond do + not MapSet.member?(allowed_positions, pos) -> sig + is_nil(tensor.data.quantization_config) -> sig + true -> Map.put(sig, pos, tensor.data.quantization_config) + end + end) + end + + defp input_refs(tensors) do + Enum.map(tensors, fn %Nx.Tensor{data: %EMLX.Backend{ref: {_dev, ref}}} -> ref end) + end + + # Builds the per-call eval closure for the flat (no-while) native path. + # `output_expr` is only used *here*, up front, to derive `output_template`, + # `real_output_count`, `quantizable_positions`, and `output_param_positions` + # (each bounded by output/parameter count, not graph size) — the returned + # closure itself never captures `output_expr` directly. The one exception + # is the rare quantized-specialization branch below, which needs the full + # expression to lower+compile a specialized program; it re-derives + # `output_expr` on demand via `fun.(vars)` instead of keeping it resident + # in every call's closure environment — see `get_or_compile_program/6`. + # `base_key` is the structural dispatch key (`dispatch_key/3`), threaded + # through so a quantized specialization is cached persistently too, not + # just the plain program. + # `output_template` (derived from `output_expr` above) serves as the + # type/shape template for reconstructing output tensors after the NIF + # returns raw resource refs. `plain_hooks` + # (from `EMLX.Native.Expr.lower/2`'s `hooks` field, see its moduledoc + # "Hooks" section) ride along as extra outputs after the real ones; the + # corresponding Elixir callbacks fire here, once, right after the single + # NIF call returns. `plain_runtime_calls` (from the `runtime_calls` field, + # see its moduledoc "Runtime calls" section) is threaded down to + # `await_worker/2`, which fires each one's real callback *during* the + # single `eval_program` NIF call, as its `:emlx_runtime_call` requests + # arrive. + defp build_native_eval_fn( + base_key, + plain_resource, + plain_hooks, + plain_runtime_calls, + output_expr, + num_inputs, + effective_device, + fun, + vars + ) do + output_template = Nx.Defn.Composite.traverse(output_expr, &Nx.to_template/1) + real_output_count = [output_template] |> Nx.Defn.Composite.flatten_list() |> length() + + # See `quant_signature/2`'s doc for why this must be restricted to + # positions actually consumed by a `:dot` — otherwise a quantized + # tensor merely passed to e.g. an `EMLX.Fast.*` fused kernel would + # fragment the dispatch cache with one dead-weight entry per + # distinct tensor identity. + quantizable_positions = EMLX.Native.Expr.quantizable_param_positions(output_expr) + + # Static (independent of quant_signature): which output leaves are a bare, + # untouched pass-through of a parameter position — e.g. a loop-invariant + # carry threaded across an Nx.Defn.Graph while-split stage boundary, never + # consumed by any op in *this* stage. Flat, parallel to out_refs below + # (both walk output_expr/output_template in the same Composite order). + output_param_positions = + [output_expr] + |> Nx.Defn.Composite.flatten_list() + |> Enum.map(fn + %Nx.Tensor{data: %Nx.Defn.Expr{op: :parameter, args: [pos]}} -> pos + _ -> nil + end) + + fn [params] -> + {worker, dev} = resolve_worker(effective_device) + tensors = materialise_input_tensors(params, dev) + quant_signature = quant_signature(tensors, quantizable_positions) + + # The plain (non-quantized) resource was already looked up/compiled + # (and cached under `{base_key, %{}}`) by our caller, so the common + # case skips a redundant cache round-trip; a quantized signature is + # looked up (or compiled once and cached) here — see + # get_or_compile_program/6. `fn -> fun.(vars) end` is only forced on + # that cache's own miss (once per distinct `{base_key, quant_signature}` + # for the process's lifetime) — never eagerly, so a cache *hit* here + # costs nothing beyond the lookup itself. + {program_resource, hooks, runtime_calls} = + if quant_signature == %{} do + {plain_resource, plain_hooks, plain_runtime_calls} + else + get_or_compile_program( + base_key, + quant_signature, + fn -> fun.(vars) end, + num_inputs, + worker, + dev + ) + end + + job_ref = + EMLX.NIF.eval_program(worker, program_resource, input_refs(tensors)) + |> unwrap!() + + all_refs = await_worker(job_ref, runtime_calls, tensors, dev) + + {out_refs, hook_refs} = Enum.split(all_refs, real_output_count) + + # Reconstruct output tensors: traverse the output expression template + # (provides type/shape/names) and attach each returned resource ref as + # a proper EMLX.Backend data value. Exception: a quantized-parameter + # pass-through leaf (see output_param_positions above) — its Nx-visible + # shape/type is a logical fiction (see quant_signature/2) that does not + # match the physical (packed) array the NIF actually returns for an + # untouched pass-through, so EMLX.Backend.to_nx/2 would raise a shape + # mismatch. Substitute back the original bound tensor (quantization_config + # and all) instead for those leaves. Checked against the tensor's own + # `quantization_config` directly rather than `quant_signature` — the + # latter is deliberately narrowed to positions a `:dot` node actually + # consumes (see `quantizable_param_positions/1`'s doc), but a + # loop-invariant quantized carry passed straight through untouched + # (this leaf) need not be consumed by any `:dot` in this stage at all. + {output_container, {[], []}} = + Nx.Defn.Composite.traverse( + output_template, + {out_refs, output_param_positions}, + fn leaf, {[ref | refs_rest], [pos | pos_rest]} -> + emlx_tensor = + case pos && Enum.at(tensors, pos) do + %Nx.Tensor{data: %EMLX.Backend{quantization_config: %EMLX.Quantization.Config{}}} = + t -> + t + + _ -> + EMLX.Backend.to_nx({dev, ref}, leaf) + end + + {emlx_tensor, {refs_rest, pos_rest}} + end + ) + + fire_hooks(hooks, hook_refs, dev) + + [output_container] + end + end + + # Looks up (or lazily lowers+compiles) the program for `{base_key, + # quant_signature}` in the process-lifetime dispatch cache. + # `base_key` (see `dispatch_key/3`) is a structural signature of the + # `Expr` — stable across `Nx.Defn.Graph.run/3`'s per-call re-tracing and + # shared across structurally-identical call sites (e.g. every one of + # Qwen3's 28 attention layers's surrounding stages), not scoped to one + # `build_native_eval_fn/6` closure the way this cache was pre-Stage-32. + # First-compile-wins under concurrent calls with a never-before-seen key: + # `:ets.insert_new/2` returning `false` means another caller raced us and + # won — we discard our (still valid, just unused) compiled resource and + # reuse theirs. Wasted work on that race, not a correctness hazard (no + # shared mutable state is corrupted). + # + # `output_expr_thunk` (a 0-arity fun, not `output_expr` itself) is only + # forced on this cache's own miss (the `[]` branch below) — callers that + # already have `output_expr` in hand (the plain, non-quantized path) pay + # nothing extra for the indirection, but the quantized-specialization + # caller (`build_native_eval_fn/9`) that re-derives `output_expr` via + # `fun.(vars)` only actually retraces when this cache misses, not on + # every call. + defp get_or_compile_program( + base_key, + quant_signature, + output_expr_thunk, + num_inputs, + worker, + dev + ) do + cache_key = {base_key, quant_signature} + table = dispatch_cache_table() + + case :ets.lookup(table, cache_key) do + [{_key, resource, hooks, runtime_calls}] -> + {resource, hooks, runtime_calls} + + [] -> + program = EMLX.Native.Expr.lower(output_expr_thunk.(), num_inputs, quant_signature) + resource = compile_native_program(worker, dev, program) + + if :ets.insert_new(table, {cache_key, resource, program.hooks, program.runtime_calls}) do + {resource, program.hooks, program.runtime_calls} + else + [{_key, winner_resource, winner_hooks, winner_runtime_calls}] = + :ets.lookup(table, cache_key) + + {winner_resource, winner_hooks, winner_runtime_calls} + end + end + end + + # Lazily creates (idempotently — races are resolved by `:ets.new/2` raising + # `ArgumentError` on the loser, which we swallow) the named, public, + # process-lifetime ETS table backing the dispatch cache. Named + # (not `:persistent_term`-stashed) so no GC-triggering `:persistent_term` + # writes are needed to publish it — the atom name is the handle. + defp dispatch_cache_table, do: ensure_named_ets_table(@native_dispatch_cache_table) + + # Lazily creates (idempotently — races are resolved by `:ets.new/2` raising + # `ArgumentError` on the loser, which we swallow) a named, public, + # process-lifetime ETS table. Named (not `:persistent_term`-stashed) so no + # GC-triggering `:persistent_term` writes are needed to publish it — the + # atom name is the handle. + defp ensure_named_ets_table(name) do + case :ets.whereis(name) do + :undefined -> + try do + :ets.new(name, [:named_table, :public, :set, read_concurrency: true]) + rescue + ArgumentError -> :ok + end + + name + + _tid -> + name + end + end + + # A structural, id-independent signature of `output_expr` — the cache key + # `get_or_compile_program/6` dispatches on (paired with a call-time + # `quant_signature`). Two `Expr`s that are the *same shape of + # computation* (same op sequence, shapes, dtypes, and static attrs) hash to + # the same key even though `Nx.Defn.Expr` assigns each a fresh `id` per + # trace — which is what lets Qwen3's 28 structurally-identical attention + # layers (or the same layer across decode steps) share one compiled + # program instead of recompiling per call. Coarser than a bit-identical + # `Expr` hash (deliberately): tensor-valued args are replaced by their position in + # `EMLX.Defn.Tree.post_order/2`'s dependency-first listing (itself + # id-independent and structurally stable across identical call sites), + # not by the referenced node's `id`. + # `sanitize_key_term/2`'s opaque-scope fallback (below) recurses into a + # `while`/`block`/`fun` sub-scope's own structural signature. Real models + # reuse the *same* shared sub-expression (e.g. a RoPE frequency table, a + # causal-mask block) across many call sites, so without memoization this + # recomputes an identical signature once per reference — the same + # unmemoized-shared-subexpression blowup `nx-graph-split-bugreport.md`'s + # Bug 1 hit in `rewrite_subtree`. `@dispatch_key_memo` is a process-local + # key => signature cache, live only for the duration of one `dispatch_key/3` + # call (cleared in the `after`), so it's safe to reuse across unrelated + # calls without stale entries leaking. + @dispatch_key_memo_pdict_key {__MODULE__, :dispatch_key_memo} + + defp dispatch_key(output_expr, num_inputs, device) do + id_key = {expr_id_fingerprint(output_expr), num_inputs, device} + table = ensure_named_ets_table(@dispatch_key_by_id_table) + + case :ets.lookup(table, id_key) do + [{_key, base_key}] -> + base_key + + [] -> + base_key = compute_dispatch_key(output_expr, num_inputs, device) + :ets.insert(table, {id_key, base_key}) + base_key + end + end + + # Cheap (O(number of output leaves), not O(nodes)) identity fingerprint — + # see `@dispatch_key_by_id_table`. Distinct traces never share an `Expr` + # node id, so matching ids here guarantee the exact same graph. + defp expr_id_fingerprint(output_expr) do + [output_expr] + |> Nx.Defn.Composite.flatten_list() + |> Enum.map(fn %Nx.Tensor{data: %Nx.Defn.Expr{id: id}} -> id end) + end + + defp compute_dispatch_key(output_expr, num_inputs, device) do + Process.put(@dispatch_key_memo_pdict_key, %{}) + + try do + {node_sigs, output_sigs} = expr_structural_signature(output_expr) + {node_sigs, output_sigs, num_inputs, device} + after + Process.delete(@dispatch_key_memo_pdict_key) + end + end + + # Computes `{node_sigs, output_sigs}` for one `EMLX.Defn.Tree.post_order/2` + # scope. Split out from `dispatch_key/3` so `sanitize_key_term/2` can + # recurse into an *opaque* scope root (a `while` condition/body, a + # `block`'s `default_expr`, a `fun` body — see `EMLX.Defn.Tree.post_order/2`'s + # "Sub-scope handling" doc) with a fresh, self-contained position map, + # instead of assuming every tensor reachable from a visited node's `args` + # was itself visited at this level. + defp expr_structural_signature(expr_or_container) do + nodes = EMLX.Defn.Tree.post_order(expr_or_container, &EMLX.Native.Expr.scope_dependencies/1) + positions = nodes |> Enum.with_index() |> Map.new(fn {t, i} -> {t.data.id, i} end) + + node_sigs = + Enum.map(nodes, fn + %Nx.Tensor{data: %Nx.Defn.Expr{op: :metadata, args: [_inner, %{__EMLX__: emlx} = meta]}} = + t + when map_size(meta) == 1 -> + # `_inner` (the plain-Nx reference formula) is deliberately excluded + # from the signature, not just skipped like `EMLX.Defn.Tree.post_order/2` + # (via `EMLX.Native.Expr.scope_dependencies/1`) does for lowering: `_inner` + # closes over real upstream operands (e.g. + # `x` in `rms_norm_reference/3`), so treating it like any other arg + # here would make `sanitize_key_term/2` fall into its "opaque + # sub-scope" fallback below and re-walk `_inner`'s *entire* transitive + # dependency graph (i.e. every prior layer) from scratch — once per + # fused-kernel call site. Signing only the real `__EMLX__` payload + # (which is exactly what `EMLX.Native.Expr.lower/2` actually lowers) + # keeps this O(1) per node instead of O(depth) per node. + {:__emlx__, t.shape, t.type, sanitize_key_term(emlx, positions)} + + %Nx.Tensor{ + data: %Nx.Defn.Expr{op: :block, args: [struct, in_args, _default_expr, _fun] = args} + } = t -> + # See `EMLX.Native.Expr.native_lowerable_block?/2`'s doc: recognized + # native blocks (e.g. `Nx.LinAlg.svd/1`'s `full_matrices?: true`) + # never consult `default_expr` when lowering, so `default_expr` is + # deliberately excluded from the signature too — walking/hashing it + # (potentially a ~100-node Jacobi-rotation fallback graph, per call) + # would be pure waste for a sub-scope that provably can't affect the + # output. Anything not recognized falls back to signing the full + # `args` (including `default_expr`, via the opaque-scope path in + # `sanitize_key_term/2` below) exactly as before — the conservative + # default. + if EMLX.Native.Expr.native_lowerable_block?(struct, in_args) do + {:block, t.shape, t.type, struct, sanitize_key_term(in_args, positions)} + else + {:block, t.shape, t.type, sanitize_key_term(args, positions)} + end + + %Nx.Tensor{data: %Nx.Defn.Expr{op: op, args: args}} = t -> + {op, t.shape, t.type, sanitize_key_term(args, positions)} + end) + + output_sigs = + [expr_or_container] + |> Nx.Defn.Composite.flatten_list() + |> Enum.map(fn %Nx.Tensor{data: %Nx.Defn.Expr{id: id}} -> Map.fetch!(positions, id) end) + + {node_sigs, output_sigs} + end + + # Recursively strips an `Nx.Defn.Expr` node's `args` down to a hashable, + # id-independent term for `dispatch_key/3`: an in-scope `%Nx.Tensor{}` + # operand becomes its `post_order/1` position (see above); a tensor that + # belongs to an opaque sub-scope (not in `positions` at all — a `while` + # condition/body, a `block`'s `default_expr`, a `fun` body) recurses into + # its own self-contained structural signature; an `Nx.TemplateBackend`-backed + # tensor riding as a static arg (e.g. a `:runtime_call`'s `out_template`) + # becomes its shape/type, since it has no `id`/scope of its own; a captured + # function (e.g. a `:runtime_call`'s callback, or a hook's callback) + # becomes its `{module, name, arity}` — identity without hashing its + # closure environment. Everything else (numbers, atoms, strings, structs + # used as plain option carriers) is kept as-is. + # + # An opaque sub-scope's signature is condensed to a `:crypto.hash/2` digest + # (`{:scope, digest}`) rather than embedded verbatim: `@dispatch_key_memo_pdict_key` + # only dedups *building* the raw signature once per `id`, but the same + # (interned, i.e. shared-by-reference) large signature term can still be + # *referenced* from many call sites — e.g. Axon's own `axon_layer:` + # metadata wraps every layer, so a late layer's real (non-`__EMLX__`) + # metadata `inner` recurses into an opaque scope covering that layer's + # entire upstream history. `:ets.lookup/2`'s key hash walks every logical + # occurrence of a subterm regardless of sharing, so embedding the raw + # signature at O(occurrences) call sites made the final cache-key hash + # O(occurrences × signature size) — cheap to *build* (reference reuse) but + # catastrophically slow to *hash* on a real 28-layer model. Digesting once + # per unique `id` and reusing the (small, fixed-size) digest everywhere + # keeps both build and hash O(unique upstream work). + defp sanitize_key_term(%Nx.Tensor{data: %Nx.Defn.Expr{id: id}} = t, positions) do + case Map.fetch(positions, id) do + {:ok, pos} -> + {:ref, pos} + + :error -> + memo = Process.get(@dispatch_key_memo_pdict_key, %{}) + + case Map.fetch(memo, id) do + {:ok, sig} -> + sig + + :error -> + digest = :crypto.hash(:sha256, :erlang.term_to_binary(expr_structural_signature(t))) + sig = {:scope, digest} + Process.put(@dispatch_key_memo_pdict_key, Map.put(memo, id, sig)) + sig + end + end + end + + defp sanitize_key_term(%Nx.Tensor{} = t, _positions) do + {:tensor_literal, t.shape, t.type} + end + + defp sanitize_key_term(fun, _positions) when is_function(fun) do + info = Function.info(fun) + {:fun, info[:module], info[:name], info[:arity]} + end + + defp sanitize_key_term(tuple, positions) when is_tuple(tuple) do + tuple |> Tuple.to_list() |> sanitize_key_term(positions) |> List.to_tuple() + end + + defp sanitize_key_term(list, positions) when is_list(list) do + Enum.map(list, &sanitize_key_term(&1, positions)) + end + + defp sanitize_key_term(map, positions) when is_map(map) and not is_struct(map) do + Map.new(map, fn {k, v} -> {k, sanitize_key_term(v, positions)} end) + end + + defp sanitize_key_term(%_struct{} = struct, positions) do + {struct.__struct__, struct |> Map.from_struct() |> sanitize_key_term(positions)} + end + + defp sanitize_key_term(other, _positions), do: other + + # Reconstructs each hook's value from its slice of `hook_refs` (in + # `hooks` order, matching `EMLX.Native.Expr.to_native/1`'s flattening) and + # invokes its callback for the side effect. Return value is discarded, + # matching `Nx.Defn.Evaluator`'s hook semantics. + defp fire_hooks([], [], _dev), do: :ok + + defp fire_hooks([%{template: template, refs: refs, callback: callback} | rest], hook_refs, dev) do + {consumed, remaining} = Enum.split(hook_refs, length(refs)) + + {value, []} = + Nx.Defn.Composite.traverse(template, consumed, fn leaf, [ref | more] -> + {EMLX.Backend.to_nx({dev, ref}, leaf), more} + end) + + callback.(value) + fire_hooks(rest, remaining, dev) + end + + # Builds the eval closure for a `while` split point surrounded by other + # computation. The expression is split on every `while` node and replayed + # by `Nx.Defn.Graph`: `compiler: EMLX` makes every stage re-enter this + # compiler, so straight-line stages compile flat (any `:runtime_call` in + # them lowers in-graph — see `build_eval_fn/4`) and the isolated `while` + # stage hits `build_while_base_eval_fn/2`. `device:` keeps stage + # compilation on the same device; the command queue (if any) is + # propagated through the process binding set by the outer wrapper. + defp build_split_chain_eval_fn(output_expr, effective_device) do + stages = Nx.Defn.Graph.split(output_expr, &if(split_point?(&1), do: :both, else: :none)) + + fn [params] -> + {_worker, dev} = resolve_worker(effective_device) + inputs = Enum.map(params, &materialise_tensor(&1, dev)) + result = Nx.Defn.Graph.run(stages, inputs, compiler: __MODULE__, device: effective_device) + [result] + end + end + + # Builds the eval closure for the base case: a bare tail `while` whose initial + # carry is exactly the stage inputs (every output leaf is the `while` node or + # an `:elem` of it). The condition and body are compiled by re-entering this + # compiler via `Nx.Defn.jit/2`, so a nested `while` in the body recurses + # through the same splitting machinery. The loop itself is driven host-side. + defp build_while_base_eval_fn(output_expr, effective_device) do + while_node = find_while_node(output_expr) + [initial, _arg, cond_expr, body_expr] = while_node.data.args + + jit_opts = [compiler: __MODULE__, device: effective_device] + cond_fn = Nx.Defn.jit(fn _ -> cond_expr end, jit_opts) + body_fn = Nx.Defn.jit(fn _ -> body_expr end, jit_opts) + + # The stage inputs arrive in stage-argument order, which need not match the + # carry (sub-scope parameter) order: each flattened `initial` leaf is the + # parameter whose position picks the stage input feeding that carry slot. + # Reorder inputs into carry order so the condition/body parameters bind to + # the right tensors. + initial_positions = + [initial] + |> Nx.Defn.Composite.flatten_list() + |> Enum.map(fn %Nx.Tensor{data: %Nx.Defn.Expr{op: :parameter, args: [pos]}} -> pos end) + + while_id = while_node.data.id + output_flat = Nx.Defn.Composite.flatten_list([output_expr]) + carry_indices = Enum.map(output_flat, &while_output_index(&1, while_id)) + output_template = Nx.Defn.Composite.traverse(output_expr, &Nx.to_template/1) + + fn [params] -> + {_worker, dev} = resolve_worker(effective_device) + inputs = Enum.map(params, &materialise_tensor(&1, dev)) + carry = Enum.map(initial_positions, &Enum.at(inputs, &1)) + final_carry = run_while_loop(carry, cond_fn, body_fn) + + output_tensors = Enum.map(carry_indices, &Enum.at(final_carry, &1)) + + {output_container, []} = + Nx.Defn.Composite.traverse(output_template, output_tensors, fn _leaf, [t | rest] -> + {t, rest} + end) + + [output_container] + end + end + + # Host-driven while loop. `carry` is the flat list of carry tensors; it is + # passed to the compiled condition/body as a single tuple argument (mirroring + # how `Nx.Defn.Graph.run` invokes a stage), which jit re-flattens to match the + # carry parameters of the sub-scope. The body result is flattened back to a + # flat list so a nested-container carry stays aligned with the leaf indices. + defp run_while_loop(carry, cond_fn, body_fn) do + if Nx.to_number(cond_fn.(List.to_tuple(carry))) == 0 do + carry else - inner + new_carry = Nx.Defn.Composite.flatten_list([body_fn.(List.to_tuple(carry))]) + run_while_loop(new_carry, cond_fn, body_fn) + end + end + + defp find_while_node(output_expr) do + output_expr + |> EMLX.Defn.Tree.post_order(&EMLX.Native.Expr.scope_dependencies/1) + |> Enum.find(&(&1.data.op == :while)) + end + + # True for the base case: the output projects exactly one `while` (each leaf + # is the `while` node or an `:elem` of it) and that `while`'s initial carry is + # made entirely of parameters — i.e. all pre-loop work has already been split + # into an earlier stage, so the carry is the stage input as-is. + defp bare_while?(output_expr) do + leaves = Nx.Defn.Composite.flatten_list([output_expr]) + + while_ids = + leaves + |> Enum.flat_map(fn + %Nx.Tensor{data: %Nx.Defn.Expr{op: :while, id: id}} -> + [id] + + %Nx.Tensor{ + data: %Nx.Defn.Expr{ + op: :elem, + args: [%Nx.Tensor{data: %Nx.Defn.Expr{op: :while, id: id}}, _] + } + } -> + [id] + + _ -> + [] + end) + |> Enum.uniq() + + case while_ids do + [wid] -> + all_project = + Enum.all?(leaves, fn + %Nx.Tensor{data: %Nx.Defn.Expr{op: :while, id: ^wid}} -> + true + + %Nx.Tensor{ + data: %Nx.Defn.Expr{op: :elem, args: [%Nx.Tensor{data: %Nx.Defn.Expr{id: ^wid}}, _]} + } -> + true + + _ -> + false + end) + + all_project and while_initial_all_params?(find_while_node(output_expr)) + + _ -> + false + end + end + + defp while_initial_all_params?(%Nx.Tensor{data: %Nx.Defn.Expr{args: [initial | _]}}) do + [initial] + |> Nx.Defn.Composite.flatten_list() + |> Enum.all?(&match?(%Nx.Tensor{data: %Nx.Defn.Expr{op: :parameter}}, &1)) + end + + # Maps a flat output leaf to the carry index it projects from `while_id`. + defp while_output_index(%Nx.Tensor{data: %Nx.Defn.Expr{op: :while, id: id}}, while_id) + when id == while_id, + do: 0 + + defp while_output_index( + %Nx.Tensor{ + data: %Nx.Defn.Expr{op: :elem, args: [%Nx.Tensor{data: %Nx.Defn.Expr{id: id}}, i]} + }, + while_id + ) + when id == while_id, + do: i + + # Materialises a defn input lazy ref to a concrete EMLX-backed tensor on `dev` + # (for use as a `Nx.Defn.Graph.run` / jitted-stage argument). + defp materialise_tensor(lazy, dev) do + case lazy.() do + %Nx.Tensor{data: %EMLX.Backend{}} = tensor -> tensor + %Nx.Tensor{} = tensor -> Nx.backend_copy(tensor, {EMLX.Backend, device: dev}) end end @@ -2000,12 +2504,6 @@ defmodule EMLX do Application.get_env(:emlx, :default_device, :gpu) end - # Splits opts into {emlx_compiler_opts, rest_opts}. The rest_opts are - # forwarded to Nx.Defn.Evaluator; EMLX-specific keys are consumed here. - defp split_compiler_opts(opts) do - Enum.split_with(opts, fn {k, _v} -> k in @valid_compiler_keys end) - end - @doc """ Returns a map with current memory usage information. @@ -2033,7 +2531,6 @@ defmodule EMLX do EMLX.clear_cache() #=> :ok """ - @spec clear_cache() :: :ok def clear_cache, do: EMLX.NIF.clear_cache() |> unwrap!() @doc """ @@ -2044,7 +2541,6 @@ defmodule EMLX do EMLX.reset_peak_memory() #=> :ok """ - @spec reset_peak_memory() :: :ok def reset_peak_memory, do: EMLX.NIF.reset_peak_memory() |> unwrap!() @doc """ diff --git a/emlx/lib/emlx/application.ex b/emlx/lib/emlx/application.ex index e839329..706fd0b 100644 --- a/emlx/lib/emlx/application.ex +++ b/emlx/lib/emlx/application.ex @@ -8,6 +8,11 @@ defmodule EMLX.Application do `EMLX.to_blob/1` for any process that has not bound its own queue via `EMLX.CommandQueue.with_queue/2`. + Also allocates one *runtime_call worker* per device — see + `runtime_call_worker/1`'s doc for why a `Nx.runtime_call/4` callback's own + eager EMLX calls must never route back to the worker that is, right now, + blocked waiting for that very callback to return. + See `clean-room-import/01-worker-thread-dispatch.md` for the rationale behind `:persistent_term` instead of a `GenServer` + `Registry`. @@ -28,6 +33,13 @@ defmodule EMLX.Application do this module silently skips the GPU worker. Subsequent `EMLX.eval/1` calls on a GPU tensor will raise at use time with the underlying `:persistent_term` `ArgumentError`. + + ## CPU JIT compilation and SIGCHLD + + Hitting `** (EMLX.NIFError) ... pclose() failed.` on the CPU backend? See + the "CPU JIT compilation and SIGCHLD" section of the `EMLX` moduledoc — + it's a BEAM/MLX interaction this application deliberately does not paper + over for you. """ use Application @@ -37,6 +49,8 @@ defmodule EMLX.Application do EMLX.Profiling.init() ensure_default_worker!(:cpu, _gpu_optional? = false) ensure_default_worker!(:gpu, _gpu_optional? = true) + ensure_worker!(:runtime_call_worker, :cpu, _gpu_optional? = false) + ensure_worker!(:runtime_call_worker, :gpu, _gpu_optional? = true) Supervisor.start_link([], strategy: :one_for_one, name: __MODULE__) end @@ -51,13 +65,46 @@ defmodule EMLX.Application do (`{EMLX, :default_worker, device}`) — overwriting a `:persistent_term` value triggers a node-wide GC. """ - @spec default_worker(:cpu | :gpu) :: reference() def default_worker(device) when device in [:cpu, :gpu] do - :persistent_term.get(persistent_term_key(device)) + :persistent_term.get(persistent_term_key(:default_worker, device)) + end + + @doc """ + Returns the `EMLX.CommandQueue` NIF resource dedicated to running + `Nx.runtime_call/4` callbacks for the given device. + + `EMLX.handle_runtime_call/5` binds *this* worker (via + `EMLX.CommandQueue.with_queue/2`'s same `Process.put(:emlx_command_queue, + _)` mechanism) for the duration of every real callback invocation, so any + eager EMLX call the callback itself makes (e.g. + `EMLX.Quantization.dequantize_callback/2` calling `EMLX.dequantize/1`) + gets worker-routed *here* instead of to the default (or bound) worker that + is, right now, blocked inside `EMLXRuntimeCall::eval_cpu`/`eval_gpu` + waiting for that very callback to finish. Reusing that same, still-busy + worker would either deadlock (its one dedicated OS thread cannot service + a new job while blocked — see `emlx_worker.hpp`) or, if the blocked wait + instead pumps that worker's own queue inline, crash (MLX's GPU `eval()` + is not reentrant on one OS thread — nesting a second `mlx::core::eval` + call while the first is still on the C++ stack corrupts Metal + command-buffer/completion-handler state). A second, otherwise-idle + worker sidesteps both failure modes: it is a distinct OS thread with its + own `mlx::core::Stream`, so its jobs never contend with (or nest inside) + the blocked worker's. + + Raises `ArgumentError` if no worker has been allocated for the device + (e.g. asking for `:gpu` on a system without Metal) — same contract as + `default_worker/1`. + """ + def runtime_call_worker(device) when device in [:cpu, :gpu] do + :persistent_term.get(persistent_term_key(:runtime_call_worker, device)) end defp ensure_default_worker!(device, gpu_optional?) do - key = persistent_term_key(device) + ensure_worker!(:default_worker, device, gpu_optional?) + end + + defp ensure_worker!(kind, device, gpu_optional?) do + key = persistent_term_key(kind, device) case :persistent_term.get(key, :unset) do :unset -> @@ -70,7 +117,7 @@ defmodule EMLX.Application do {:error, reason} -> raise EMLX.NIFError, - "EMLX.Application could not allocate default #{device} worker: " <> + "EMLX.Application could not allocate #{kind} #{device} worker: " <> List.to_string(reason) end @@ -79,5 +126,5 @@ defmodule EMLX.Application do end end - defp persistent_term_key(device), do: {EMLX, :default_worker, device} + defp persistent_term_key(kind, device), do: {EMLX, kind, device} end diff --git a/emlx/lib/emlx/backend.ex b/emlx/lib/emlx/backend.ex index f558ba3..2fe9173 100644 --- a/emlx/lib/emlx/backend.ex +++ b/emlx/lib/emlx/backend.ex @@ -4,15 +4,15 @@ defmodule EMLX.Backend do alias Nx.Tensor, as: T alias EMLX.Backend, as: Backend - require Logger + require EMLX.Debug + import EMLX.Debug, only: [assert_no_nan_inf!: 2] - # Compile-time debug flags. Both default to false (zero runtime cost when off). + # Compile-time debug flag. Defaults to false (zero runtime cost when off). # Enable only in development via config/dev.exs: # config :emlx, enable_bounds_check: true - # config :emlx, detect_non_finites: true - # After flipping a flag run: mix compile --force + # After flipping run: mix compile --force + # (:detect_non_finites lives in EMLX.Debug, shared with EMLX.Fast.) @enable_bounds_check Application.compile_env(:emlx, :enable_bounds_check, false) - @detect_non_finites Application.compile_env(:emlx, :detect_non_finites, false) # Each macro expands to the assertion body when its flag is true, or to nil # when false — leaving no trace in the BEAM opcodes (no call instruction, @@ -79,25 +79,6 @@ defmodule EMLX.Backend do end end - # Checks a raw MLX ref for NaN or Inf. Forces two eval syncs — development only. - defmacrop assert_no_nan_inf!(tensor_ref, op) do - if @detect_non_finites do - quote do - has_nan = unquote(tensor_ref) |> EMLX.is_nan() |> EMLX.any([], false) - has_inf = unquote(tensor_ref) |> EMLX.is_infinity() |> EMLX.any([], false) - EMLX.eval(has_nan) - EMLX.eval(has_inf) - - if EMLX.item(has_nan) == 1 or EMLX.item(has_inf) == 1 do - raise ArgumentError, - "#{unquote(op)} produced NaN or Inf. Disable :detect_non_finites for production." - end - end - else - quote do: nil - end - end - @moduledoc """ EMLX backend for Nx tensors, providing GPU acceleration via Apple's MLX. @@ -381,8 +362,10 @@ defmodule EMLX.Backend do @impl true def to_binary(tensor, limit) do - EMLX.to_blob(from_nx(tensor), limit) - |> maybe_modify_binary(to_nx_type(to_mlx_type(tensor.type)), tensor.type) + EMLX.Telemetry.span_to_binary(tensor, fn -> + EMLX.to_blob(from_nx(tensor), limit) + |> maybe_modify_binary(to_nx_type(to_mlx_type(tensor.type)), tensor.type) + end) end @impl true @@ -450,10 +433,16 @@ defmodule EMLX.Backend do end end - defp maybe_modify_binary(binary, {:u, size}, {:u, 8}) when size in [2, 4] do - for <>, into: <<>> do - <> - end + defp maybe_modify_binary(binary, {:u, 2}, {:u, 8}) do + for <>, + into: <<>>, + do: <> + end + + defp maybe_modify_binary(binary, {:u, 4}, {:u, 8}) do + for <>, + into: <<>>, + do: <> end defp maybe_modify_binary(binary, {:u, 8}, {:u, size}) when size in [2, 4] do @@ -701,26 +690,7 @@ defmodule EMLX.Backend do defp needs_type_conversion?({:u, 8}, :bool), do: true defp needs_type_conversion?(_, _), do: false - defp to_mlx_type({:u, 2}), do: :uint8 - defp to_mlx_type({:u, 4}), do: :uint8 - defp to_mlx_type({:u, 8}), do: :uint8 - defp to_mlx_type({:u, 16}), do: :uint16 - defp to_mlx_type({:u, 32}), do: :uint32 - defp to_mlx_type({:u, 64}), do: :uint64 - defp to_mlx_type({:s, 2}), do: :int8 - defp to_mlx_type({:s, 4}), do: :int8 - defp to_mlx_type({:s, 8}), do: :int8 - defp to_mlx_type({:s, 16}), do: :int16 - defp to_mlx_type({:s, 32}), do: :int32 - defp to_mlx_type({:s, 64}), do: :int64 - defp to_mlx_type({:f, 8}), do: :float16 - defp to_mlx_type({:f, 16}), do: :float16 - defp to_mlx_type({:f, 32}), do: :float32 - defp to_mlx_type({:f, 64}), do: :float32 - defp to_mlx_type({:bf, 16}), do: :bfloat16 - defp to_mlx_type({:c, 64}), do: :complex64 - defp to_mlx_type({:c, 128}), do: :complex64 - defp to_mlx_type(:bool), do: :bool + defp to_mlx_type(type), do: EMLX.Native.to_mlx_type(type) defp to_nx_type(:uint8), do: {:u, 8} defp to_nx_type(:uint16), do: {:u, 16} @@ -865,7 +835,6 @@ defmodule EMLX.Backend do # Get the actual shape after summation actual_shape = EMLX.shape(result) # FIXME: MLX returns whatever the original type is, but Nx expects u8 -> u32 - # scalar_type = EMLX.scalar_type(result) # Create a new output tensor with the correct shape %{out | shape: actual_shape} @@ -916,7 +885,7 @@ defmodule EMLX.Backend do # Apply reversal for tie_break after NaN check t_mx = if opts[:tie_break] == :high do - reverse_mlx(t_mx, tensor.shape, [axis] || Nx.axes(tensor)) + reverse_mlx(t_mx, tensor.shape, if(axis, do: [axis], else: Nx.axes(tensor))) else t_mx end @@ -1119,19 +1088,23 @@ defmodule EMLX.Backend do |> Enum.sort() |> Enum.map(&elem(&1, 1)) - input_mx - |> EMLX.conv_general( - kernel_mx, - strides, - padding_low, - padding_high, - kernel_dilation, - input_dilation, - feature_group_count - ) - |> EMLX.transpose(permute_channels_first) - |> EMLX.transpose(output_permutation) - |> to_nx(out) + result_ref = + input_mx + |> EMLX.conv_general( + kernel_mx, + strides, + padding_low, + padding_high, + kernel_dilation, + input_dilation, + feature_group_count + ) + |> EMLX.transpose(permute_channels_first) + |> EMLX.transpose(output_permutation) + + assert_no_nan_inf!(result_ref, :conv) + + to_nx(result_ref, out) end defp dot_spec_to_einsum_spec( @@ -1250,8 +1223,13 @@ defmodule EMLX.Backend do defp quantized_dot(out, activation, qw_tensor, right_axes) do %Backend{ref: weight_ref, quantization_config: cfg} = qw_tensor.data - %EMLX.Quantization.Config{scales: scales_nx, biases: biases_nx, group_size: gs, bits: bits} = - cfg + %EMLX.Quantization.Config{ + scales: scales_nx, + biases: biases_nx, + group_size: gs, + bits: bits, + mode: mode + } = cfg weight_rank = tuple_size(qw_tensor.shape) last_dim = weight_rank - 1 @@ -1265,15 +1243,20 @@ defmodule EMLX.Backend do explicit -> explicit end + # biases_nx is nil for microscaled modes (mxfp4/mxfp8/nvfp4) — mx::fp_quantize + # doesn't emit them. + biases_ref = biases_nx && from_nx(biases_nx) + result = EMLX.quantized_matmul( from_nx(activation), weight_ref, from_nx(scales_nx), - from_nx(biases_nx), + biases_ref, transpose, gs, - bits + bits, + mode ) to_nx(result, out) @@ -1310,8 +1293,10 @@ defmodule EMLX.Backend do ) EMLX.einsum( - to_typed_ref(left_mx, left_type, computation_out_type), - to_typed_ref(right_mx, right_type, computation_out_type), + [ + to_typed_ref(left_mx, left_type, computation_out_type), + to_typed_ref(right_mx, right_type, computation_out_type) + ], einsum_spec ) |> to_typed_ref(computation_out_type, out_type) @@ -2054,13 +2039,23 @@ defmodule EMLX.Backend do b_mx = to_typed_ref(from_nx(b), b.type, {:f, 32}) |> EMLX.to_device(:cpu) upper = !opts[:lower] + a_axes_swap = swap_last_two_axes(Nx.rank(a)) # Apply transform_a: transposing flips the upper/lower triangularity. - # For real types, :conjugate is equivalent to :transpose. + # :conjugate means "conjugate the entries" (a no-op for the real dtypes + # EMLX supports here — complex is rejected above), NOT a transpose. It is + # only equivalent to :transpose under LAPACK's conjugate-*transpose* + # convention, which is not what Nx.LinAlg.triangular_solve documents (see + # Nx.BinaryBackend.Matrix.ts/8: :conjugate pre-conjugates and is treated + # as :none). + # + # NB: axes must be a non-negative full permutation (e.g. [1, 0] for rank 2) + # here — a literal [-2, -1] normalizes to [0, 1] for rank 2, which is a + # silent no-op, not a swap. {a_mx, effective_upper} = case opts[:transform_a] do - :none -> {a_typed, upper} - t when t in [:transpose, :conjugate] -> {EMLX.transpose(a_typed, [-2, -1]), !upper} + :transpose -> {EMLX.transpose(a_typed, a_axes_swap), !upper} + _ -> {a_typed, upper} end out_mx = @@ -2069,16 +2064,17 @@ defmodule EMLX.Backend do EMLX.linalg_solve_triangular(a_mx, b_mx, effective_upper) else # Solve XA = B → A^T x = b (works for both 1D and 2D b) - a_t = EMLX.transpose(a_mx, [-2, -1]) + a_t = EMLX.transpose(a_mx, a_axes_swap) if Nx.rank(b) == 1 do # b is 1D: solve_triangular handles A^T x = b directly EMLX.linalg_solve_triangular(a_t, b_mx, !effective_upper) else - b_t = EMLX.transpose(b_mx, [-2, -1]) + b_axes_swap = swap_last_two_axes(Nx.rank(b)) + b_t = EMLX.transpose(b_mx, b_axes_swap) EMLX.linalg_solve_triangular(a_t, b_t, !effective_upper) - |> EMLX.transpose([-2, -1]) + |> EMLX.transpose(b_axes_swap) end end @@ -2088,6 +2084,13 @@ defmodule EMLX.Backend do |> to_nx(out) end + # Non-negative permutation swapping the last two axes, identity elsewhere. + # rank 2 -> [1, 0]; rank 3 -> [0, 2, 1]; etc. + defp swap_last_two_axes(rank) do + {front, [x, y]} = Enum.split(Enum.to_list(0..(rank - 1)), rank - 2) + front ++ [y, x] + end + cumulative_blocks = [ {Nx.Block.CumulativeSum, :cumulative_sum}, {Nx.Block.CumulativeProduct, :cumulative_product}, @@ -2180,7 +2183,7 @@ defmodule EMLX.Backend do to_nx(cholesky, out) :gpu -> - fun.(struct, tensor) + eager_fallback(device, struct, [tensor], fun) end end @@ -2204,14 +2207,14 @@ defmodule EMLX.Backend do {to_nx(q, out_q), to_nx(r, out_r)} :gpu -> - fun.(struct, tensor) + eager_fallback(device, struct, [tensor], fun) end end @impl true def block(%Nx.Block.LinAlg.QR{mode: :complete} = struct, _output, args, fun) do # MLX only supports reduced QR; fall back to defn for :complete mode - apply(fun, [struct | args]) + eager_fallback(struct, args, fun) end @impl true @@ -2225,7 +2228,7 @@ defmodule EMLX.Backend do {to_nx(eigenvalues, out_eigenvals), to_nx(eigenvectors, out_eigenvecs)} :gpu -> - fun.(struct, tensor) + eager_fallback(device, struct, [tensor], fun) end end @@ -2245,7 +2248,7 @@ defmodule EMLX.Backend do {to_nx(u, out_u), to_nx(s, out_s), to_nx(vt, out_v)} :gpu -> - fun.(struct, tensor) + eager_fallback(device, struct, [tensor], fun) end end @@ -2265,7 +2268,32 @@ defmodule EMLX.Backend do @impl true def block(struct, _output, args, fun) do - apply(fun, [struct | args]) + eager_fallback(struct, args, fun) + end + + # Runs a block's plain-Nx composite `fun` (the traced `default_expr` + # implementation) with `Nx.default_backend/1` pinned to `tensor`'s own + # backend/device for the duration of the call. + # + # Without this, a literal tensor `fun` creates internally without an + # explicit `:backend` (e.g. `Nx.LinAlg.SVD.svd/2`'s `qdwh/2` helper does + # `Nx.iota({}, type: :u8, ...) + 1` to seed its `while` loop's condition + # state) is allocated on `Nx.default_backend()` — `Nx.BinaryBackend` unless + # configured otherwise — regardless of which backend/device the real + # operand (`tensor`) lives on. Mixing that literal with `tensor` inside the + # composite's own control flow (e.g. the `while` loop's per-iteration + # state) then crashes some `Nx.BinaryBackend` ops (e.g. `put_slice/5`, + # which assumes all its operands are already `Nx.BinaryBackend`) with a + # `FunctionClauseError` instead of raising a clean cross-backend error. + # Pinning the default backend to `tensor`'s keeps every literal `fun` + # creates on the same backend/device, avoiding the mismatch entirely. + defp eager_fallback(device, struct, args, fun) when is_atom(device) do + Nx.with_default_backend({Backend, device: device}, fn -> apply(fun, [struct | args]) end) + end + + defp eager_fallback(struct, [tensor | _] = args, fun) do + {device, _ref} = from_nx(tensor) + eager_fallback(device, struct, args, fun) end for {op, arity} <- [ diff --git a/emlx/lib/emlx/command_queue.ex b/emlx/lib/emlx/command_queue.ex index dd67cef..0adf76f 100644 --- a/emlx/lib/emlx/command_queue.ex +++ b/emlx/lib/emlx/command_queue.ex @@ -63,7 +63,6 @@ defmodule EMLX.CommandQueue do `{:error, reason}` if the NIF cannot allocate the stream (e.g. `:gpu` on a system without Metal). """ - @spec new(:cpu | :gpu) :: {:ok, t()} | {:error, list()} def new(device) do case EMLX.NIF.command_queue_new(device) do {:ok, ref} -> {:ok, %__MODULE__{ref: ref, device: device}} @@ -78,7 +77,6 @@ defmodule EMLX.CommandQueue do Useful in one-liner contexts (e.g. `__partitions_options__/1`) where the caller has no reasonable recovery path if the device is unavailable. """ - @spec new!(:cpu | :gpu) :: t() def new!(device) do case new(device) do {:ok, q} -> q @@ -93,7 +91,6 @@ defmodule EMLX.CommandQueue do Internally posts a barrier job that calls `mlx::core::synchronize(stream)` on the worker thread, then blocks in `receive` until the reply arrives. """ - @spec synchronize(t()) :: :ok def synchronize(%__MODULE__{ref: ref}) do case EMLX.NIF.command_queue_synchronize(ref) do {:ok, job_ref} -> @@ -115,7 +112,6 @@ defmodule EMLX.CommandQueue do `after` block. Nested calls are safe — each level saves and restores independently. """ - @spec with_queue(t(), (-> result)) :: result when result: term() def with_queue(%__MODULE__{ref: ref, device: device}, fun) when is_function(fun, 0) do previous = Process.get(:emlx_command_queue) Process.put(:emlx_command_queue, {ref, device}) diff --git a/emlx/lib/emlx/debug.ex b/emlx/lib/emlx/debug.ex new file mode 100644 index 0000000..c95a699 --- /dev/null +++ b/emlx/lib/emlx/debug.ex @@ -0,0 +1,24 @@ +defmodule EMLX.Debug do + @moduledoc false + + @detect_non_finites Application.compile_env(:emlx, :detect_non_finites, false) + + @doc false + defmacro assert_no_nan_inf!(tensor_ref, op) do + if @detect_non_finites do + quote do + has_nan = unquote(tensor_ref) |> EMLX.is_nan() |> EMLX.any([], false) + has_inf = unquote(tensor_ref) |> EMLX.is_infinity() |> EMLX.any([], false) + EMLX.eval(has_nan) + EMLX.eval(has_inf) + + if EMLX.item(has_nan) == 1 or EMLX.item(has_inf) == 1 do + raise ArgumentError, + "#{unquote(op)} produced NaN or Inf. Disable :detect_non_finites for production." + end + end + else + quote do: nil + end + end +end diff --git a/emlx/lib/emlx/defn/tree.ex b/emlx/lib/emlx/defn/tree.ex new file mode 100644 index 0000000..792f5ce --- /dev/null +++ b/emlx/lib/emlx/defn/tree.ex @@ -0,0 +1,103 @@ +defmodule EMLX.Defn.Tree do + @moduledoc """ + Topological traversal of `Nx.Defn.Expr` DAGs. + + This module is an upstream candidate for `Nx.Defn.Tree`; the only change + needed to upstream it is a rename. It has **zero EMLX dependencies** — it + only uses `Nx.Defn.{Tree, Composite}` and `Nx.Tensor`. Callers needing + compiler-specific scope handling (e.g. EMLX's `:__EMLX__` metadata nodes, + see `EMLX.Native.Expr.scope_dependencies/1`) inject it via `post_order/2`'s + `scope_dependencies` argument instead of this module knowing about it. + + ## Scope model + + An `Nx.Defn.Expr` DAG may contain scope boundaries: `while`/`fun`/`block` + nodes each introduce an inner scope whose sub-graphs must not be mixed with + the parent ordering. `cond` clauses share the parent scope and are traversed + normally. + + `post_order/2` respects these boundaries: `while`, `fun`, and `block` nodes + are treated as **opaque leaves** in the parent ordering — their inner-scope + sub-graphs do not appear in the result. + """ + + alias Nx.Defn.Composite + alias Nx.Defn.Tree + alias Nx.Tensor, as: T + + @doc """ + Returns the `%Nx.Tensor{}` nodes reachable from `output` in post-order. + + `output` may be any `Nx.Container.t()` (a tensor, a tuple, or any struct + implementing `Nx.Container`). The result is a flat list of the same + `%Nx.Tensor{}` structs that make up the DAG (no rewriting), ordered so that + every node's same-scope operands appear at a strictly smaller index than the + node itself. + + Each node appears **exactly once**, deduplicated by `node.data.id`. + + `scope_dependencies` lets a caller override which nodes count as a given node's + same-scope dependencies, without this module needing to know about + compiler-specific node shapes. For each visited node it is called as + `scope_dependencies.(node)` and must return either: + + * `{:ok, deps}` — treat `deps` (a list of `%Nx.Tensor{}`) as the node's + *only* same-scope dependencies, recursing into them instead of the + node's own `args`. Use this to redirect traversal away from a node's + literal `args` (e.g. a node wrapping an inner sub-expression that + should stay out of the ordering entirely). + * `:default` — fall back to the default traversal below. + + ## Sub-scope handling + + `while`, `fun`, and `block` nodes are returned **opaque**: their inner-scope + sub-graphs are not traversed and do not appear in the result. The caller is + responsible for recurring into sub-scopes as needed. For each such node, the + relevant inner scope can be accessed from its `args`: + + * `:while` — `args = [initial, arg, condition, body]`; `arg`, `condition`, + and `body` belong to the while scope. Call `post_order/2` on `condition` + or `body` to obtain their orderings independently. + * `:fun` — `args = [params, body, mfa]`; `body` is the function scope. + * `:block` — `args = [struct, in_args, default_expr, callback]`; + `default_expr` is the block's inner scope. + """ + def post_order(output, scope_dependencies \\ fn _node -> :default end) do + roots = Composite.flatten_list([output]) + {_visited, rev} = Enum.reduce(roots, {MapSet.new(), []}, &visit(&1, &2, scope_dependencies)) + Enum.reverse(rev) + end + + # --- recursive post-order DFS --- + + defp visit(%T{data: %Nx.Defn.Expr{id: id}} = node, {visited, output}, scope_dependencies) do + if MapSet.member?(visited, id) do + {visited, output} + else + # Mark visited before recursing to handle shared subexpressions correctly. + visited = MapSet.put(visited, id) + {visited, output} = visit_scope_deps(node, {visited, output}, scope_dependencies) + {visited, [node | output]} + end + end + + # fun's args[0] contains inner-scope parameter templates — no parent-scope deps. + defp visit_scope_deps(%T{data: %Nx.Defn.Expr{op: :fun}}, acc, _scope_dependencies), do: acc + + defp visit_scope_deps(node, acc, scope_dependencies) do + case scope_dependencies.(node) do + {:ok, deps} -> + Enum.reduce(deps, acc, &visit(&1, &2, scope_dependencies)) + + :default -> + # For all other ops (including while and block, which expose parent-scope + # deps via apply_args :scope), recurse into same-scope operands before emitting. + {_, acc} = + Tree.apply_args(node, :scope, acc, fn dep, a -> + {dep, visit(dep, a, scope_dependencies)} + end) + + acc + end + end +end diff --git a/emlx/lib/emlx/fast.ex b/emlx/lib/emlx/fast.ex index ed2e59e..973149f 100644 --- a/emlx/lib/emlx/fast.ex +++ b/emlx/lib/emlx/fast.ex @@ -1,10 +1,43 @@ defmodule EMLX.Fast do @moduledoc """ Single-kernel Metal shaders from `mlx::fast`, exposed as `deftransform` - functions backed by `Nx.runtime_call`. + functions. Every function is defn-safe: call inside `defn`, `Nx.Defn.jit`, or from - `Axon.rewrite_nodes/2` rewrite callbacks without restriction. + `Axon.rewrite_nodes/2` rewrite callbacks without restriction — **except + `einsum/2`**, which is eager-only (see its docs). + + ## Two execution paths + + Each `deftransform` here dispatches on whether its tensor arguments are + concrete (eager — a real `EMLX.Backend`-backed tensor) or expression-backed + (traced — building an `Nx.Defn.Expr` graph): + + - **Eager**: the fused NIF (e.g. `EMLX.fast_rms_norm/3`) runs immediately. + - **Traced**: the function returns an `Nx.Defn.Expr.metadata/2` node + carrying a `:__EMLX__` key — `%{op: opcode, operands: [...], attrs: [...]}` + — naming the native EMLX opcode, its operand tensors, and its int-encoded + attributes (see `EMLX.Native.Expr.f64_bits/1`). `EMLX.Native.Expr`'s + `:metadata` `expand_node` clause recognizes this key and lowers straight + to the native op — no graph split. The metadata node wraps an + `Nx.runtime_call/4` of the *same* eager callback used above (`inner`) — + this compiler never evaluates it (it's discarded in favor of the + `:__EMLX__` payload); it exists only so (a) the operand tensors are + ordinary reachable dependencies for `EMLX.Defn.Tree.post_order/2` to + visit, and (b) any other `Nx.Defn.Compiler` — notably the default + `Nx.Defn.Evaluator` and `Nx.Defn.Grad` — still gets a correct fallback: + `runtime_call` just runs the real NIF against concrete tensors, so it's + both exact (not a slower plain-`Nx` approximation) and free to build + (unlike a full composite reference formula, `Nx.runtime_call/4` is a + single lightweight node — no per-op sub-expression tracing cost). The + `:__EMLX__` node is itself wrapped once more in a + `Nx.Defn.Kernel.custom_grad/3` annotation (`with_reference_grad/3`) so + `Nx.Defn.grad` differentiates through each op's plain-`Nx` `*_reference/N` + formula (VJP via `Nx.Defn.Grad.transform/3` — see `with_reference_grad/3`) + instead of hitting the non-differentiable `runtime_call` forward pass. + + A bare `Nx.runtime_call` (anything *not* wrapped in `:__EMLX__` metadata) + always forces a graph split — see `EMLX.split_point?/1`. ## Functions @@ -15,10 +48,16 @@ defmodule EMLX.Fast do - `rope_with_positions/6` — fused RoPE accepting a `position_ids` tensor - `rope_with_freqs/6` — fused RoPE with precomputed inv-frequency tensor (for `:llama3` scaling) - `scaled_dot_product_attention/4` — flash-attention SDPA (no mask) - - `scaled_dot_product_attention/5` — flash-attention SDPA (additive/bool mask) + - `scaled_dot_product_attention/5` — flash-attention SDPA, either an additive/bool + `mask` tensor, or an `opts` keyword list (`:sinks`) when there's no mask + - `scaled_dot_product_attention/6` — flash-attention SDPA with `mask` and `opts` (`:sinks`) - `scaled_dot_product_attention_causal/4` — flash-attention SDPA with built-in causal mask + - `scaled_dot_product_attention_causal/5` — same, plus `opts` (`:sinks`) - `scaled_dot_product_attention_causal_key_masked/5` — causal SDPA; checks key_mask at C++ level, fast-paths to pure causal when all-ones + - `scaled_dot_product_attention_causal_key_masked/6` — same, plus `opts` (`:sinks`) - `swiglu/2` — fused SwiGLU: `silu(gate) * up` + - `einsum/2` — variadic-operand Einstein summation (`mlx::core::einsum`); + **eager-only, not defn-safe** (see its docs) ## Axon graph rewrite example @@ -31,6 +70,65 @@ defmodule EMLX.Fast do import Nx.Defn + require EMLX.Debug + import EMLX.Debug, only: [assert_no_nan_inf!: 2] + + alias EMLX.Native.Expr, as: NativeExpr + + # The `*_reference/N` plain-`Nx` formulas below (and their private helpers: + # `position_freqs/6`, `rope_broadcast_shape/2`, `rope_rotate/5`, + # `rope_split_half/5`, `rope_interleaved/5`, `repeat_kv_heads/2`, + # `apply_causal_mask/4`, `apply_causal_key_mask/5`, `apply_generic_mask/2`, + # `iota_bin/2`, `neg_inf_like/1`) are no longer used for the traced + # *forward* dispatch path (see moduledoc) — each is instead called from a + # `with_reference_grad/3` closure to compute `Nx.Defn.grad`'s backward pass. + + # ── Traced/eager dispatch helpers ─────────────────────────────────────────── + + # True when any of `tensors` is expression-backed (i.e. we're being traced + # by `Nx.Defn`, as opposed to running eagerly on concrete tensors). + defp traced?(tensors) do + Enum.any?(List.wrap(tensors), &match?(%Nx.Tensor{data: %Nx.Defn.Expr{}}, &1)) + end + + # Wraps `reference` (a plain-Nx formula, never evaluated by EMLX's + # compiler) with the `:__EMLX__` metadata naming the real native opcode, + # its operand tensors, and its int-encoded attrs. + defp emlx_metadata(reference, opcode, operands, attrs) do + Nx.Defn.Expr.metadata(reference, %{__EMLX__: %{op: opcode, operands: operands, attrs: attrs}}) + end + + # Wraps `fused` (an `emlx_metadata/4` node) with a `custom_grad/3` + # annotation so `Nx.Defn.grad` differentiates through the op's plain-`Nx` + # `*_reference/N` formula instead of hitting the opaque, non-differentiable + # `Nx.runtime_call` forward pass. `reference_fn` is applied positionally to + # `inputs` (same order as the `emlx_metadata/4` `operands` list) and must + # recompute the same (single-tensor) value as `fused`. + # + # Reuses `Nx.Defn.Grad` itself rather than hand-deriving each op's backward + # formula: for any (possibly multi-output-shaped) `y = reference_fn(inputs)` + # and upstream cotangent `g` (same shape as `y`), the VJP w.r.t. `inputs` is + # exactly `grad(inputs, sum(g * y))` — a standard trick for building a VJP + # out of a scalar-output `grad`. + defp with_reference_grad(fused, inputs, reference_fn) when is_function(reference_fn) do + Nx.Defn.Kernel.custom_grad(fused, inputs, fn g -> + {_value, grad} = + Nx.Defn.Grad.transform( + List.to_tuple(inputs), + fn args -> + args + |> Tuple.to_list() + |> then(&apply(reference_fn, &1)) + |> Nx.multiply(g) + |> Nx.sum() + end, + & &1 + ) + + Tuple.to_list(grad) + end) + end + # ── RMS Norm ──────────────────────────────────────────────────────────────── @doc """ @@ -43,14 +141,41 @@ defmodule EMLX.Fast do Output shape and type match `x`. """ deftransform rms_norm(x, weight, eps) do - out = Nx.template(Nx.shape(x), Nx.type(x)) - Nx.runtime_call(out, {x, weight}, [eps: eps], &__MODULE__.rms_norm_callback/2) + if traced?([x, weight]) do + fused = + emlx_metadata( + Nx.runtime_call(Nx.to_template(x), {x, weight}, [eps: eps], &rms_norm_callback/2), + :fast_rms_norm, + [x, weight], + [NativeExpr.f64_bits(eps)] + ) + + with_reference_grad(fused, [x, weight], fn x, weight -> + rms_norm_reference(x, weight, eps) + end) + else + rms_norm_callback({x, weight}, eps: eps) + end + end + + defp rms_norm_reference(x, weight, eps) do + x + |> Nx.pow(2) + |> Nx.mean(axes: [-1], keep_axes: true) + |> Nx.add(eps) + |> Nx.sqrt() + |> then(&Nx.divide(x, &1)) + |> Nx.multiply(weight) + |> Nx.as_type(Nx.type(x)) end @doc false def rms_norm_callback({%Nx.Tensor{} = x, %Nx.Tensor{} = weight}, opts) do - EMLX.fast_rms_norm(EMLX.Backend.from_nx(x), EMLX.Backend.from_nx(weight), opts[:eps]) - |> EMLX.Backend.to_nx() + result_ref = + EMLX.fast_rms_norm(EMLX.Backend.from_nx(x), EMLX.Backend.from_nx(weight), opts[:eps]) + + assert_no_nan_inf!(result_ref, :rms_norm) + EMLX.Backend.to_nx(result_ref) end # ── Layer Norm ─────────────────────────────────────────────────────────────── @@ -66,19 +191,52 @@ defmodule EMLX.Fast do Output shape and type match `x`. """ deftransform layer_norm(x, weight, bias, eps) do - out = Nx.template(Nx.shape(x), Nx.type(x)) - Nx.runtime_call(out, {x, weight, bias}, [eps: eps], &__MODULE__.layer_norm_callback/2) + if traced?([x, weight, bias]) do + fused = + emlx_metadata( + Nx.runtime_call( + Nx.to_template(x), + {x, weight, bias}, + [eps: eps], + &layer_norm_callback/2 + ), + :fast_layer_norm, + [x, weight, bias], + [NativeExpr.f64_bits(eps)] + ) + + with_reference_grad(fused, [x, weight, bias], fn x, weight, bias -> + layer_norm_reference(x, weight, bias, eps) + end) + else + layer_norm_callback({x, weight, bias}, eps: eps) + end + end + + defp layer_norm_reference(x, weight, bias, eps) do + mean = Nx.mean(x, axes: [-1], keep_axes: true) + centered = Nx.subtract(x, mean) + variance = Nx.mean(Nx.pow(centered, 2), axes: [-1], keep_axes: true) + + centered + |> Nx.divide(Nx.sqrt(Nx.add(variance, eps))) + |> Nx.multiply(weight) + |> Nx.add(bias) + |> Nx.as_type(Nx.type(x)) end @doc false def layer_norm_callback({%Nx.Tensor{} = x, %Nx.Tensor{} = weight, %Nx.Tensor{} = bias}, opts) do - EMLX.fast_layer_norm( - EMLX.Backend.from_nx(x), - EMLX.Backend.from_nx(weight), - EMLX.Backend.from_nx(bias), - opts[:eps] - ) - |> EMLX.Backend.to_nx() + result_ref = + EMLX.fast_layer_norm( + EMLX.Backend.from_nx(x), + EMLX.Backend.from_nx(weight), + EMLX.Backend.from_nx(bias), + opts[:eps] + ) + + assert_no_nan_inf!(result_ref, :layer_norm) + EMLX.Backend.to_nx(result_ref) end @doc """ @@ -91,24 +249,57 @@ defmodule EMLX.Fast do Output shape and type match `x`. """ deftransform layer_norm(x, weight, eps) do - out = Nx.template(Nx.shape(x), Nx.type(x)) - Nx.runtime_call(out, {x, weight}, [eps: eps], &__MODULE__.layer_norm_no_bias_callback/2) + if traced?([x, weight]) do + fused = + emlx_metadata( + Nx.runtime_call( + Nx.to_template(x), + {x, weight}, + [eps: eps], + &layer_norm_no_bias_callback/2 + ), + :fast_layer_norm_no_bias, + [x, weight], + [NativeExpr.f64_bits(eps)] + ) + + with_reference_grad(fused, [x, weight], fn x, weight -> + layer_norm_no_bias_reference(x, weight, eps) + end) + else + layer_norm_no_bias_callback({x, weight}, eps: eps) + end + end + + defp layer_norm_no_bias_reference(x, weight, eps) do + mean = Nx.mean(x, axes: [-1], keep_axes: true) + centered = Nx.subtract(x, mean) + variance = Nx.mean(Nx.pow(centered, 2), axes: [-1], keep_axes: true) + + centered + |> Nx.divide(Nx.sqrt(Nx.add(variance, eps))) + |> Nx.multiply(weight) + |> Nx.as_type(Nx.type(x)) end @doc false def layer_norm_no_bias_callback({%Nx.Tensor{} = x, %Nx.Tensor{} = weight}, opts) do - EMLX.fast_layer_norm_no_bias( - EMLX.Backend.from_nx(x), - EMLX.Backend.from_nx(weight), - opts[:eps] - ) - |> EMLX.Backend.to_nx() + result_ref = + EMLX.fast_layer_norm_no_bias( + EMLX.Backend.from_nx(x), + EMLX.Backend.from_nx(weight), + opts[:eps] + ) + + assert_no_nan_inf!(result_ref, :layer_norm) + EMLX.Backend.to_nx(result_ref) end @doc """ - Causal SDPA with the key_mask check delegated to the C++ NIF. + Causal SDPA with the key_mask check delegated to the C++ NIF (eager) or + folded directly into the compiled graph (traced). - At runtime the NIF evaluates `all(key_mask == 1)`: + When eager, the NIF evaluates `all(key_mask == 1)`: - **true** (no padding, e.g. single-sequence decode) → pure causal SDPA, no mask tensor allocated. - **false** (padded batch or multi-sequence) → builds a combined @@ -117,6 +308,10 @@ defmodule EMLX.Fast do This avoids the `Nx.cond` double-evaluation problem: the NIF forces eval of only the small `{B, T_kv}` key_mask subgraph, then branches in C++. + When traced, the compiled `:fast_sdpa_causal_key_masked*` opcode always + builds the combined causal+key_mask additive mask in-graph (a compiled + program can't branch on a runtime `all(key_mask)` check). + Input/output layout matches `scaled_dot_product_attention_causal/4`: - `q` — `{B, N_q, T_q, D}` - `k` — `{B, N_kv, T_kv, D}` @@ -124,8 +319,16 @@ defmodule EMLX.Fast do - `scale` — pre-computed scalar - `key_mask` — `{B, T_kv}` boolean/int tensor (1 = attend, 0 = masked) - Output — `{B, N_q, T_q, D}`, same dtype as `q` + + An optional 6th `opts` keyword list accepts `:sinks` (see + `scaled_dot_product_attention/5`). """ deftransform scaled_dot_product_attention_causal_key_masked(q, k, v, scale, key_mask) do + scaled_dot_product_attention_causal_key_masked(q, k, v, scale, key_mask, []) + end + + @doc false + deftransform scaled_dot_product_attention_causal_key_masked(q, k, v, scale, key_mask, opts) do # Compute kv_offset at JIT/deftransform time: shapes are compile-time constants here. # - Decode (T_q = 1): single query attends to all prior positions; # key_mask filters which positions are valid. kv_offset = T_kv - 1. @@ -134,15 +337,67 @@ defmodule EMLX.Fast do t_q = elem(Nx.shape(q), 2) t_kv = elem(Nx.shape(k), 2) kv_offset = if t_q == 1, do: t_kv - 1, else: 0 - - out = Nx.template(Nx.shape(q), Nx.type(q)) - - Nx.runtime_call( - out, - {q, k, v, key_mask}, - [scale: scale, kv_offset: kv_offset], - &__MODULE__.sdpa_causal_key_masked_callback/2 - ) + sinks = Keyword.get(opts, :sinks) + + if traced?([q, k, v, key_mask, sinks]) do + out = Nx.to_template(q) + + if sinks do + inner = + Nx.runtime_call( + out, + {q, k, v, key_mask, sinks}, + [scale: scale, kv_offset: kv_offset], + &sdpa_causal_key_masked_sinks_callback/2 + ) + + fused = + emlx_metadata( + inner, + :fast_sdpa_causal_key_masked_sinks, + [q, k, v, key_mask, sinks], + [NativeExpr.f64_bits(scale), kv_offset] + ) + + with_reference_grad(fused, [q, k, v, key_mask, sinks], fn q, k, v, key_mask, sinks -> + sdpa_reference(q, k, v, scale, + causal: true, + key_mask: key_mask, + kv_offset: kv_offset, + sinks: sinks + ) + end) + else + inner = + Nx.runtime_call( + out, + {q, k, v, key_mask}, + [scale: scale, kv_offset: kv_offset], + &sdpa_causal_key_masked_callback/2 + ) + + fused = + emlx_metadata( + inner, + :fast_sdpa_causal_key_masked, + [q, k, v, key_mask], + [NativeExpr.f64_bits(scale), kv_offset] + ) + + with_reference_grad(fused, [q, k, v, key_mask], fn q, k, v, key_mask -> + sdpa_reference(q, k, v, scale, causal: true, key_mask: key_mask, kv_offset: kv_offset) + end) + end + else + if sinks do + sdpa_causal_key_masked_sinks_callback({q, k, v, key_mask, sinks}, + scale: scale, + kv_offset: kv_offset + ) + else + sdpa_causal_key_masked_callback({q, k, v, key_mask}, scale: scale, kv_offset: kv_offset) + end + end end @doc false @@ -157,7 +412,7 @@ defmodule EMLX.Fast do scale = opts[:scale] kv_offset = opts[:kv_offset] - out = + result_ref = EMLX.fast_sdpa_causal_key_masked( EMLX.Backend.from_nx(q), EMLX.Backend.from_nx(k), @@ -166,7 +421,36 @@ defmodule EMLX.Fast do EMLX.Backend.from_nx(key_mask), kv_offset ) - |> EMLX.Backend.to_nx() + + assert_no_nan_inf!(result_ref, :sdpa) + out = EMLX.Backend.to_nx(result_ref) + + dtype = Nx.type(q) + if Nx.type(out) != dtype, do: Nx.as_type(out, dtype), else: out + end + + @doc false + def sdpa_causal_key_masked_sinks_callback( + {%Nx.Tensor{} = q, %Nx.Tensor{} = k, %Nx.Tensor{} = v, %Nx.Tensor{} = key_mask, + %Nx.Tensor{} = sinks}, + opts + ) do + scale = opts[:scale] + kv_offset = opts[:kv_offset] + + result_ref = + EMLX.fast_sdpa_causal_key_masked( + EMLX.Backend.from_nx(q), + EMLX.Backend.from_nx(k), + EMLX.Backend.from_nx(v), + scale, + EMLX.Backend.from_nx(key_mask), + kv_offset, + EMLX.Backend.from_nx(sinks) + ) + + assert_no_nan_inf!(result_ref, :sdpa) + out = EMLX.Backend.to_nx(result_ref) dtype = Nx.type(q) if Nx.type(out) != dtype, do: Nx.as_type(out, dtype), else: out @@ -190,14 +474,30 @@ defmodule EMLX.Fast do Output shape and type match `a`. """ deftransform rope(a, dims, traditional, base, scale, offset) do - out = Nx.template(Nx.shape(a), Nx.type(a)) - - Nx.runtime_call( - out, - a, - [dims: dims, traditional: traditional, base: base, scale: scale, offset: offset], - &__MODULE__.rope_callback/2 - ) + if traced?(a) do + traditional_int = if traditional, do: 1, else: 0 + opts = [dims: dims, traditional: traditional, base: base, scale: scale, offset: offset] + + fused = + emlx_metadata( + Nx.runtime_call(Nx.to_template(a), a, opts, &rope_callback/2), + :fast_rope, + [a], + [dims, traditional_int, NativeExpr.f64_bits(base), NativeExpr.f64_bits(scale), offset] + ) + + with_reference_grad(fused, [a], fn a -> + rope_reference(a, dims, traditional, base, scale, offset) + end) + else + rope_callback(a, + dims: dims, + traditional: traditional, + base: base, + scale: scale, + offset: offset + ) + end end @doc false @@ -233,18 +533,17 @@ defmodule EMLX.Fast do Output shape and type match `a`. > ### Sequential positions only (fast T=1 path) {: .warning} - > For **decode** with `T = 1` and `base` below about `1.0e5`, the `fast_rope_ids` NIF - > is used; it assumes sequential positions from `position_ids[b, 0]`. For **larger** - > `base` (e.g. Qwen3 `rope_theta` 1M) or **prefill** (`T > 1`), the Nx per-token - > path is used, matching Bumblebee for arbitrary per-token `position_ids`. + > For **decode** with `T = 1` and `base` below about `1.0e5`, the `fast_rope_ids` + > opcode is used; it assumes sequential positions from `position_ids[b, 0]`. For + > **larger** `base` (e.g. Qwen3 `rope_theta` 1M) or **prefill** (`T > 1`), the + > per-token `fast_rope_positions` opcode is used, matching Bumblebee for + > arbitrary per-token `position_ids`. """ deftransform rope_with_positions(a, position_ids, dims, traditional, base, scale) do - out = Nx.template(Nx.shape(a), Nx.type(a)) - # Branch at JIT/deftransform time on T (index 1 in Bumblebee {B, T, N, D} layout). # T is a compile-time constant when Bumblebee uses static sequence_length compilation. - # - Decode (T = 1): sequential positions — fast_rope_ids NIF (1 Metal dispatch). - # - Prefill (T > 1): arbitrary per-token positions — fast_rope_positions NIF. + # - Decode (T = 1): sequential positions — fast_rope_ids (1 Metal dispatch). + # - Prefill (T > 1): arbitrary per-token positions — fast_rope_positions. t = elem(Nx.shape(a), 1) base = base * 1.0 # `fast_rope_ids` uses a scalar `base` in `mlx::fast::rope` that does not @@ -252,21 +551,58 @@ defmodule EMLX.Fast do # `rope_theta` (Qwen2/3 use 1e6+). `fast_rope_positions` matches # Bumblebee; use it for T=1 when `base` is in that regime (A13). t1_use_fast? = t == 1 and base < 1.0e5 - - if t1_use_fast? do - Nx.runtime_call( - out, - {a, position_ids}, - [dims: dims, traditional: traditional, base: base, scale: scale], - &__MODULE__.rope_with_positions_fast_callback/2 - ) + traditional_int = if traditional, do: 1, else: 0 + + if traced?([a, position_ids]) do + out = Nx.to_template(a) + opts = [dims: dims, traditional: traditional, base: base, scale: scale] + + reference_fn = fn a, position_ids -> + rope_with_positions_reference(a, position_ids, dims, traditional, base, scale) + end + + if t1_use_fast? do + inner = + Nx.runtime_call(out, {a, position_ids}, opts, &rope_with_positions_fast_callback/2) + + fused = + emlx_metadata( + inner, + :fast_rope_ids, + [a, position_ids], + [dims, traditional_int, NativeExpr.f64_bits(base), NativeExpr.f64_bits(scale)] + ) + + with_reference_grad(fused, [a, position_ids], reference_fn) + else + inner = Nx.runtime_call(out, {a, position_ids}, opts, &rope_with_positions_callback/2) + + fused = + emlx_metadata( + inner, + :fast_rope_positions, + [a, position_ids], + [dims, traditional_int, NativeExpr.f64_bits(base), NativeExpr.f64_bits(scale)] + ) + + with_reference_grad(fused, [a, position_ids], reference_fn) + end else - Nx.runtime_call( - out, - {a, position_ids}, - [dims: dims, traditional: traditional, base: base, scale: scale], - &__MODULE__.rope_with_positions_callback/2 - ) + if t1_use_fast? do + rope_with_positions_fast_callback({a, position_ids}, + dims: dims, + traditional: traditional, + base: base, + scale: scale + ) + else + rope_with_positions_callback({a, position_ids}, + dims: dims, + traditional: traditional, + base: base, + scale: scale + ) + end end end @@ -316,8 +652,8 @@ defmodule EMLX.Fast do - `position_ids` — `{B, T}` integer tensor. For **decode** (`T = 1`) the fast path uses `position_ids[b,0]` as the per-batch offset into `freqs` (same contract as `mlx::fast::rope` with a scalar offset per batch). For **prefill** (`T > 1`) a - per-token Nx path runs so arbitrary positions (e.g. left-padded - `[0,…,0,1,2,…]`) are correct; the NIF’s offset-only entry point cannot represent + per-token path runs so arbitrary positions (e.g. left-padded + `[0,…,0,1,2,…]`) are correct; the offset-only entry point cannot represent that. - `dims` — number of feature dims to rotate. - `traditional` — `false` for split-half (Bumblebee / Qwen3); `true` for interleaved. @@ -327,23 +663,57 @@ defmodule EMLX.Fast do Output shape and type match `a`. """ deftransform rope_with_freqs(a, position_ids, dims, traditional, scale, freqs) do - out = Nx.template(Nx.shape(a), Nx.type(a)) t = elem(Nx.shape(a), 1) - - if t == 1 do - Nx.runtime_call( - out, - {a, position_ids, freqs}, - [dims: dims, traditional: traditional, scale: scale], - &__MODULE__.rope_with_freqs_fast_callback/2 - ) + traditional_int = if traditional, do: 1, else: 0 + + if traced?([a, position_ids, freqs]) do + out = Nx.to_template(a) + opts = [dims: dims, traditional: traditional, scale: scale] + + reference_fn = fn a, position_ids, freqs -> + rope_with_freqs_reference(a, position_ids, dims, traditional, scale, freqs) + end + + if t == 1 do + inner = + Nx.runtime_call(out, {a, position_ids, freqs}, opts, &rope_with_freqs_fast_callback/2) + + fused = + emlx_metadata( + inner, + :fast_rope_with_freqs, + [a, position_ids, freqs], + [dims, traditional_int, NativeExpr.f64_bits(scale)] + ) + + with_reference_grad(fused, [a, position_ids, freqs], reference_fn) + else + inner = Nx.runtime_call(out, {a, position_ids, freqs}, opts, &rope_with_freqs_callback/2) + + fused = + emlx_metadata( + inner, + :fast_rope_with_freqs_positions, + [a, position_ids, freqs], + [dims, traditional_int, NativeExpr.f64_bits(scale)] + ) + + with_reference_grad(fused, [a, position_ids, freqs], reference_fn) + end else - Nx.runtime_call( - out, - {a, position_ids, freqs}, - [dims: dims, traditional: traditional, scale: scale], - &__MODULE__.rope_with_freqs_callback/2 - ) + if t == 1 do + rope_with_freqs_fast_callback({a, position_ids, freqs}, + dims: dims, + traditional: traditional, + scale: scale + ) + else + rope_with_freqs_callback({a, position_ids, freqs}, + dims: dims, + traditional: traditional, + scale: scale + ) + end end end @@ -390,6 +760,153 @@ defmodule EMLX.Fast do end end + # ── RoPE reference formulas ───────────────────────────────────────────────── + # + # Plain-Nx formulas used only from a `with_reference_grad/3` closure (see + # moduledoc) to compute `Nx.Defn.grad`'s backward pass — the fused forward + # pass never evaluates them. + + # `a` is `{B, ..., T, D}` (T second-to-last axis) — matches `rope/6`. + defp rope_reference(a, dims, traditional, base, scale, offset) do + rank = Nx.rank(a) + t_axis = rank - 2 + d_axis = rank - 1 + t = elem(Nx.shape(a), t_axis) + half = div(dims, 2) + + inv_freq = + Nx.iota({half}, type: :f32, backend: Nx.BinaryBackend) + |> Nx.multiply(2.0 / dims) + |> then(&Nx.pow(base * 1.0, &1)) + |> then(&Nx.divide(1.0, &1)) + + positions = + Nx.iota({t}, type: :f32, backend: Nx.BinaryBackend) + |> Nx.add(offset * 1.0) + |> Nx.multiply(scale * 1.0) + + freqs = Nx.outer(positions, inv_freq) + freqs_bcast = Nx.reshape(freqs, rope_broadcast_shape(rank, [{t_axis, t}, {d_axis, half}])) + + rope_rotate(a, dims, traditional, freqs_bcast, d_axis) + end + + # `a` is `{B, T, ..., D}` (Bumblebee convention, T at axis 1) — matches + # `rope_with_positions/6`. `position_ids` is `{B, T}`. + defp rope_with_positions_reference(a, position_ids, dims, traditional, base, scale) do + rank = Nx.rank(a) + d_axis = rank - 1 + half = div(dims, 2) + {b, t} = {elem(Nx.shape(position_ids), 0), elem(Nx.shape(position_ids), 1)} + + inv_freq = + Nx.iota({half}, type: :f32, backend: Nx.BinaryBackend) + |> Nx.multiply(2.0 / dims) + |> then(&Nx.pow(base * 1.0, &1)) + |> then(&Nx.divide(1.0, &1)) + + freqs_bt = position_freqs(position_ids, inv_freq, scale, b, t, half) + + freqs_bcast = + Nx.reshape(freqs_bt, rope_broadcast_shape(rank, [{0, b}, {1, t}, {d_axis, half}])) + + rope_rotate(a, dims, traditional, freqs_bcast, d_axis) + end + + # Same layout as `rope_with_positions_reference/6`, but `inv_freq` is + # supplied directly as `reciprocal(freqs)` (matches `mlx::fast::rope`'s + # freqs overload — see the `fast_rope_with_freqs*` opcodes). + defp rope_with_freqs_reference(a, position_ids, dims, traditional, scale, freqs) do + rank = Nx.rank(a) + d_axis = rank - 1 + half = div(dims, 2) + {b, t} = {elem(Nx.shape(position_ids), 0), elem(Nx.shape(position_ids), 1)} + + inv_freq = Nx.divide(1.0, Nx.as_type(freqs, :f32)) + + freqs_bt = position_freqs(position_ids, inv_freq, scale, b, t, half) + + freqs_bcast = + Nx.reshape(freqs_bt, rope_broadcast_shape(rank, [{0, b}, {1, t}, {d_axis, half}])) + + rope_rotate(a, dims, traditional, freqs_bcast, d_axis) + end + + defp position_freqs(position_ids, inv_freq, scale, b, t, half) do + pos_bt1 = + position_ids + |> Nx.as_type(:f32) + |> Nx.multiply(scale * 1.0) + |> Nx.reshape({b, t, 1}) + + Nx.multiply(pos_bt1, Nx.reshape(inv_freq, {1, 1, half})) + end + + # Builds a reshape target with `1`s everywhere except the given + # `{axis, size}` pairs — used to broadcast a `{..., half}`-shaped + # cos/sin table against `a`'s full rank without touching axes it + # passes through untouched (e.g. the heads axis). + defp rope_broadcast_shape(rank, axis_sizes) do + for i <- 0..(rank - 1) do + case List.keyfind(axis_sizes, i, 0) do + {^i, size} -> size + nil -> 1 + end + end + |> List.to_tuple() + end + + # Rotates the first `dims` elements of `a`'s last axis by `freqs_bcast` + # (already reshaped to broadcast against `a`), passing the remainder + # through unchanged. `traditional` selects interleaved-pair vs + # split-half rotation (must match the model checkpoint's convention). + defp rope_rotate(a, dims, traditional, freqs_bcast, d_axis) do + full_d = elem(Nx.shape(a), d_axis) + half = div(dims, 2) + rotate_part = Nx.slice_along_axis(a, 0, dims, axis: d_axis) + cos = Nx.cos(freqs_bcast) + sin = Nx.sin(freqs_bcast) + + rotated = + if traditional do + rope_interleaved(rotate_part, cos, sin, d_axis, half) + else + rope_split_half(rotate_part, cos, sin, d_axis, half) + end + + result = + if dims == full_d do + rotated + else + pass_part = Nx.slice_along_axis(a, dims, full_d - dims, axis: d_axis) + Nx.concatenate([rotated, pass_part], axis: d_axis) + end + + Nx.as_type(result, Nx.type(a)) + end + + defp rope_split_half(rotate_part, cos, sin, d_axis, half) do + x1 = Nx.slice_along_axis(rotate_part, 0, half, axis: d_axis) + x2 = Nx.slice_along_axis(rotate_part, half, half, axis: d_axis) + rotated_half = Nx.concatenate([Nx.negate(x2), x1], axis: d_axis) + cos2 = Nx.concatenate([cos, cos], axis: d_axis) + sin2 = Nx.concatenate([sin, sin], axis: d_axis) + Nx.add(Nx.multiply(rotate_part, cos2), Nx.multiply(rotated_half, sin2)) + end + + defp rope_interleaved(rotate_part, cos, sin, d_axis, half) do + front = rotate_part |> Nx.shape() |> Tuple.to_list() |> Enum.take(d_axis) + paired = Nx.reshape(rotate_part, List.to_tuple(front ++ [half, 2])) + + x1 = paired |> Nx.slice_along_axis(0, 1, axis: d_axis + 1) |> Nx.squeeze(axes: [d_axis + 1]) + x2 = paired |> Nx.slice_along_axis(1, 1, axis: d_axis + 1) |> Nx.squeeze(axes: [d_axis + 1]) + + out1 = Nx.subtract(Nx.multiply(x1, cos), Nx.multiply(x2, sin)) + out2 = Nx.add(Nx.multiply(x1, sin), Nx.multiply(x2, cos)) + + Nx.stack([out1, out2], axis: d_axis + 1) |> Nx.reshape(Nx.shape(rotate_part)) + end + # ── SwiGLU ────────────────────────────────────────────────────────────────── @doc """ @@ -404,8 +921,20 @@ defmodule EMLX.Fast do Output has the same shape and dtype as `gate`. """ deftransform swiglu(gate, up) do - out = Nx.template(Nx.shape(gate), Nx.type(gate)) - Nx.runtime_call(out, {gate, up}, [], &__MODULE__.swiglu_callback/2) + if traced?([gate, up]) do + inner = Nx.runtime_call(Nx.to_template(gate), {gate, up}, [], &swiglu_callback/2) + fused = emlx_metadata(inner, :fast_swiglu, [gate, up], []) + with_reference_grad(fused, [gate, up], &swiglu_reference/2) + else + swiglu_callback({gate, up}, []) + end + end + + defp swiglu_reference(gate, up) do + gate + |> Nx.multiply(Nx.sigmoid(gate)) + |> Nx.multiply(up) + |> Nx.as_type(Nx.type(gate)) end @doc false @@ -430,38 +959,142 @@ defmodule EMLX.Fast do Softmax accumulates in float32 internally regardless of input dtype. """ deftransform scaled_dot_product_attention(q, k, v, scale) do - out = Nx.template(Nx.shape(q), Nx.type(q)) - Nx.runtime_call(out, {q, k, v}, [scale: scale], &__MODULE__.sdpa_callback/2) + scaled_dot_product_attention(q, k, v, scale, []) + end + + @doc """ + Flash-attention SDPA — either an additive/boolean `mask` tensor (5th arg), + or, with no mask, an `opts` keyword list (disambiguated by `is_list/1`): + + * `:sinks` — optional learned per-head attention-sink logits tensor, + appended as an extra key/value pair the softmax normalises against (see + `mlx::fast::scaled_dot_product_attention`'s `sinks` parameter). Shape + must broadcast against `{B, N_q}` (typically `{N_q}`). + + `mask` must be broadcast-compatible with `{B, N_q, T_q, T_kv}`. + Boolean `false` entries are masked out (`-∞`); float entries are added to + the pre-softmax scores. + + For causal masking in decode (single query token), prefer the no-mask arity + since `T_q=1` is always trivially causal. + """ + deftransform scaled_dot_product_attention(q, k, v, scale, opts) when is_list(opts) do + sinks = Keyword.get(opts, :sinks) + + if traced?([q, k, v, sinks]) do + out = Nx.to_template(q) + + if sinks do + inner = Nx.runtime_call(out, {q, k, v, sinks}, [scale: scale], &sdpa_sinks_callback/2) + + fused = + emlx_metadata(inner, :fast_sdpa_sinks, [q, k, v, sinks], [NativeExpr.f64_bits(scale)]) + + with_reference_grad(fused, [q, k, v, sinks], fn q, k, v, sinks -> + sdpa_reference(q, k, v, scale, sinks: sinks) + end) + else + inner = Nx.runtime_call(out, {q, k, v}, [scale: scale], &sdpa_callback/2) + fused = emlx_metadata(inner, :fast_sdpa, [q, k, v], [NativeExpr.f64_bits(scale)]) + + with_reference_grad(fused, [q, k, v], fn q, k, v -> + sdpa_reference(q, k, v, scale, []) + end) + end + else + if sinks do + sdpa_sinks_callback({q, k, v, sinks}, scale: scale) + else + sdpa_callback({q, k, v}, scale: scale) + end + end + end + + deftransform scaled_dot_product_attention(q, k, v, scale, mask) do + scaled_dot_product_attention(q, k, v, scale, mask, []) + end + + @doc """ + Flash-attention SDPA with an additive/boolean `mask` and `opts` (`:sinks` + — see `scaled_dot_product_attention/5`). + """ + deftransform scaled_dot_product_attention(q, k, v, scale, mask, opts) when is_list(opts) do + sinks = Keyword.get(opts, :sinks) + + if traced?([q, k, v, mask, sinks]) do + out = Nx.to_template(q) + + if sinks do + inner = + Nx.runtime_call( + out, + {q, k, v, mask, sinks}, + [scale: scale], + &sdpa_masked_sinks_callback/2 + ) + + fused = + emlx_metadata(inner, :fast_sdpa_masked_sinks, [q, k, v, mask, sinks], [ + NativeExpr.f64_bits(scale) + ]) + + with_reference_grad(fused, [q, k, v, mask, sinks], fn q, k, v, mask, sinks -> + sdpa_reference(q, k, v, scale, mask: mask, sinks: sinks) + end) + else + inner = Nx.runtime_call(out, {q, k, v, mask}, [scale: scale], &sdpa_masked_callback/2) + + fused = + emlx_metadata(inner, :fast_sdpa_masked, [q, k, v, mask], [NativeExpr.f64_bits(scale)]) + + with_reference_grad(fused, [q, k, v, mask], fn q, k, v, mask -> + sdpa_reference(q, k, v, scale, mask: mask) + end) + end + else + if sinks do + sdpa_masked_sinks_callback({q, k, v, mask, sinks}, scale: scale) + else + sdpa_masked_callback({q, k, v, mask}, scale: scale) + end + end end @doc false def sdpa_callback({%Nx.Tensor{} = q, %Nx.Tensor{} = k, %Nx.Tensor{} = v}, opts) do - out = + result_ref = EMLX.fast_sdpa( EMLX.Backend.from_nx(q), EMLX.Backend.from_nx(k), EMLX.Backend.from_nx(v), opts[:scale] ) - |> EMLX.Backend.to_nx() + + assert_no_nan_inf!(result_ref, :sdpa) + out = EMLX.Backend.to_nx(result_ref) # mlx::fast::sdpa may upcast to f32 internally; cast back to q's dtype if Nx.type(out) != Nx.type(q), do: Nx.as_type(out, Nx.type(q)), else: out end - @doc """ - Flash-attention SDPA with an additive or boolean `mask`. + @doc false + def sdpa_sinks_callback( + {%Nx.Tensor{} = q, %Nx.Tensor{} = k, %Nx.Tensor{} = v, %Nx.Tensor{} = sinks}, + opts + ) do + result_ref = + EMLX.fast_sdpa( + EMLX.Backend.from_nx(q), + EMLX.Backend.from_nx(k), + EMLX.Backend.from_nx(v), + opts[:scale], + EMLX.Backend.from_nx(sinks) + ) - `mask` must be broadcast-compatible with `{B, N_q, T_q, T_kv}`. - Boolean `false` entries are masked out (`-∞`); float entries are added to - the pre-softmax scores. + assert_no_nan_inf!(result_ref, :sdpa) + out = EMLX.Backend.to_nx(result_ref) - For causal masking in decode (single query token), prefer the no-mask arity - since `T_q=1` is always trivially causal. - """ - deftransform scaled_dot_product_attention(q, k, v, scale, mask) do - out = Nx.template(Nx.shape(q), Nx.type(q)) - Nx.runtime_call(out, {q, k, v, mask}, [scale: scale], &__MODULE__.sdpa_masked_callback/2) + if Nx.type(out) != Nx.type(q), do: Nx.as_type(out, Nx.type(q)), else: out end @doc false @@ -469,7 +1102,7 @@ defmodule EMLX.Fast do {%Nx.Tensor{} = q, %Nx.Tensor{} = k, %Nx.Tensor{} = v, %Nx.Tensor{} = mask}, opts ) do - out = + result_ref = EMLX.fast_sdpa_masked( EMLX.Backend.from_nx(q), EMLX.Backend.from_nx(k), @@ -477,7 +1110,31 @@ defmodule EMLX.Fast do EMLX.Backend.from_nx(mask), opts[:scale] ) - |> EMLX.Backend.to_nx() + + assert_no_nan_inf!(result_ref, :sdpa) + out = EMLX.Backend.to_nx(result_ref) + + if Nx.type(out) != Nx.type(q), do: Nx.as_type(out, Nx.type(q)), else: out + end + + @doc false + def sdpa_masked_sinks_callback( + {%Nx.Tensor{} = q, %Nx.Tensor{} = k, %Nx.Tensor{} = v, %Nx.Tensor{} = mask, + %Nx.Tensor{} = sinks}, + opts + ) do + result_ref = + EMLX.fast_sdpa_masked( + EMLX.Backend.from_nx(q), + EMLX.Backend.from_nx(k), + EMLX.Backend.from_nx(v), + EMLX.Backend.from_nx(mask), + opts[:scale], + EMLX.Backend.from_nx(sinks) + ) + + assert_no_nan_inf!(result_ref, :sdpa) + out = EMLX.Backend.to_nx(result_ref) if Nx.type(out) != Nx.type(q), do: Nx.as_type(out, Nx.type(q)), else: out end @@ -500,21 +1157,257 @@ defmodule EMLX.Fast do - Output — `{B, N_q, T_q, D}`, same dtype as `q` """ deftransform scaled_dot_product_attention_causal(q, k, v, scale) do - out = Nx.template(Nx.shape(q), Nx.type(q)) - Nx.runtime_call(out, {q, k, v}, [scale: scale], &__MODULE__.sdpa_causal_callback/2) + scaled_dot_product_attention_causal(q, k, v, scale, []) + end + + @doc """ + Causal flash-attention SDPA with `opts` (`:sinks` — see + `scaled_dot_product_attention/5`). + """ + deftransform scaled_dot_product_attention_causal(q, k, v, scale, opts) when is_list(opts) do + sinks = Keyword.get(opts, :sinks) + + if traced?([q, k, v, sinks]) do + out = Nx.to_template(q) + + if sinks do + inner = + Nx.runtime_call(out, {q, k, v, sinks}, [scale: scale], &sdpa_causal_sinks_callback/2) + + fused = + emlx_metadata(inner, :fast_sdpa_causal_sinks, [q, k, v, sinks], [ + NativeExpr.f64_bits(scale) + ]) + + with_reference_grad(fused, [q, k, v, sinks], fn q, k, v, sinks -> + sdpa_reference(q, k, v, scale, causal: true, sinks: sinks) + end) + else + inner = Nx.runtime_call(out, {q, k, v}, [scale: scale], &sdpa_causal_callback/2) + fused = emlx_metadata(inner, :fast_sdpa_causal, [q, k, v], [NativeExpr.f64_bits(scale)]) + + with_reference_grad(fused, [q, k, v], fn q, k, v -> + sdpa_reference(q, k, v, scale, causal: true) + end) + end + else + if sinks do + sdpa_causal_sinks_callback({q, k, v, sinks}, scale: scale) + else + sdpa_causal_callback({q, k, v}, scale: scale) + end + end end @doc false def sdpa_causal_callback({%Nx.Tensor{} = q, %Nx.Tensor{} = k, %Nx.Tensor{} = v}, opts) do - out = + result_ref = EMLX.fast_sdpa_causal( EMLX.Backend.from_nx(q), EMLX.Backend.from_nx(k), EMLX.Backend.from_nx(v), opts[:scale] ) - |> EMLX.Backend.to_nx() + + assert_no_nan_inf!(result_ref, :sdpa) + out = EMLX.Backend.to_nx(result_ref) if Nx.type(out) != Nx.type(q), do: Nx.as_type(out, Nx.type(q)), else: out end + + @doc false + def sdpa_causal_sinks_callback( + {%Nx.Tensor{} = q, %Nx.Tensor{} = k, %Nx.Tensor{} = v, %Nx.Tensor{} = sinks}, + opts + ) do + result_ref = + EMLX.fast_sdpa_causal( + EMLX.Backend.from_nx(q), + EMLX.Backend.from_nx(k), + EMLX.Backend.from_nx(v), + opts[:scale], + EMLX.Backend.from_nx(sinks) + ) + + assert_no_nan_inf!(result_ref, :sdpa) + out = EMLX.Backend.to_nx(result_ref) + + if Nx.type(out) != Nx.type(q), do: Nx.as_type(out, Nx.type(q)), else: out + end + + # ── SDPA reference formula ─────────────────────────────────────────────────── + # + # Plain-Nx formula used only from a `with_reference_grad/3` closure (see + # moduledoc) to compute `Nx.Defn.grad`'s backward pass. q/k/v are + # `{B, N, T, D}`; GQA-repeats k/v to N_q heads before a broadcasted + # dot-product (no batched-matmul axis juggling). + + defp sdpa_reference(q, k, v, scale, opts) do + causal = Keyword.get(opts, :causal, false) + mask = Keyword.get(opts, :mask) + key_mask = Keyword.get(opts, :key_mask) + sinks = Keyword.get(opts, :sinks) + kv_offset = Keyword.get(opts, :kv_offset, 0) + + {_b, n_q, t_q, _d} = Nx.shape(q) + {_b, n_kv, t_kv, _d} = Nx.shape(k) + groups = div(n_q, n_kv) + + k_g = repeat_kv_heads(k, groups) + v_g = repeat_kv_heads(v, groups) + + # Batched contraction over D (no explicit {B,N,T_q,T_kv,D} intermediate — + # keeps this reference cheap to *trace* even at prefill sequence lengths, + # since it's rebuilt (never evaluated) on every EMLX-compiled call). + scores = + q + |> Nx.dot([3], [0, 1], k_g, [3], [0, 1]) + |> Nx.multiply(scale * 1.0) + + scores = + cond do + causal and key_mask != nil -> + apply_causal_key_mask(scores, kv_offset, t_q, t_kv, key_mask) + + causal -> + apply_causal_mask(scores, kv_offset, t_q, t_kv) + + mask != nil -> + apply_generic_mask(scores, mask) + + true -> + scores + end + + scores = + if sinks do + b = elem(Nx.shape(scores), 0) + + sinks_col = + sinks + |> Nx.as_type(Nx.type(scores)) + |> Nx.reshape({1, n_q, 1, 1}) + |> Nx.broadcast({b, n_q, t_q, 1}) + + Nx.concatenate([scores, sinks_col], axis: 3) + else + scores + end + + probs = + scores + |> Nx.subtract(Nx.reduce_max(scores, axes: [3], keep_axes: true)) + |> Nx.exp() + + probs = Nx.divide(probs, Nx.sum(probs, axes: [3], keep_axes: true)) + probs = Nx.slice_along_axis(probs, 0, t_kv, axis: 3) + + probs + |> Nx.dot([3], [0, 1], v_g, [2], [0, 1]) + |> Nx.as_type(Nx.type(q)) + end + + defp repeat_kv_heads(kv, 1), do: kv + + defp repeat_kv_heads(kv, groups) do + {b, n_kv, t, d} = Nx.shape(kv) + + kv + |> Nx.new_axis(2) + |> Nx.broadcast({b, n_kv, groups, t, d}) + |> Nx.reshape({b, n_kv * groups, t, d}) + end + + # `Nx.select/3` uses its *predicate*'s shape as the output shape verbatim + # (it does not pick the pairwise-broadcast-compatible union like `add`/ + # `multiply` do), so `mask`/`keep` must already be pre-broadcast to + # `Nx.shape(scores)` here. That broadcast runs on `Nx.iota`'s constant + # backend (see `iota_bin/2`'s note) — `Nx.BinaryBackend`, plain CPU memory — + # rather than on `scores`'s own (traced) backend, so it never touches a + # real GPU/MLX allocation. + defp apply_causal_mask(scores, kv_offset, t_q, t_kv) do + query_positions = iota_bin({t_q}, :s32) |> Nx.add(kv_offset) |> Nx.reshape({1, 1, t_q, 1}) + key_positions = iota_bin({t_kv}, :s32) |> Nx.reshape({1, 1, 1, t_kv}) + + mask = + key_positions |> Nx.less_equal(query_positions) |> Nx.broadcast(Nx.shape(scores)) + + Nx.select(mask, scores, neg_inf_like(scores)) + end + + defp apply_causal_key_mask(scores, kv_offset, t_q, t_kv, key_mask) do + b = elem(Nx.shape(key_mask), 0) + km = key_mask |> Nx.not_equal(0) |> Nx.reshape({b, 1, 1, t_kv}) + + query_positions = iota_bin({t_q}, :s32) |> Nx.add(kv_offset) |> Nx.reshape({1, 1, t_q, 1}) + key_positions = iota_bin({t_kv}, :s32) |> Nx.reshape({1, 1, 1, t_kv}) + causal_bool = Nx.less_equal(key_positions, query_positions) + + keep = km |> Nx.logical_and(causal_bool) |> Nx.broadcast(Nx.shape(scores)) + Nx.select(keep, scores, neg_inf_like(scores)) + end + + defp apply_generic_mask(scores, mask) do + if Nx.Type.integer?(Nx.type(mask)) do + bool_mask = mask |> Nx.not_equal(0) |> Nx.broadcast(Nx.shape(scores)) + Nx.select(bool_mask, scores, neg_inf_like(scores)) + else + Nx.add(scores, mask) + end + end + + # Pinned to `Nx.BinaryBackend` (cheap, CPU-only) rather than + # `Nx.default_backend/0` — `deftransform` code (unlike `defn` bodies) isn't + # auto-rewritten into `Nx.Defn.Expr` construction, so plain `Nx.iota/2` and + # `Nx.tensor/2` calls here would otherwise allocate *real* tensors on + # whatever backend the caller has configured (typically `EMLX.Backend`, + # i.e. real, uncomputed MLX graph nodes) merely to be embedded as constants + # in this never-evaluated reference formula. + defp iota_bin(shape, type), do: Nx.iota(shape, type: type, backend: Nx.BinaryBackend) + + defp neg_inf_like(scores), + do: Nx.tensor(-1.0e9, type: Nx.type(scores), backend: Nx.BinaryBackend) + + # ── Einsum ──────────────────────────────────────────────────────────────── + + @doc """ + Variadic-operand einsum computed by MLX's path-optimised + `mlx::core::einsum` kernel. + + `subscripts` is a standard Einstein-summation equation (e.g. + `"ij,jk->ik"`, `"bij,bjk->bik"`, `"bhid,bhjd->bhij"`, + `"ij,jk,kl->il"`). `operands` is the corresponding list of 2+ tensors. + + ## Eager-only, not defn-callable + + Unlike the other helpers in this module, `einsum/2` does **not** emit an + `Nx.Defn.Expr` node — it takes refs directly off `EMLX.Backend`-backed + tensors and calls the NIF eagerly, in the same "direct-call helper" style + as `EMLX.quantized_matmul/2`. Every operand must live on `EMLX.Backend`; + anything else raises `ArgumentError`. + + ## Examples + + iex> a = Nx.iota({2, 3}, backend: EMLX.Backend, type: :f32) + iex> b = Nx.iota({3, 4}, backend: EMLX.Backend, type: :f32) + iex> y = EMLX.Fast.einsum("ij,jk->ik", [a, b]) + iex> Nx.shape(y) + {2, 4} + + """ + def einsum(subscripts, operands) when is_binary(subscripts) and is_list(operands) do + refs = Enum.map(operands, &einsum_operand_ref!/1) + + EMLX.einsum(refs, subscripts) + |> EMLX.Backend.to_nx() + end + + defp einsum_operand_ref!(%Nx.Tensor{data: %EMLX.Backend{ref: ref}}), do: ref + + defp einsum_operand_ref!(%Nx.Tensor{data: %other_backend{}}) do + raise ArgumentError, + "EMLX.Fast.einsum/2: every operand must live on EMLX.Backend, got a " <> + "#{inspect(other_backend)}-backed tensor. Transfer with " <> + "Nx.backend_transfer/2 first." + end end diff --git a/emlx/lib/emlx/native.ex b/emlx/lib/emlx/native.ex new file mode 100644 index 0000000..12af191 --- /dev/null +++ b/emlx/lib/emlx/native.ex @@ -0,0 +1,31 @@ +defmodule EMLX.Native do + @moduledoc """ + Shared utilities for the EMLX native compiler layer. + """ + + @doc """ + Maps an `Nx.Type.t()` to the corresponding MLX type atom understood by the + C++ NIFs. Sub-byte integer types (`{:u, 2}`, `{:u, 4}`, `{:s, 2}`, `{:s, 4}`) + are widened to the smallest MLX integer type that can hold them. + """ + def to_mlx_type({:u, 2}), do: :uint8 + def to_mlx_type({:u, 4}), do: :uint8 + def to_mlx_type({:u, 8}), do: :uint8 + def to_mlx_type({:u, 16}), do: :uint16 + def to_mlx_type({:u, 32}), do: :uint32 + def to_mlx_type({:u, 64}), do: :uint64 + def to_mlx_type({:s, 2}), do: :int8 + def to_mlx_type({:s, 4}), do: :int8 + def to_mlx_type({:s, 8}), do: :int8 + def to_mlx_type({:s, 16}), do: :int16 + def to_mlx_type({:s, 32}), do: :int32 + def to_mlx_type({:s, 64}), do: :int64 + def to_mlx_type({:f, 8}), do: :float16 + def to_mlx_type({:f, 16}), do: :float16 + def to_mlx_type({:f, 32}), do: :float32 + def to_mlx_type({:f, 64}), do: :float32 + def to_mlx_type({:bf, 16}), do: :bfloat16 + def to_mlx_type({:c, 64}), do: :complex64 + def to_mlx_type({:c, 128}), do: :complex64 + def to_mlx_type(:bool), do: :bool +end diff --git a/emlx/lib/emlx/native/expr.ex b/emlx/lib/emlx/native/expr.ex new file mode 100644 index 0000000..064f6b1 --- /dev/null +++ b/emlx/lib/emlx/native/expr.ex @@ -0,0 +1,2480 @@ +defmodule EMLX.Native.Expr do + @moduledoc false + import Bitwise + + alias Nx.Defn.Composite + alias Nx.Tensor, as: T + + @compiler_debug Application.compile_env(:emlx, :compiler_debug, false) + defmacrop maybe_debug_check(do: block) do + if @compiler_debug do + block + else + :ok + end + end + + @enforce_keys [:inputs, :captures, :constants, :instructions, :outputs] + defstruct [ + :inputs, + :captures, + :constants, + :instructions, + :outputs, + hooks: [], + runtime_calls: [] + ] + + @type node_ref :: reference() + @type hook :: %{ + name: atom(), + callback: (Nx.Container.t() -> term()), + template: Nx.Container.t(), + refs: [node_ref()] + } + @type runtime_call :: %{ + index: non_neg_integer(), + callback: (Nx.Container.t(), keyword() -> Nx.Container.t()), + args_template: Nx.Container.t(), + arg_param_positions: [non_neg_integer() | nil], + opts: keyword() + } + @type t :: %__MODULE__{ + inputs: [node_ref()], + captures: [{node_ref(), Nx.Tensor.t()}], + constants: [{node_ref(), number(), Nx.Type.t()}], + instructions: [{node_ref(), atom(), [node_ref()], [integer()]}], + outputs: [node_ref()], + hooks: [hook()], + runtime_calls: [runtime_call()] + } + + # ── lowering ────────────────────────────────────────────────────────────── + + @doc false + def lower(output, num_inputs \\ nil, quant_signature \\ %{}) do + ordered = EMLX.Defn.Tree.post_order(output, &scope_dependencies/1) + + # inputs is a map of pos → ref during lowering; densified to a list at the end. + state = %{ + inputs: %{}, + captures: [], + constants: [], + instructions: [], + node_to_ref: %{}, + hooks: [], + runtime_calls: [], + top_scope_ids: output |> Nx.Defn.Tree.scope_ids() |> Map.keys() |> MapSet.new(), + quant_signature: quant_signature + } + + state = Enum.reduce(ordered, state, &expand_node/2) + + max_referenced_pos = state.inputs |> Map.keys() |> Enum.max(fn -> -1 end) + arity = max(num_inputs || 0, max_referenced_pos + 1) + + inputs_list = + for pos <- 0..(arity - 1)//1 do + Map.get_lazy(state.inputs, pos, &make_ref/0) + end + + flat_outputs = Composite.flatten_list([output]) + output_refs = Enum.map(flat_outputs, &Map.fetch!(state.node_to_ref, &1.data.id)) + + %__MODULE__{ + inputs: inputs_list, + captures: Enum.reverse(state.captures), + constants: Enum.reverse(state.constants), + instructions: Enum.reverse(state.instructions), + outputs: output_refs, + hooks: Enum.reverse(state.hooks), + runtime_calls: Enum.reverse(state.runtime_calls) + } + end + + # ── node expansion ──────────────────────────────────────────────────────── + + @doc false + def scope_dependencies(%T{ + data: %Nx.Defn.Expr{op: :metadata, args: [_inner, %{__EMLX__: %{operands: operands}}]} + }) do + {:ok, operands} + end + + def scope_dependencies(_node), do: :default + + defp expand_node(%T{data: %Nx.Defn.Expr{id: id, op: :parameter, args: [pos]}}, state) do + ref = make_ref() + + %{ + state + | inputs: Map.put(state.inputs, pos, ref), + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node(%T{data: %Nx.Defn.Expr{id: id, op: :constant, args: [number]}} = node, state) do + ref = make_ref() + + state = %{ + state + | constants: [{ref, number, node.type} | state.constants] + } + + if node.shape == {} do + %{state | node_to_ref: Map.put(state.node_to_ref, id, ref)} + else + broadcast_ref = make_ref() + shape_list = Tuple.to_list(node.shape) + attrs = [length(shape_list) | shape_list] ++ [0] + + %{ + state + | instructions: [{broadcast_ref, :broadcast, [ref], attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, broadcast_ref) + } + end + end + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :tensor, args: [backend_tensor]}}, + state + ) do + ref = make_ref() + + %{ + state + | captures: [{ref, backend_tensor} | state.captures], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :metadata, + args: [_inner, %{__EMLX__: %{op: opcode, operands: operands, attrs: attrs}}] + } + }, + state + ) do + ref = make_ref() + operand_refs = Enum.map(operands, &Map.fetch!(state.node_to_ref, &1.data.id)) + + %{ + state + | instructions: [{ref, opcode, operand_refs, attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :metadata, args: [inner, _meta]}}, + state + ) do + inner_ref = + if is_tuple(inner) do + inner |> Tuple.to_list() |> Enum.map(&Map.fetch!(state.node_to_ref, &1.data.id)) + else + Map.fetch!(state.node_to_ref, inner.data.id) + end + + %{state | node_to_ref: Map.put(state.node_to_ref, id, inner_ref)} + end + + # ── unary elementwise ops ───────────────────────────────────────────────── + + # Direct unary ops — no coercion; MLX infers result dtype from input. + @unary_direct_ops [ + :abs, + :ceil, + :floor, + :negate, + :round, + :sign, + :real, + :imag, + :is_nan, + :is_infinity, + :bitwise_not, + :conjugate, + :logical_not, + :sigmoid, + :asin, + :asinh, + :acos, + :acosh, + :atan, + :atanh, + :cos, + :cosh, + :erf, + :erf_inv, + :exp, + :expm1, + :log, + :log1p, + :rsqrt, + :sin, + :sinh, + :sqrt, + :tan, + :tanh, + :cbrt, + :erfc + ] + + for op <- @unary_direct_ops do + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: unquote(op), args: [operand]}}, + state + ) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, operand.data.id) + + %{ + state + | instructions: [{ref, unquote(op), [operand_ref], []} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + end + + defp expand_node(%T{data: %Nx.Defn.Expr{op: :count_leading_zeros}}, _state) do + raise ArgumentError, "count_leading_zeros is not supported by EMLX" + end + + defp expand_node(%T{data: %Nx.Defn.Expr{op: :population_count}}, _state) do + raise ArgumentError, "population_count is not supported by EMLX" + end + + # block: dispatch on struct — handles Nx.Block.LogicalNot. + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.LogicalNot{}, [operand], _default, _fun] + } + }, + state + ) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, operand.data.id) + + %{ + state + | instructions: [{ref, :logical_not, [operand_ref], []} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # ── binary elementwise ops ──────────────────────────────────────────────── + + @binary_arithmetic_ops [:add, :subtract, :multiply, :pow, :left_shift] + + for op <- @binary_arithmetic_ops do + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: unquote(op), args: [left, right]}}, + state + ) do + expand_binary_node(id, unquote(op), out_type, left, right, state) + end + end + + @binary_generic_ops [ + :divide, + :quotient, + :atan2, + :right_shift, + :bitwise_and, + :bitwise_or, + :bitwise_xor, + :equal, + :not_equal, + :greater, + :less, + :greater_equal, + :less_equal, + :logical_and, + :logical_or, + :logical_xor + ] + + for op <- @binary_generic_ops do + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: unquote(op), args: [left, right]}}, + state + ) do + expand_binary_node(id, unquote(op), out_type, left, right, state) + end + end + + # min → minimum, max → maximum (mapped in C++ registry) + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: :min, args: [left, right]}}, + state + ) do + expand_binary_node(id, :min, out_type, left, right, state) + end + + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: :max, args: [left, right]}}, + state + ) do + expand_binary_node(id, :max, out_type, left, right, state) + end + + # remainder: composite sign-fix is handled entirely in the C++ registry. + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: :remainder, args: [left, right]}}, + state + ) do + expand_binary_node(id, :remainder, out_type, left, right, state) + end + + # ── shape / movement ops ────────────────────────────────────────────────────── + + # reshape: attrs = new shape dims (flat list); shape from the output tensor. + defp expand_node( + %T{shape: out_shape, data: %Nx.Defn.Expr{id: id, op: :reshape, args: [tensor]}}, + state + ) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + shape_attrs = Tuple.to_list(out_shape) + + %{ + state + | instructions: [{ref, :reshape, [operand_ref], shape_attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # squeeze: attrs = axes to remove (non-negative). + defp expand_node(%T{data: %Nx.Defn.Expr{id: id, op: :squeeze, args: [tensor, axes]}}, state) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + norm_axes = normalize_axes(axes, tuple_size(tensor.shape)) + + %{ + state + | instructions: [{ref, :squeeze, [operand_ref], norm_axes} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # transpose: attrs = axis permutation (non-negative). + defp expand_node(%T{data: %Nx.Defn.Expr{id: id, op: :transpose, args: [tensor, axes]}}, state) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + norm_axes = normalize_axes(axes, tuple_size(tensor.shape)) + + %{ + state + | instructions: [{ref, :transpose, [operand_ref], norm_axes} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # as_type: reuse existing :astype opcode; always emit the cast. + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: :as_type, args: [tensor]}}, + state + ) do + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + {result_ref, state} = emit_cast_to(operand_ref, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + # bitcast: attrs = [target_dtype]. Target type from the output tensor. + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: :bitcast, args: [tensor]}}, + state + ) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + mlx_type = EMLX.Native.to_mlx_type(out_type) + + %{ + state + | instructions: [{ref, :bitcast, [operand_ref], [mlx_type]} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # broadcast: attrs = [n_shape, d0…, n_axes, a0…] (both shape and axes, length-delimited). + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :broadcast, args: [tensor, shape, axes]}}, + state + ) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + shape_list = Tuple.to_list(shape) + n_shape = length(shape_list) + n_axes = length(axes) + attrs = [n_shape | shape_list] ++ [n_axes | axes] + + %{ + state + | instructions: [{ref, :broadcast, [operand_ref], attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :pad, args: [tensor, pad_value, config]}}, + state + ) do + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + pad_value_ref = Map.fetch!(state.node_to_ref, pad_value.data.id) + + {result_ref, state} = + if Enum.any?(config, fn {lo, hi, interior} -> lo < 0 or hi < 0 or interior > 0 end) do + expand_pad_general(tensor_ref, pad_value_ref, Tuple.to_list(tensor.shape), config, state) + else + ref = make_ref() + n_dims = length(config) + attrs = [n_dims | Enum.flat_map(config, fn {lo, hi, interior} -> [lo, hi, interior] end)] + + {ref, + %{ + state + | instructions: [{ref, :pad, [tensor_ref, pad_value_ref], attrs} | state.instructions] + }} + end + + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + # reverse: attrs = axes to flip (non-negative). + defp expand_node(%T{data: %Nx.Defn.Expr{id: id, op: :reverse, args: [tensor, axes]}}, state) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + norm_axes = normalize_axes(axes, tuple_size(tensor.shape)) + + %{ + state + | instructions: [{ref, :reverse, [operand_ref], norm_axes} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :concatenate, args: [tensors, axis]}}, + state + ) do + ref = make_ref() + operand_refs = Enum.map(tensors, &Map.fetch!(state.node_to_ref, &1.data.id)) + norm_axis = if axis < 0, do: tuple_size(hd(tensors).shape) + axis, else: axis + + %{ + state + | instructions: [{ref, :concatenate, operand_refs, [norm_axis]} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :stack, args: [tensors, axis]}}, + state + ) do + ref = make_ref() + operand_refs = Enum.map(tensors, &Map.fetch!(state.node_to_ref, &1.data.id)) + # stack output rank = input rank + 1; normalise axis against output rank + out_rank = tuple_size(hd(tensors).shape) + 1 + norm_axis = if axis < 0, do: out_rank + axis, else: axis + + %{ + state + | instructions: [{ref, :stack, operand_refs, [norm_axis]} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # ── reductions ────────────────────────────────────────────────────────────── + + @reduction_cast_ops [:sum, :product, :all, :any] + + for op <- @reduction_cast_ops do + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: unquote(op), args: [tensor, opts]}}, + state + ) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + axes = opts[:axes] || Nx.axes(tensor) + keep_axes = if opts[:keep_axes], do: 1, else: 0 + attrs = [keep_axes | normalize_axes(axes, tuple_size(tensor.shape))] + + state = %{ + state + | instructions: [{ref, unquote(op), [operand_ref], attrs} | state.instructions] + } + + {result_ref, state} = emit_cast_to(ref, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + end + + # reduce_max, reduce_min: MLX preserves input dtype, no cast needed. + @reduction_nocast_ops [:reduce_max, :reduce_min] + + for op <- @reduction_nocast_ops do + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: unquote(op), args: [tensor, opts]}}, + state + ) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + axes = opts[:axes] || Nx.axes(tensor) + keep_axes = if opts[:keep_axes], do: 1, else: 0 + attrs = [keep_axes | normalize_axes(axes, tuple_size(tensor.shape))] + + %{ + state + | instructions: [{ref, unquote(op), [operand_ref], attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + end + + @argreduce_ops [:argmax, :argmin] + + for op <- @argreduce_ops do + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: unquote(op), args: [tensor, opts]}}, + state + ) do + ref = make_ref() + operand_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + axis = opts[:axis] + keep_axis = if opts[:keep_axis], do: 1, else: 0 + + norm_axis = + cond do + is_nil(axis) -> -1 + axis < 0 -> tuple_size(tensor.shape) + axis + true -> axis + end + + state = %{ + state + | instructions: [ + {ref, unquote(op), [operand_ref], [norm_axis, keep_axis]} | state.instructions + ] + } + + {result_ref, state} = emit_cast_to(ref, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + end + + defp expand_node( + %T{ + type: out_type, + shape: out_shape, + data: %Nx.Defn.Expr{id: id, op: :reduce, args: [tensor, acc, opts, fun]} + }, + state + ) do + expand_reduce_unroll(id, out_type, out_shape, tensor, acc, opts, fun, state) + end + + # ── dot ───────────────────────────────────────────────────────────────────── + + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{ + id: id, + op: :dot, + args: [left, c_left, b_left, right, c_right, b_right] + } + }, + state + ) do + if quantized_param_config(left, state.quant_signature) do + raise ArgumentError, + "does not yet lower op :dot with a quantized left operand. " <> + "Dequantize it first with EMLX.dequantize/1." + end + + case quantized_param_config(right, state.quant_signature) do + nil -> + expand_plain_dot(id, out_type, left, c_left, right, c_right, b_left, b_right, state) + + cfg -> + expand_quantized_dot(id, out_type, left, c_left, b_left, right, c_right, cfg, state) + end + end + + # ── conv ───────────────────────────────────────────────────────────────────── + + defp expand_node( + %T{ + type: out_type, + shape: out_shape, + data: %Nx.Defn.Expr{id: id, op: :conv, args: [input, kernel, opts]} + }, + state + ) do + batch_group_size = opts[:batch_group_size] + + if batch_group_size != 1 do + raise ArgumentError, "does not yet lower op :conv with batch_group_size != 1" + end + + input_permutation = opts[:input_permutation] + kernel_permutation = opts[:kernel_permutation] + output_permutation = opts[:output_permutation] + strides = opts[:strides] + padding = opts[:padding] + input_dilation = opts[:input_dilation] + kernel_dilation = opts[:kernel_dilation] + feature_group_count = opts[:feature_group_size] + + input_ref = Map.fetch!(state.node_to_ref, input.data.id) + kernel_ref = Map.fetch!(state.node_to_ref, kernel.data.id) + + # 1. Cast to out_type. + {input_casted, state} = emit_cast_if_needed(input_ref, input.type, out_type, state) + {kernel_casted, state} = emit_cast_if_needed(kernel_ref, kernel.type, out_type, state) + + # 2. Transpose input: user permutation then channels-last. + input_rank = tuple_size(input.shape) + {input_perm1, state} = emit_transpose_instr(input_casted, input_permutation, state) + + {input_processed, state} = + emit_transpose_instr( + input_perm1, + move_channels_last(Enum.to_list(0..(input_rank - 1))), + state + ) + + # 3. Transpose kernel: user permutation then channels-last. + kernel_rank = tuple_size(kernel.shape) + {kernel_perm1, state} = emit_transpose_instr(kernel_casted, kernel_permutation, state) + + {kernel_processed, state} = + emit_transpose_instr( + kernel_perm1, + move_channels_last(Enum.to_list(0..(kernel_rank - 1))), + state + ) + + # 4. :conv_general — attrs = [n_dims, s…, pl0,ph0,…, kd…, id…, fgs] + n_dims = input_rank - 2 + {padding_low, padding_high} = Enum.unzip(padding) + + conv_attrs = + [n_dims | strides] ++ + Enum.flat_map(Enum.zip(padding_low, padding_high), fn {lo, hi} -> [lo, hi] end) ++ + kernel_dilation ++ + input_dilation ++ + [feature_group_count] + + conv_ref = make_ref() + + state = %{ + state + | instructions: [ + {conv_ref, :conv_general, [input_processed, kernel_processed], conv_attrs} + | state.instructions + ] + } + + # 5. Transpose output: channels-first then inverse of output_permutation. + out_rank = tuple_size(out_shape) + [batch | spatial_and_channels] = Enum.to_list(0..(out_rank - 1)) + {channels, spatial} = List.pop_at(spatial_and_channels, -1) + permute_channels_first = [batch, channels | spatial] + + output_perm_inverse = + output_permutation + |> Enum.with_index() + |> Enum.sort() + |> Enum.map(&elem(&1, 1)) + + {conv_perm1, state} = emit_transpose_instr(conv_ref, permute_channels_first, state) + {result_ref, state} = emit_transpose_instr(conv_perm1, output_perm_inverse, state) + + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + # ── indexing / selection ops ────────────────────────────────────────────── + + # select: cast on_true and on_false to out_type, then emit :select. + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{id: id, op: :select, args: [pred, on_true, on_false]} + }, + state + ) do + pred_ref = Map.fetch!(state.node_to_ref, pred.data.id) + true_ref0 = Map.fetch!(state.node_to_ref, on_true.data.id) + false_ref0 = Map.fetch!(state.node_to_ref, on_false.data.id) + {true_ref, state} = emit_cast_if_needed(true_ref0, on_true.type, out_type, state) + {false_ref, state} = emit_cast_if_needed(false_ref0, on_false.type, out_type, state) + ref = make_ref() + + %{ + state + | instructions: [{ref, :select, [pred_ref, true_ref, false_ref], []} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # clip: operands = [tensor, min, max]. + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :clip, args: [tensor, min_t, max_t]}}, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + min_ref = Map.fetch!(state.node_to_ref, min_t.data.id) + max_ref = Map.fetch!(state.node_to_ref, max_t.data.id) + + %{ + state + | instructions: [{ref, :clip, [tensor_ref, min_ref, max_ref], []} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :slice, + args: [tensor, start_indices, lengths, strides] + } + }, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + n_dims = tuple_size(tensor.shape) + input_shape = Tuple.to_list(tensor.shape) + + # Partition start_indices into dynamic (tensor refs) and static (integers). + {dynamic_mask, static_vals, dyn_operand_refs, state} = + Enum.reduce(Enum.with_index(start_indices), {0, [], [], state}, fn + {idx, _i}, {mask, statics, dyn_refs, st} when is_integer(idx) -> + {mask, statics ++ [idx], dyn_refs, st} + + {%T{} = idx_tensor, i}, {mask, statics, dyn_refs, st} -> + dyn_ref = Map.fetch!(st.node_to_ref, idx_tensor.data.id) + {mask ||| 1 <<< i, statics ++ [0], dyn_refs ++ [dyn_ref], st} + end) + + attrs = + [n_dims, dynamic_mask] ++ + input_shape ++ + lengths ++ + strides ++ + static_vals + + operands = [tensor_ref | dyn_operand_refs] + + %{ + state + | instructions: [{ref, :slice, operands, attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{id: id, op: :put_slice, args: [input, start_indices, slice]} + }, + state + ) do + ref = make_ref() + input_ref0 = Map.fetch!(state.node_to_ref, input.data.id) + slice_ref0 = Map.fetch!(state.node_to_ref, slice.data.id) + + # Cast both to out_type. + {input_ref, state} = emit_cast_if_needed(input_ref0, input.type, out_type, state) + {slice_ref, state} = emit_cast_if_needed(slice_ref0, slice.type, out_type, state) + + n_dims = tuple_size(input.shape) + input_shape = Tuple.to_list(input.shape) + lengths = Tuple.to_list(slice.shape) + + {dynamic_mask, static_vals, dyn_operand_refs, state} = + Enum.reduce(Enum.with_index(start_indices), {0, [], [], state}, fn + {idx, _i}, {mask, statics, dyn_refs, st} when is_integer(idx) -> + {mask, statics ++ [idx], dyn_refs, st} + + {%T{} = idx_tensor, i}, {mask, statics, dyn_refs, st} -> + dyn_ref = Map.fetch!(st.node_to_ref, idx_tensor.data.id) + {mask ||| 1 <<< i, statics ++ [0], dyn_refs ++ [dyn_ref], st} + end) + + attrs = [n_dims, dynamic_mask] ++ input_shape ++ lengths ++ static_vals + operands = [input_ref, slice_ref | dyn_operand_refs] + + %{ + state + | instructions: [{ref, :put_slice, operands, attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{ + shape: out_shape, + data: %Nx.Defn.Expr{id: id, op: :gather, args: [tensor, indices, opts]} + }, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + indices_ref = Map.fetch!(state.node_to_ref, indices.data.id) + axes = opts[:axes] + n_gather_axes = length(axes) + n_tensor_dims = tuple_size(tensor.shape) + + slice_sizes = + Enum.map(Nx.axes(tensor), fn axis -> + if axis in axes, do: 1, else: elem(tensor.shape, axis) + end) + + out_shape_list = Tuple.to_list(out_shape) + + attrs = + [n_gather_axes | axes] ++ + [n_tensor_dims | slice_sizes] ++ + [length(out_shape_list) | out_shape_list] + + %{ + state + | instructions: [{ref, :gather, [tensor_ref, indices_ref], attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # block: Nx.Block.Take — take(tensor, indices, axis). + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.Take{axis: axis}, [tensor, indices], _default, _fun] + } + }, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + indices_ref = Map.fetch!(state.node_to_ref, indices.data.id) + norm_axis = if axis < 0, do: tuple_size(tensor.shape) + axis, else: axis + + %{ + state + | instructions: [{ref, :take, [tensor_ref, indices_ref], [norm_axis]} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # block: Nx.Block.TakeAlongAxis — take_along_axis(tensor, indices, axis). + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.TakeAlongAxis{axis: axis}, [tensor, indices], _default, _fun] + } + }, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + indices_ref = Map.fetch!(state.node_to_ref, indices.data.id) + norm_axis = if axis < 0, do: tuple_size(tensor.shape) + axis, else: axis + + %{ + state + | instructions: [ + {ref, :take_along_axis, [tensor_ref, indices_ref], [norm_axis]} | state.instructions + ], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{id: id, op: :indexed_add, args: [target, indices, updates, opts]} + }, + state + ) do + expand_indexed_node(id, :indexed_add, out_type, target, indices, updates, opts, state) + end + + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{id: id, op: :indexed_put, args: [target, indices, updates, opts]} + }, + state + ) do + expand_indexed_node(id, :indexed_put, out_type, target, indices, updates, opts, state) + end + + # ── sort / argsort ──────────────────────────────────────────────────────── + + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: :sort, args: [tensor, opts]}}, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + axis = opts[:axis] + norm_axis = if axis < 0, do: tuple_size(tensor.shape) + axis, else: axis + asc_int = if opts[:direction] == :asc, do: 1, else: 0 + + state = %{ + state + | instructions: [{ref, :sort, [tensor_ref], [norm_axis, asc_int]} | state.instructions] + } + + {result_ref, state} = emit_cast_if_needed(ref, tensor.type, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + defp expand_node( + %T{type: out_type, data: %Nx.Defn.Expr{id: id, op: :argsort, args: [tensor, opts]}}, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + axis = opts[:axis] + norm_axis = if axis < 0, do: tuple_size(tensor.shape) + axis, else: axis + asc_int = if opts[:direction] == :asc, do: 1, else: 0 + + state = %{ + state + | instructions: [{ref, :argsort, [tensor_ref], [norm_axis, asc_int]} | state.instructions] + } + + {result_ref, state} = emit_cast_if_needed(ref, {:u, 32}, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + # ── window reductions ───────────────────────────────────────────────────── + + @window_op_int %{window_sum: 0, window_product: 1, window_max: 2, window_min: 3} + + for op <- [:window_sum, :window_product, :window_max, :window_min] do + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{ + id: id, + op: unquote(op), + args: [tensor, window_dims_tuple, opts] + } + }, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + n_dims = tuple_size(window_dims_tuple) + window_dims = Tuple.to_list(window_dims_tuple) + {low_pads, high_pads} = Enum.unzip(opts[:padding]) + strides = opts[:strides] || List.duplicate(1, n_dims) + window_dilations = opts[:window_dilations] || List.duplicate(1, n_dims) + op_int = @window_op_int[unquote(op)] + + attrs = + [n_dims, op_int] ++ + Enum.flat_map(0..(n_dims - 1), fn i -> + [Enum.at(low_pads, i), Enum.at(high_pads, i)] + end) ++ + strides ++ window_dims ++ window_dilations + + state = %{ + state + | instructions: [{ref, unquote(op), [tensor_ref], attrs} | state.instructions] + } + + {result_ref, state} = emit_cast_if_needed(ref, tensor.type, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + end + + for op <- [:window_scatter_max, :window_scatter_min] do + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: unquote(op), + args: [tensor_t, source, init_value, window_dims_tuple, opts] + } + }, + state + ) do + ref = make_ref() + t_ref = Map.fetch!(state.node_to_ref, tensor_t.data.id) + src_ref = Map.fetch!(state.node_to_ref, source.data.id) + init_ref = Map.fetch!(state.node_to_ref, init_value.data.id) + + n_dims = tuple_size(window_dims_tuple) + window_dims = Tuple.to_list(window_dims_tuple) + {low_pads, high_pads} = Enum.unzip(opts[:padding]) + strides = opts[:strides] || List.duplicate(1, n_dims) + + attrs = + [n_dims] ++ + Enum.flat_map(0..(n_dims - 1), fn i -> + [Enum.at(low_pads, i), Enum.at(high_pads, i)] + end) ++ + strides ++ window_dims + + %{ + state + | instructions: [ + {ref, unquote(op), [t_ref, src_ref, init_ref], attrs} | state.instructions + ], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + end + + defp expand_node( + %T{ + type: out_type, + shape: out_shape, + data: %Nx.Defn.Expr{ + id: id, + op: :window_reduce, + args: [tensor, acc, window_dims_tuple, opts, fun] + } + }, + state + ) do + n_dims = tuple_size(window_dims_tuple) + window_dims = Tuple.to_list(window_dims_tuple) + {low_pads, high_pads} = Enum.unzip(opts[:padding]) + strides = opts[:strides] || List.duplicate(1, n_dims) + dilations = opts[:window_dilations] || List.duplicate(1, n_dims) + + if Enum.any?(low_pads ++ high_pads, &(&1 < 0)) do + raise ArgumentError, "does not yet lower op :window_reduce with negative padding" + end + + [params, body, _mfa] = fun.data.args + + unless length(params) == 2 do + raise ArgumentError, "does not yet lower op :window_reduce with a non-binary reducer" + end + + # Cast input + acc to the reducer/output type before padding/folding. + tensor_ref0 = Map.fetch!(state.node_to_ref, tensor.data.id) + {tensor_ref, state} = emit_cast_if_needed(tensor_ref0, tensor.type, out_type, state) + acc_ref0 = Map.fetch!(state.node_to_ref, acc.data.id) + {acc_scalar_ref, state} = emit_cast_if_needed(acc_ref0, acc.type, out_type, state) + + # Pad with acc (interior 0). Padded shape drives the slice input-shape iattr. + in_dims = Tuple.to_list(tensor.shape) + + padded_shape = + [in_dims, low_pads, high_pads] + |> Enum.zip() + |> Enum.map(fn {d, lo, hi} -> d + lo + hi end) + + {padded_ref, state} = + if Enum.all?(low_pads ++ high_pads, &(&1 == 0)) do + {tensor_ref, state} + else + emit_pad_with(tensor_ref, acc_scalar_ref, low_pads, high_pads, state) + end + + out_dims = Tuple.to_list(out_shape) + + {acc_ref, state} = emit_broadcast_to(acc_scalar_ref, out_dims, state) + extent = Enum.product(window_dims) + + # :slice takes a span (stop = start + length); with a stride it yields + spans = Enum.zip_with(out_dims, strides, fn d, s -> (d - 1) * s + 1 end) + + {final_ref, state} = + Enum.reduce(0..(extent - 1)//1, {acc_ref, state}, fn k, {acc_k, st} -> + offsets = window_offsets(k, window_dims) + starts = Enum.zip_with(offsets, dilations, &(&1 * &2)) + {slice_ref, st} = emit_static_slice(padded_ref, padded_shape, starts, spans, strides, st) + lower_fun_body(body, %{0 => slice_ref, 1 => acc_k}, st) + end) + + %{state | node_to_ref: Map.put(state.node_to_ref, id, final_ref)} + end + + # ── Nx.Block.Cumulative* — recognize-struct path ───────────────────────── + for {block_mod, op} <- [ + {Nx.Block.CumulativeSum, :cumulative_sum}, + {Nx.Block.CumulativeProduct, :cumulative_product}, + {Nx.Block.CumulativeMin, :cumulative_min}, + {Nx.Block.CumulativeMax, :cumulative_max} + ] do + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%unquote(block_mod){axis: axis, reverse: reverse}, [tensor], _default, _fun] + } + }, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + norm_axis = if axis < 0, do: tuple_size(tensor.shape) + axis, else: axis + reverse_int = if reverse, do: 1, else: 0 + + state = %{ + state + | instructions: [ + {ref, unquote(op), [tensor_ref], [norm_axis, reverse_int]} | state.instructions + ] + } + + {result_ref, state} = emit_cast_if_needed(ref, tensor.type, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + end + + # ── fft / ifft ──────────────────────────────────────────────────────────── + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :fft, args: [tensor, opts]}}, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + axis = opts[:axis] + n = opts[:length] + + %{ + state + | instructions: [{ref, :fft, [tensor_ref], [axis, n]} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :ifft, args: [tensor, opts]}}, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + axis = opts[:axis] + n = opts[:length] + + %{ + state + | instructions: [{ref, :ifft, [tensor_ref], [axis, n]} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # ── Nx.Block.FFT2 / IFFT2 — recognize-struct path ───────────────────────── + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.FFT2{lengths: lengths, axes: axes}, [tensor], _default, _fun] + } + }, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + [ax0, ax1] = axes || [-2, -1] + rank = tuple_size(tensor.shape) + ax0 = if ax0 < 0, do: rank + ax0, else: ax0 + ax1 = if ax1 < 0, do: rank + ax1, else: ax1 + [n0, n1] = lengths || [elem(tensor.shape, ax0), elem(tensor.shape, ax1)] + + %{ + state + | instructions: [{ref, :fft2, [tensor_ref], [ax0, ax1, n0, n1]} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.IFFT2{lengths: lengths, axes: axes}, [tensor], _default, _fun] + } + }, + state + ) do + ref = make_ref() + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + [ax0, ax1] = axes || [-2, -1] + rank = tuple_size(tensor.shape) + ax0 = if ax0 < 0, do: rank + ax0, else: ax0 + ax1 = if ax1 < 0, do: rank + ax1, else: ax1 + [n0, n1] = lengths || [elem(tensor.shape, ax0), elem(tensor.shape, ax1)] + + %{ + state + | instructions: [{ref, :ifft2, [tensor_ref], [ax0, ax1, n0, n1]} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # ── Nx.Block.LinAlg.* — recognize-struct native path ───────────────────── + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.LinAlg.Cholesky{}, [tensor], _default, _fun] + } + }, + state + ) do + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + {f32_ref, state} = emit_cast_if_needed(tensor_ref, tensor.type, {:f, 32}, state) + + ref = make_ref() + state = %{state | instructions: [{ref, :cholesky, [f32_ref], []} | state.instructions]} + + {result_ref, state} = emit_cast_if_needed(ref, {:f, 32}, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + # solve: single-output. operands = [a, b]; no attrs. Solves A x = b. + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.LinAlg.Solve{}, [a, b], _default, _fun] + } + }, + state + ) do + a_ref = Map.fetch!(state.node_to_ref, a.data.id) + b_ref = Map.fetch!(state.node_to_ref, b.data.id) + {a_f, state} = emit_cast_if_needed(a_ref, a.type, {:f, 32}, state) + {b_f, state} = emit_cast_if_needed(b_ref, b.type, {:f, 32}, state) + + ref = make_ref() + state = %{state | instructions: [{ref, :solve, [a_f, b_f], []} | state.instructions]} + + {result_ref, state} = emit_cast_if_needed(ref, {:f, 32}, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.LinAlg.QR{mode: :reduced}, [tensor], _default, _fun] + } + }, + state + ) do + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + {f32_ref, state} = emit_cast_if_needed(tensor_ref, tensor.type, {:f, 32}, state) + + q_ref = make_ref() + r_ref = make_ref() + + state = %{ + state + | instructions: [{[q_ref, r_ref], :qr, [f32_ref], []} | state.instructions] + } + + %{state | node_to_ref: Map.put(state.node_to_ref, id, [q_ref, r_ref])} + end + + # eigh: multi-output [eigenvalues, eigenvectors]. operands = [a]; lower triangle. + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.LinAlg.Eigh{}, [tensor], _default, _fun] + } + }, + state + ) do + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + {f32_ref, state} = emit_cast_if_needed(tensor_ref, tensor.type, {:f, 32}, state) + + w_ref = make_ref() + v_ref = make_ref() + + state = %{ + state + | instructions: [{[w_ref, v_ref], :eigh, [f32_ref], []} | state.instructions] + } + + %{state | node_to_ref: Map.put(state.node_to_ref, id, [w_ref, v_ref])} + end + + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.LinAlg.SVD{full_matrices?: true}, [tensor], _default, _fun] + } + }, + state + ) do + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + {f32_ref, state} = emit_cast_if_needed(tensor_ref, tensor.type, {:f, 32}, state) + + u_ref = make_ref() + s_ref = make_ref() + vt_ref = make_ref() + + state = %{ + state + | instructions: [{[u_ref, s_ref, vt_ref], :svd, [f32_ref], []} | state.instructions] + } + + %{state | node_to_ref: Map.put(state.node_to_ref, id, [u_ref, s_ref, vt_ref])} + end + + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [%Nx.Block.LinAlg.LU{}, [tensor], _default, _fun] + } + }, + state + ) do + tensor_ref = Map.fetch!(state.node_to_ref, tensor.data.id) + {f32_ref, state} = emit_cast_if_needed(tensor_ref, tensor.type, {:f, 32}, state) + + piv_ref = make_ref() + l_ref = make_ref() + u_ref = make_ref() + + state = %{ + state + | instructions: [{[piv_ref, l_ref, u_ref], :lu, [f32_ref], []} | state.instructions] + } + + # P = take(eye(n), pivots, axis: 0). n is the trailing matrix dimension. + n = elem(tensor.shape, tuple_size(tensor.shape) - 1) + f32_int = EMLX.Native.to_mlx_type({:f, 32}) + + eye_ref = make_ref() + p_ref = make_ref() + + state = %{ + state + | instructions: [ + {p_ref, :take, [eye_ref, piv_ref], [0]}, + {eye_ref, :eye, [], [f32_int, n, n]} + | state.instructions + ] + } + + %{state | node_to_ref: Map.put(state.node_to_ref, id, [p_ref, l_ref, u_ref])} + end + + defp expand_node( + %T{ + type: out_type, + data: %Nx.Defn.Expr{id: id, op: :triangular_solve, args: [a, b, opts]} + }, + state + ) do + left_side = Keyword.get(opts, :left_side, true) + transform_a = Keyword.get(opts, :transform_a, :none) + lower = Keyword.get(opts, :lower, true) + upper = not lower + + a_ref = Map.fetch!(state.node_to_ref, a.data.id) + b_ref = Map.fetch!(state.node_to_ref, b.data.id) + {a_f, state} = emit_cast_if_needed(a_ref, a.type, {:f, 32}, state) + {b_f, state} = emit_cast_if_needed(b_ref, b.type, {:f, 32}, state) + + a_rank = tuple_size(a.shape) + + {a_op, effective_upper, state} = + case transform_a do + :transpose -> + {a_t, state} = emit_transpose_instr(a_f, swap_last_two_axes(a_rank), state) + {a_t, not upper, state} + + _ -> + {a_f, upper, state} + end + + {ref, state} = + if left_side do + emit_solve_triangular_instr(a_op, b_f, effective_upper, state) + else + # Solve XA = B → A^T x = b (works for both 1D and 2D b). + {a_t, state} = emit_transpose_instr(a_op, swap_last_two_axes(a_rank), state) + b_rank = tuple_size(b.shape) + + if b_rank == 1 do + emit_solve_triangular_instr(a_t, b_f, not effective_upper, state) + else + {b_t, state} = emit_transpose_instr(b_f, swap_last_two_axes(b_rank), state) + {out_t, state} = emit_solve_triangular_instr(a_t, b_t, not effective_upper, state) + emit_transpose_instr(out_t, swap_last_two_axes(b_rank), state) + end + end + + {result_ref, state} = emit_cast_if_needed(ref, {:f, 32}, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + # ── block fallback: descend into default_expr ───────────────────────────── + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :block, + args: [_struct, in_args, default_expr, _fun] + } + }, + state + ) do + expand_block_via_default(id, in_args, default_expr, state) + end + + # ── cond: lower as nested :select ops ──────────────────────────────────── + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :cond, args: [clauses, last]}}, + state + ) do + last_refs = flat_refs(last, state) + + clause_ref_pairs = + Enum.map(clauses, fn {pred, body} -> + {Map.fetch!(state.node_to_ref, pred.data.id), flat_refs(body, state)} + end) + + n = length(last_refs) + + {per_elem_refs, state} = + Enum.reduce(0..(n - 1), {[], state}, fn i, {elem_results, st} -> + last_ref_i = Enum.at(last_refs, i) + + # Right-fold: most-priority clause wraps the least-priority accumulator. + {result_ref, st} = + Enum.reduce(Enum.reverse(clause_ref_pairs), {last_ref_i, st}, fn {pred_ref, body_refs}, + {acc_ref, st2} -> + body_ref_i = Enum.at(body_refs, i) + ref = make_ref() + + st2 = %{ + st2 + | instructions: [ + {ref, :select, [pred_ref, body_ref_i, acc_ref], []} | st2.instructions + ] + } + + {ref, st2} + end) + + {elem_results ++ [result_ref], st} + end) + + node_val = if n == 1, do: hd(per_elem_refs), else: per_elem_refs + %{state | node_to_ref: Map.put(state.node_to_ref, id, node_val)} + end + + # ── elem: extract element from a tuple-output op (cond/while) ───────────── + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :elem, args: [tuple_node, pos]}}, + state + ) do + case Map.fetch!(state.node_to_ref, tuple_node.data.id) do + refs when is_list(refs) -> + elem_ref = Enum.at(refs, pos) + %{state | node_to_ref: Map.put(state.node_to_ref, id, elem_ref)} + + single_ref when pos == 0 -> + # Degenerate single-element tuple — shouldn't normally appear. + %{state | node_to_ref: Map.put(state.node_to_ref, id, single_ref)} + end + end + + # ── creation ops ───────────────────────────────────────────────────────── + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :iota, args: [axis]}} = node, + state + ) do + ref = make_ref() + shape = Tuple.to_list(node.shape) + n_dims = length(shape) + dtype_int = EMLX.Native.to_mlx_type(node.type) + axis_int = if axis == nil, do: -1, else: axis + + %{ + state + | instructions: [ + {ref, :iota, [], [dtype_int, n_dims, axis_int | shape]} | state.instructions + ], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :eye, args: []}} = node, + state + ) do + ref = make_ref() + shape_list = Tuple.to_list(node.shape) + n_dims = length(shape_list) + [m, n] = Enum.take(shape_list, -2) + dtype_int = EMLX.Native.to_mlx_type(node.type) + + state = %{ + state + | instructions: [{ref, :eye, [], [dtype_int, m, n]} | state.instructions] + } + + if n_dims == 2 do + %{state | node_to_ref: Map.put(state.node_to_ref, id, ref)} + else + broadcast_ref = make_ref() + axes = [n_dims - 2, n_dims - 1] + attrs = [n_dims | shape_list] ++ [length(axes) | axes] + + %{ + state + | instructions: [{broadcast_ref, :broadcast, [ref], attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, broadcast_ref) + } + end + end + + defp expand_node(%T{data: %Nx.Defn.Expr{op: :fun}}, state), do: state + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :token, args: [%Nx.Defn.Token{hooks: hooks}]}}, + state + ) do + unless MapSet.member?(state.top_scope_ids, id) do + raise ArgumentError, + "cannot lower a hook nested inside a cond branch: EMLX's cond compiles by " <> + "evaluating every branch unconditionally (:select), which would fire this " <> + "hook on every call regardless of which branch is actually taken -- a " <> + "behavior divergence from Nx.Defn.Evaluator (which only fires the selected " <> + "branch's hook). Move the hook outside the cond." + end + + Enum.reduce(hooks, state, fn + %{callback: nil}, state -> + state + + %{callback: callback, expr: expr, name: name}, state -> + refs = + [expr] + |> Composite.flatten_list() + |> Enum.map(&Map.fetch!(state.node_to_ref, &1.data.id)) + + template = Composite.traverse(expr, &Nx.to_template/1) + hook = %{name: name, callback: callback, template: template, refs: refs} + %{state | hooks: [hook | state.hooks]} + end) + end + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :attach_token, args: [_token, expr]}}, + state + ) do + ref = Map.fetch!(state.node_to_ref, expr.data.id) + %{state | node_to_ref: Map.put(state.node_to_ref, id, ref)} + end + + defp expand_node( + %T{ + data: %Nx.Defn.Expr{ + id: id, + op: :runtime_call, + args: [tensor_expr, callback, out_template, opts] + } + }, + state + ) do + unless MapSet.member?(state.top_scope_ids, id) do + raise ArgumentError, + "cannot lower a runtime_call nested inside a cond branch: EMLX's cond compiles " <> + "by evaluating every branch unconditionally (:select), which would fire this " <> + "runtime_call's callback on every call regardless of which branch is actually " <> + "taken -- a behavior divergence from Nx.Defn.Evaluator (which only fires the " <> + "selected branch's callback). Move the runtime_call outside the cond." + end + + operand_leaves = Composite.flatten_list([tensor_expr]) + + operand_refs = Enum.map(operand_leaves, &Map.fetch!(state.node_to_ref, &1.data.id)) + + arg_param_positions = + Enum.map(operand_leaves, fn + %T{data: %Nx.Defn.Expr{op: :parameter, args: [pos]}} -> pos + _ -> nil + end) + + output_templates = + case out_template do + %Nx.Tensor{} = t -> [t] + container -> Composite.flatten_list([container]) + end + + callback_index = length(state.runtime_calls) + + attrs = + [callback_index, length(output_templates)] ++ + Enum.flat_map(output_templates, fn t -> + dtype_int = EMLX.Native.to_mlx_type(t.type) + shape = Tuple.to_list(t.shape) + [dtype_int, length(shape) | shape] + end) + + runtime_call = %{ + index: callback_index, + callback: callback, + args_template: Composite.traverse(tensor_expr, &Nx.to_template/1), + arg_param_positions: arg_param_positions, + opts: opts + } + + result_ref = + case out_template do + %Nx.Tensor{} -> make_ref() + _ -> Enum.map(output_templates, fn _ -> make_ref() end) + end + + %{ + state + | instructions: [{result_ref, :runtime_call, operand_refs, attrs} | state.instructions], + node_to_ref: Map.put(state.node_to_ref, id, result_ref), + runtime_calls: [runtime_call | state.runtime_calls] + } + end + + defp expand_node( + %T{data: %Nx.Defn.Expr{id: id, op: :while, args: [initial, arg, condition, body]}}, + state + ) do + case detect_static_while_trip_count(initial, arg, condition, body) do + {:ok, count} -> expand_while_unroll(id, initial, arg, body, count, state) + :error -> raise ArgumentError, "does not yet lower op :while" + end + end + + defp expand_node(%T{data: %Nx.Defn.Expr{op: op}}, _state) do + raise ArgumentError, "does not yet lower op #{inspect(op)}" + end + + @doc false + def native_lowerable_block?(%Nx.Block.LogicalNot{}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.Take{}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.TakeAlongAxis{}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.FFT2{}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.IFFT2{}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.LinAlg.Cholesky{}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.LinAlg.Solve{}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.LinAlg.QR{mode: :reduced}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.LinAlg.Eigh{}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.LinAlg.SVD{full_matrices?: true}, _in_args), do: true + def native_lowerable_block?(%Nx.Block.LinAlg.LU{}, _in_args), do: true + def native_lowerable_block?(_struct, _in_args), do: false + + # ── dot helpers ──────────────────────────────────────────────────────────── + + defp quantized_param_config( + %T{data: %Nx.Defn.Expr{op: :parameter, args: [pos]}}, + quant_signature + ), + do: Map.get(quant_signature, pos) + + defp quantized_param_config(_node, _quant_signature), do: nil + + defp expand_plain_dot(id, out_type, left, c_left, right, c_right, b_left, b_right, state) do + left_ref0 = Map.fetch!(state.node_to_ref, left.data.id) + right_ref0 = Map.fetch!(state.node_to_ref, right.data.id) + + computation_type = + if Nx.Type.integer?(out_type), do: Nx.Type.to_floating(out_type), else: out_type + + {left_ref, state} = emit_cast_if_needed(left_ref0, left.type, computation_type, state) + {right_ref, state} = emit_cast_if_needed(right_ref0, right.type, computation_type, state) + + attrs = + [length(c_left) | c_left] ++ + [length(c_right) | c_right] ++ + [length(b_left) | b_left] ++ + [length(b_right) | b_right] + + dot_ref = make_ref() + + state = %{ + state + | instructions: [{dot_ref, :dot, [left_ref, right_ref], attrs} | state.instructions] + } + + {result_ref, state} = emit_cast_if_needed(dot_ref, computation_type, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + defp expand_quantized_dot( + id, + out_type, + left, + c_left, + b_left, + right, + c_right, + %EMLX.Quantization.Config{} = cfg, + state + ) do + unless b_left == [] do + raise ArgumentError, + "does not yet lower op :dot with a quantized right operand and batch axes " <> + "(mx::quantized_matmul does not support batching)" + end + + unless c_left == [tuple_size(left.shape) - 1] do + raise ArgumentError, + "does not yet lower op :dot with a quantized right operand contracted " <> + "on a non-last left axis" + end + + unless match?([_], c_right) do + raise ArgumentError, + "does not yet lower op :dot with a quantized right operand contracted " <> + "on more than one axis" + end + + last_dim = tuple_size(right.shape) - 1 + + transpose = + case cfg.transpose do + nil -> c_right == [last_dim] + explicit -> explicit + end + + left_ref = Map.fetch!(state.node_to_ref, left.data.id) + right_ref = Map.fetch!(state.node_to_ref, right.data.id) + + {scales_ref, state} = emit_capture(cfg.scales, state) + + {operands, has_bias, state} = + if cfg.biases do + {biases_ref, state} = emit_capture(cfg.biases, state) + {[left_ref, right_ref, scales_ref, biases_ref], 1, state} + else + {[left_ref, right_ref, scales_ref], 0, state} + end + + mode_int = String.to_atom(cfg.mode) + transpose_int = if transpose, do: 1, else: 0 + attrs = [cfg.group_size, cfg.bits, transpose_int, mode_int, has_bias] + + qmm_ref = make_ref() + + state = %{ + state + | instructions: [{qmm_ref, :quantized_matmul, operands, attrs} | state.instructions] + } + + # mx::quantized_matmul returns the activation's dtype (matching the eager + {result_ref, state} = emit_cast_if_needed(qmm_ref, left.type, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + defp emit_capture(%Nx.Tensor{} = tensor, state) do + ref = make_ref() + {ref, %{state | captures: [{ref, tensor} | state.captures]}} + end + + # ── indexing helpers ────────────────────────────────────────────────────── + + defp expand_indexed_node(id, op, out_type, target, indices, updates, opts, state) do + ref = make_ref() + axes = opts[:axes] || Nx.axes(target) + num_axes = elem(indices.shape, tuple_size(indices.shape) - 1) + + # Mirror EMLX.Backend.indexed_op: compute reshape of updates. + insert_index = + axes + |> Enum.scan(&(&1 - &2)) + |> Enum.find_index(&(&1 > 1)) + |> then(&(&1 || num_axes)) + + [num_updates | updates_inner_shape] = Tuple.to_list(updates.shape) + + updates_shape = + [num_updates | List.duplicate(1, num_axes)] + |> List.insert_at(insert_index + 1, updates_inner_shape) + |> List.flatten() + + target_ref0 = Map.fetch!(state.node_to_ref, target.data.id) + indices_ref = Map.fetch!(state.node_to_ref, indices.data.id) + updates_ref0 = Map.fetch!(state.node_to_ref, updates.data.id) + + {target_ref, state} = emit_cast_if_needed(target_ref0, target.type, out_type, state) + {updates_ref, state} = emit_cast_if_needed(updates_ref0, updates.type, out_type, state) + + attrs = [length(axes) | axes] ++ [length(updates_shape) | updates_shape] + + %{ + state + | instructions: [ + {ref, op, [target_ref, indices_ref, updates_ref], attrs} | state.instructions + ], + node_to_ref: Map.put(state.node_to_ref, id, ref) + } + end + + # ── block-descent helper ────────────────────────────────────────────────── + + defp expand_block_via_default(id, in_args, default_expr, state) do + inner_ordered = EMLX.Defn.Tree.post_order(default_expr, &scope_dependencies/1) + + # Collect inner :parameter nodes; sort by position (args[0]). + inner_params = + inner_ordered + |> Enum.filter(&(&1.data.op == :parameter)) + |> Enum.sort_by(fn t -> hd(t.data.args) end) + + # Map each inner param id → the corresponding parent-scope arg ref. + parent_arg_refs = Enum.map(in_args, &Map.fetch!(state.node_to_ref, &1.data.id)) + + inner_param_id_set = + inner_params + |> Enum.zip(parent_arg_refs) + |> Enum.reduce(MapSet.new(), fn {param, _}, acc -> + MapSet.put(acc, param.data.id) + end) + + inner_param_ref_map = + inner_params + |> Enum.zip(parent_arg_refs) + |> Map.new(fn {param, ref} -> {param.data.id, ref} end) + + # Extend node_to_ref with inner param → parent ref mappings. + merged_node_to_ref = Map.merge(state.node_to_ref, inner_param_ref_map) + inner_state = %{state | node_to_ref: merged_node_to_ref} + + # Expand inner scope, skipping inner :parameter nodes (already mapped). + inner_state = + Enum.reduce(inner_ordered, inner_state, fn node, st -> + if MapSet.member?(inner_param_id_set, node.data.id) do + st + else + expand_node(node, st) + end + end) + + result_ref = + if is_tuple(default_expr) do + flat_refs(default_expr, inner_state) + else + Map.fetch!(inner_state.node_to_ref, default_expr.data.id) + end + + %{inner_state | node_to_ref: Map.put(inner_state.node_to_ref, id, result_ref)} + end + + # ── while (static unroll for counted range loops) ────────────────────────── + defp detect_static_while_trip_count(initial, arg, condition, body) when is_tuple(arg) do + index_param = elem(arg, 0) + + with {:ok, start} <- constant_value(elem(initial, 0)), + {:ok, bound, le?} <- while_condition_bound(condition, index_param), + {:ok, step} <- while_body_step(elem(body, 0), index_param), + {:ok, count} <- static_trip_count(start, bound, step, le?) do + {:ok, count} + else + _ -> :error + end + end + + defp detect_static_while_trip_count(_initial, _arg, _condition, _body), do: :error + + defp static_trip_count(start, bound, step, true) when step > 0 and bound >= start, + do: {:ok, div(bound - start, step) + 1} + + defp static_trip_count(start, bound, step, false) when step < 0 and bound <= start, + do: {:ok, div(start - bound, -step) + 1} + + defp static_trip_count(_start, _bound, _step, _le?), do: :error + + defp while_condition_bound( + %T{data: %Nx.Defn.Expr{op: op, args: [%T{data: %Nx.Defn.Expr{id: pid}}, bound_node]}}, + %T{data: %Nx.Defn.Expr{id: pid}} + ) + when op in [:less_equal, :greater_equal] do + case constant_value(bound_node) do + {:ok, bound} -> {:ok, bound, op == :less_equal} + :error -> :error + end + end + + defp while_condition_bound(_condition, _index_param), do: :error + + defp while_body_step( + %T{data: %Nx.Defn.Expr{op: :add, args: [a, b]}}, + %T{data: %Nx.Defn.Expr{id: pid}} + ) do + case {a, b} do + {%T{data: %Nx.Defn.Expr{id: ^pid}}, step_node} -> constant_value(step_node) + {step_node, %T{data: %Nx.Defn.Expr{id: ^pid}}} -> constant_value(step_node) + _ -> :error + end + end + + defp while_body_step(_next_index, _index_param), do: :error + + defp constant_value(%T{data: %Nx.Defn.Expr{op: :constant, args: [n]}}) when is_integer(n), + do: {:ok, n} + + defp constant_value(_node), do: :error + + defp expand_while_unroll(id, initial, arg, body, count, state) do + initial_list = Tuple.to_list(initial) + arg_list = Tuple.to_list(arg) + body_list = Tuple.to_list(body) + + init_refs = Enum.map(initial_list, &Map.fetch!(state.node_to_ref, &1.data.id)) + + {final_refs, state} = + Enum.reduce(1..count//1, {init_refs, state}, fn _iteration, {carry_refs, acc_state} -> + param_ref_by_pos = + arg_list + |> Enum.zip(carry_refs) + |> Map.new(fn {param, ref} -> {hd(param.data.args), ref} end) + + lower_tuple_body(body_list, param_ref_by_pos, acc_state) + end) + + %{state | node_to_ref: Map.put(state.node_to_ref, id, final_refs)} + end + + defp lower_tuple_body(body_list, param_ref_by_pos, state) do + body_tuple = List.to_tuple(body_list) + state = merge_scope_ids(state, body_tuple) + inner_ordered = EMLX.Defn.Tree.post_order(body_tuple, &scope_dependencies/1) + + param_id_to_ref = + inner_ordered + |> Enum.filter(&(&1.data.op == :parameter)) + |> Map.new(fn p -> {p.data.id, Map.fetch!(param_ref_by_pos, hd(p.data.args))} end) + + param_id_set = MapSet.new(Map.keys(param_id_to_ref)) + local_base = Map.merge(state.node_to_ref, param_id_to_ref) + + inner_state = + Enum.reduce(inner_ordered, %{state | node_to_ref: local_base}, fn node, st -> + if MapSet.member?(param_id_set, node.data.id), do: st, else: expand_node(node, st) + end) + + result_refs = Enum.map(body_list, &Map.fetch!(inner_state.node_to_ref, &1.data.id)) + + {result_refs, + %{ + state + | instructions: inner_state.instructions, + captures: inner_state.captures, + constants: inner_state.constants, + inputs: inner_state.inputs, + hooks: inner_state.hooks + }} + end + + # ── custom-fun reduce (static unroll) ────────────────────────────────────── + defp expand_reduce_unroll(id, out_type, out_shape, tensor, acc, opts, fun, state) do + in_rank = tuple_size(tensor.shape) + + reduce_axes = + case opts[:axes] do + nil -> Enum.to_list(0..(in_rank - 1)//1) + axes -> Enum.sort(normalize_axes(axes, in_rank)) + end + + if in_rank == 0 or reduce_axes == [] do + raise ArgumentError, "does not yet lower op :reduce with no reduction axes" + end + + [params, body, _mfa] = fun.data.args + + unless length(params) == 2 do + raise ArgumentError, "does not yet lower op :reduce with a non-binary reducer" + end + + reduce_set = MapSet.new(reduce_axes) + kept_axes = Enum.reject(0..(in_rank - 1)//1, &MapSet.member?(reduce_set, &1)) + in_dims = Tuple.to_list(tensor.shape) + kept_shape = Enum.map(kept_axes, &Enum.at(in_dims, &1)) + reduce_extent = reduce_axes |> Enum.map(&Enum.at(in_dims, &1)) |> Enum.product() + + # Cast input + initial acc to the reducer/output type. + tensor_ref0 = Map.fetch!(state.node_to_ref, tensor.data.id) + {tensor_ref, state} = emit_cast_if_needed(tensor_ref0, tensor.type, out_type, state) + acc_ref0 = Map.fetch!(state.node_to_ref, acc.data.id) + {acc_scalar_ref, state} = emit_cast_if_needed(acc_ref0, acc.type, out_type, state) + + # Move reduce axes last, then collapse them into a single trailing axis. + perm = kept_axes ++ reduce_axes + + {perm_ref, state} = + if perm == Enum.to_list(0..(in_rank - 1)//1) do + {tensor_ref, state} + else + emit_transpose_instr(tensor_ref, perm, state) + end + + combined_shape = kept_shape ++ [reduce_extent] + {combined_ref, state} = emit_reshape_instr(perm_ref, combined_shape, state) + + # Seed acc broadcast to the kept shape, then fold over the extent. + {acc_ref, state} = emit_broadcast_to(acc_scalar_ref, kept_shape, state) + + {final_ref, state} = + Enum.reduce(0..(reduce_extent - 1)//1, {acc_ref, state}, fn i, {acc_i, st} -> + {elem_ref, st} = emit_reduce_slice(combined_ref, combined_shape, kept_shape, i, st) + lower_fun_body(body, %{0 => elem_ref, 1 => acc_i}, st) + end) + + # Reshape to the declared output shape (restores keep_axes 1-dims). + {out_ref, state} = emit_reshape_instr(final_ref, Tuple.to_list(out_shape), state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, out_ref)} + end + + # Inline fun/while bodies sit outside top-level scope_ids; extend before lowering hooks. + defp merge_scope_ids(state, body) do + extra = body |> Nx.Defn.Tree.scope_ids() |> Map.keys() |> MapSet.new() + %{state | top_scope_ids: MapSet.union(state.top_scope_ids, extra)} + end + + defp lower_fun_body(body, param_ref_by_pos, state) do + state = merge_scope_ids(state, body) + inner_ordered = EMLX.Defn.Tree.post_order(body, &scope_dependencies/1) + + param_id_to_ref = + inner_ordered + |> Enum.filter(&(&1.data.op == :parameter)) + |> Map.new(fn p -> {p.data.id, Map.fetch!(param_ref_by_pos, hd(p.data.args))} end) + + param_id_set = MapSet.new(Map.keys(param_id_to_ref)) + local_base = Map.merge(state.node_to_ref, param_id_to_ref) + + inner_state = + Enum.reduce(inner_ordered, %{state | node_to_ref: local_base}, fn node, st -> + if MapSet.member?(param_id_set, node.data.id), do: st, else: expand_node(node, st) + end) + + result_ref = Map.fetch!(inner_state.node_to_ref, body.data.id) + + # Carry forward everything the body produced, but discard its node_to_ref. + {result_ref, + %{ + state + | instructions: inner_state.instructions, + captures: inner_state.captures, + constants: inner_state.constants, + inputs: inner_state.inputs, + hooks: inner_state.hooks + }} + end + + # Emit a :reshape instruction; returns {new_ref, state}. + defp emit_reshape_instr(ref, shape_list, state) do + new_ref = make_ref() + + {new_ref, + %{state | instructions: [{new_ref, :reshape, [ref], shape_list} | state.instructions]}} + end + + # Broadcast a scalar ref to `shape_list` (no-op when the target is scalar). + defp emit_broadcast_to(ref, [], state), do: {ref, state} + + defp emit_broadcast_to(ref, shape_list, state) do + new_ref = make_ref() + attrs = [length(shape_list) | shape_list] ++ [0] + + {new_ref, %{state | instructions: [{new_ref, :broadcast, [ref], attrs} | state.instructions]}} + end + + # Slice element `i` along the collapsed trailing axis then squeeze it away, + defp emit_reduce_slice(ref, combined_shape, kept_shape, i, state) do + n_dims = length(combined_shape) + last_axis = n_dims - 1 + lengths = kept_shape ++ [1] + strides = List.duplicate(1, n_dims) + starts = List.duplicate(0, length(kept_shape)) ++ [i] + attrs = [n_dims, 0] ++ combined_shape ++ lengths ++ strides ++ starts + + slice_ref = make_ref() + state = %{state | instructions: [{slice_ref, :slice, [ref], attrs} | state.instructions]} + squeeze_ref = make_ref() + + {squeeze_ref, + %{ + state + | instructions: [{squeeze_ref, :squeeze, [slice_ref], [last_axis]} | state.instructions] + }} + end + + defp emit_pad_with(ref, pad_value_ref, low_pads, high_pads, state) do + new_ref = make_ref() + n_dims = length(low_pads) + + attrs = + [n_dims | Enum.flat_map(Enum.zip(low_pads, high_pads), fn {lo, hi} -> [lo, hi, 0] end)] + + {new_ref, + %{ + state + | instructions: [{new_ref, :pad, [ref, pad_value_ref], attrs} | state.instructions] + }} + end + + defp emit_static_slice(ref, input_shape, starts, lengths, strides, state) do + new_ref = make_ref() + n_dims = length(input_shape) + attrs = [n_dims, 0] ++ input_shape ++ lengths ++ strides ++ starts + + {new_ref, %{state | instructions: [{new_ref, :slice, [ref], attrs} | state.instructions]}} + end + + defp expand_pad_general(tensor_ref, pad_value_ref, in_dims, config, state) do + interior_list = Enum.map(config, fn {_lo, _hi, interior} -> interior end) + + {interior_ref, state} = + if Enum.all?(interior_list, &(&1 == 0)) do + {tensor_ref, state} + else + emit_interior_padding(tensor_ref, pad_value_ref, in_dims, interior_list, state) + end + + interior_shape = + Enum.zip(in_dims, interior_list) + |> Enum.map(fn {d, interior} -> d + max(d - 1, 0) * interior end) + + {cropped_ref, _cropped_shape, state} = + if Enum.any?(config, fn {lo, hi, _interior} -> lo < 0 or hi < 0 end) do + emit_negative_crop(interior_ref, interior_shape, config, state) + else + {interior_ref, interior_shape, state} + end + + low_pads = Enum.map(config, fn {lo, _hi, _interior} -> max(lo, 0) end) + high_pads = Enum.map(config, fn {_lo, hi, _interior} -> max(hi, 0) end) + + if Enum.all?(low_pads ++ high_pads, &(&1 == 0)) do + {cropped_ref, state} + else + emit_pad_with(cropped_ref, pad_value_ref, low_pads, high_pads, state) + end + end + + defp emit_interior_padding(ref, pad_value_ref, in_dims, interior_list, state) do + rank = length(in_dims) + shape0 = in_dims ++ [1] + {ref, state} = emit_reshape_instr(ref, shape0, state) + + {final_ref, _final_shape, state} = + interior_list + |> Enum.with_index() + |> Enum.reduce({ref, shape0, state}, fn + {0, _axis_index}, {acc_ref, shape, st} -> + {acc_ref, shape, st} + + {interior, axis_index}, {acc_ref, shape, st} -> + next_axis = axis_index + 1 + axis_size = Enum.at(shape, axis_index) + next_axis_size = Enum.at(shape, next_axis) + + pad_lows = List.duplicate(0, rank + 1) + pad_highs = List.replace_at(pad_lows, next_axis, next_axis_size * interior) + {padded_ref, st} = emit_pad_with(acc_ref, pad_value_ref, pad_lows, pad_highs, st) + + new_axis_size = axis_size + axis_size * interior + + reshaped_shape = + shape + |> List.replace_at(axis_index, new_axis_size) + |> List.replace_at(next_axis, next_axis_size) + + {reshaped_ref, st} = emit_reshape_instr(padded_ref, reshaped_shape, st) + + sliced_shape = List.replace_at(reshaped_shape, axis_index, new_axis_size - interior) + starts = List.duplicate(0, rank + 1) + strides = List.duplicate(1, rank + 1) + + {sliced_ref, st} = + emit_static_slice(reshaped_ref, reshaped_shape, starts, sliced_shape, strides, st) + + {sliced_ref, sliced_shape, st} + end) + + squeeze_ref = make_ref() + + {squeeze_ref, + %{state | instructions: [{squeeze_ref, :squeeze, [final_ref], [rank]} | state.instructions]}} + end + + defp emit_negative_crop(ref, shape, config, state) do + starts = Enum.map(config, fn {lo, _hi, _interior} -> max(-lo, 0) end) + + lengths = + [shape, config, starts] + |> Enum.zip_with(fn [d, {_lo, hi, _interior}, start] -> + stop = if hi < 0, do: d + hi, else: d + stop - start + end) + + strides = List.duplicate(1, length(shape)) + {new_ref, state} = emit_static_slice(ref, shape, starts, lengths, strides, state) + {new_ref, lengths, state} + end + + defp window_offsets(k, dims) do + {digits, _} = + dims + |> Enum.reverse() + |> Enum.reduce({[], k}, fn d, {acc, n} -> {[rem(n, d) | acc], div(n, d)} end) + + digits + end + + @doc false + def quantizable_param_positions(output) do + output + |> EMLX.Defn.Tree.post_order(&scope_dependencies/1) + |> Enum.reduce(MapSet.new(), fn + %T{data: %Nx.Defn.Expr{op: :dot, args: [left, _c_left, _b_left, right, _c_right, _b_right]}}, + acc -> + acc + |> maybe_put_param_position(left) + |> maybe_put_param_position(right) + + _, acc -> + acc + end) + end + + defp maybe_put_param_position(acc, %T{data: %Nx.Defn.Expr{op: :parameter, args: [pos]}}), + do: MapSet.put(acc, pos) + + defp maybe_put_param_position(acc, _node), do: acc + + @doc false + def f64_bits(v) when is_number(v) do + <> = <> + bits + end + + @doc false + def bits_to_f64(bits) when is_integer(bits) do + <> = <> + v + end + + # ── cond helper ─────────────────────────────────────────────────────────── + + defp flat_refs(composite, state) do + Composite.flatten_list([composite]) + |> Enum.map(&Map.fetch!(state.node_to_ref, &1.data.id)) + end + + # ── binary lowering helpers ──────────────────────────────────────────────── + + defp expand_binary_node(id, op, out_type, left, right, state) do + merge_type = Nx.Type.merge(left.type, right.type) + left_ref0 = Map.fetch!(state.node_to_ref, left.data.id) + right_ref0 = Map.fetch!(state.node_to_ref, right.data.id) + + {left_ref, state} = emit_cast_if_needed(left_ref0, left.type, merge_type, state) + {right_ref, state} = emit_cast_if_needed(right_ref0, right.type, merge_type, state) + + op_ref = make_ref() + + state = %{ + state + | instructions: [{op_ref, op, [left_ref, right_ref], []} | state.instructions] + } + + {result_ref, state} = emit_cast_if_needed(op_ref, merge_type, out_type, state) + %{state | node_to_ref: Map.put(state.node_to_ref, id, result_ref)} + end + + # Emit an :astype instruction; returns {new_ref, updated_state}. + defp emit_cast_to(ref, nx_type, state) do + cast_ref = make_ref() + mlx_type = EMLX.Native.to_mlx_type(nx_type) + instr = {cast_ref, :astype, [ref], [mlx_type]} + {cast_ref, %{state | instructions: [instr | state.instructions]}} + end + + # Emit :astype only when the MLX type representation of from_type differs from to_type. + defp emit_cast_if_needed(ref, from_type, to_type, state) do + if EMLX.Native.to_mlx_type(from_type) == EMLX.Native.to_mlx_type(to_type) do + {ref, state} + else + emit_cast_to(ref, to_type, state) + end + end + + # Normalise negative axis values to non-negative, given the input tensor rank. + defp normalize_axes(axes, rank) do + Enum.map(axes, fn ax -> if ax < 0, do: rank + ax, else: ax end) + end + + # Emit a :transpose instruction; returns {result_ref, updated_state}. + defp emit_transpose_instr(operand_ref, perm, state) do + ref = make_ref() + {ref, %{state | instructions: [{ref, :transpose, [operand_ref], perm} | state.instructions]}} + end + + defp swap_last_two_axes(rank) do + {front, [x, y]} = Enum.split(Enum.to_list(0..(rank - 1)), rank - 2) + front ++ [y, x] + end + + # Emit a :solve_triangular instruction; returns {result_ref, updated_state}. + defp emit_solve_triangular_instr(a_ref, b_ref, upper, state) do + ref = make_ref() + upper_int = if upper, do: 1, else: 0 + + {ref, + %{ + state + | instructions: [ + {ref, :solve_triangular, [a_ref, b_ref], [upper_int]} | state.instructions + ] + }} + end + + # Move the second element (channels) to the last position. + defp move_channels_last([head | [second | rest]]) do + [head | rest] ++ [second] + end + + # ── wire serialisation ──────────────────────────────────────────────────── + + @doc false + def to_native(%__MODULE__{} = prog) do + # Build ref → wire-ref map for all non-instruction nodes. A wire ref is a + # tagged tuple ({:input, i} / {:capture, i} / {:const, i} / {:result, i}) + # instead of a bit-packed int — see EMLX.Native.Instruction.ref/0. + input_map = + prog.inputs + |> Enum.with_index() + |> Map.new(fn {ref, i} -> {ref, {:input, i}} end) + + capture_map = + prog.captures + |> Enum.with_index() + |> Map.new(fn {{ref, _t}, i} -> {ref, {:capture, i}} end) + + constant_map = + prog.constants + |> Enum.with_index() + |> Map.new(fn {{ref, _v, _t}, i} -> {ref, {:const, i}} end) + + ref_to_wire = Map.merge(input_map, Map.merge(capture_map, constant_map)) + + maybe_debug_check do + expected_size = map_size(input_map) + map_size(capture_map) + map_size(constant_map) + + if map_size(ref_to_wire) != expected_size do + raise ArgumentError, + "EMLX.Native.Expr.to_native: ref id collision across inputs/captures/constants -- " <> + "#{map_size(input_map)} input(s), #{map_size(capture_map)} capture(s), " <> + "#{map_size(constant_map)} constant(s) should merge to #{expected_size} distinct " <> + "refs, but only #{map_size(ref_to_wire)} survived Map.merge/2. This means two " <> + "refs of different categories share the same id, silently dropping one from the " <> + "wire map -- inputs: #{inspect(Map.keys(input_map))}, " <> + "captures: #{inspect(Map.keys(capture_map))}, " <> + "constants: #{inspect(Map.keys(constant_map))}" + end + end + + {instructions, ref_to_wire, _flat} = + prog.instructions + |> Enum.reduce({[], ref_to_wire, 0}, fn {id, op, operand_refs, attrs}, + {instrs, rmap, flat} -> + wire_operands = Enum.map(operand_refs, &Map.fetch!(rmap, &1)) + + maybe_debug_check do + for {:result, idx} <- wire_operands, idx >= flat do + raise ArgumentError, + "EMLX.Native.Expr.to_native: instruction #{inspect(op)} (id=#{inspect(id)}) " <> + "references result index #{idx} of the flat results accumulator, but only " <> + "#{flat} result(s) have been produced so far -- this is a forward/self " <> + "reference bug in program lowering, not a valid program. Full instruction " <> + "list: #{inspect(prog.instructions)}" + end + end + + maybe_debug_check do + for one <- List.wrap(id), Map.has_key?(rmap, one) do + raise ArgumentError, + "EMLX.Native.Expr.to_native: instruction #{inspect(op)} produces result ref " <> + "#{inspect(one)}, but that ref is already bound (to " <> + "#{inspect(Map.fetch!(rmap, one))}) -- the same node id was lowered twice, " <> + "silently overwriting its earlier binding for every prior instruction that " <> + "already referenced it. Full instruction list: #{inspect(prog.instructions)}" + end + end + + {rmap2, flat2} = + case id do + ids when is_list(ids) -> + Enum.reduce(ids, {rmap, flat}, fn one, {m, f} -> + {Map.put(m, one, {:result, f}), f + 1} + end) + + one -> + {Map.put(rmap, one, {:result, flat}), flat + 1} + end + + instr = %EMLX.Native.Instruction{op: op, operands: wire_operands, attrs: attrs} + {[instr | instrs], rmap2, flat2} + end) + + hook_refs = Enum.flat_map(prog.hooks, & &1.refs) + wire_outputs = Enum.map(prog.outputs ++ hook_refs, &Map.fetch!(ref_to_wire, &1)) + + capture_nif_refs = + Enum.map(prog.captures, fn {_ref, %Nx.Tensor{data: %EMLX.Backend{ref: {_, nif_ref}}}} -> + nif_ref + end) + + wire_constants = + Enum.map(prog.constants, fn {_, v, t} -> {v * 1.0, EMLX.Native.to_mlx_type(t)} end) + + %EMLX.Native.Program{ + num_inputs: length(prog.inputs), + captures: capture_nif_refs, + constants: wire_constants, + instructions: Enum.reverse(instructions), + outputs: wire_outputs + } + end +end diff --git a/emlx/lib/emlx/native/instruction.ex b/emlx/lib/emlx/native/instruction.ex new file mode 100644 index 0000000..feb0458 --- /dev/null +++ b/emlx/lib/emlx/native/instruction.ex @@ -0,0 +1,28 @@ +defmodule EMLX.Native.Instruction do + @moduledoc false + + # One compiled-program instruction, as passed to the `compile_program` NIF. + # Decoded directly on the C++ side by Fine (see emlx_compiler.hpp's + # `Instruction` struct + `fine::Decoder`), replacing the old + # to_native/1 format's parallel `op_names`/`operands`/`attrs` lists. + + @enforce_keys [:op, :operands, :attrs] + defstruct [:op, :operands, :attrs] + + # A reference to an already-produced value, resolved on the C++ side against + # the runtime inputs / closed-over captures / closed-over constants / the + # flat per-eval results accumulator (in that order). + @type ref :: + {:input, non_neg_integer()} + | {:capture, non_neg_integer()} + | {:const, non_neg_integer()} + | {:result, non_neg_integer()} + + # Most attrs are plain integers (shapes, axes, flags, f64_bits-encoded + # floats); a few are MLX dtype atoms or quantized_matmul mode atoms (see + # emlx_compiler.hpp's `Attr` type) — no int<->meaning lookup table needs to + # be kept in sync between Elixir and C++ for those. + @type attr :: integer() | atom() + + @type t :: %__MODULE__{op: atom(), operands: [ref()], attrs: [attr()]} +end diff --git a/emlx/lib/emlx/native/program.ex b/emlx/lib/emlx/native/program.ex new file mode 100644 index 0000000..d58c013 --- /dev/null +++ b/emlx/lib/emlx/native/program.ex @@ -0,0 +1,15 @@ +defmodule EMLX.Native.Program do + @moduledoc false + @enforce_keys [:num_inputs, :captures, :constants, :instructions, :outputs] + defstruct [:num_inputs, :captures, :constants, :instructions, :outputs] + + @type constant :: {float(), atom()} + + @type t :: %__MODULE__{ + num_inputs: non_neg_integer(), + captures: [reference()], + constants: [constant()], + instructions: [EMLX.Native.Instruction.t()], + outputs: [EMLX.Native.Instruction.ref()] + } +end diff --git a/emlx/lib/emlx/native/qwen3.ex b/emlx/lib/emlx/native/qwen3.ex new file mode 100644 index 0000000..fcb75b6 --- /dev/null +++ b/emlx/lib/emlx/native/qwen3.ex @@ -0,0 +1,762 @@ +defmodule EMLX.Native.Qwen3 do + @moduledoc """ + Fused native NIF wrappers for Qwen3 transformer inference. + + These are hand-written (non-`@mlx_function`) wrappers around the + `qwen3_*` NIFs registered directly in `EMLX.NIF` (see that module's + `qwen3_*` stubs) — the underlying NIF/C++ names keep the `qwen3_` prefix + (`emlx/c_src/emlx_fast/qwen3.cpp`), only this module's public API drops it + since the module name itself already disambiguates. + """ + + import EMLX, only: [is_tensor: 2] + + @doc """ + Qwen3 fused RoPE + KV cache update + SDPA for the native decode path. + + Accepts `query`, `new_key`, and `new_value` in projection layout + `{B, T, N, D}`. The NIF transposes Q/K/V, applies Qwen3 RoPE to Q/K, updates + the owned KV cache, runs SDPA, and returns flattened attention output + `{B, T, N * D}` ready for projection, plus updated cache refs. + """ + def kv_cache_attention( + {dev_q, ref_q}, + {_dev_k, ref_k}, + {_dev_v, ref_v}, + {_dev_kc, ref_kc}, + {_dev_vc, ref_vc}, + offset, + scale, + head_dim, + theta + ) + when is_tensor(dev_q, ref_q) and is_integer(offset) and is_float(scale) and + is_integer(head_dim) and is_number(theta) do + device = dev_q + {worker, effective_device} = EMLX.resolve_worker(device) + + {attn_ref, k_upd_ref, v_upd_ref} = + EMLX.NIF.qwen3_kv_cache_attention( + worker, + ref_q, + ref_k, + ref_v, + ref_kc, + ref_vc, + offset, + scale, + head_dim, + theta * 1.0, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {{effective_device, attn_ref}, {effective_device, k_upd_ref}, {effective_device, v_upd_ref}} + end + + @doc """ + Qwen3 dense MLP helper. + + Accepts hidden states `{B, T, H}`, RMSNorm weight after attention `{H}`, + dense gate/up projections `{H, I}`, dense down projection `{I, H}`, and RMSNorm + epsilon. Returns `hidden + mlp_output` as `{B, T, H}`. + """ + def mlp( + {dev_h, ref_h}, + {_dev_norm, ref_norm}, + {_dev_gate, ref_gate}, + {_dev_up, ref_up}, + {_dev_down, ref_down}, + eps + ) + when is_tensor(dev_h, ref_h) and is_float(eps) do + device = dev_h + {worker, effective_device} = EMLX.resolve_worker(device) + + out_ref = + EMLX.NIF.qwen3_mlp( + worker, + ref_h, + ref_norm, + ref_gate, + ref_up, + ref_down, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {effective_device, out_ref} + end + + @doc """ + Qwen3 dense attention output projection plus residual add. + + Accepts residual hidden state `{B, T, H}`, flattened attention output + `{B, T, I}`, and dense output projection `{I, H}`. Returns + `hidden + projected_attention`. + """ + def attention_residual( + {dev_h, ref_h}, + {_dev_attn, ref_attn}, + {_dev_o, ref_o} + ) + when is_tensor(dev_h, ref_h) do + device = dev_h + {worker, effective_device} = EMLX.resolve_worker(device) + + out_ref = + EMLX.NIF.qwen3_attention_residual( + worker, + ref_h, + ref_attn, + ref_o, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {effective_device, out_ref} + end + + @doc """ + Qwen3 dense transformer layer helper. + + Accepts hidden states `{B, T, H}`, input RMSNorm and RMSNorm weights after attention, + dense attention projections, Q/K RMSNorm weights, owned KV cache refs, dense + MLP projections, offset, scale, RoPE parameters, and RMSNorm epsilon. Returns + `{hidden_out, k_cache, v_cache}`. + """ + def layer( + {dev_h, ref_h}, + {_dev_norm1, ref_norm1}, + {_dev_q, ref_q}, + {_dev_k, ref_k}, + {_dev_v, ref_v}, + {_dev_o, ref_o}, + {_dev_qn, ref_qn}, + {_dev_kn, ref_kn}, + {_dev_kc, ref_kc}, + {_dev_vc, ref_vc}, + {_dev_norm2, ref_norm2}, + {_dev_gate, ref_gate}, + {_dev_up, ref_up}, + {_dev_down, ref_down}, + offset, + scale, + head_dim, + theta, + eps + ) + when is_tensor(dev_h, ref_h) and is_integer(offset) and is_float(scale) and + is_integer(head_dim) and is_number(theta) and is_float(eps) do + device = dev_h + {worker, effective_device} = EMLX.resolve_worker(device) + + {out_ref, k_upd_ref, v_upd_ref} = + EMLX.NIF.qwen3_layer( + worker, + ref_h, + ref_norm1, + ref_q, + ref_k, + ref_v, + ref_o, + ref_qn, + ref_kn, + ref_kc, + ref_vc, + ref_norm2, + ref_gate, + ref_up, + ref_down, + offset, + scale, + head_dim, + theta * 1.0, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {{effective_device, out_ref}, {effective_device, k_upd_ref}, {effective_device, v_upd_ref}} + end + + @doc """ + Qwen3 generalized transformer layer helper: same fusion as `layer/18`, + but each of the 7 projections (q/k/v/o/gate/up/down) independently accepts + a dense or quantized (`EMLX.quantize/2`) `Nx.Tensor`, collapsing the + quantized native lane's per-op NIF calls down to one fused call. + + Accepts hidden states `{B, T, H}`, input RMSNorm weight, q/k/v/o projections + (dense or quantized), Q/K RMSNorm weights, owned KV cache refs, post-attention + RMSNorm weight, gate/up/down projections (dense or quantized), offset, scale, + RoPE parameters, and RMSNorm epsilon. Returns `{hidden_out, k_cache, v_cache}`. + """ + def layer_quantized( + {dev_h, ref_h}, + norm1, + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + {_dev_kc, ref_kc}, + {_dev_vc, ref_vc}, + norm2, + gate_proj, + up_proj, + down_proj, + offset, + scale, + head_dim, + theta, + eps + ) + when is_tensor(dev_h, ref_h) and is_integer(offset) and is_float(scale) and + is_integer(head_dim) and is_number(theta) and is_float(eps) do + device = dev_h + {worker, effective_device} = EMLX.resolve_worker(device) + + {out_ref, k_upd_ref, v_upd_ref} = + EMLX.NIF.qwen3_layer_quantized( + worker, + ref_h, + tensor_ref!(norm1), + linear_weight_term(q_proj), + linear_weight_term(k_proj), + linear_weight_term(v_proj), + linear_weight_term(o_proj), + tensor_ref!(q_norm), + tensor_ref!(k_norm), + ref_kc, + ref_vc, + tensor_ref!(norm2), + linear_weight_term(gate_proj), + linear_weight_term(up_proj), + linear_weight_term(down_proj), + offset, + scale, + head_dim, + theta * 1.0, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {{effective_device, out_ref}, {effective_device, k_upd_ref}, {effective_device, v_upd_ref}} + end + + @doc """ + Qwen3 embedding lookup plus dense forward through all layers and greedy token. + + This accepts token ids and embedding weights directly, so the dense greedy + path can avoid constructing a separate embedding lookup graph before entering + the native Qwen3 worker call. + """ + def forward_greedy_ids( + {dev_ids, ref_ids}, + {_dev_embed, ref_embed}, + layers, + kv_cache, + {_dev_norm, ref_norm}, + {_dev_lm_head, ref_lm_head}, + offset, + scale, + head_dim, + theta, + eps + ) + when is_tensor(dev_ids, ref_ids) and is_list(layers) and is_list(kv_cache) and + is_integer(offset) and is_float(scale) and is_integer(head_dim) and + is_number(theta) and is_float(eps) do + device = dev_ids + {worker, effective_device} = EMLX.resolve_worker(device) + + layer_refs = Enum.map(layers, &layer_refs!/1) + kv_refs = kv_refs(kv_cache) + + {token_ref, kv_updated_refs} = + EMLX.NIF.qwen3_forward_greedy_ids( + worker, + ref_ids, + ref_embed, + layer_refs, + kv_refs, + ref_norm, + ref_lm_head, + offset, + scale, + head_dim, + theta * 1.0, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {{effective_device, token_ref}, + Enum.map(kv_updated_refs, fn {k_ref, v_ref} -> + {{effective_device, k_ref}, {effective_device, v_ref}} + end)} + end + + @doc """ + Qwen3 greedy decode chunk helper. + + Starting from a `{1, 1}` token id tensor, runs `count` greedy decode steps + without returning to Elixir between steps. Returns `{token_refs, kv_cache}`, + where `token_refs` is a list of raw EMLX token refs in generation order and + `kv_cache` is the final updated raw cache. + """ + def forward_greedy_ids_chunk( + {dev_ids, ref_ids}, + {_dev_embed, ref_embed}, + layers, + kv_cache, + {_dev_norm, ref_norm}, + {_dev_lm_head, ref_lm_head}, + offset, + count, + scale, + head_dim, + theta, + eps + ) + when is_tensor(dev_ids, ref_ids) and is_list(layers) and is_list(kv_cache) and + is_integer(offset) and is_integer(count) and count > 0 and is_float(scale) and + is_integer(head_dim) and is_number(theta) and is_float(eps) do + device = dev_ids + {worker, effective_device} = EMLX.resolve_worker(device) + + assert_decode_ids!({dev_ids, ref_ids}, "forward_greedy_ids_chunk") + + layer_refs = Enum.map(layers, &layer_refs!/1) + kv_refs = kv_refs(kv_cache) + + {token_refs, kv_updated_refs} = + EMLX.NIF.qwen3_forward_greedy_ids_chunk( + worker, + ref_ids, + ref_embed, + layer_refs, + kv_refs, + ref_norm, + ref_lm_head, + offset, + count, + scale, + head_dim, + theta * 1.0, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + tokens = Enum.map(token_refs, &{effective_device, &1}) + + kv_cache = + Enum.map(kv_updated_refs, fn {k_ref, v_ref} -> + {{effective_device, k_ref}, {effective_device, v_ref}} + end) + + {tokens, kv_cache} + end + + @doc """ + Qwen3 generalized greedy decode chunk helper: same fusion as + `forward_greedy_ids_chunk/12`, but every layer's 7 projections and the + final `lm_head` each independently accept a dense or quantized + (`EMLX.quantize/2`) `Nx.Tensor`. This lets the quantized native lane fuse a + whole multi-token decode chunk into 1 NIF call, matching dense's chunk fusion. + + Starting from a `{1, 1}` token id tensor, runs `count` greedy decode steps + without returning to Elixir between steps. Returns `{token_refs, kv_cache}`, + where `token_refs` is a list of raw EMLX token refs in generation order and + `kv_cache` is the final updated raw cache. + """ + def forward_greedy_ids_chunk_quantized( + {dev_ids, ref_ids}, + {_dev_embed, ref_embed}, + layers, + kv_cache, + {_dev_norm, ref_norm}, + lm_head, + offset, + count, + scale, + head_dim, + theta, + eps + ) + when is_tensor(dev_ids, ref_ids) and is_list(layers) and is_list(kv_cache) and + is_integer(offset) and is_integer(count) and count > 0 and is_float(scale) and + is_integer(head_dim) and is_number(theta) and is_float(eps) do + device = dev_ids + {worker, effective_device} = EMLX.resolve_worker(device) + + assert_decode_ids!({dev_ids, ref_ids}, "forward_greedy_ids_chunk_quantized") + + layer_terms = Enum.map(layers, &layer_weight_terms!/1) + kv_refs = kv_refs(kv_cache) + + {token_refs, kv_updated_refs} = + EMLX.NIF.qwen3_forward_greedy_ids_chunk_quantized( + worker, + ref_ids, + ref_embed, + layer_terms, + kv_refs, + ref_norm, + linear_weight_term(lm_head), + offset, + count, + scale, + head_dim, + theta * 1.0, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + tokens = Enum.map(token_refs, &{effective_device, &1}) + + kv_cache = + Enum.map(kv_updated_refs, fn {k_ref, v_ref} -> + {{effective_device, k_ref}, {effective_device, v_ref}} + end) + + {tokens, kv_cache} + end + + @doc """ + Qwen3 embedding lookup plus dense forward through all layers and greedy token. + + This variant returns the selected token id as a BEAM integer while keeping the + updated KV cache as raw EMLX refs. It is intended for streaming decode paths + that need the token on the host anyway. + """ + def forward_greedy_ids_token_id( + {dev_ids, ref_ids}, + {_dev_embed, ref_embed}, + layers, + kv_cache, + {_dev_norm, ref_norm}, + {_dev_lm_head, ref_lm_head}, + offset, + scale, + head_dim, + theta, + eps + ) + when is_tensor(dev_ids, ref_ids) and is_list(layers) and is_list(kv_cache) and + is_integer(offset) and is_float(scale) and is_integer(head_dim) and + is_number(theta) and is_float(eps) do + device = dev_ids + {worker, effective_device} = EMLX.resolve_worker(device) + + assert_batch_size_one!({dev_ids, ref_ids}, "forward_greedy_ids_token_id") + + layer_refs = Enum.map(layers, &layer_refs!/1) + kv_refs = kv_refs(kv_cache) + + {token_id, kv_updated_refs} = + EMLX.NIF.qwen3_forward_greedy_ids_token_id( + worker, + ref_ids, + ref_embed, + layer_refs, + kv_refs, + ref_norm, + ref_lm_head, + offset, + scale, + head_dim, + theta * 1.0, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {token_id, + Enum.map(kv_updated_refs, fn {k_ref, v_ref} -> + {{effective_device, k_ref}, {effective_device, v_ref}} + end)} + end + + @doc """ + Qwen3 dense forward and greedy decode for one token. + + This variant accepts the previous token id as a BEAM integer and returns the + selected token id as a BEAM integer. It is intended for decode loops that + already synchronize once per token for streaming or EOS checks. By default it + runs on the embedding tensor's device; pass `device` explicitly to override. + """ + def forward_greedy_token_id( + token_id, + {dev_embed, ref_embed}, + layers, + kv_cache, + {_dev_norm, ref_norm}, + {_dev_lm_head, ref_lm_head}, + offset, + scale, + head_dim, + theta, + eps, + device \\ nil + ) + when is_integer(token_id) and is_list(layers) and is_list(kv_cache) and + is_integer(offset) and is_float(scale) and is_integer(head_dim) and + is_number(theta) and is_float(eps) do + device = device || dev_embed + {worker, effective_device} = EMLX.resolve_worker(device) + + layer_refs = Enum.map(layers, &layer_refs!/1) + kv_refs = kv_refs(kv_cache) + + {next_token_id, kv_updated_refs} = + EMLX.NIF.qwen3_forward_greedy_token_id( + worker, + token_id, + ref_embed, + layer_refs, + kv_refs, + ref_norm, + ref_lm_head, + offset, + scale, + head_dim, + theta * 1.0, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {next_token_id, + Enum.map(kv_updated_refs, fn {k_ref, v_ref} -> + {{effective_device, k_ref}, {effective_device, v_ref}} + end)} + end + + @doc """ + Qwen3 dense final RMSNorm + lm_head + greedy argmax helper. + + Accepts hidden states `{B, T, H}`, final RMSNorm weight `{H}`, dense lm_head + `{V, H}`, and RMSNorm epsilon. Returns token ids as `{B}`. + """ + def final_greedy( + {dev_h, ref_h}, + {_dev_norm, ref_norm}, + {_dev_lm_head, ref_lm_head}, + eps + ) + when is_tensor(dev_h, ref_h) and is_float(eps) do + device = dev_h + {worker, effective_device} = EMLX.resolve_worker(device) + + out_ref = + EMLX.NIF.qwen3_final_greedy( + worker, + ref_h, + ref_norm, + ref_lm_head, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {effective_device, out_ref} + end + + @doc """ + Qwen3 dense attention block helper. + + Accepts residual hidden state before attention `{B, T, H}`, input RMSNorm weight + `{H}`, dense Q/K/V/O projections, Q/K RMSNorm weights, owned KV cache refs, + offset, scale, RoPE parameters, and RMSNorm epsilon. Returns + `{hidden_out, k_cache, v_cache}`. + """ + def attention_block( + {dev_h, ref_h}, + {_dev_norm, ref_norm}, + {_dev_q, ref_q}, + {_dev_k, ref_k}, + {_dev_v, ref_v}, + {_dev_o, ref_o}, + {_dev_qn, ref_qn}, + {_dev_kn, ref_kn}, + {_dev_kc, ref_kc}, + {_dev_vc, ref_vc}, + offset, + scale, + head_dim, + theta, + eps + ) + when is_tensor(dev_h, ref_h) and is_integer(offset) and is_float(scale) and + is_integer(head_dim) and is_number(theta) and is_float(eps) do + device = dev_h + {worker, effective_device} = EMLX.resolve_worker(device) + + {out_ref, k_upd_ref, v_upd_ref} = + EMLX.NIF.qwen3_attention_block( + worker, + ref_h, + ref_norm, + ref_q, + ref_k, + ref_v, + ref_o, + ref_qn, + ref_kn, + ref_kc, + ref_vc, + offset, + scale, + head_dim, + theta * 1.0, + eps, + effective_device + ) + |> EMLX.unwrap!() + |> EMLX.await_worker() + + {{effective_device, out_ref}, {effective_device, k_upd_ref}, {effective_device, v_upd_ref}} + end + + defp layer_refs!( + {norm1, norm2, q_norm, k_norm, q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, + down_proj} + ) do + { + tensor_ref!(norm1), + tensor_ref!(norm2), + tensor_ref!(q_norm), + tensor_ref!(k_norm), + tensor_ref!(q_proj), + tensor_ref!(k_proj), + tensor_ref!(v_proj), + tensor_ref!(o_proj), + tensor_ref!(gate_proj), + tensor_ref!(up_proj), + tensor_ref!(down_proj) + } + end + + # Generalized variant of `layer_refs!/1`: q/k/v/o/gate/up/down each + # become a `linear_weight_term/1` (dense or quantized) instead of a + # plain ref, for `layer_quantized`/`forward_greedy_ids_chunk_quantized`. + defp layer_weight_terms!( + {norm1, norm2, q_norm, k_norm, q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, + down_proj} + ) do + { + tensor_ref!(norm1), + tensor_ref!(norm2), + tensor_ref!(q_norm), + tensor_ref!(k_norm), + linear_weight_term(q_proj), + linear_weight_term(k_proj), + linear_weight_term(v_proj), + linear_weight_term(o_proj), + linear_weight_term(gate_proj), + linear_weight_term(up_proj), + linear_weight_term(down_proj) + } + end + + defp kv_refs!({{_device_k, ref_k}, {_device_v, ref_v}}) + when is_reference(ref_k) and is_reference(ref_v), + do: {ref_k, ref_v} + + defp kv_refs!({k_cache, v_cache}), do: {tensor_ref!(k_cache), tensor_ref!(v_cache)} + + defp kv_refs([{{_device_k, ref_k}, {_device_v, ref_v}} | _rest] = kv_cache) + when is_reference(ref_k) and is_reference(ref_v), + do: kv_cache + + defp kv_refs(kv_cache), do: Enum.map(kv_cache, &kv_refs!/1) + + defp assert_batch_size_one!(tensor, function_name) do + case EMLX.shape(tensor) do + {1, _seq_len} -> + :ok + + {batch_size, _seq_len} -> + raise ArgumentError, + "#{function_name} requires batch size 1, got batch size #{batch_size}" + + shape -> + raise ArgumentError, + "#{function_name} expects rank-2 input ids, got shape #{inspect(shape)}" + end + end + + defp assert_decode_ids!(tensor, function_name) do + case EMLX.shape(tensor) do + {1, 1} -> + :ok + + {1, seq_len} -> + raise ArgumentError, + "#{function_name} requires sequence length 1, got sequence length #{seq_len}" + + {batch_size, _seq_len} -> + raise ArgumentError, + "#{function_name} requires batch size 1, got batch size #{batch_size}" + + shape -> + raise ArgumentError, + "#{function_name} expects rank-2 input ids, got shape #{inspect(shape)}" + end + end + + # Builds the weight-term tuple accepted by `qwen3_get_linear_weight` in + # emlx_fast/qwen3.cpp: `{:dense, ref}` for a plain tensor, or + # `{:quantized, weight_ref, scales_ref, biases_ref_or_nil, group_size, bits, + # mode, true}` for a tensor produced by `EMLX.quantize/2`. `transpose` is + # always `true` for quantized terms here — every Qwen3 projection and + # lm_head uses the {out,in} physical layout (same convention hardcoded by + # `EMLX.quantized_matmul/2`); dense orientation is instead selected + # by the C++ call site (`dense_transpose` arg of `qwen3_get_linear_weight`). + defp linear_weight_term(%Nx.Tensor{ + data: %EMLX.Backend{ref: {_device, ref}, quantization_config: nil} + }) do + {:dense, ref} + end + + defp linear_weight_term(%Nx.Tensor{ + data: %EMLX.Backend{ + ref: {_device, ref}, + quantization_config: %EMLX.Quantization.Config{} = cfg + } + }) do + {_scales_device, scales_ref} = EMLX.Backend.from_nx(cfg.scales) + biases_ref = cfg.biases && elem(EMLX.Backend.from_nx(cfg.biases), 1) + + {:quantized, ref, scales_ref, biases_ref, cfg.group_size, cfg.bits, cfg.mode, true} + end + + defp linear_weight_term(tensor) do + {_device, ref} = EMLX.Backend.from_nx(tensor) + {:dense, ref} + end + + defp tensor_ref!(%Nx.Tensor{data: %EMLX.Backend{ref: {_device, ref}}}), do: ref + + defp tensor_ref!(tensor) do + {_device, ref} = EMLX.Backend.from_nx(tensor) + ref + end +end diff --git a/emlx/lib/emlx/nif.ex b/emlx/lib/emlx/nif.ex index 565e56a..cd44498 100644 --- a/emlx/lib/emlx/nif.ex +++ b/emlx/lib/emlx/nif.ex @@ -33,6 +33,11 @@ defmodule EMLX.NIF do :erlang.nif_error(:nif_not_loaded) end + # See the comment on the C++ side (emlx_nif.cpp). + def eval_many(_worker, _tensor_refs) do + :erlang.nif_error(:nif_not_loaded) + end + def to_device(_worker, _tensor, _device) do :erlang.nif_error(:nif_not_loaded) end @@ -102,19 +107,244 @@ defmodule EMLX.NIF do :erlang.nif_error(:nif_not_loaded) end - # ── Graph capture / replay ───────────────────────────────────────────────── - # graph_capture/3 — walks the lazy MLX DAG from `outputs` back to `inputs`, - # builds and simplifies a replayable tape, and returns an opaque compiled_ref. - # Must be called while the arrays are still lazy (before eval). - # - # graph_replay/2 — substitutes `new_inputs` into the stored tape and returns - # new lazy output arrays without re-dispatching any Nx ops. + # ── Native.Expr compiler NIFs ───────────────────────────────────────────── + # compile_program — builds an opaque ProgramResource from the Native.Expr IR + # and a set of captured arrays. Worker-routed (argv[0] = worker). + # `program` is an `EMLX.Native.Program.t()` (see + # EMLX.Native.Expr.to_native/1), decoded directly by `fine` on the C++ side + # (emlx_compiler.hpp's `Program`/`Instruction` structs) instead of manually + # parsed positional args. + # Arity = 1 (worker) + 1 = 2 registered. + def compile_program(_worker, _program) do + :erlang.nif_error(:nif_not_loaded) + end + + # eval_program — replays a compiled ProgramResource against runtime inputs. + # Worker-routed (argv[0] = worker). Returns a list of output MLX array refs. + # Arity = 1 (worker) + 2 = 3 registered. + def eval_program(_worker, _program_ref, _input_refs) do + :erlang.nif_error(:nif_not_loaded) + end + + # resolve_runtime_call — delivers the real Elixir callback's reply for one + # in-flight Nx.runtime_call/4 round trip (see EMLX.await_worker/2 and + # EMLX.Native.Expr's moduledoc "Runtime calls" section), waking the worker + # thread blocked inside EMLXRuntimeCall::eval_cpu/eval_gpu. NOT + # worker-routed — no argv[0] worker ref; runs directly on the calling + # scheduler thread. + def resolve_runtime_call(_pending, _status, _result) do + :erlang.nif_error(:nif_not_loaded) + end + + # ── Qwen3 fused native NIFs ──────────────────────────────────────────────── + # Elixir wrappers live in `EMLX.Native.Qwen3` (without the `qwen3_` prefix); + # the NIF/C++ names below keep it (`emlx/c_src/emlx_fast/qwen3.cpp`). + # Worker-routed (argv[0] = worker), same pattern as the `@mlx_function` + # macro-generated stubs above. + + def qwen3_kv_cache_attention( + _worker, + _q, + _k, + _v, + _k_cache, + _v_cache, + _offset, + _scale, + _head_dim, + _theta, + _device + ) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_mlp(_worker, _hidden, _norm, _gate, _up, _down, _eps, _device) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_attention_residual(_worker, _hidden, _attn, _o_proj, _device) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_layer( + _worker, + _hidden, + _norm1, + _q, + _k, + _v, + _o, + _q_norm, + _k_norm, + _k_cache, + _v_cache, + _norm2, + _gate, + _up, + _down, + _offset, + _scale, + _head_dim, + _theta, + _eps, + _device + ) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_layer_quantized( + _worker, + _hidden, + _norm1, + _q, + _k, + _v, + _o, + _q_norm, + _k_norm, + _k_cache, + _v_cache, + _norm2, + _gate, + _up, + _down, + _offset, + _scale, + _head_dim, + _theta, + _eps, + _device + ) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_forward_greedy_ids( + _worker, + _ids, + _embed, + _layers, + _kv_cache, + _norm, + _lm_head, + _offset, + _scale, + _head_dim, + _theta, + _eps, + _device + ) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_forward_greedy_ids_chunk( + _worker, + _ids, + _embed, + _layers, + _kv_cache, + _norm, + _lm_head, + _offset, + _count, + _scale, + _head_dim, + _theta, + _eps, + _device + ) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_forward_greedy_ids_chunk_quantized( + _worker, + _ids, + _embed, + _layers, + _kv_cache, + _norm, + _lm_head, + _offset, + _count, + _scale, + _head_dim, + _theta, + _eps, + _device + ) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_forward_greedy_ids_token_id( + _worker, + _ids, + _embed, + _layers, + _kv_cache, + _norm, + _lm_head, + _offset, + _scale, + _head_dim, + _theta, + _eps, + _device + ) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_forward_greedy_token_id( + _worker, + _token_id, + _embed, + _layers, + _kv_cache, + _norm, + _lm_head, + _offset, + _scale, + _head_dim, + _theta, + _eps, + _device + ) do + :erlang.nif_error(:nif_not_loaded) + end + + def qwen3_final_greedy(_worker, _hidden, _norm, _lm_head, _eps, _device) do + :erlang.nif_error(:nif_not_loaded) + end - def graph_capture(_inputs, _outputs, _shapeless) do + def qwen3_attention_block( + _worker, + _hidden, + _norm, + _q, + _k, + _v, + _o, + _q_norm, + _k_norm, + _k_cache, + _v_cache, + _offset, + _scale, + _head_dim, + _theta, + _eps, + _device + ) do :erlang.nif_error(:nif_not_loaded) end - def graph_replay(_compiled_ref, _new_inputs) do + # load_plugin — `dlopen`s a standalone, name-keyed native plugin (no + # erl_nif dependency) and caches its vtable under `name` (see + # emlx_plugin_registry.hpp). Callers that decode/dispatch into a + # specific plugin's ABI (e.g. `qwen3_*` NIFs above, which fetch the + # "qwen3" plugin) error with `{:error, _}` until this has been called + # successfully for that name — for qwen3, see `EMLXAxon.Application`, + # which calls it eagerly at boot. Not worker-routed (no argv[0] worker + # ref): `dlopen` does no MLX graph work. + def load_plugin(_name, _path) do :erlang.nif_error(:nif_not_loaded) end end diff --git a/emlx/lib/emlx/quantization.ex b/emlx/lib/emlx/quantization.ex index d4ba122..2c525fc 100644 --- a/emlx/lib/emlx/quantization.ex +++ b/emlx/lib/emlx/quantization.ex @@ -1,11 +1,15 @@ defmodule EMLX.Quantization do @moduledoc """ - Affine group-wise int2/int4/int8 quantization for Apple Silicon inference. + Affine group-wise int2/int4/int8 quantization for Apple Silicon inference, + plus microscaled floating-point modes (`"mxfp4"`/`"mxfp8"`/`"nvfp4"`, via + `:mode` on `quantize/2` — see `EMLX.quantize/2`). Quantized weights are represented as annotated `Nx.Tensor` values — the tensor carries the original logical shape and type (e.g. `{:s, 4}` for 4-bit), while the `EMLX.Backend` struct stores the packed uint32 data and - a `EMLX.Quantization.Config` with scales, biases, group_size, and bits. + a `EMLX.Quantization.Config` with scales, biases, group_size, bits, and + mode. Microscaled modes have no biases (`Config.biases` is `nil` — MLX's + `fp_quantize` returns only `(wq, scales)` for them). `Nx.dot` automatically dispatches to `mx::quantized_matmul` when it detects a quantized operand on `EMLX.Backend` — no explicit call site changes needed. @@ -57,6 +61,8 @@ defmodule EMLX.Quantization do * `:type` — Nx storage type: `{:s, 2}`, `{:s, 4}` (default), or `{:s, 8}`. * `:group_size` — 32, 64, or 128 (default 64). Must evenly divide the last dimension of `tensor`. + * `:mode` — `"affine"` (default), or a microscaled mode (`"mxfp4"`, + `"mxfp8"`, `"nvfp4"`) — see `EMLX.quantize/2`. """ deftransform quantize(tensor, opts \\ []) do type = Keyword.get(opts, :type, {:s, 4}) @@ -80,7 +86,6 @@ defmodule EMLX.Quantization do * `:type` — Nx storage type: `{:s, 2}`, `{:s, 4}` (default), or `{:s, 8}`. * `:group_size` — quantization group size (default 64). """ - @spec quantized_tensor(term(), term(), term(), tuple(), keyword()) :: Nx.Tensor.t() def quantized_tensor(weight_ref, scales_ref, biases_ref, original_shape, opts \\ []) do type = Keyword.get(opts, :type, {:s, 4}) {_, bits} = type @@ -91,7 +96,7 @@ defmodule EMLX.Quantization do config = %Config{scales: scales, biases: biases, group_size: group_size, bits: bits} weight_shape = EMLX.shape(weight_ref) - template = Nx.template(original_shape, type) + %Nx.Tensor{} = template = Nx.template(original_shape, type) %Nx.Tensor{ template @@ -112,15 +117,26 @@ defmodule EMLX.Quantization do `Nx.Defn.jit`-traced forward passes. The output has the same shape as the input (the quantized tensor's logical - shape equals the dense shape). + shape equals the dense shape). Supports every mode `EMLX.quantize/2` + accepts. """ deftransform dequantize(qw) do # Infer float type from scales when we have the real backend (eager / evaluator). # Falls back to :f32 during JIT tracing where qw.data is Nx.Defn.Expr. + # Microscaled modes store scales as :u8 (a packed exponent byte, not a + # float dtype), so `Nx.type(s)` doesn't apply there — MLX's `dequantize` + # reconstructs a float array regardless of mode; :bf16 matches the + # convention used elsewhere for microscaled outputs. out_type = case qw do - %Nx.Tensor{data: %EMLX.Backend{quantization_config: %Config{scales: s}}} -> Nx.type(s) - _ -> :f32 + %Nx.Tensor{data: %EMLX.Backend{quantization_config: %Config{mode: "affine", scales: s}}} -> + Nx.type(s) + + %Nx.Tensor{data: %EMLX.Backend{quantization_config: %Config{}}} -> + {:bf, 16} + + _ -> + :f32 end out = Nx.template(Nx.shape(qw), out_type) @@ -140,7 +156,10 @@ defmodule EMLX.Quantization do forward passes. Output shape is `{batch_dims..., out_features}` where `out_features` is the - first dimension of `qw`. Output type is `:bf16`. + first dimension of `qw`. Output type matches the scales dtype for + `"affine"` (historically the activation's own float type); for microscaled + modes scales are `:u8` (a packed exponent byte), so the output type + follows the activation's dtype instead. """ deftransform quantized_matmul(activation, qw) do {out_features, _} = Nx.shape(qw) @@ -149,8 +168,14 @@ defmodule EMLX.Quantization do out_type = case qw do - %Nx.Tensor{data: %EMLX.Backend{quantization_config: %Config{scales: s}}} -> Nx.type(s) - _ -> Nx.Type.merge(Nx.type(activation), Nx.type(qw)) + %Nx.Tensor{data: %EMLX.Backend{quantization_config: %Config{mode: "affine", scales: s}}} -> + Nx.type(s) + + %Nx.Tensor{data: %EMLX.Backend{quantization_config: %Config{}}} -> + Nx.type(activation) + + _ -> + Nx.Type.merge(Nx.type(activation), Nx.type(qw)) end out = Nx.template(out_shape, out_type) @@ -165,7 +190,6 @@ defmodule EMLX.Quantization do @doc """ Returns `true` if the tensor has quantization metadata on its backend. """ - @spec quantized?(term()) :: boolean() def quantized?(%Nx.Tensor{data: %EMLX.Backend{quantization_config: cfg}}) when not is_nil(cfg), do: true diff --git a/emlx/lib/emlx/quantization/config.ex b/emlx/lib/emlx/quantization/config.ex index 72ee238..d861c70 100644 --- a/emlx/lib/emlx/quantization/config.ex +++ b/emlx/lib/emlx/quantization/config.ex @@ -2,15 +2,18 @@ defmodule EMLX.Quantization.Config do @moduledoc false @enforce_keys [:scales, :biases, :group_size, :bits] - defstruct [:scales, :biases, :group_size, :bits, transpose: nil] + defstruct [:scales, :biases, :group_size, :bits, transpose: nil, mode: "affine"] @type t :: %__MODULE__{ scales: Nx.Tensor.t(), - biases: Nx.Tensor.t(), + # `nil` for microscaled modes ("mxfp4"/"mxfp8"/"nvfp4") — MLX's + # `fp_quantize` doesn't emit biases for them. + biases: Nx.Tensor.t() | nil, group_size: pos_integer(), bits: 2 | 4 | 8, # When set, overrides the auto-detection of the transpose flag in # quantized_dot (based on right_axes). nil = auto-detect. - transpose: boolean() | nil + transpose: boolean() | nil, + mode: String.t() } end diff --git a/emlx/lib/emlx/telemetry.ex b/emlx/lib/emlx/telemetry.ex new file mode 100644 index 0000000..ea88d24 --- /dev/null +++ b/emlx/lib/emlx/telemetry.ex @@ -0,0 +1,82 @@ +defmodule EMLX.Telemetry do + @moduledoc """ + `:telemetry` events emitted by EMLX. + + All span-style events use `:telemetry.span/3` semantics, so attaching to + `*:start`, `*:stop`, and `*:exception` is sufficient for histograms and + error tracking. + + ## Events + + ### Evaluation boundaries + + `[:emlx, :eval, :start | :stop | :exception]` — `EMLX.eval/1`. Spans the + full round-trip through the resolved `EMLX.CommandQueue` worker, including + the blocking wait for the worker to finish `mlx::core::eval/1` (the real + latency, not just NIF dispatch). The `:stop` event carries `:duration` + (monotonic native units). + + `[:emlx, :to_binary, :start | :stop | :exception]` — `EMLX.Backend.to_binary/2` + (the `Nx.to_binary/1` path). Spans the sync-forcing blob copy. Metadata: + `:shape`, `:dtype`, `:byte_size` (byte size of the binary actually + returned to the caller). + + ### Memory stats (poll-driven) + + `[:emlx, :memory, :stats]` — discrete event, not a span. Call + `EMLX.Telemetry.memory_stats/0` to sample; measurements: + + * `:active_memory` — bytes currently allocated and in use + * `:peak_memory` — highest active memory since the last + `EMLX.reset_peak_memory/0` + * `:cache_memory` — bytes in the allocator cache (freed but not + returned to the OS) + + Wire this into a periodic task (e.g. `Process.send_after/3` loop) to + graph memory drift in a long-running serving. + + ## Attaching a handler + + :telemetry.attach( + "emlx-eval-log", + [:emlx, :eval, :stop], + fn _event, measurements, _metadata, _config -> + IO.inspect(measurements.duration) + end, + nil + ) + """ + + @doc false + def span_eval(fun) do + :telemetry.span([:emlx, :eval], %{}, fn -> {fun.(), %{}} end) + end + + @doc false + def span_to_binary(%Nx.Tensor{shape: shape, type: type}, fun) do + start_metadata = %{shape: shape, dtype: type} + + :telemetry.span([:emlx, :to_binary], start_metadata, fn -> + binary = fun.() + {binary, %{shape: shape, dtype: type, byte_size: byte_size(binary)}} + end) + end + + @doc """ + Sample the MLX allocator and emit `[:emlx, :memory, :stats]`. + + Returns the measurements map so callers can also log or plot inline. + + ## Examples + + iex> stats = EMLX.Telemetry.memory_stats() + iex> Map.keys(stats) |> Enum.sort() + [:active_memory, :cache_memory, :peak_memory] + + """ + def memory_stats do + stats = EMLX.memory_info() + :telemetry.execute([:emlx, :memory, :stats], stats, %{}) + stats + end +end diff --git a/emlx/mix.exs b/emlx/mix.exs index c6c7b9a..d7fc9ac 100644 --- a/emlx/mix.exs +++ b/emlx/mix.exs @@ -20,26 +20,30 @@ defmodule EMLX.MixProject do start_permanent: Mix.env() == :prod, deps: deps(), docs: docs(), - # elixir_make - make_env: %{ - "MLX_DIR" => libmlx_config.dir, - "MLX_VERSION" => libmlx_config.version, - "MLX_BUILD" => to_string(libmlx_config.features.build?), - "MLX_INCLUDE_DIR" => - Path.join( - libmlx_config.dir, - if(libmlx_config.features.build?, do: "usr/include", else: "include") - ), - "MLX_LIB_DIR" => - Path.join( - libmlx_config.dir, - if(libmlx_config.features.build?, do: "usr/lib", else: "lib") - ), - "MLX_VARIANT" => libmlx_config.variant, - "EMLX_CACHE_DIR" => libmlx_config.cache_dir, - "EMLX_VERSION" => @version, - "LIBMLX_ENABLE_DEBUG" => to_string(libmlx_config.features.debug?) - }, + # elixir_make — a function ref so Fine.include_dir/0 (needs the :fine + # dep compiled) isn't called while this project/0 map is being built. + make_env: fn -> + %{ + "MLX_DIR" => libmlx_config.dir, + "MLX_VERSION" => libmlx_config.version, + "MLX_BUILD" => to_string(libmlx_config.features.build?), + "MLX_INCLUDE_DIR" => + Path.join( + libmlx_config.dir, + if(libmlx_config.features.build?, do: "usr/include", else: "include") + ), + "MLX_LIB_DIR" => + Path.join( + libmlx_config.dir, + if(libmlx_config.features.build?, do: "usr/lib", else: "lib") + ), + "MLX_VARIANT" => libmlx_config.variant, + "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() + } + end, # Compilers compilers: compilers(libmlx_config), @@ -65,7 +69,9 @@ defmodule EMLX.MixProject do defp deps do [ {:elixir_make, "~> 0.6"}, - {:nx, "~> 0.12"}, + {:fine, "~> 0.1", runtime: false}, + {:nx, github: "elixir-nx/nx", branch: "main", sparse: "nx"}, + {:telemetry, "~> 1.0"}, {:ex_doc, "~> 0.34", only: :docs} ] end @@ -259,10 +265,6 @@ defmodule EMLX.MixProject do defp to_boolean(nil), do: false - defp to_boolean(var) when is_boolean(var) do - var - end - defp to_boolean(var) do String.downcase(to_string(var)) in ["1", "true", "on", "yes", "y"] end diff --git a/emlx/mix.lock b/emlx/mix.lock index 5526e91..587c7bc 100644 --- a/emlx/mix.lock +++ b/emlx/mix.lock @@ -3,10 +3,11 @@ "earmark_parser": {:hex, :earmark_parser, "1.4.44", "f20830dd6b5c77afe2b063777ddbbff09f9759396500cdbe7523efd58d7a339c", [:mix], [], "hexpm", "4778ac752b4701a5599215f7030989c989ffdc4f6df457c5f36938cc2d2a2750"}, "elixir_make": {:hex, :elixir_make, "0.9.0", "6484b3cd8c0cee58f09f05ecaf1a140a8c97670671a6a0e7ab4dc326c3109726", [:mix], [], "hexpm", "db23d4fd8b757462ad02f8aa73431a426fe6671c80b200d9710caf3d1dd0ffdb"}, "ex_doc": {:hex, :ex_doc, "0.40.2", "f50edec428c4b0a457a167de42414c461122a3585a99515a69d09fff19e5597e", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "4fa426e2beb47854a162e2c488727fdec51cd4692e319b23810c2804cb1a40fe"}, + "fine": {:hex, :fine, "0.1.6", "4bf7151493443c454aac9f2fa2f34f5fefd0346a83fb5586a016c4a135c63247", [:mix], [], "hexpm", "5638eb4495488e885ebec167fa57973e5c35e1a50c344eb7666c90ec1c4e3b12"}, "makeup": {:hex, :makeup, "1.2.1", "e90ac1c65589ef354378def3ba19d401e739ee7ee06fb47f94c687016e3713d1", [:mix], [{:nimble_parsec, "~> 1.4", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "d36484867b0bae0fea568d10131197a4c2e47056a6fbe84922bf6ba71c8d17ce"}, "makeup_elixir": {:hex, :makeup_elixir, "1.0.1", "e928a4f984e795e41e3abd27bfc09f51db16ab8ba1aebdba2b3a575437efafc2", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "7284900d412a3e5cfd97fdaed4f5ed389b8f2b4cb49efc0eb3bd10e2febf9507"}, "makeup_erlang": {:hex, :makeup_erlang, "1.1.0", "835f7e60792e08824cda445639555d7bf1bbbddb1b60b306e33cb6f6db24dc74", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "1cd6780fb1dd1a03979abaed0fe82712b0625118fd5257d3ebbf73f960c73c3c"}, "nimble_parsec": {:hex, :nimble_parsec, "1.4.2", "8efba0122db06df95bfaa78f791344a89352ba04baedd3849593bfce4d0dc1c6", [:mix], [], "hexpm", "4b21398942dda052b403bbe1da991ccd03a053668d147d53fb8c4e0efe09c973"}, - "nx": {:hex, :nx, "0.12.0", "32bc205bab5486d73892132d17a11ea113e97427a29bb70606a544724b95e193", [:mix], [{:complex, "~> 0.7", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "7d022a33ea3c900eb6e2e91b4e0793759459c886f482be61978004b5e4843b5e"}, - "telemetry": {:hex, :telemetry, "1.4.1", "ab6de178e2b29b58e8256b92b382ea3f590a47152ca3651ea857a6cae05ac423", [:rebar3], [], "hexpm", "2172e05a27531d3d31dd9782841065c50dd5c3c7699d95266b2edd54c2dafa1c"}, + "nx": {:git, "https://github.com/elixir-nx/nx.git", "da2c15c9beb03ed3d7edd79ea8ce6d9377775631", [branch: "main", sparse: "nx"]}, + "telemetry": {:hex, :telemetry, "1.4.2", "a0cb522801dffb1c49fe6e30561badffc7b6d0e180db1300df759faa22062855", [:rebar3], [], "hexpm", "928f6495066506077862c0d1646609eed891a4326bee3126ba54b60af61febb1"}, } diff --git a/emlx/test/emlx/compiler_test.exs b/emlx/test/emlx/compiler_test.exs index a00483a..b79f574 100644 --- a/emlx/test/emlx/compiler_test.exs +++ b/emlx/test/emlx/compiler_test.exs @@ -72,6 +72,7 @@ defmodule EMLX.CompilerTest do end end + @tag :metal test "valid opts :device and :max_concurrency do not raise" do defmodule IdentFn2 do import Nx.Defn diff --git a/emlx/test/emlx/conv_pool_training_canary_test.exs b/emlx/test/emlx/conv_pool_training_canary_test.exs new file mode 100644 index 0000000..0fdb2c6 --- /dev/null +++ b/emlx/test/emlx/conv_pool_training_canary_test.exs @@ -0,0 +1,173 @@ +defmodule EMLX.ConvPoolTrainingCanaryTest do + @moduledoc """ + Reference convention matches `grad_equivalence_test.exs`: scope the + process-global default backend to `Nx.BinaryBackend` around + `Nx.Defn.Evaluator` so any tensor the Evaluator synthesizes internally + (e.g. a window op's padding fill) doesn't silently pull in `EMLX.Backend`. + """ + + use EMLX.Case, async: false + import Nx.Defn + + @batch 4 + @in_hw 8 + @kernel_hw 3 + @out_channels 4 + @pooled_hw 3 + @flat_size @out_channels * @pooled_hw * @pooled_hw + @num_classes 3 + @lr 0.05 + @steps 20 + @num_batches 3 + + # ── deterministic, backend-agnostic data/param generation ────────────────── + # Avoids Nx.Random entirely so the two backends under comparison start from + # bit-identical inputs without relying on cross-backend RNG parity (not + # this stage's concern — see grad_equivalence_test.exs for that). + defp seeded(shape, seed) do + size = Tuple.product(shape) + + for i <- 0..(size - 1) do + :math.sin((i + seed) * 0.7) * 0.5 + end + |> Nx.tensor(type: {:f, 32}, backend: Nx.BinaryBackend) + |> Nx.reshape(shape) + end + + defp one_hot_targets(offset) do + for i <- 0..(@batch - 1) do + label = rem(i + offset, @num_classes) + for c <- 0..(@num_classes - 1), do: if(c == label, do: 1.0, else: 0.0) + end + |> Nx.tensor(type: {:f, 32}, backend: Nx.BinaryBackend) + end + + defp init_params do + %{ + w1: seeded({@out_channels, 1, @kernel_hw, @kernel_hw}, 1), + w2: seeded({@flat_size, @num_classes}, 2), + b2: seeded({@num_classes}, 3) + } + end + + defp dataset do + for b <- 0..(@num_batches - 1) do + {seeded({@batch, 1, @in_hw, @in_hw}, 10 + b), one_hot_targets(b)} + end + end + + # ── model ──────────────────────────────────────────────────────────────── + + defn forward(params, x) do + x + |> Nx.conv(params.w1, padding: :valid) + |> Nx.max(0.0) + |> Nx.window_max({1, 1, 2, 2}, strides: [1, 1, 2, 2], padding: :valid) + |> Nx.reshape({@batch, @flat_size}) + |> Nx.dot(params.w2) + |> Nx.add(params.b2) + end + + defn loss_fn(params, x, y) do + preds = forward(params, x) + Nx.mean(Nx.pow(Nx.subtract(preds, y), 2)) + end + + defn train_step(params, x, y) do + {loss, grads} = value_and_grad(params, &loss_fn(&1, x, y)) + + new_params = %{ + w1: Nx.subtract(params.w1, Nx.multiply(@lr, grads.w1)), + w2: Nx.subtract(params.w2, Nx.multiply(@lr, grads.w2)), + b2: Nx.subtract(params.b2, Nx.multiply(@lr, grads.b2)) + } + + {loss, new_params} + end + + # ── runners ────────────────────────────────────────────────────────────── + + defp transfer_params(params, backend) do + %{ + w1: Nx.backend_transfer(params.w1, backend), + w2: Nx.backend_transfer(params.w2, backend), + b2: Nx.backend_transfer(params.b2, backend) + } + end + + defp run_steps(params0, dataset, jit_fun) do + {_final_params, losses} = + Enum.reduce(1..@steps, {params0, []}, fn step, {params, losses} -> + {x, y} = Enum.at(dataset, rem(step - 1, length(dataset))) + {loss, new_params} = jit_fun.(params, x, y) + {new_params, [Nx.to_number(loss) | losses]} + end) + + Enum.reverse(losses) + end + + defp reference_curve(params0, dataset) do + previous = Nx.default_backend() + Nx.default_backend(Nx.BinaryBackend) + + try do + run_steps(params0, dataset, fn params, x, y -> + Nx.Defn.jit_apply(&train_step/3, [params, x, y], compiler: Nx.Defn.Evaluator) + end) + after + Nx.default_backend(previous) + end + end + + defp eager_emlx_curve(params0, dataset) do + params = transfer_params(params0, EMLX.Backend) + + dataset = + Enum.map(dataset, fn {x, y} -> + {Nx.backend_transfer(x, EMLX.Backend), Nx.backend_transfer(y, EMLX.Backend)} + end) + + run_steps(params, dataset, fn params, x, y -> + Nx.Defn.jit_apply(&train_step/3, [params, x, y], compiler: Nx.Defn.Evaluator) + end) + end + + defp native_curve(params0, dataset) do + run_steps(params0, dataset, fn params, x, y -> + Nx.Defn.jit_apply(&train_step/3, [params, x, y], compiler: EMLX) + end) + end + + defp assert_curves_match(curve, reference, tol) do + curve + |> Enum.zip(reference) + |> Enum.with_index() + |> Enum.each(fn {{a, b}, i} -> + assert_in_delta(a, b, tol, "loss curves diverge at step #{i}: #{a} vs #{b}") + end) + end + + describe "conv+max_pool small-CNN training curve" do + test "eager EMLX.Backend matches the Nx.BinaryBackend/Evaluator reference" do + params0 = init_params() + data = dataset() + + reference = reference_curve(params0, data) + eager = eager_emlx_curve(params0, data) + + assert_curves_match(eager, reference, 1.0e-3) + assert List.last(reference) < List.first(reference) * 0.5 + end + + test "native compiler: EMLX matches the Nx.BinaryBackend/Evaluator reference" do + params0 = init_params() + data = dataset() + + reference = reference_curve(params0, data) + native = native_curve(params0, data) + + assert_curves_match(native, reference, 1.0e-3) + assert List.last(reference) < List.first(reference) * 0.5 + end + end +end diff --git a/emlx/test/emlx/debug_flags_functional_test.exs b/emlx/test/emlx/debug_flags_functional_test.exs new file mode 100644 index 0000000..644d352 --- /dev/null +++ b/emlx/test/emlx/debug_flags_functional_test.exs @@ -0,0 +1,114 @@ +defmodule EMLX.DebugFlagsFunctionalTest do + use EMLX.Case, async: false + + @moduledoc """ + Functional (raise-on-violation) counterpart to `debug_flags_test.exs`'s + zero-cost-when-off opcode checks. + + `:enable_bounds_check`, `:detect_non_finites`, and `:compiler_debug` are all + `Application.compile_env/3`-gated, so none can be toggled at runtime — + these tests only pass against a build compiled with all three flags on. Run: + + EMLX_DEBUG_FLAGS=1 mix test --force --include debug_flags_functional + + Excluded by default (see `test/test_helper.exs`) so a normal `mix test` + run — compiled with every debug flag off, per production defaults — + doesn't fail here. + """ + + @moduletag :debug_flags_functional + + alias EMLX.Native.Expr + + setup_all do + # `Application.get_env/3` here (not `compile_env/3`): the flags below are + # genuinely compile-time-baked into the *lib* code under test, but this + # guard is just a friendly diagnostic — reading via `get_env` avoids the + # type-checker constant-folding the whole `unless` away (both flags are + # literal `false` in a default build) into a "warnings-as-errors" hit. + flags_on? = + Application.get_env(:emlx, :enable_bounds_check, false) and + Application.get_env(:emlx, :detect_non_finites, false) and + Application.get_env(:emlx, :compiler_debug, false) + + unless flags_on? do + flunk(""" + compiled with a debug flag off — these tests always fail here. + Run: EMLX_DEBUG_FLAGS=1 mix test --force --include debug_flags_functional + """) + end + + :ok + end + + defp gpu(tensor), do: Nx.backend_transfer(tensor, {EMLX.Backend, device: :gpu}) + + test "take raises on an out-of-bounds index (:enable_bounds_check, pre-existing coverage, verified not regressed)" do + t = Nx.tensor([1, 2, 3], backend: EMLX.Backend) + idx = Nx.tensor([5], type: :s64, backend: EMLX.Backend) + + assert_raise ArgumentError, ~r/index out of bounds/, fn -> + Nx.take(t, idx) |> Nx.backend_transfer() + end + end + + test "dot raises on a NaN-producing matmul (pre-existing coverage, unmoved by the EMLX.Debug extraction)" do + nan = Nx.tensor([:nan, 1.0, 1.0], type: :f32) + ones = Nx.tensor([1.0, 1.0, 1.0], type: :f32) + + assert_raise ArgumentError, ~r/dot produced NaN or Inf/, fn -> + Nx.dot(nan, ones) |> Nx.backend_transfer() + end + end + + test "conv raises on a NaN-producing convolution" do + input = Nx.tensor([[[[:nan, 1.0, 1.0, 1.0]]]], type: :f32) + kernel = Nx.tensor([[[[1.0, 1.0]]]], type: :f32) + + assert_raise ArgumentError, ~r/conv produced NaN or Inf/, fn -> + Nx.conv(input, kernel, padding: :valid) |> Nx.backend_transfer() + end + end + + @tag :metal + test "EMLX.Fast.rms_norm raises on a NaN-producing input" do + x = Nx.tensor([[:nan, 1.0, 1.0, 1.0]], type: :f32) |> gpu() + w = Nx.tensor([1.0, 1.0, 1.0, 1.0], type: :f32) |> gpu() + + assert_raise ArgumentError, ~r/rms_norm produced NaN or Inf/, fn -> + EMLX.Fast.rms_norm(x, w, 1.0e-5) + end + end + + test "EMLX.Native.Expr.to_native raises on a ref id collision across categories (:compiler_debug)" do + ref = make_ref() + + prog = %Expr{ + inputs: [ref], + captures: [{ref, Nx.tensor(1.0)}], + constants: [], + instructions: [], + outputs: [ref] + } + + assert_raise ArgumentError, ~r/ref id collision across inputs\/captures\/constants/, fn -> + Expr.to_native(prog) + end + end + + test "EMLX.Native.Expr.to_native raises when an instruction's result ref is already bound (:compiler_debug)" do + ref = make_ref() + + prog = %Expr{ + inputs: [ref], + captures: [], + constants: [], + instructions: [{ref, :add, [], []}], + outputs: [ref] + } + + assert_raise ArgumentError, ~r/that ref is already bound/, fn -> + Expr.to_native(prog) + end + end +end diff --git a/emlx/test/emlx/debug_flags_test.exs b/emlx/test/emlx/debug_flags_test.exs index b340405..14fa2fc 100644 --- a/emlx/test/emlx/debug_flags_test.exs +++ b/emlx/test/emlx/debug_flags_test.exs @@ -5,21 +5,28 @@ defmodule EMLX.DebugFlagsTest do Verifies that debug flags default to false and that call-site branches are dead-code-eliminated in the compiled BEAM when the flags are off. - The private helpers (`assert_in_bounds!`, `assert_no_nan_inf!`) are defined - unconditionally so the compiler can resolve them at the abstract-code stage. - Call sites live inside `if @flag do … end` blocks; the Erlang compiler - eliminates the unreachable branch from the final BEAM opcodes. - - When `EMLX.Fast` is implemented (task 05), add a parallel opcode test for - each fast-kernel `*_impl/2` function here. + The assertion macros (`EMLX.Backend`'s private `assert_in_bounds!`, and + `EMLX.Debug`'s public `assert_no_nan_inf!`, shared with `EMLX.Fast`) are + defined unconditionally so the compiler can resolve them at the + abstract-code stage. Each one branches on its flag *inside the macro body* + at the macro's own compile time, so a call site either inlines the real + assertion or inlines nothing at all — never a runtime `if`. The + `dot`/`gather` tests below check this the way the macro's own definition + would suggest (no call named after the macro ever appears, on or off, + since macros don't compile to calls); the `assert_no_nan_inf!` extension + tests (`conv`, `EMLX.Fast`) check the more meaningful thing instead: that + the *expanded* NaN/Inf-check calls (`EMLX.is_nan/1`, `EMLX.is_infinity/1`) + are themselves absent from the compiled function when the flag is off. """ @enable_bounds_check Application.compile_env(:emlx, :enable_bounds_check, false) @detect_non_finites Application.compile_env(:emlx, :detect_non_finites, false) + @compiler_debug Application.compile_env(:emlx, :compiler_debug, false) test "debug flags default to false" do assert @enable_bounds_check == false assert @detect_non_finites == false + assert @compiler_debug == false end test "assert_in_bounds! call is absent from BEAM opcodes of gather/4 when flag is off" do @@ -63,6 +70,27 @@ defmodule EMLX.DebugFlagsTest do "assert_no_nan_inf! must not appear in dot/7 opcodes when :detect_non_finites is false" end + test "EMLX.is_nan/EMLX.is_infinity calls are absent from EMLX.Backend's conv/4 opcodes when flag is off" do + refute finite_check_calls(EMLX.Backend, :conv, 4) != [], + "EMLX.is_nan/EMLX.is_infinity must not appear in conv/4 opcodes when :detect_non_finites is false" + end + + test "EMLX.is_nan/EMLX.is_infinity calls are absent from EMLX.Fast's rms_norm/layer_norm/sdpa callbacks when flag is off" do + for {fun, arity} <- [ + {:rms_norm_callback, 2}, + {:layer_norm_callback, 2}, + {:layer_norm_no_bias_callback, 2}, + {:sdpa_callback, 2}, + {:sdpa_masked_callback, 2}, + {:sdpa_causal_callback, 2}, + {:sdpa_causal_key_masked_callback, 2} + ] do + refute finite_check_calls(EMLX.Fast, fun, arity) != [], + "EMLX.is_nan/EMLX.is_infinity must not appear in #{fun}/#{arity} opcodes " <> + "when :detect_non_finites is false" + end + end + # Collect all local (same-module) function names referenced in a BEAM opcode list. defp local_calls(instructions) do for instr <- instructions, @@ -72,4 +100,30 @@ defmodule EMLX.DebugFlagsTest do name end end + + # `assert_no_nan_inf!` (EMLX.Debug) is a macro: it never compiles to a call + # named after itself, on or off (see moduledoc). What actually distinguishes + # "compiled in" from "dead-code-eliminated" is whether the *expanded* body's + # own calls — `EMLX.is_nan/1` and `EMLX.is_infinity/1` — appear in the + # target function's opcode list at all. + defp finite_check_calls(module, fun, arity) do + beam_path = :code.which(module) + {:beam_file, _mod, _exports, _attrs, _info, functions} = :beam_disasm.file(beam_path) + + target_fn = + Enum.find(functions, fn + {:function, ^fun, ^arity, _, _} -> true + _ -> false + end) + + refute target_fn == nil, "#{inspect(module)}.#{fun}/#{arity} must exist" + + {:function, ^fun, ^arity, _, instructions} = target_fn + + for instr <- instructions, + match?({:call_ext, _arity, {:extfunc, EMLX, :is_nan, _}}, instr) or + match?({:call_ext, _arity, {:extfunc, EMLX, :is_infinity, _}}, instr) do + instr + end + end end diff --git a/emlx/test/emlx/defn/tree_test.exs b/emlx/test/emlx/defn/tree_test.exs new file mode 100644 index 0000000..7230321 --- /dev/null +++ b/emlx/test/emlx/defn/tree_test.exs @@ -0,0 +1,173 @@ +defmodule EMLX.Defn.TreeTest do + use ExUnit.Case, async: true + import Nx.Defn + + alias EMLX.Defn.Tree + + alias Nx.Tensor, as: T + alias Nx.Defn.Expr + + # A defn function that produces a :while node; used to verify scope-boundary. + defn count_up(n) do + while n, Nx.less(n, 10) do + Nx.add(n, 1) + end + end + + # Verifies the core post-order invariant: + # For every node in `order`, each of its same-scope tensor operands (if + # present in the ordering) must appear at a strictly smaller index. + # + # Deps that are *absent* from the index are inner-scope nodes (e.g. fun + # parameter templates) and are expected — they are silently skipped. + defp assert_post_order(order) do + id_to_idx = Map.new(Enum.with_index(order), fn {t, i} -> {t.data.id, i} end) + + for {node, node_idx} <- Enum.with_index(order) do + Nx.Defn.Tree.apply_args(node, :scope, :ok, fn dep, :ok -> + case Map.fetch(id_to_idx, dep.data.id) do + {:ok, dep_idx} -> + assert dep_idx < node_idx, + "#{inspect(dep.data.op)} (idx #{dep_idx}) is not before " <> + "its consumer #{inspect(node.data.op)} (idx #{node_idx})" + + :error -> + # inner-scope dep (e.g. fun parameter template) — not in order, expected + :ok + end + + {dep, :ok} + end) + end + end + + describe "post_order/1" do + test "linear chain: nodes appear in dependency-first order" do + expr = + Nx.Defn.debug_expr_apply( + fn x -> + y = Nx.negate(x) + Nx.abs(y) + end, + [Nx.tensor(1.0)] + ) + + order = Tree.post_order(expr) + + assert length(order) == 3 + ops = Enum.map(order, & &1.data.op) + assert [:parameter, :negate, :abs] == ops + assert_post_order(order) + end + + test "diamond: shared subexpression appears exactly once, after its operands" do + expr = + Nx.Defn.debug_expr_apply( + fn x -> + y = Nx.negate(x) + # y is shared: should appear once + Nx.add(y, y) + end, + [Nx.tensor(1.0)] + ) + + order = Tree.post_order(expr) + + assert length(order) == 3 + ops = Enum.map(order, & &1.data.op) + assert [:parameter, :negate, :add] == ops + assert_post_order(order) + end + + test "multi-output container: all roots are collected, shared param appears once" do + {e1, e2} = + Nx.Defn.debug_expr_apply( + fn x -> {Nx.negate(x), Nx.abs(x)} end, + [Nx.tensor(1.0)] + ) + + order = Tree.post_order({e1, e2}) + ops = Enum.map(order, & &1.data.op) + assert [:parameter, :negate, :abs] == ops + assert_post_order(order) + + # reordering the output tuple should cause the order to change + order = Tree.post_order({e2, e1}) + ops = Enum.map(order, & &1.data.op) + assert [:parameter, :abs, :negate] == ops + assert_post_order(order) + end + + test "constant and parameter leaves appear before their consumers" do + expr = + Nx.Defn.debug_expr_apply( + fn x -> Nx.add(x, Nx.tensor(1.0)) end, + [Nx.tensor(2.0)] + ) + + order = Tree.post_order(expr) + + # It doesn't really matter if constants appear closer to invocation + # than parameters as they are decoupled, so we don't assert on them here. + ops = Enum.map(order, & &1.data.op) + assert :parameter in ops + assert :constant in ops + # add must be last + assert Enum.find_index(order, &(&1.data.op == :add)) == 2 + assert_post_order(order) + end + + test "while node is opaque: appears as a single node; inner body does not leak" do + expr = Nx.Defn.debug_expr_apply(&count_up/1, [Nx.tensor(0)]) + + order = Tree.post_order(expr) + + # There is exactly one while node + assert {while, not_while} = Enum.split_with(order, &(&1.data.op == :while)) + + assert [%T{data: %Expr{op: :while, args: [_, _, pred, body]}}] = while + + # Inner-body ops (:less, :add at inner scope and hoisted :constant) must NOT appear + not_while = Enum.map(not_while, & &1.data.op) + + assert [:parameter] == not_while, + "inner-scope nodes leaked into parent ordering: #{inspect(not_while)}" + + assert_post_order(order) + + pred_order = Tree.post_order(pred) + assert [:parameter, :constant, :less] == Enum.map(pred_order, & &1.data.op) + assert_post_order(pred_order) + + body_order = Tree.post_order(body) + assert [:constant, :parameter, :add] == Enum.map(body_order, & &1.data.op) + assert_post_order(body_order) + end + + test "post-order invariant holds on a diamond through a shared interior node" do + expr = + Nx.Defn.debug_expr_apply( + fn x, y -> + a = Nx.add(x, y) + b = Nx.multiply(a, x) + c = Nx.subtract(b, y) + # 'a' is shared between b and this add + Nx.add(a, c) + end, + [Nx.tensor(1.0), Nx.tensor(2.0)] + ) + + order = Tree.post_order(expr) + + # :multiply and :subtract could appear in different others, this doesn't really impact the test + assert [:parameter, :parameter, :add, :multiply, :subtract, :add] == + Enum.map(order, & &1.data.op) + + # Every operand must precede its consumer + assert_post_order(order) + # Sanity-check that both :parameter nodes are distinct as well as both :add nodes + ids = Enum.map(order, & &1.data.id) + assert ids == Enum.uniq(ids) + end + end +end diff --git a/emlx/test/emlx/fast/einsum_test.exs b/emlx/test/emlx/fast/einsum_test.exs new file mode 100644 index 0000000..f9422d2 --- /dev/null +++ b/emlx/test/emlx/fast/einsum_test.exs @@ -0,0 +1,91 @@ +defmodule EMLX.Fast.EinsumTest do + @moduledoc """ + Tests for `EMLX.Fast.einsum/2` — the eager-only helper that dispatches + to `mlx::core::einsum`. + + Covered: + + * Two-operand contractions (matmul, batched matmul, attention-style + QKᵀ fold). + * Three-operand chain — sanity check that MLX's internal path + optimisation produces the same result as both left-to-right and + right-to-left hand contractions. + * Error path — a non-EMLX backend raises a clear transfer-first + message. + """ + + use EMLX.Case, async: true + + doctest EMLX.Fast, only: [einsum: 2] + + @f32_tol 1.0e-5 + + describe "two-operand contractions" do + test "\"ij,jk->ik\" matches Nx.dot" do + a = Nx.iota({3, 4}, backend: EMLX.Backend, type: :f32) + b = Nx.iota({4, 5}, backend: EMLX.Backend, type: :f32) + + got = EMLX.Fast.einsum("ij,jk->ik", [a, b]) + expected = Nx.dot(a, b) + + assert Nx.shape(got) == {3, 5} + assert_all_close(got, expected, atol: @f32_tol, rtol: @f32_tol) + end + + test "\"bij,bjk->bik\" matches Nx.dot with explicit batch axes" do + a = Nx.iota({2, 3, 4}, backend: EMLX.Backend, type: :f32) + b = Nx.iota({2, 4, 5}, backend: EMLX.Backend, type: :f32) + + got = EMLX.Fast.einsum("bij,bjk->bik", [a, b]) + # Contract a's last axis with b's second-to-last; keep batch 0 as + # a batch dim on both sides. + expected = Nx.dot(a, [2], [0], b, [1], [0]) + + assert Nx.shape(got) == {2, 3, 5} + assert_all_close(got, expected, atol: @f32_tol, rtol: @f32_tol) + end + + test "attention-style \"bhid,bhjd->bhij\"" do + q = Nx.iota({2, 3, 4, 5}, backend: EMLX.Backend, type: :f32) + k = Nx.iota({2, 3, 6, 5}, backend: EMLX.Backend, type: :f32) + + got = EMLX.Fast.einsum("bhid,bhjd->bhij", [q, k]) + # Contract d (axis 3 on both); keep b, h as batch dims on both. + expected = Nx.dot(q, [3], [0, 1], k, [3], [0, 1]) + + assert Nx.shape(got) == {2, 3, 4, 6} + assert_all_close(got, expected, atol: @f32_tol, rtol: @f32_tol) + end + end + + describe "three-operand contractions" do + test "\"ij,jk,kl->il\" matches both hand-chosen contraction orders" do + a = Nx.iota({2, 3}, backend: EMLX.Backend, type: :f32) + b = Nx.iota({3, 4}, backend: EMLX.Backend, type: :f32) + c = Nx.iota({4, 5}, backend: EMLX.Backend, type: :f32) + + got = EMLX.Fast.einsum("ij,jk,kl->il", [a, b, c]) + + left_first = Nx.dot(Nx.dot(a, b), c) + right_first = Nx.dot(a, Nx.dot(b, c)) + + assert Nx.shape(got) == {2, 5} + # Associativity means both orders are mathematically identical; + # any FP reordering across MLX's chosen path is still within + # f32 tolerance on these small shapes. + assert_all_close(got, left_first, atol: @f32_tol, rtol: @f32_tol) + assert_all_close(got, right_first, atol: @f32_tol, rtol: @f32_tol) + end + end + + describe "error paths" do + test "raises a transfer-first ArgumentError on non-EMLX operands" do + a = Nx.iota({3, 4}, backend: Nx.BinaryBackend, type: :f32) + b = Nx.iota({4, 5}, backend: EMLX.Backend, type: :f32) + + assert_raise ArgumentError, ~r/Nx\.BinaryBackend.*Nx\.backend_transfer/s, fn -> + EMLX.Fast.einsum("ij,jk->ik", [a, b]) + end + end + end +end diff --git a/emlx/test/emlx/fast_test.exs b/emlx/test/emlx/fast_test.exs index 39cafcf..5a4444a 100644 --- a/emlx/test/emlx/fast_test.exs +++ b/emlx/test/emlx/fast_test.exs @@ -479,1111 +479,6 @@ defmodule EMLX.FastTest do end end - describe "EMLX.qwen3_kv_cache_attention/9 validation" do - test "rejects malformed Q rank before reading shapes" do - q = Nx.broadcast(0.0, {1, 1, 4}) |> Nx.as_type(:f16) |> gpu() - k = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() - v = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() - k_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() - v_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() - - assert_raise EMLX.NIFError, ~r/q expects rank 4/, fn -> - EMLX.qwen3_kv_cache_attention( - EMLX.Backend.from_nx(q), - EMLX.Backend.from_nx(k), - EMLX.Backend.from_nx(v), - EMLX.Backend.from_nx(k_cache), - EMLX.Backend.from_nx(v_cache), - 0, - 1.0 / :math.sqrt(4), - 4, - 10_000.0 - ) - end - end - - test "rejects negative offsets" do - q = Nx.broadcast(0.0, {1, 1, 2, 4}) |> Nx.as_type(:f16) |> gpu() - k = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() - v = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() - k_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() - v_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() - - assert_raise EMLX.NIFError, ~r/offset must be non-negative/, fn -> - EMLX.qwen3_kv_cache_attention( - EMLX.Backend.from_nx(q), - EMLX.Backend.from_nx(k), - EMLX.Backend.from_nx(v), - EMLX.Backend.from_nx(k_cache), - EMLX.Backend.from_nx(v_cache), - -1, - 1.0 / :math.sqrt(4), - 4, - 10_000.0 - ) - end - end - - test "rejects cache capacity overflow" do - q = Nx.broadcast(0.0, {1, 1, 2, 4}) |> Nx.as_type(:f16) |> gpu() - k = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() - v = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() - k_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() - v_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() - - assert_raise EMLX.NIFError, ~r/KV cache capacity 4 is smaller than required length 5/, fn -> - EMLX.qwen3_kv_cache_attention( - EMLX.Backend.from_nx(q), - EMLX.Backend.from_nx(k), - EMLX.Backend.from_nx(v), - EMLX.Backend.from_nx(k_cache), - EMLX.Backend.from_nx(v_cache), - 4, - 1.0 / :math.sqrt(4), - 4, - 10_000.0 - ) - end - end - - test "matches pure Nx reference for GQA prefill and updates caches" do - {q_cpu, k_cpu, v_cpu, k_cache_cpu, v_cache_cpu} = - Nx.with_default_backend(Nx.BinaryBackend, fn -> - { - Nx.iota({1, 2, 4, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(100), - Nx.iota({1, 2, 2, 4}, type: :f32) |> Nx.add(11) |> Nx.divide(110), - Nx.iota({1, 2, 2, 4}, type: :f32) |> Nx.add(17) |> Nx.divide(120), - Nx.iota({1, 2, 5, 4}, type: :f32) |> Nx.divide(1_000), - Nx.iota({1, 2, 5, 4}, type: :f32) |> Nx.add(50) |> Nx.divide(1_000) - } - end) - - offset = 1 - head_dim = 4 - theta = 10_000.0 - scale = 1.0 / :math.sqrt(head_dim) - - {attn_ref, k_ref, v_ref} = - EMLX.qwen3_kv_cache_attention( - EMLX.Backend.from_nx(gpu(q_cpu)), - EMLX.Backend.from_nx(gpu(k_cpu)), - EMLX.Backend.from_nx(gpu(v_cpu)), - EMLX.Backend.from_nx(gpu(k_cache_cpu)), - EMLX.Backend.from_nx(gpu(v_cache_cpu)), - offset, - scale, - head_dim, - theta - ) - - {expected_attn, expected_k, expected_v} = - qwen3_kv_cache_attention_reference( - q_cpu, - k_cpu, - v_cpu, - k_cache_cpu, - v_cache_cpu, - offset, - scale, - head_dim, - theta - ) - - assert_all_close( - attn_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_attn, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_k, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_v, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - end - - test "matches pure Nx reference for batched GQA prefill with nonzero offset" do - {q_cpu, k_cpu, v_cpu, k_cache_cpu, v_cache_cpu} = - Nx.with_default_backend(Nx.BinaryBackend, fn -> - { - Nx.iota({2, 2, 4, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(100), - Nx.iota({2, 2, 2, 4}, type: :f32) |> Nx.add(23) |> Nx.divide(130), - Nx.iota({2, 2, 2, 4}, type: :f32) |> Nx.add(37) |> Nx.divide(140), - Nx.iota({2, 2, 6, 4}, type: :f32) |> Nx.add(5) |> Nx.divide(1_000), - Nx.iota({2, 2, 6, 4}, type: :f32) |> Nx.add(95) |> Nx.divide(1_000) - } - end) - - offset = 2 - head_dim = 4 - theta = 10_000.0 - scale = 1.0 / :math.sqrt(head_dim) - - {attn_ref, k_ref, v_ref} = - EMLX.qwen3_kv_cache_attention( - EMLX.Backend.from_nx(gpu(q_cpu)), - EMLX.Backend.from_nx(gpu(k_cpu)), - EMLX.Backend.from_nx(gpu(v_cpu)), - EMLX.Backend.from_nx(gpu(k_cache_cpu)), - EMLX.Backend.from_nx(gpu(v_cache_cpu)), - offset, - scale, - head_dim, - theta - ) - - {expected_attn, expected_k, expected_v} = - qwen3_kv_cache_attention_reference( - q_cpu, - k_cpu, - v_cpu, - k_cache_cpu, - v_cache_cpu, - offset, - scale, - head_dim, - theta - ) - - assert Nx.shape(expected_attn) == {2, 2, 16} - - assert_all_close( - attn_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_attn, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_k, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_v, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - end - end - - describe "dense Qwen3 native helpers" do - test "qwen3_mlp matches pure Nx reference" do - {hidden_cpu, norm_cpu, gate_cpu, up_cpu, down_cpu} = - Nx.with_default_backend(Nx.BinaryBackend, fn -> - { - Nx.iota({1, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(10), - Nx.tensor([1.0, 1.1, 0.9, 1.2], type: :f32), - Nx.iota({4, 6}, type: :f32) |> Nx.add(1) |> Nx.divide(50), - Nx.iota({4, 6}, type: :f32) |> Nx.add(7) |> Nx.divide(60), - Nx.iota({6, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(70) - } - end) - - eps = 1.0e-6 - - out_ref = - EMLX.qwen3_mlp( - EMLX.Backend.from_nx(gpu(hidden_cpu)), - EMLX.Backend.from_nx(gpu(norm_cpu)), - EMLX.Backend.from_nx(gpu(gate_cpu)), - EMLX.Backend.from_nx(gpu(up_cpu)), - EMLX.Backend.from_nx(gpu(down_cpu)), - eps - ) - - expected = qwen3_mlp_reference(hidden_cpu, norm_cpu, gate_cpu, up_cpu, down_cpu, eps) - - assert_all_close( - out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - end - - test "qwen3_attention_block matches pure Nx reference" do - fixtures = qwen3_dense_attention_fixtures() - scale = 1.0 / :math.sqrt(fixtures.head_dim) - - {out_ref, k_ref, v_ref} = - EMLX.qwen3_attention_block( - EMLX.Backend.from_nx(gpu(fixtures.hidden)), - EMLX.Backend.from_nx(gpu(fixtures.norm1)), - EMLX.Backend.from_nx(gpu(fixtures.q_proj)), - EMLX.Backend.from_nx(gpu(fixtures.k_proj)), - EMLX.Backend.from_nx(gpu(fixtures.v_proj)), - EMLX.Backend.from_nx(gpu(fixtures.o_proj)), - EMLX.Backend.from_nx(gpu(fixtures.q_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_cache)), - EMLX.Backend.from_nx(gpu(fixtures.v_cache)), - fixtures.offset, - scale, - fixtures.head_dim, - fixtures.theta, - fixtures.eps - ) - - {expected, expected_k, expected_v} = qwen3_attention_block_reference(fixtures) - - assert_all_close( - out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_k, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_v, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - end - - test "qwen3_attention_block matches pure Nx reference for batched GQA" do - fixtures = qwen3_batched_dense_attention_fixtures() - scale = 1.0 / :math.sqrt(fixtures.head_dim) - - {out_ref, k_ref, v_ref} = - EMLX.qwen3_attention_block( - EMLX.Backend.from_nx(gpu(fixtures.hidden)), - EMLX.Backend.from_nx(gpu(fixtures.norm1)), - EMLX.Backend.from_nx(gpu(fixtures.q_proj)), - EMLX.Backend.from_nx(gpu(fixtures.k_proj)), - EMLX.Backend.from_nx(gpu(fixtures.v_proj)), - EMLX.Backend.from_nx(gpu(fixtures.o_proj)), - EMLX.Backend.from_nx(gpu(fixtures.q_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_cache)), - EMLX.Backend.from_nx(gpu(fixtures.v_cache)), - fixtures.offset, - scale, - fixtures.head_dim, - fixtures.theta, - fixtures.eps - ) - - {expected, expected_k, expected_v} = qwen3_attention_block_reference(fixtures) - - assert Nx.shape(expected) == {2, 2, 4} - - assert_all_close( - out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_k, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_v, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - end - - test "qwen3_layer matches pure Nx reference" do - fixtures = qwen3_dense_attention_fixtures() - scale = 1.0 / :math.sqrt(fixtures.head_dim) - - {out_ref, k_ref, v_ref} = - EMLX.qwen3_layer( - EMLX.Backend.from_nx(gpu(fixtures.hidden)), - EMLX.Backend.from_nx(gpu(fixtures.norm1)), - EMLX.Backend.from_nx(gpu(fixtures.q_proj)), - EMLX.Backend.from_nx(gpu(fixtures.k_proj)), - EMLX.Backend.from_nx(gpu(fixtures.v_proj)), - EMLX.Backend.from_nx(gpu(fixtures.o_proj)), - EMLX.Backend.from_nx(gpu(fixtures.q_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_cache)), - EMLX.Backend.from_nx(gpu(fixtures.v_cache)), - EMLX.Backend.from_nx(gpu(fixtures.norm2)), - EMLX.Backend.from_nx(gpu(fixtures.gate_proj)), - EMLX.Backend.from_nx(gpu(fixtures.up_proj)), - EMLX.Backend.from_nx(gpu(fixtures.down_proj)), - fixtures.offset, - scale, - fixtures.head_dim, - fixtures.theta, - fixtures.eps - ) - - {attn_expected, expected_k, expected_v} = qwen3_attention_block_reference(fixtures) - - expected = - qwen3_mlp_reference( - attn_expected, - fixtures.norm2, - fixtures.gate_proj, - fixtures.up_proj, - fixtures.down_proj, - fixtures.eps - ) - - assert_all_close( - out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_k, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_v, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - end - - test "qwen3_layer matches pure Nx reference for batched GQA" do - fixtures = qwen3_batched_dense_attention_fixtures() - scale = 1.0 / :math.sqrt(fixtures.head_dim) - - {out_ref, k_ref, v_ref} = - EMLX.qwen3_layer( - EMLX.Backend.from_nx(gpu(fixtures.hidden)), - EMLX.Backend.from_nx(gpu(fixtures.norm1)), - EMLX.Backend.from_nx(gpu(fixtures.q_proj)), - EMLX.Backend.from_nx(gpu(fixtures.k_proj)), - EMLX.Backend.from_nx(gpu(fixtures.v_proj)), - EMLX.Backend.from_nx(gpu(fixtures.o_proj)), - EMLX.Backend.from_nx(gpu(fixtures.q_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_cache)), - EMLX.Backend.from_nx(gpu(fixtures.v_cache)), - EMLX.Backend.from_nx(gpu(fixtures.norm2)), - EMLX.Backend.from_nx(gpu(fixtures.gate_proj)), - EMLX.Backend.from_nx(gpu(fixtures.up_proj)), - EMLX.Backend.from_nx(gpu(fixtures.down_proj)), - fixtures.offset, - scale, - fixtures.head_dim, - fixtures.theta, - fixtures.eps - ) - - {attn_expected, expected_k, expected_v} = qwen3_attention_block_reference(fixtures) - - expected = - qwen3_mlp_reference( - attn_expected, - fixtures.norm2, - fixtures.gate_proj, - fixtures.up_proj, - fixtures.down_proj, - fixtures.eps - ) - - assert Nx.shape(expected) == {2, 2, 4} - - assert_all_close( - out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_k, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - - assert_all_close( - v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected_v, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - end - - test "qwen3_attention_residual matches pure Nx reference" do - {hidden, attn_out, o_proj} = - Nx.with_default_backend(Nx.BinaryBackend, fn -> - { - Nx.iota({1, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(10), - Nx.iota({1, 2, 3}, type: :f32) |> Nx.add(5) |> Nx.divide(20), - Nx.iota({3, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(30) - } - end) - - out_ref = - EMLX.qwen3_attention_residual( - EMLX.Backend.from_nx(gpu(hidden)), - EMLX.Backend.from_nx(gpu(attn_out)), - EMLX.Backend.from_nx(gpu(o_proj)) - ) - - expected = - hidden - |> Nx.add(Nx.dot(attn_out, [2], o_proj, [0])) - |> Nx.backend_transfer(Nx.BinaryBackend) - - assert_all_close( - out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected, - atol: 1.0e-5, - rtol: 1.0e-5 - ) - end - - test "qwen3_final_greedy matches pure Nx reference" do - {hidden, norm, lm_head} = - Nx.with_default_backend(Nx.BinaryBackend, fn -> - { - Nx.iota({1, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(10), - Nx.tensor([1.0, 1.1, 0.9, 1.2], type: :f32), - Nx.tensor( - [ - [0.2, 0.1, 0.0, 0.3], - [0.0, 0.5, 0.2, 0.1], - [0.3, 0.1, 0.4, 0.2] - ], - type: :f32 - ) - } - end) - - eps = 1.0e-6 - - token_ref = - EMLX.qwen3_final_greedy( - EMLX.Backend.from_nx(gpu(hidden)), - EMLX.Backend.from_nx(gpu(norm)), - EMLX.Backend.from_nx(gpu(lm_head)), - eps - ) - - expected = - hidden - |> qwen3_final_logits_reference(norm, lm_head, eps) - |> Nx.argmax(axis: 1) - - assert_all_close( - token_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), - expected - ) - end - - test "qwen3_forward_greedy_token_id defaults to embedding tensor device" do - fixtures = qwen3_dense_attention_fixtures() - cpu = fn tensor -> Nx.backend_transfer(tensor, {EMLX.Backend, device: :cpu}) end - - layer = { - cpu.(fixtures.norm1), - cpu.(fixtures.norm2), - cpu.(fixtures.q_norm), - cpu.(fixtures.k_norm), - cpu.(fixtures.q_proj), - cpu.(fixtures.k_proj), - cpu.(fixtures.v_proj), - cpu.(fixtures.o_proj), - cpu.(fixtures.gate_proj), - cpu.(fixtures.up_proj), - cpu.(fixtures.down_proj) - } - - kv_cache = - {EMLX.Backend.from_nx(cpu.(fixtures.k_cache)), - EMLX.Backend.from_nx(cpu.(fixtures.v_cache))} - - embed_tokens = - Nx.tensor(List.duplicate(List.duplicate(0.1, 4), 4), - type: :f32, - backend: {EMLX.Backend, device: :cpu} - ) - - norm = Nx.tensor(List.duplicate(1.0, 4), type: :f32, backend: {EMLX.Backend, device: :cpu}) - - lm_head = - Nx.tensor(List.duplicate(List.duplicate(0.1, 4), 4), - type: :f32, - backend: {EMLX.Backend, device: :cpu} - ) - - assert {:cpu, _ref} = EMLX.Backend.from_nx(embed_tokens) - - {_token_id, [{{k_device, _k_ref}, {v_device, _v_ref}}]} = - EMLX.qwen3_forward_greedy_token_id( - 0, - EMLX.Backend.from_nx(embed_tokens), - [layer], - [kv_cache], - EMLX.Backend.from_nx(norm), - EMLX.Backend.from_nx(lm_head), - 0, - 1.0 / :math.sqrt(fixtures.head_dim), - fixtures.head_dim, - fixtures.theta, - fixtures.eps - ) - - assert k_device == :cpu - assert v_device == :cpu - end - - test "qwen3_forward_greedy_ids returns tensor token matching token id path" do - fixtures = qwen3_dense_attention_fixtures() - input_ids = Nx.tensor([[0]], type: :s64) |> gpu() - - embed_tokens = - Nx.tensor(List.duplicate(List.duplicate(0.1, 4), 4), type: :f32) - |> gpu() - - norm = Nx.tensor(List.duplicate(1.0, 4), type: :f32) |> gpu() - - lm_head = - Nx.tensor(List.duplicate(List.duplicate(0.1, 4), 4), type: :f32) - |> gpu() - - layer = { - gpu(fixtures.norm1), - gpu(fixtures.norm2), - gpu(fixtures.q_norm), - gpu(fixtures.k_norm), - gpu(fixtures.q_proj), - gpu(fixtures.k_proj), - gpu(fixtures.v_proj), - gpu(fixtures.o_proj), - gpu(fixtures.gate_proj), - gpu(fixtures.up_proj), - gpu(fixtures.down_proj) - } - - {token_ref, _kv_cache} = - EMLX.qwen3_forward_greedy_ids( - EMLX.Backend.from_nx(input_ids), - EMLX.Backend.from_nx(embed_tokens), - [layer], - [{gpu(fixtures.k_cache), gpu(fixtures.v_cache)}], - EMLX.Backend.from_nx(norm), - EMLX.Backend.from_nx(lm_head), - 0, - 1.0 / :math.sqrt(fixtures.head_dim), - fixtures.head_dim, - fixtures.theta, - fixtures.eps - ) - - {token_id, _kv_cache} = - EMLX.qwen3_forward_greedy_token_id( - 0, - EMLX.Backend.from_nx(embed_tokens), - [layer], - [{gpu(fixtures.k_cache), gpu(fixtures.v_cache)}], - EMLX.Backend.from_nx(norm), - EMLX.Backend.from_nx(lm_head), - 0, - 1.0 / :math.sqrt(fixtures.head_dim), - fixtures.head_dim, - fixtures.theta, - fixtures.eps - ) - - token = - token_ref - |> EMLX.Backend.to_nx() - |> Nx.backend_transfer(Nx.BinaryBackend) - |> Nx.to_flat_list() - |> hd() - - assert token == token_id - end - end - - describe "EMLX.qwen3_attention_block/15 validation" do - test "qwen3_kv_cache_attention rejects required cache length overflow before graph construction" do - q = Nx.broadcast(0.0, {1, 1, 2, 2}) |> Nx.as_type(:f32) |> gpu() - k = Nx.broadcast(0.0, {1, 1, 1, 2}) |> Nx.as_type(:f32) |> gpu() - v = Nx.broadcast(0.0, {1, 1, 1, 2}) |> Nx.as_type(:f32) |> gpu() - k_cache = Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f32) |> gpu() - v_cache = Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f32) |> gpu() - - assert_raise EMLX.NIFError, - ~r/KV cache capacity 4 is smaller than required length 2147483648/, - fn -> - EMLX.qwen3_kv_cache_attention( - EMLX.Backend.from_nx(q), - EMLX.Backend.from_nx(k), - EMLX.Backend.from_nx(v), - EMLX.Backend.from_nx(k_cache), - EMLX.Backend.from_nx(v_cache), - 2_147_483_647, - 1.0 / :math.sqrt(2), - 2, - 10_000.0 - ) - end - end - - test "qwen3_layer rejects required cache length overflow before graph construction" do - fixtures = qwen3_dense_attention_fixtures() - - assert_raise EMLX.NIFError, - ~r/KV cache capacity 4 is smaller than required length 2147483649/, - fn -> - EMLX.qwen3_layer( - EMLX.Backend.from_nx(gpu(fixtures.hidden)), - EMLX.Backend.from_nx(gpu(fixtures.norm1)), - EMLX.Backend.from_nx(gpu(fixtures.q_proj)), - EMLX.Backend.from_nx(gpu(fixtures.k_proj)), - EMLX.Backend.from_nx(gpu(fixtures.v_proj)), - EMLX.Backend.from_nx(gpu(fixtures.o_proj)), - EMLX.Backend.from_nx(gpu(fixtures.q_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_cache)), - EMLX.Backend.from_nx(gpu(fixtures.v_cache)), - EMLX.Backend.from_nx(gpu(fixtures.norm2)), - EMLX.Backend.from_nx(gpu(fixtures.gate_proj)), - EMLX.Backend.from_nx(gpu(fixtures.up_proj)), - EMLX.Backend.from_nx(gpu(fixtures.down_proj)), - 2_147_483_647, - 1.0 / :math.sqrt(fixtures.head_dim), - fixtures.head_dim, - fixtures.theta, - fixtures.eps - ) - end - end - - test "qwen3_attention_block rejects required cache length overflow before graph construction" do - fixtures = qwen3_dense_attention_fixtures() - - assert_raise EMLX.NIFError, - ~r/KV cache capacity 4 is smaller than required length 2147483649/, - fn -> - EMLX.qwen3_attention_block( - EMLX.Backend.from_nx(gpu(fixtures.hidden)), - EMLX.Backend.from_nx(gpu(fixtures.norm1)), - EMLX.Backend.from_nx(gpu(fixtures.q_proj)), - EMLX.Backend.from_nx(gpu(fixtures.k_proj)), - EMLX.Backend.from_nx(gpu(fixtures.v_proj)), - EMLX.Backend.from_nx(gpu(fixtures.o_proj)), - EMLX.Backend.from_nx(gpu(fixtures.q_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_norm)), - EMLX.Backend.from_nx(gpu(fixtures.k_cache)), - EMLX.Backend.from_nx(gpu(fixtures.v_cache)), - 2_147_483_647, - 1.0 / :math.sqrt(fixtures.head_dim), - fixtures.head_dim, - fixtures.theta, - fixtures.eps - ) - end - end - - test "qwen3_forward_greedy_ids_token_id rejects batches before layer/cache normalization" do - input_ids = Nx.tensor([[0], [1]], type: :s64) |> gpu() - embed_tokens = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f32) |> gpu() - norm = Nx.broadcast(1.0, {4}) |> Nx.as_type(:f32) |> gpu() - lm_head = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f32) |> gpu() - - assert_raise ArgumentError, - ~r/qwen3_forward_greedy_ids_token_id requires batch size 1, got batch size 2/, - fn -> - EMLX.qwen3_forward_greedy_ids_token_id( - EMLX.Backend.from_nx(input_ids), - EMLX.Backend.from_nx(embed_tokens), - [:invalid_layer], - [:invalid_cache], - EMLX.Backend.from_nx(norm), - EMLX.Backend.from_nx(lm_head), - 0, - 1.0 / :math.sqrt(2), - 2, - 10_000.0, - 1.0e-6 - ) - end - end - - test "qwen3_forward_greedy_ids_chunk rejects input that is not a decode step before layer/cache normalization" do - input_ids = Nx.tensor([[0, 1]], type: :s64) |> gpu() - embed_tokens = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f32) |> gpu() - norm = Nx.broadcast(1.0, {4}) |> Nx.as_type(:f32) |> gpu() - lm_head = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f32) |> gpu() - - assert_raise ArgumentError, - ~r/qwen3_forward_greedy_ids_chunk requires sequence length 1, got sequence length 2/, - fn -> - EMLX.qwen3_forward_greedy_ids_chunk( - EMLX.Backend.from_nx(input_ids), - EMLX.Backend.from_nx(embed_tokens), - [:invalid_layer], - [:invalid_cache], - EMLX.Backend.from_nx(norm), - EMLX.Backend.from_nx(lm_head), - 0, - 1, - 1.0 / :math.sqrt(2), - 2, - 10_000.0, - 1.0e-6 - ) - end - end - - test "rejects projection widths before deriving head counts" do - hidden = Nx.broadcast(0.0, {1, 1, 4}) |> Nx.as_type(:f16) |> gpu() - norm = Nx.broadcast(1.0, {4}) |> Nx.as_type(:f16) |> gpu() - q_proj = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f16) |> gpu() - k_proj = Nx.broadcast(0.1, {4, 1}) |> Nx.as_type(:f16) |> gpu() - v_proj = Nx.broadcast(0.1, {4, 1}) |> Nx.as_type(:f16) |> gpu() - o_proj = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f16) |> gpu() - q_norm = Nx.broadcast(1.0, {2}) |> Nx.as_type(:f16) |> gpu() - k_norm = Nx.broadcast(1.0, {2}) |> Nx.as_type(:f16) |> gpu() - k_cache = Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f16) |> gpu() - v_cache = Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f16) |> gpu() - - assert_raise EMLX.NIFError, - ~r/projection output widths must be divisible by head_dim/, - fn -> - EMLX.qwen3_attention_block( - EMLX.Backend.from_nx(hidden), - EMLX.Backend.from_nx(norm), - EMLX.Backend.from_nx(q_proj), - EMLX.Backend.from_nx(k_proj), - EMLX.Backend.from_nx(v_proj), - EMLX.Backend.from_nx(o_proj), - EMLX.Backend.from_nx(q_norm), - EMLX.Backend.from_nx(k_norm), - EMLX.Backend.from_nx(k_cache), - EMLX.Backend.from_nx(v_cache), - 0, - 1.0 / :math.sqrt(2), - 2, - 10_000.0, - 1.0e-6 - ) - end - end - end - - defp qwen3_dense_attention_fixtures do - Nx.with_default_backend(Nx.BinaryBackend, fn -> - %{ - hidden: Nx.iota({1, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(10), - norm1: Nx.tensor([1.0, 1.1, 0.9, 1.2], type: :f32), - norm2: Nx.tensor([0.9, 1.0, 1.1, 1.2], type: :f32), - q_norm: Nx.tensor([1.0, 1.1], type: :f32), - k_norm: Nx.tensor([0.9, 1.2], type: :f32), - q_proj: Nx.iota({4, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(50), - k_proj: Nx.iota({4, 2}, type: :f32) |> Nx.add(2) |> Nx.divide(60), - v_proj: Nx.iota({4, 2}, type: :f32) |> Nx.add(3) |> Nx.divide(70), - o_proj: Nx.iota({4, 4}, type: :f32) |> Nx.add(4) |> Nx.divide(80), - gate_proj: Nx.iota({4, 6}, type: :f32) |> Nx.add(1) |> Nx.divide(50), - up_proj: Nx.iota({4, 6}, type: :f32) |> Nx.add(7) |> Nx.divide(60), - down_proj: Nx.iota({6, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(70), - k_cache: Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f32), - v_cache: Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f32), - offset: 0, - head_dim: 2, - theta: 10_000.0, - eps: 1.0e-6 - } - end) - end - - defp qwen3_batched_dense_attention_fixtures do - Nx.with_default_backend(Nx.BinaryBackend, fn -> - %{ - hidden: Nx.iota({2, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(20), - norm1: Nx.tensor([1.0, 1.1, 0.9, 1.2], type: :f32), - norm2: Nx.tensor([1.2, 0.8, 1.1, 0.95], type: :f32), - q_norm: Nx.tensor([1.0, 1.1], type: :f32), - k_norm: Nx.tensor([0.9, 1.2], type: :f32), - q_proj: Nx.iota({4, 8}, type: :f32) |> Nx.add(1) |> Nx.divide(80), - k_proj: Nx.iota({4, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(90), - v_proj: Nx.iota({4, 4}, type: :f32) |> Nx.add(5) |> Nx.divide(100), - o_proj: Nx.iota({8, 4}, type: :f32) |> Nx.add(7) |> Nx.divide(110), - gate_proj: Nx.iota({4, 6}, type: :f32) |> Nx.add(13) |> Nx.divide(120), - up_proj: Nx.iota({4, 6}, type: :f32) |> Nx.add(29) |> Nx.divide(130), - down_proj: Nx.iota({6, 4}, type: :f32) |> Nx.add(41) |> Nx.divide(140), - k_cache: Nx.iota({2, 2, 6, 2}, type: :f32) |> Nx.add(11) |> Nx.divide(1_000), - v_cache: Nx.iota({2, 2, 6, 2}, type: :f32) |> Nx.add(101) |> Nx.divide(1_000), - offset: 2, - head_dim: 2, - theta: 10_000.0, - eps: 1.0e-6 - } - end) - end - - defp qwen3_kv_cache_attention_reference( - q, - new_k, - new_v, - k_cache, - v_cache, - offset, - scale, - head_dim, - theta - ) do - {batch, seq_len, q_heads, _head_dim} = Nx.shape(q) - {_batch, _seq_len, kv_heads, _head_dim} = Nx.shape(new_k) - - q_bn = - q - |> Nx.transpose(axes: [0, 2, 1, 3]) - |> qwen3_rope_reference(head_dim, theta, offset) - |> Nx.backend_transfer(Nx.BinaryBackend) - - k_bn = - new_k - |> Nx.transpose(axes: [0, 2, 1, 3]) - |> qwen3_rope_reference(head_dim, theta, offset) - |> Nx.backend_transfer(Nx.BinaryBackend) - - v_bn = new_v |> Nx.transpose(axes: [0, 2, 1, 3]) |> Nx.backend_transfer(Nx.BinaryBackend) - - k_cache = Nx.put_slice(k_cache, [0, 0, offset, 0], k_bn) - v_cache = Nx.put_slice(v_cache, [0, 0, offset, 0], v_bn) - - valid_len = offset + seq_len - k_valid = Nx.slice_along_axis(k_cache, 0, valid_len, axis: 2) - v_valid = Nx.slice_along_axis(v_cache, 0, valid_len, axis: 2) - - groups = div(q_heads, kv_heads) - k_repeated = repeat_kv_heads_bn(k_valid, groups) - v_repeated = repeat_kv_heads_bn(v_valid, groups) - - scores = - q_bn - |> Nx.new_axis(3) - |> Nx.multiply(Nx.new_axis(k_repeated, 2)) - |> Nx.sum(axes: [4]) - |> Nx.multiply(scale) - |> apply_causal_mask(offset, seq_len, valid_len) - - weights = softmax_reference(scores, 3) - - attn = - weights - |> Nx.new_axis(4) - |> Nx.multiply(Nx.new_axis(v_repeated, 2)) - |> Nx.sum(axes: [3]) - - attn_out = - attn - |> Nx.transpose(axes: [0, 2, 1, 3]) - |> Nx.reshape({batch, seq_len, q_heads * head_dim}) - - {attn_out, k_cache, v_cache} - end - - defp qwen3_mlp_reference(hidden, norm, gate_proj, up_proj, down_proj, eps) do - xn = rms_norm_reference(hidden, norm, eps) - gate = Nx.dot(xn, [2], gate_proj, [0]) - up = Nx.dot(xn, [2], up_proj, [0]) - mlp = Nx.multiply(Nx.multiply(gate, Nx.sigmoid(gate)), up) - out = Nx.dot(mlp, [2], down_proj, [0]) - Nx.add(hidden, out) - end - - defp qwen3_attention_block_reference(fixtures) do - hidden = fixtures.hidden - offset = fixtures.offset - head_dim = fixtures.head_dim - scale = 1.0 / :math.sqrt(head_dim) - - {batch, seq_len, _hidden_size} = Nx.shape(hidden) - - xn = rms_norm_reference(hidden, fixtures.norm1, fixtures.eps) - - q_flat = Nx.dot(xn, [2], fixtures.q_proj, [0]) - k_flat = Nx.dot(xn, [2], fixtures.k_proj, [0]) - v_flat = Nx.dot(xn, [2], fixtures.v_proj, [0]) - - q_heads = div(elem(Nx.shape(fixtures.q_proj), 1), head_dim) - kv_heads = div(elem(Nx.shape(fixtures.k_proj), 1), head_dim) - - q = Nx.reshape(q_flat, {batch, seq_len, q_heads, head_dim}) - k = Nx.reshape(k_flat, {batch, seq_len, kv_heads, head_dim}) - v = Nx.reshape(v_flat, {batch, seq_len, kv_heads, head_dim}) - - q = rms_norm_reference(q, fixtures.q_norm, fixtures.eps) - k = rms_norm_reference(k, fixtures.k_norm, fixtures.eps) - - q_bn = - q - |> Nx.transpose(axes: [0, 2, 1, 3]) - |> qwen3_rope_reference(head_dim, fixtures.theta, offset) - - k_bn = - k - |> Nx.transpose(axes: [0, 2, 1, 3]) - |> qwen3_rope_reference(head_dim, fixtures.theta, offset) - - q_bn = Nx.backend_transfer(q_bn, Nx.BinaryBackend) - k_bn = Nx.backend_transfer(k_bn, Nx.BinaryBackend) - v_bn = v |> Nx.transpose(axes: [0, 2, 1, 3]) |> Nx.backend_transfer(Nx.BinaryBackend) - - k_cache = Nx.put_slice(fixtures.k_cache, [0, 0, offset, 0], k_bn) - v_cache = Nx.put_slice(fixtures.v_cache, [0, 0, offset, 0], v_bn) - - valid_len = offset + seq_len - k_valid = Nx.slice_along_axis(k_cache, 0, valid_len, axis: 2) - v_valid = Nx.slice_along_axis(v_cache, 0, valid_len, axis: 2) - - groups = div(q_heads, kv_heads) - k_repeated = repeat_kv_heads_bn(k_valid, groups) - v_repeated = repeat_kv_heads_bn(v_valid, groups) - - scores = - q_bn - |> Nx.new_axis(3) - |> Nx.multiply(Nx.new_axis(k_repeated, 2)) - |> Nx.sum(axes: [4]) - |> Nx.multiply(scale) - - scores = apply_causal_mask(scores, offset, seq_len, valid_len) - weights = softmax_reference(scores, 3) - - attn = - weights - |> Nx.new_axis(4) - |> Nx.multiply(Nx.new_axis(v_repeated, 2)) - |> Nx.sum(axes: [3]) - - attn_out = - attn - |> Nx.transpose(axes: [0, 2, 1, 3]) - |> Nx.reshape({batch, seq_len, q_heads * head_dim}) - - projected = Nx.dot(attn_out, [2], fixtures.o_proj, [0]) - - {Nx.add(hidden, projected), k_cache, v_cache} - end - - defp qwen3_final_logits_reference(hidden, norm, lm_head, eps) do - {_batch, seq_len, hidden_size} = Nx.shape(hidden) - - last = - hidden - |> Nx.slice([0, seq_len - 1, 0], [1, 1, hidden_size]) - |> Nx.reshape({1, hidden_size}) - - last - |> rms_norm_reference(norm, eps) - |> Nx.dot([1], lm_head, [1]) - end - - defp rms_norm_reference(tensor, weight, eps) do - tensor - |> Nx.pow(2) - |> Nx.mean(axes: [-1], keep_axes: true) - |> Nx.add(eps) - |> Nx.sqrt() - |> then(&Nx.divide(tensor, &1)) - |> Nx.multiply(weight) - end - - defp qwen3_rope_reference(tensor, dims, theta, offset) do - {_batch, _heads, seq_len, _dims} = Nx.shape(tensor) - half = div(dims, 2) - - inv_freq = - Nx.iota({half}, type: :f32) - |> Nx.multiply(2) - |> Nx.divide(dims) - |> then(&Nx.pow(theta, &1)) - |> then(&Nx.divide(1.0, &1)) - - positions = - Nx.iota({seq_len}, type: :f32) - |> Nx.add(offset) - |> Nx.reshape({seq_len, 1}) - - freqs = Nx.multiply(positions, Nx.reshape(inv_freq, {1, half})) - - cos = - Nx.concatenate([Nx.cos(freqs), Nx.cos(freqs)], axis: 1) - |> Nx.reshape({1, 1, seq_len, dims}) - - sin = - Nx.concatenate([Nx.sin(freqs), Nx.sin(freqs)], axis: 1) - |> Nx.reshape({1, 1, seq_len, dims}) - - first = tensor[[.., .., .., 0..(half - 1)//1]] - second = tensor[[.., .., .., half..(dims - 1)//1]] - rotated = Nx.concatenate([Nx.negate(second), first], axis: 3) - - Nx.add(Nx.multiply(tensor, cos), Nx.multiply(rotated, sin)) - end - - defp repeat_kv_heads_bn(tensor, 1), do: tensor - - defp repeat_kv_heads_bn(tensor, groups) do - {batch, n_kv, t_kv, head_dim} = Nx.shape(tensor) - - tensor - |> Nx.new_axis(2) - |> Nx.broadcast({batch, n_kv, groups, t_kv, head_dim}) - |> Nx.reshape({batch, n_kv * groups, t_kv, head_dim}) - end - - defp apply_causal_mask(scores, offset, seq_len, valid_len) do - query_positions = - Nx.iota({seq_len}, type: :s32, backend: Nx.BinaryBackend) - |> Nx.add(offset) - |> Nx.reshape({1, 1, seq_len, 1}) - - key_positions = - Nx.iota({valid_len}, type: :s32, backend: Nx.BinaryBackend) - |> Nx.reshape({1, 1, 1, valid_len}) - - mask = - key_positions - |> Nx.less_equal(query_positions) - |> Nx.broadcast(Nx.shape(scores)) - - Nx.select(mask, scores, Nx.broadcast(-1.0e9, Nx.shape(scores))) - end - - defp softmax_reference(tensor, axis) do - shifted = Nx.subtract(tensor, Nx.reduce_max(tensor, axes: [axis], keep_axes: true)) - exp = Nx.exp(shifted) - Nx.divide(exp, Nx.sum(exp, axes: [axis], keep_axes: true)) - end - defp causal_kv_attention_reference(q, new_k, new_v, k_cache, v_cache, offset, scale, key_mask) do q = Nx.backend_transfer(q, Nx.BinaryBackend) new_k = Nx.backend_transfer(new_k, Nx.BinaryBackend) diff --git a/emlx/test/emlx/grad_equivalence_test.exs b/emlx/test/emlx/grad_equivalence_test.exs new file mode 100644 index 0000000..693ec8a --- /dev/null +++ b/emlx/test/emlx/grad_equivalence_test.exs @@ -0,0 +1,445 @@ +defmodule EMLX.GradEquivalenceTest do + @moduledoc """ + Permanent grad-equivalence regression suite, widening + `grad_triage_test.exs`'s 8-scenario triage into materially more op-class + combinations, shapes, and dtypes. + + Two adjustments to the original plan (user directive + advisor sign-off): + + * **No `StreamData`.** Breadth comes from table-driven fixed + scenario × shape × dtype combinations (looped at test-run time), not + generated inputs. + * **Non-differentiable ops are included, not excluded.** + `Nx.Defn.Grad`'s `@constants` list (`deps/nx/nx/lib/nx/defn/grad.ex`) + already treats `argmax`/`argmin`/`floor`/`sign`/comparisons as a + stop-gradient boundary when they appear as an operand of a + differentiable op — grad is always well-defined for them (typically + zero contribution through that operand). The real EMLX-specific + question this suite checks is whether the *native* compiled backward + pass applies the exact same stop-gradient rule as the Evaluator, not + whether grad exists at all. + + Reference convention (same as `grad_triage_test.exs`): `reference/2` runs + `Nx.Defn.jit_apply/3` with `compiler: Nx.Defn.Evaluator` on + `Nx.BinaryBackend` tensors; `native/2` runs the same with `compiler: EMLX`. + """ + + use EMLX.Case, async: false + import Nx.Defn + + # ── shared helpers ───────────────────────────────────────────────────────── + + # `EMLX.Case`'s setup sets the *process-global* default backend to + # `EMLX.Backend`. `Nx.Defn.Evaluator` uses that default backend for any + # tensor it synthesizes internally (e.g. `window_max`'s padding fill value) + # rather than matching the explicit-backend args passed in — so without + # this scoping, the reference silently mixes in `EMLX.Backend` and stops being + # a pure BinaryBackend/Evaluator reference. Found via this stage's + # `window_max`-with-explicit-padding scenario (see Results below). + defp reference(fun, args) do + previous = Nx.default_backend() + Nx.default_backend(Nx.BinaryBackend) + + try do + Nx.Defn.jit_apply(fun, args, compiler: Nx.Defn.Evaluator) + after + Nx.default_backend(previous) + end + end + + defp native(fun, args), do: Nx.Defn.jit_apply(fun, args, compiler: EMLX) + + # Unique-valued, zero-free, tie-free tensor generator — safe for argmax + # (no ties), sign/floor (no zero/no exact-integer boundary), and division + # (no zero denominator) all at once, across any shape/dtype combination. + defp bin(shape, dtype) do + size = Tuple.product(shape) + + values = + for i <- 0..(size - 1) do + (i - size / 2) * 0.37 + 0.6 + end + + values + |> Nx.tensor(type: dtype, backend: Nx.BinaryBackend) + |> Nx.reshape(shape) + end + + defp assert_grad_equivalent(fun, args, tol \\ [atol: 1.0e-3, rtol: 1.0e-3]) do + assert_all_close(native(fun, args), reference(fun, args), tol) + end + + @shapes [{}, {4}, {2, 3}, {2, 2, 3}] + @rank1plus_shapes [{4}, {2, 3}, {2, 2, 3}] + @dtypes [{:f, 32}, {:f, 64}] + + # ── 1. smooth elementwise chain (broader op mix than sin*cos) ───────────── + + defn smooth_elementwise_loss(x) do + x + |> Nx.sin() + |> Nx.multiply(Nx.cos(x)) + |> Nx.add(Nx.log1p(Nx.abs(x))) + |> Nx.subtract(Nx.tanh(x)) + |> Nx.multiply(Nx.sigmoid(x)) + |> Nx.add(Nx.sqrt(Nx.abs(x))) + |> Nx.sum() + end + + defn smooth_elementwise_grad(x), do: grad(x, &smooth_elementwise_loss/1) + + describe "smooth elementwise chain grad (sin/cos/log1p/tanh/sigmoid/sqrt/abs composed)" do + test "matches the Evaluator reference across shapes and dtypes" do + for shape <- @shapes, dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&smooth_elementwise_grad/1, [x]) + end + end + end + + # ── 2. reduction chain (sum/mean/reduce_max/reduce_min composed) ─────────── + + defn reduction_zoo_loss(x) do + x + |> Nx.abs() + |> Nx.add(1) + |> Nx.log() + |> Nx.sum(axes: [-1]) + |> Nx.mean() + |> Nx.add(Nx.multiply(Nx.reduce_max(x), 0.1)) + |> Nx.add(Nx.multiply(Nx.reduce_min(x), 0.1)) + end + + defn reduction_zoo_grad(x), do: grad(x, &reduction_zoo_loss/1) + + describe "reduction chain grad (sum/mean/reduce_max/reduce_min composed)" do + test "matches the Evaluator reference across rank>=1 shapes and dtypes" do + for shape <- @rank1plus_shapes, dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&reduction_zoo_grad/1, [x]) + end + end + end + + # ── 3. dot chain (matmul backward through tanh + sum) ────────────────────── + + defn dot_chain_loss(a, b) do + a |> Nx.dot(b) |> Nx.tanh() |> Nx.sum() + end + + defn dot_chain_grad(a, b), do: grad(a, fn a -> dot_chain_loss(a, b) end) + + describe "dot chain grad (dot -> tanh -> sum)" do + test "matches the Evaluator reference across shape pairs and dtypes" do + for {shape_a, shape_b} <- [{{2, 3}, {3, 2}}, {{4, 4}, {4, 4}}, {{3, 4}, {4, 2}}], + dtype <- @dtypes do + a = bin(shape_a, dtype) + b = bin(shape_b, dtype) + assert_grad_equivalent(&dot_chain_grad/2, [a, b]) + end + end + end + + # ── 4. nested cond (cond inside cond, all four leaf branches exercised) ──── + + defn nested_cond_loss(x) do + cond do + Nx.all(Nx.greater(x, 0)) -> + cond do + Nx.all(Nx.greater(Nx.sum(x), 10)) -> Nx.sum(Nx.multiply(x, x)) + true -> Nx.sum(Nx.exp(x)) + end + + true -> + cond do + Nx.all(Nx.less(x, -10)) -> Nx.sum(Nx.abs(x)) + true -> Nx.sum(Nx.multiply(x, 3)) + end + end + end + + defn nested_cond_grad(x), do: grad(x, &nested_cond_loss/1) + + describe "nested cond grad (cond inside cond, all four leaf branches)" do + test "matches the Evaluator reference for each of the four branch combinations" do + # {all positive?, sum > 10 or all < -10?} + cases = [ + Nx.tensor([1.0, 2.0, 8.0]), + Nx.tensor([0.1, 0.2, 0.3]), + Nx.tensor([-11.0, -12.0, -13.0]), + Nx.tensor([-1.0, 2.0, -3.0]) + ] + + for x <- cases do + assert_grad_equivalent(&nested_cond_grad/1, [Nx.backend_transfer(x, Nx.BinaryBackend)]) + end + end + end + + # ── 5. while whose body itself contains a cond ───────────────────────────── + + defn while_with_cond_loss(x) do + {out, _i} = + while {out = x, i = 0}, Nx.less(i, 4) do + out = + cond do + Nx.all(Nx.greater(Nx.sum(out), 0)) -> Nx.multiply(out, 1.1) + true -> Nx.add(out, 0.05) + end + + {out, i + 1} + end + + Nx.sum(out) + end + + defn while_with_cond_grad(x), do: grad(x, &while_with_cond_loss/1) + + describe "while-body-contains-cond grad" do + # Not checked against the `Nx.Defn.Evaluator` reference here — a genuine + # `Nx.Defn.Grad` bug (not an EMLX bug) makes that reference itself wrong for + # this exact scenario (backward `:while` + nested data-dependent `cond`). + # See `workdir/native-compiler/nx-grad-while-cond-bugreport.md`: EMLX's + # native result is finite-difference-correct; `Nx.Defn.Evaluator`'s + # (pure-BinaryBackend, no EMLX involved) is off by up to 20 orders of + # magnitude on some elements. So this scenario is checked against a + # finite-difference reference instead. + test "matches a finite-difference reference (not the known-broken Evaluator path)" do + eps = 1.0e-4 + + for dtype <- @dtypes do + x = bin({3}, dtype) + native_grad = native(&while_with_cond_grad/1, [x]) + + fd_grad = + for i <- 0..2 do + plus = Nx.indexed_add(x, Nx.tensor([[i]]), Nx.tensor([eps], type: dtype)) + minus = Nx.indexed_add(x, Nx.tensor([[i]]), Nx.tensor([-eps], type: dtype)) + + (Nx.to_number(native(&while_with_cond_loss/1, [plus])) - + Nx.to_number(native(&while_with_cond_loss/1, [minus]))) / (2 * eps) + end + |> Nx.tensor(type: dtype) + + assert_all_close(native_grad, fd_grad, atol: 5.0e-3, rtol: 5.0e-3) + end + end + end + + # ── 6. multi-output while carries (3 carried tensors) ────────────────────── + + defn while_multi_carry_loss(x) do + {acc1, acc2, _i} = + while {acc1 = x, acc2 = Nx.multiply(x, 2), i = 0}, Nx.less(i, 3) do + {Nx.add(acc1, acc2), Nx.multiply(acc2, 1.05), i + 1} + end + + Nx.sum(Nx.add(acc1, acc2)) + end + + defn while_multi_carry_grad(x), do: grad(x, &while_multi_carry_loss/1) + + describe "multi-output while carries grad (3 carried tensors)" do + test "matches the Evaluator reference" do + for dtype <- @dtypes do + x = bin({4}, dtype) + assert_grad_equivalent(&while_multi_carry_grad/1, [x]) + end + end + end + + # ── 7. nested while (while inside while) ─────────────────────────────────── + + defn nested_while_loss(x) do + {out, _i} = + while {out = x, i = 0}, Nx.less(i, 2) do + {inner, _j} = + while {inner = out, j = 0}, Nx.less(j, 2) do + {Nx.add(inner, 0.1), j + 1} + end + + {inner, i + 1} + end + + Nx.sum(out) + end + + defn nested_while_grad(x), do: grad(x, &nested_while_loss/1) + + describe "nested while grad (while inside while)" do + test "matches the Evaluator reference" do + for dtype <- @dtypes do + x = bin({3}, dtype) + assert_grad_equivalent(&nested_while_grad/1, [x]) + end + end + end + + # ── 8. windowed ops with non-default strides/padding ─────────────────────── + + defn window_sum_strided_loss(x), do: Nx.sum(Nx.window_sum(x, {2, 2}, strides: [2, 1])) + defn window_sum_strided_grad(x), do: grad(x, &window_sum_strided_loss/1) + + defn window_max_padded_loss(x), + do: Nx.sum(Nx.window_max(x, {2, 2}, padding: [{1, 1}, {1, 1}])) + + defn window_max_padded_grad(x), do: grad(x, &window_max_padded_loss/1) + + defn window_sum_strided_3d_loss(x), + do: Nx.sum(Nx.window_sum(x, {2, 2, 2}, strides: [2, 1, 2])) + + defn window_sum_strided_3d_grad(x), do: grad(x, &window_sum_strided_3d_loss/1) + + describe "windowed ops grad with non-default strides/padding" do + # `window_sum`'s backward (grad.ex's `grad(:window_sum, …)`) un-strides + # the cotangent via `Nx.pad` with *interior* padding whenever + # `strides != 1`. `:pad` with interior padding (and negative lo/hi) lowers + # natively (see `EMLX.Native.Expr.expand_pad_general/5`). The default + # (unit-stride) `window_sum` grad scenario elsewhere in this suite never + # exercises this path. + test "window_sum with non-unit strides matches the Evaluator reference" do + for shape <- [{4, 4}, {3, 5}], dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&window_sum_strided_grad/1, [x]) + end + end + + test "window_sum with non-unit strides matches the Evaluator reference (3D)" do + for dtype <- @dtypes do + x = bin({4, 3, 5}, dtype) + assert_grad_equivalent(&window_sum_strided_3d_grad/1, [x]) + end + end + + test "window_max with padding matches the Evaluator reference" do + for shape <- [{4, 4}, {3, 5}], dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&window_max_padded_grad/1, [x]) + end + end + end + + # ── 8b. direct :pad / :slice grad ──────────────────────────────────────── + # + # Broader than window ops: grad.ex's + # own `grad(:pad, …)` un-pads the cotangent via *negative* lo/hi whenever the + # forward `Nx.pad` had positive lo/hi, unconditionally (no strides needed); + # `grad(:slice, …)` re-inserts *interior* padding into the cotangent whenever + # the forward `Nx.slice` used non-unit strides. Both are more common than the + # window_sum path above and exercise the same `:pad` decomposition from a + # different direction. + + defn pad_positive_loss(x), do: Nx.sum(Nx.pad(x, 0.0, [{2, 1, 0}, {0, 3, 0}])) + defn pad_positive_grad(x), do: grad(x, &pad_positive_loss/1) + + defn slice_strided_loss(x), do: Nx.sum(Nx.slice(x, [0, 0], [2, 3], strides: [2, 3])) + defn slice_strided_grad(x), do: grad(x, &slice_strided_loss/1) + + describe "direct :pad / :slice grad" do + test "pad with positive lo/hi (negative-lo/hi backward) matches the Evaluator reference" do + for shape <- [{4, 4}, {3, 5}], dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&pad_positive_grad/1, [x]) + end + end + + test "slice with non-unit strides (interior-pad backward) matches the Evaluator reference" do + for shape <- [{5, 8}], dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&slice_strided_grad/1, [x]) + end + end + end + + # ── 9. non-differentiable ops used as an operand (stop-gradient boundary) ── + # + # Per this stage's plan adjustment: these aren't excluded — `Nx.Defn.Grad`'s + # `@constants` list makes them act as a stop-gradient boundary when used as + # an operand of a differentiable op, so grad is well-defined. What's tested + # here is that EMLX's native backward lowering applies that exact same + # boundary as the Evaluator, at points chosen away from each op's + # discontinuity (no exact zero for `sign`, no argmax ties, no integer + # boundary for `floor`). + + defn sign_operand_loss(x), do: Nx.sum(Nx.multiply(x, Nx.sign(x))) + defn sign_operand_grad(x), do: grad(x, &sign_operand_loss/1) + + defn floor_operand_loss(x), do: Nx.sum(Nx.multiply(x, Nx.floor(Nx.abs(x) |> Nx.add(1)))) + defn floor_operand_grad(x), do: grad(x, &floor_operand_loss/1) + + defn comparison_select_loss(x) do + Nx.sum(Nx.select(Nx.greater(x, 0), Nx.multiply(x, x), Nx.multiply(x, -1))) + end + + defn comparison_select_grad(x), do: grad(x, &comparison_select_loss/1) + + defn argmax_gather_loss(x) do + idx = Nx.argmax(x, axis: -1, keep_axis: true) + Nx.sum(Nx.take_along_axis(x, idx, axis: -1)) + end + + defn argmax_gather_grad(x), do: grad(x, &argmax_gather_loss/1) + + describe "non-differentiable-op-as-operand grad (stop-gradient boundary parity)" do + test "sign(x) used as a multiplicative operand matches the Evaluator reference" do + for shape <- @shapes, dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&sign_operand_grad/1, [x]) + end + end + + test "floor(x) used as a multiplicative operand matches the Evaluator reference" do + for shape <- @shapes, dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&floor_operand_grad/1, [x]) + end + end + + test "a comparison feeding select matches the Evaluator reference" do + for shape <- @shapes, dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&comparison_select_grad/1, [x]) + end + end + + test "argmax used to gather (max-pooling-style pattern) matches the Evaluator reference" do + for shape <- @rank1plus_shapes, dtype <- @dtypes do + x = bin(shape, dtype) + assert_grad_equivalent(&argmax_gather_grad/1, [x]) + end + end + end + + # ── 10. finite-difference reference (smooth ops only) ───────────────────────── + # + # FD is meaningless exactly at a non-differentiable op's discontinuity, so + # this is restricted to the smooth subset (per advisor guidance) and run at + # points chosen away from any domain edge (all positive, away from zero). + # A vector shape is used (not just scalar): `Nx.Defn.grad/2` treats extra + # dims as batch dims, so this also validates the gradient is correct + # *per element*, not just in aggregate. + + @smooth_unary_ops [:sin, :cos, :exp, :tanh, :sigmoid, :sqrt, :log, :cbrt, :expm1, :log1p] + @fd_eps 1.0e-3 + # f32 central differences bottom out ~1e-3 relative — use a matching + # tolerance. + @fd_tol [atol: 5.0e-3, rtol: 5.0e-3] + + describe "finite-difference reference (smooth unary ops, points away from discontinuities)" do + test "native grad matches central-difference FD for each smooth unary op" do + x = Nx.tensor([0.3, 0.8, 1.4, 2.1], type: {:f, 64}, backend: Nx.BinaryBackend) + + for op <- @smooth_unary_ops do + grad_fn = fn x -> Nx.Defn.grad(x, &apply(Nx, op, [&1])) end + native_grad = Nx.Defn.jit_apply(grad_fn, [x], compiler: EMLX) + + fd_grad = + apply(Nx, op, [Nx.add(x, @fd_eps)]) + |> Nx.subtract(apply(Nx, op, [Nx.subtract(x, @fd_eps)])) + |> Nx.divide(2 * @fd_eps) + + assert_all_close(native_grad, fd_grad, @fd_tol) + end + end + end +end diff --git a/emlx/test/emlx/grad_triage_test.exs b/emlx/test/emlx/grad_triage_test.exs new file mode 100644 index 0000000..c2942f2 --- /dev/null +++ b/emlx/test/emlx/grad_triage_test.exs @@ -0,0 +1,146 @@ +defmodule EMLX.GradTriageTest do + use EMLX.Case, async: false + import Nx.Defn + + # ── forward functions (the zoo) ─────────────────────────────────────────── + + defn elementwise_loss(x), do: Nx.sum(Nx.multiply(Nx.sin(x), Nx.cos(x))) + defn elementwise_grad(x), do: grad(x, &elementwise_loss/1) + + defn reduction_loss(x), do: Nx.mean(Nx.sum(x, axes: [1])) + defn reduction_grad(x), do: grad(x, &reduction_loss/1) + + defn dot_loss(a, b), do: Nx.sum(Nx.dot(a, b)) + defn dot_grad(a, b), do: grad(a, fn a -> dot_loss(a, b) end) + + defn cond_loss(x) do + cond do + Nx.all(Nx.greater(x, 0)) -> Nx.sum(Nx.multiply(x, x)) + true -> Nx.sum(Nx.abs(x)) + end + end + + defn cond_grad(x), do: grad(x, &cond_loss/1) + + defn while_loss(x) do + {out, _i} = + while {out = x, i = 0}, Nx.less(i, 3) do + {Nx.multiply(out, out), Nx.add(i, 1)} + end + + Nx.sum(out) + end + + defn while_grad(x), do: grad(x, &while_loss/1) + + defn window_loss(x), do: Nx.sum(Nx.window_sum(x, {2, 2})) + defn window_grad(x), do: grad(x, &window_loss/1) + + # window_max's backward hits :window_scatter_max (not just window_sum + # again, unlike window_sum's own backward) — a distinct opcode, tested + # separately. + defn window_max_loss(x), do: Nx.sum(Nx.window_max(x, {2, 2})) + defn window_max_grad(x), do: grad(x, &window_max_loss/1) + + # ── reference helper ────────────────────────────────────────────────────────── + + # Reference runs the same defn through the plain Evaluator on BinaryBackend + # tensors (no EMLX involved at all) — independent of whatever the process + # default backend/compiler happens to be. + defp reference(fun, args) do + Nx.Defn.jit_apply(fun, args, compiler: Nx.Defn.Evaluator) + end + + defp native(fun, args) do + Nx.Defn.jit_apply(fun, args, compiler: EMLX) + end + + defp bin(list_or_tensor, opts \\ []) do + Nx.tensor(list_or_tensor, [backend: Nx.BinaryBackend] ++ opts) + end + + describe "elementwise grad" do + test "matches the Evaluator reference under compiler: EMLX" do + x = bin([0.1, 0.5, -0.3, 1.2]) + + assert_all_close( + native(&elementwise_grad/1, [x]), + reference(&elementwise_grad/1, [x]) + ) + end + end + + describe "reduction grad" do + test "matches the Evaluator reference under compiler: EMLX" do + x = bin([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + + assert_all_close( + native(&reduction_grad/1, [x]), + reference(&reduction_grad/1, [x]) + ) + end + end + + describe "dot grad" do + test "matches the Evaluator reference under compiler: EMLX" do + a = bin([[1.0, 2.0], [3.0, 4.0]]) + b = bin([[5.0, 6.0], [7.0, 8.0]]) + + assert_all_close( + native(&dot_grad/2, [a, b]), + reference(&dot_grad/2, [a, b]) + ) + end + end + + describe "cond grad" do + test "matches the Evaluator reference under compiler: EMLX (branch taken: true)" do + x = bin([1.0, 2.0, 3.0]) + + assert_all_close( + native(&cond_grad/1, [x]), + reference(&cond_grad/1, [x]) + ) + end + + test "matches the Evaluator reference under compiler: EMLX (branch taken: false)" do + x = bin([-1.0, -2.0, 3.0]) + + assert_all_close( + native(&cond_grad/1, [x]), + reference(&cond_grad/1, [x]) + ) + end + end + + describe "while grad (backward pass builds its own :while node — Nx.Defn.Graph splits it same as a forward while)" do + test "matches the Evaluator reference under compiler: EMLX" do + x = bin([0.5, 0.6, 0.7]) + + assert_all_close( + native(&while_grad/1, [x]), + reference(&while_grad/1, [x]) + ) + end + end + + describe "windowed-reduce grad (eager EMLX.Backend.window_reduce/6 hard-raises; window_sum forward is native-only in the compiler's IR)" do + test "window_sum grad matches the Evaluator reference under compiler: EMLX" do + x = bin([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) + + assert_all_close( + native(&window_grad/1, [x]), + reference(&window_grad/1, [x]) + ) + end + + test "window_max grad (backward hits :window_scatter_max) matches the Evaluator reference under compiler: EMLX" do + x = bin([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) + + assert_all_close( + native(&window_max_grad/1, [x]), + reference(&window_max_grad/1, [x]) + ) + end + end +end diff --git a/emlx/test/emlx/microscaled_quantization_test.exs b/emlx/test/emlx/microscaled_quantization_test.exs new file mode 100644 index 0000000..bebbeb5 --- /dev/null +++ b/emlx/test/emlx/microscaled_quantization_test.exs @@ -0,0 +1,164 @@ +defmodule EMLX.MicroscaledQuantizationTest do + use EMLX.Case, async: true + + @moduletag :metal + + alias EMLX.Quantization + + # {mode, group_size, bits}. Each microscaled mode pins an exact + # {group_size, bits} pair (mlx::core::quantize's mode-specific constraint). + @modes [ + {"affine", 64, 4}, + {"mxfp4", 32, 4}, + {"mxfp8", 32, 8}, + {"nvfp4", 16, 4} + ] + + describe "EMLX.quantize/2 per-mode" do + for {mode, group_size, bits} <- @modes do + test "#{mode}: round-trips through EMLX.quantize/EMLX.dequantize" do + mode = unquote(mode) + group_size = unquote(group_size) + bits = unquote(bits) + + weight = Nx.iota({64, 64}, type: :f32) |> Nx.divide(100) + weight = Nx.backend_transfer(weight, {EMLX.Backend, device: :gpu}) + + qw = EMLX.quantize(weight, type: {:s, bits}, group_size: group_size, mode: mode) + + %EMLX.Backend{quantization_config: cfg} = qw.data + assert cfg.mode == mode + assert cfg.group_size == group_size + assert cfg.bits == bits + + if mode == "affine" do + assert %Nx.Tensor{} = cfg.biases + else + assert cfg.biases == nil + end + + dense = EMLX.dequantize(qw) + assert Nx.shape(dense) == {64, 64} + + original = Nx.backend_transfer(weight, Nx.BinaryBackend) + recovered = Nx.backend_transfer(dense, Nx.BinaryBackend) + + original_mean = Nx.mean(original) |> Nx.to_number() + recovered_mean = Nx.mean(recovered) |> Nx.to_number() + + # Lossy quantization (esp. 4-bit) — assert same ballpark, not exact match. + assert abs(original_mean - recovered_mean) / original_mean < 0.5 + end + end + + test "raises on an unknown :mode" do + weight = Nx.iota({64, 64}, type: :f32) + weight = Nx.backend_transfer(weight, {EMLX.Backend, device: :gpu}) + + assert_raise ArgumentError, ~r/:mode must be one of/, fn -> + EMLX.quantize(weight, mode: "not_a_real_mode") + end + end + + for {mode, expected_gs, expected_bits} <- Enum.reject(@modes, &(elem(&1, 0) == "affine")) do + test "#{mode}: raises when group_size doesn't match the mode's fixed constraint" do + mode = unquote(mode) + expected_gs = unquote(expected_gs) + expected_bits = unquote(expected_bits) + wrong_gs = if expected_gs == 64, do: 32, else: 64 + + weight = Nx.iota({64, 128}, type: :f32) + weight = Nx.backend_transfer(weight, {EMLX.Backend, device: :gpu}) + + assert_raise ArgumentError, ~r/requires group_size=#{expected_gs}/, fn -> + EMLX.quantize(weight, mode: mode, group_size: wrong_gs, type: {:s, expected_bits}) + end + end + end + end + + describe "EMLX.quantized_matmul/2 per-mode vs dense dot" do + for {mode, group_size, bits} <- @modes do + test "#{mode}: matches Nx.dot(x, transpose(dense)) within quantization tolerance" do + mode = unquote(mode) + group_size = unquote(group_size) + bits = unquote(bits) + + x = Nx.iota({1, 4, 128}, type: :f32) |> Nx.divide(100) + x = Nx.backend_transfer(x, {EMLX.Backend, device: :gpu}) + + w = Nx.iota({64, 128}, type: :f32) |> Nx.divide(1000) + w = Nx.backend_transfer(w, {EMLX.Backend, device: :gpu}) + + qw = EMLX.quantize(w, type: {:s, bits}, group_size: group_size, mode: mode) + dense = EMLX.dequantize(qw) + + # dense (and qw) is {out_features, in_features}; quantized_matmul's + # transpose: true means x @ dense^T, i.e. contract x's last axis + # directly against dense's in_features axis (no explicit transpose). + expected = Nx.dot(x, [2], dense, [1]) + result = EMLX.quantized_matmul(x, qw) + + assert Nx.shape(result) == {1, 4, 64} + + # quantized_matmul must agree with dequantize-then-dense-dot using the + # *same* quantized weight (not the pre-quantization original) — this + # isolates matmul correctness from quantization lossiness. + assert Nx.all_close(result, expected, atol: 1.0e-2, rtol: 1.0e-2) |> Nx.to_number() == 1 + end + end + end + + describe "transparent Nx.dot dispatch per-mode" do + for {mode, group_size, bits} <- @modes do + test "#{mode}: Nx.dot dispatches to quantized_matmul for a quantized right operand" do + mode = unquote(mode) + group_size = unquote(group_size) + bits = unquote(bits) + + x = Nx.iota({1, 4, 128}, type: :f32) |> Nx.divide(100) + x = Nx.backend_transfer(x, {EMLX.Backend, device: :gpu}) + + w = Nx.iota({64, 128}, type: :f32) |> Nx.divide(1000) + w = Nx.backend_transfer(w, {EMLX.Backend, device: :gpu}) + + qw = EMLX.quantize(w, type: {:s, bits}, group_size: group_size, mode: mode) + + result_dot = Nx.dot(x, [2], qw, [1]) + result_explicit = EMLX.quantized_matmul(x, qw) + + assert Nx.shape(result_dot) == {1, 4, 64} + assert Nx.all_close(result_dot, result_explicit, atol: 1.0e-4) |> Nx.to_number() == 1 + end + end + end + + describe "EMLX.Quantization (runtime_call deftransform) per-mode" do + for {mode, group_size, bits} <- @modes do + test "#{mode}: Quantization.quantize/2 + dequantize/1 + quantized_matmul/2 agree with the eager path" do + mode = unquote(mode) + group_size = unquote(group_size) + bits = unquote(bits) + + x = Nx.iota({1, 4, 128}, type: :f32) |> Nx.divide(100) + x = Nx.backend_transfer(x, {EMLX.Backend, device: :gpu}) + + w = Nx.iota({64, 128}, type: :f32) |> Nx.divide(1000) + w = Nx.backend_transfer(w, {EMLX.Backend, device: :gpu}) + + qw_eager = EMLX.quantize(w, type: {:s, bits}, group_size: group_size, mode: mode) + qw_defn = Quantization.quantize(w, type: {:s, bits}, group_size: group_size, mode: mode) + + assert Quantization.quantized?(qw_defn) + + dequant_eager = EMLX.dequantize(qw_eager) + dequant_defn = Quantization.dequantize(qw_defn) + assert Nx.all_close(dequant_eager, dequant_defn, atol: 1.0e-4) |> Nx.to_number() == 1 + + matmul_eager = EMLX.quantized_matmul(x, qw_eager) + matmul_defn = Quantization.quantized_matmul(x, qw_defn) + assert Nx.all_close(matmul_eager, matmul_defn, atol: 1.0e-4) |> Nx.to_number() == 1 + end + end + end +end diff --git a/emlx/test/emlx/native/expr_test.exs b/emlx/test/emlx/native/expr_test.exs new file mode 100644 index 0000000..fa45edc --- /dev/null +++ b/emlx/test/emlx/native/expr_test.exs @@ -0,0 +1,3582 @@ +defmodule EMLX.Native.ExprTest do + @moduledoc """ + Tests for the EMLX native defn compiler: + - EMLX.Native.Expr struct shape (refs as node IDs, atom opcodes, attrs) + - lower/1 (parameter, constant, tensor/capture, add, identity) + - compile_program / eval_program NIFs via to_native/1 (C++ replay) + - Compiler seam: Nx.Defn.compile(..., compiler: EMLX) via single-NIF replay + - Unary + binary + compare/logical equivalence vs EMLX.Backend + """ + use ExUnit.Case, async: false + import Nx.Defn + + alias EMLX.Native.Expr + + # ── module-level defn helpers ───────────────────────────────────────────── + defn add_two(a, b), do: Nx.add(a, b) + defn add_one(x), do: Nx.add(x, 1) + defn identity(x), do: x + + defn gt_f32(a, b), do: Nx.greater(a, b) + defn cmp_mixed(a, b), do: Nx.greater(a, b) + defn mixed_add(a, b), do: Nx.add(a, b) + + defn two_quantized_dots(x, w1, w2) do + Nx.concatenate([Nx.dot(x, w1), Nx.dot(x, w2)], axis: 1) + end + + defn dequant_surrounded(x, qw) do + dense = EMLX.Quantization.dequantize(qw) + Nx.add(x, dense) |> Nx.multiply(2.0) + end + + defn dequant_only(qw), do: EMLX.Quantization.dequantize(qw) + + defn two_runtime_calls(x, qw1, qw2) do + Nx.add(EMLX.Quantization.dequantize(qw1), EMLX.Quantization.dequantize(qw2)) |> Nx.add(x) + end + + defn runtime_call_inside_while(x0, n, qw) do + {result, _, _, _} = + while {acc = x0, i = 0, n, qw}, Nx.less(i, n) do + {Nx.add(acc, EMLX.Quantization.dequantize(qw)), i + 1, n, qw} + end + + result + end + + defn quantized_matmul_surrounded(x, qw) do + EMLX.Quantization.quantized_matmul(x, qw) |> Nx.add(1.0) |> Nx.multiply(2.0) + end + + defn dequant_surrounded_other_site(x, qw) do + dense = EMLX.Quantization.dequantize(qw) + Nx.add(x, dense) |> Nx.multiply(2.0) + end + + defn hook_top_level(a, b) do + hooked = hook(Nx.multiply(Nx.add(a, b), 2), :mid, fn t -> send(self(), {:mid, t}) end) + Nx.subtract(hooked, 1) + end + + defn hook_unused_value(a, b) do + _ = hook(Nx.multiply(a, b), :dbg, fn t -> send(self(), {:dbg, t}) end) + Nx.add(a, b) + end + + defn hook_name_only(a, b) do + _ = hook(Nx.multiply(a, b), :no_fn) + Nx.add(a, b) + end + + defn hook_tuple_payload(a, b) do + hook({Nx.add(a, b), Nx.subtract(a, b)}, :pair, fn {s, d} -> send(self(), {:pair, s, d}) end) + end + + defn hook_in_cond_branch(a, b) do + pred = Nx.any(Nx.greater(a, 0)) + + cond do + pred -> hook(Nx.add(a, b), :branch_true, fn t -> send(self(), {:branch_true, t}) end) + true -> Nx.subtract(a, b) + end + end + + defn hook_in_while_body(a) do + {result, _} = + while {acc = Nx.tensor(0), i = a}, Nx.less(0, i) do + acc = hook(Nx.add(acc, i), :iter, fn t -> send(self(), {:iter, t}) end) + {acc, i - 1} + end + + result + end + + # Exercises the Nx.Defn.Graph.split-chain path (hook before AND after a + # non-bare `while`, i.e. the while has surrounding work on both sides) -- + # this is the shape that surfaced the `Nx.Defn.Graph` `:token` rewrite gap + # (see Results). + defn hook_around_while(a) do + seed = hook(Nx.multiply(a, 2), :seed, fn t -> send(self(), {:seed, t}) end) + + {result, _} = + while {acc = Nx.tensor(0), i = seed}, Nx.less(0, i) do + {Nx.add(acc, i), i - 1} + end + + hook(Nx.add(result, 1), :final, fn t -> send(self(), {:final, t}) end) + end + + defn hook_in_reduce_body(t) do + Nx.reduce(t, Nx.tensor(0), fn x, acc -> + hook(Nx.add(acc, x), :step, fn v -> send(self(), {:step, v}) end) + end) + end + + # A hook genuinely nested inside a `cond` that is itself inside a reduce + # body must still raise -- the always-executes exemption above must not + # blanket-exempt a real cond-branch hook one level deeper. + defn hook_in_cond_in_reduce_body(t) do + Nx.reduce(t, Nx.tensor(0), fn x, acc -> + cond do + Nx.any(Nx.greater(x, 0)) -> + hook(Nx.add(acc, x), :pos, fn v -> send(self(), {:pos, v}) end) + + true -> + acc + end + end) + end + + # ── IR shape ───────────────────────────────────────────────────────────── + + describe "program shape" do + test "node IDs are Erlang refs" do + expr = Nx.Defn.debug_expr_apply(&add_two/2, [Nx.template({}, :f32), Nx.template({}, :f32)]) + prog = Expr.lower(expr) + + assert Enum.all?(prog.inputs, &is_reference/1) + assert Enum.all?(prog.outputs, &is_reference/1) + + for {id, op, operands, attrs} <- prog.instructions do + assert is_reference(id) + assert is_atom(op) + assert Enum.all?(operands, &is_reference/1) + assert is_list(attrs) + end + end + + test "each input ref is unique" do + expr = Nx.Defn.debug_expr_apply(&add_two/2, [Nx.template({}, :f32), Nx.template({}, :f32)]) + prog = Expr.lower(expr) + + assert length(prog.inputs) == length(Enum.uniq(prog.inputs)) + end + + test "operand refs point to known nodes" do + expr = Nx.Defn.debug_expr_apply(&add_two/2, [Nx.template({}, :f32), Nx.template({}, :f32)]) + prog = Expr.lower(expr) + + known = + MapSet.new(prog.inputs) + |> MapSet.union(MapSet.new(prog.captures, fn {r, _} -> r end)) + |> MapSet.union(MapSet.new(prog.constants, fn {r, _, _} -> r end)) + |> MapSet.union(MapSet.new(prog.instructions, fn {r, _, _, _} -> r end)) + + for {_id, _op, operands, _} <- prog.instructions do + for ref <- operands do + assert MapSet.member?(known, ref), + "operand ref #{inspect(ref)} not in known node set" + end + end + end + end + + # ── lower/1 ────────────────────────────────────────────────────────────── + + describe "lower/1" do + test "identity: one input, no instructions, output = input ref" do + expr = Nx.Defn.debug_expr_apply(&identity/1, [Nx.template({3}, :f32)]) + prog = Expr.lower(expr) + + assert length(prog.inputs) == 1 + assert prog.captures == [] + assert prog.constants == [] + assert prog.instructions == [] + assert prog.outputs == prog.inputs + end + + test "add two parameters: one :add instruction, operands are the input refs" do + expr = Nx.Defn.debug_expr_apply(&add_two/2, [Nx.template({}, :f32), Nx.template({}, :f32)]) + prog = Expr.lower(expr) + + assert length(prog.inputs) == 2 + assert prog.captures == [] + assert prog.constants == [] + assert [{result_ref, :add, [left_ref, right_ref], []}] = prog.instructions + assert left_ref == Enum.at(prog.inputs, 0) + assert right_ref == Enum.at(prog.inputs, 1) + assert prog.outputs == [result_ref] + end + + test "add parameter + scalar literal: one const entry, one :add instruction" do + expr = Nx.Defn.debug_expr_apply(&add_one/1, [Nx.template({}, :f32)]) + prog = Expr.lower(expr) + + assert length(prog.inputs) == 1 + assert prog.captures == [] + assert [{const_ref, 1, _int_type}] = prog.constants + # add_one: f32 input + integer constant. Lowerer may emit an :astype to promote + # the integer constant to f32 before :add. Assert the :add instruction is present + # and references both the input and the constant (possibly via a cast ref). + add_instr = Enum.find(prog.instructions, fn {_, op, _, _} -> op == :add end) + assert {result_ref, :add, [left_ref, right_ref], []} = add_instr + # Either a direct ref or a cast of the const ref must be in the add operands. + all_refs = MapSet.new(prog.inputs ++ [const_ref]) + + cast_or_direct = fn r -> + r in prog.inputs or r == const_ref or + Enum.any?(prog.instructions, fn {id, :astype, [src], _} -> + id == r and (src in prog.inputs or src == const_ref) + end) + end + + assert cast_or_direct.(left_ref) or cast_or_direct.(right_ref) + # suppress unused warning + _ = all_refs + assert prog.outputs == [result_ref] + end + + test "closed-over backend tensor becomes a capture" do + weight_tensor = Nx.tensor(1.0, backend: EMLX.Backend) + + # Construct the :tensor Expr node manually — Nx.Defn.Expr.tensor/1 rejects + # EMLX backend tensors (intentional guard in Nx to prevent accidental capture). + weight_expr = %Nx.Tensor{ + data: %Nx.Defn.Expr{id: make_ref(), op: :tensor, args: [weight_tensor], context: nil}, + type: weight_tensor.type, + shape: weight_tensor.shape, + names: weight_tensor.names + } + + param_a = Nx.Defn.Expr.parameter(:root, {:f, 32}, {}, 0) + output = Nx.add(param_a, weight_expr) + prog = Expr.lower(output) + + assert length(prog.inputs) == 1 + assert [{capture_ref, ^weight_tensor}] = prog.captures + assert prog.constants == [] + assert [{_result, :add, [_input_ref, ^capture_ref], []}] = prog.instructions + end + + test "closed-over tensor executes correctly (capture, not just lowering shape)" do + # Same manually-built capture Expr as the lowering-shape test above -- + # `Nx.Defn.jit`/`check_equiv` can't exercise this naturally: Nx itself + # guards against a real `defn` closure embedding an EMLX.Backend tensor + # (see the comment above), so this can only be reached by hand-building + # the Expr tree. Runs it through the real to_native/NIF path (not just + # inspecting `prog`) to prove the capture actually evaluates correctly. + weight_tensor = Nx.tensor(10.0, backend: EMLX.Backend) + + weight_expr = %Nx.Tensor{ + data: %Nx.Defn.Expr{id: make_ref(), op: :tensor, args: [weight_tensor], context: nil}, + type: weight_tensor.type, + shape: weight_tensor.shape, + names: weight_tensor.names + } + + param_a = Nx.Defn.Expr.parameter(:root, {:f, 32}, {}, 0) + prog = Expr.lower(Nx.add(param_a, weight_expr)) + + x = Nx.tensor(5.0, backend: EMLX.Backend) + [out] = run_nif(prog, [x]) + + assert_in_delta Nx.to_number(out), 15.0, 1.0e-6 + end + + test "unknown op raises ArgumentError with 'does not yet lower op'" do + expr = %Nx.Tensor{ + data: %Nx.Defn.Expr{id: make_ref(), op: :this_op_does_not_exist, args: [], context: nil}, + type: {:f, 32}, + shape: {}, + names: [] + } + + assert_raise ArgumentError, ~r/does not yet lower op/, fn -> Expr.lower(expr) end + end + end + + # ── to_native/1 ────────────────────────────────────────────────────────── + + describe "to_native/1" do + test "identity program: n_inputs=1, no instructions, output encodes input ref" do + expr = Nx.Defn.debug_expr_apply(&identity/1, [Nx.template({3}, :f32)]) + prog = Expr.lower(expr) + program = Expr.to_native(prog) + + assert %EMLX.Native.Program{} = program + assert program.num_inputs == 1 + assert program.captures == [] + assert program.constants == [] + assert program.instructions == [] + assert program.outputs == [{:input, 0}] + end + + test "add program: op_name is :add, operands encode the two inputs" do + expr = Nx.Defn.debug_expr_apply(&add_two/2, [Nx.template({}, :f32), Nx.template({}, :f32)]) + prog = Expr.lower(expr) + %EMLX.Native.Program{instructions: [instr], outputs: [output]} = Expr.to_native(prog) + + assert instr.op == :add + assert instr.operands == [{:input, 0}, {:input, 1}] + assert output == {:result, 0} + end + end + + # ── compile_program / eval_program NIFs ────────────────────────────────── + + describe "compile_program / eval_program NIFs" do + setup do + device = EMLX.default_device() + {worker, _} = EMLX.resolve_worker(device) + %{worker: worker, device: device} + end + + test "identity program via to_native: output equals input", %{worker: worker, device: device} do + expr = Nx.Defn.debug_expr_apply(&identity/1, [Nx.template({3}, :f32)]) + prog = Expr.lower(expr) + wire = Expr.to_native(prog) + + prog_ref = compile_nif!(worker, wire) + + x = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + %EMLX.Backend{ref: {_, input_ref}} = x.data + + [out_ref] = eval_nif!(worker, prog_ref, [input_ref]) + out = EMLX.Backend.to_nx({device, out_ref}, x) + + assert_all_close(out, x) + end + + test "add program via to_native: correct result", %{worker: worker, device: device} do + expr = Nx.Defn.debug_expr_apply(&add_two/2, [Nx.template({}, :f32), Nx.template({}, :f32)]) + prog = Expr.lower(expr) + wire = Expr.to_native(prog) + + prog_ref = compile_nif!(worker, wire) + + a = Nx.tensor(3.0, backend: EMLX.Backend) + b = Nx.tensor(4.0, backend: EMLX.Backend) + %EMLX.Backend{ref: {_, ref_a}} = a.data + %EMLX.Backend{ref: {_, ref_b}} = b.data + + [out_ref] = eval_nif!(worker, prog_ref, [ref_a, ref_b]) + out = EMLX.Backend.to_nx({device, out_ref}, a) + + assert_in_delta Nx.to_number(out), 7.0, 1.0e-6 + end + end + + # ── compiler seam (end-to-end) ──────────────────────────────────────────── + + describe "compiler seam E2E" do + test "add_one defn via EMLX compiler returns correct result" do + compiled = Nx.Defn.compile(&add_one/1, [Nx.template({3}, :f32)], compiler: EMLX) + x = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + result = compiled.(x) + + assert_all_close(result, Nx.tensor([2.0, 3.0, 4.0])) + end + + test "identity defn routes through native path (zero-instruction program)" do + compiled = Nx.Defn.compile(&identity/1, [Nx.template({3}, :f32)], compiler: EMLX) + x = Nx.tensor([7.0, 8.0, 9.0], backend: EMLX.Backend) + + assert_all_close(compiled.(x), x) + end + + test "add_two defn with two parameters" do + compiled = + Nx.Defn.compile(&add_two/2, [Nx.template({3}, :f32), Nx.template({3}, :f32)], + compiler: EMLX + ) + + a = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + b = Nx.tensor([4.0, 5.0, 6.0], backend: EMLX.Backend) + + assert_all_close(compiled.(a, b), Nx.tensor([5.0, 7.0, 9.0])) + end + + test "unsupported op raises through the compiler seam (no Evaluator fallback)" do + f = fn a -> Nx.population_count(a) end + + a = Nx.tensor([1, 2, 3], type: :s32, backend: EMLX.Backend) + + assert_raise ArgumentError, ~r/population_count is not supported by EMLX/, fn -> + Nx.Defn.jit(f, compiler: EMLX).(a) + end + end + + test "result matches eager EMLX.Backend within tolerance" do + compiled = + Nx.Defn.compile(&add_two/2, [Nx.template({3}, :f32), Nx.template({3}, :f32)], + compiler: EMLX + ) + + a = Nx.tensor([1.5, 2.5, 3.5], backend: EMLX.Backend) + b = Nx.tensor([0.5, 1.5, 2.5], backend: EMLX.Backend) + + eager_result = EMLX.add(EMLX.Backend.from_nx(a), EMLX.Backend.from_nx(b)) + assert_all_close(compiled.(a, b), EMLX.Backend.to_nx(eager_result, a)) + end + end + + # ── perf gate ───────────────────────────────────────────────────────────── + + defn chain_10(x, y) do + x + |> Nx.add(y) + |> Nx.add(y) + |> Nx.add(y) + |> Nx.add(y) + |> Nx.add(y) + |> Nx.add(y) + |> Nx.add(y) + |> Nx.add(y) + |> Nx.add(y) + |> Nx.add(y) + end + + # Helper: compile + eval the defn via the native path, compare to eager backend. + defp check_equiv(fun, inputs_eager, opts \\ []) do + tol = Keyword.get(opts, :tol, 1.0e-4) + templates = Enum.map(inputs_eager, &Nx.template(&1.shape, &1.type)) + compiled = Nx.Defn.compile(fun, templates, compiler: EMLX) + result = apply(compiled, inputs_eager) + eager = apply(Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator), inputs_eager) + assert_close(result, eager, tol) + end + + defp check_reduce_equiv(fun, inputs_eager, opts \\ []) do + tol = Keyword.get(opts, :tol, 1.0e-4) + templates = Enum.map(inputs_eager, &Nx.template(&1.shape, &1.type)) + compiled = Nx.Defn.compile(fun, templates, compiler: EMLX) + result = apply(compiled, inputs_eager) + + bin_inputs = Enum.map(inputs_eager, &Nx.backend_copy(&1, Nx.BinaryBackend)) + prev = Nx.default_backend(Nx.BinaryBackend) + + eager = + try do + apply(Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator), bin_inputs) + after + Nx.default_backend(prev) + end + + assert result.shape == eager.shape + assert result.type == eager.type + assert_close(result, eager, tol) + end + + defp assert_close(a, b, tol) do + # For complex tensors, compare real and imaginary parts separately. + case a.type do + {:c, _} -> + assert_close(Nx.real(a), Nx.real(b), tol) + assert_close(Nx.imag(a), Nx.imag(b), tol) + + _ -> + a_vals = Nx.to_flat_list(a) + b_vals = Nx.to_flat_list(b) + + Enum.zip(a_vals, b_vals) + |> Enum.each(fn {av, bv} -> + case {av, bv} do + {:nan, :nan} -> + :ok + + {:infinity, :infinity} -> + :ok + + {:neg_infinity, :neg_infinity} -> + :ok + + {a_num, b_num} when is_number(a_num) and is_number(b_num) -> + assert_in_delta(a_num * 1.0, b_num * 1.0, tol) + + _ -> + flunk("Values differ: #{inspect(av)} vs #{inspect(bv)}") + end + end) + end + end + + describe "unary elementwise" do + # Sample unary ops over f32 and bf16 using representative positive values + # to avoid NaN from log/sqrt/etc. + test "abs/ceil/floor/negate/round/sign — f32" do + x = Nx.tensor([1.7, -2.3, 0.0, -0.5], backend: EMLX.Backend) + + for fun <- [&Nx.abs/1, &Nx.ceil/1, &Nx.floor/1, &Nx.negate/1, &Nx.round/1, &Nx.sign/1] do + check_equiv(fun, [x]) + end + end + + test "exp/log/sqrt/tanh/sigmoid — f32" do + x = Nx.tensor([0.5, 1.0, 2.0, 4.0], backend: EMLX.Backend) + + for fun <- [&Nx.exp/1, &Nx.log/1, &Nx.sqrt/1, &Nx.tanh/1, &Nx.sigmoid/1] do + check_equiv(fun, [x]) + end + end + + test "sin/cos/tan/asin/acos/atan — f32" do + x = Nx.tensor([0.1, 0.5, 0.9, -0.5], backend: EMLX.Backend) + + for fun <- [&Nx.sin/1, &Nx.cos/1, &Nx.tan/1, &Nx.asin/1, &Nx.acos/1, &Nx.atan/1] do + check_equiv(fun, [x]) + end + end + + test "sinh/cosh/tanh/asinh/acosh/atanh — f32" do + x = Nx.tensor([0.1, 0.5, 1.0, 1.5], backend: EMLX.Backend) + + for fun <- [&Nx.sinh/1, &Nx.cosh/1, &Nx.tanh/1, &Nx.asinh/1, &Nx.acosh/1] do + check_equiv(fun, [x]) + end + + x2 = Nx.tensor([0.1, 0.5, 0.9, -0.5], backend: EMLX.Backend) + check_equiv(&Nx.atanh/1, [x2]) + end + + test "erf/erf_inv/erfc/rsqrt/expm1/log1p — f32" do + x = Nx.tensor([0.1, 0.5, 1.0, 2.0], backend: EMLX.Backend) + + for fun <- [&Nx.erf/1, &Nx.erf_inv/1, &Nx.erfc/1, &Nx.rsqrt/1, &Nx.expm1/1, &Nx.log1p/1] do + check_equiv(fun, [x]) + end + end + + test "cbrt — f32 positive values" do + x = Nx.tensor([1.0, 8.0, 27.0, 0.125], backend: EMLX.Backend) + check_equiv(&Nx.cbrt/1, [x], tol: 1.0e-3) + end + + test "is_nan/is_infinity — f32" do + x = Nx.tensor([1.0, :nan, :infinity, :neg_infinity], type: :f32, backend: EMLX.Backend) + check_equiv(&Nx.is_nan/1, [x]) + check_equiv(&Nx.is_infinity/1, [x]) + end + + test "bitwise_not — s32" do + x = Nx.tensor([0, 1, -1, 255], type: :s32, backend: EMLX.Backend) + check_equiv(&Nx.bitwise_not/1, [x]) + end + + test "logical_not — u8" do + x = Nx.tensor([0, 1, 1, 0], type: :u8, backend: EMLX.Backend) + check_equiv(&Nx.logical_not/1, [x]) + end + + test "real/imag/conjugate — c64" do + # Complex tensor: values are reals with zero imaginary parts. + c = Nx.tensor([1.5, 2.5, -1.0], type: {:c, 64}, backend: EMLX.Backend) + check_equiv(&Nx.real/1, [c]) + check_equiv(&Nx.imag/1, [c]) + check_equiv(&Nx.conjugate/1, [c]) + end + + test "unary ops — bf16" do + x = Nx.tensor([0.5, 1.0, 2.0], type: :bf16, backend: EMLX.Backend) + + for fun <- [&Nx.abs/1, &Nx.exp/1, &Nx.tanh/1, &Nx.sqrt/1] do + check_equiv(fun, [x], tol: 1.0e-2) + end + end + end + + describe "binary arithmetic + bitwise" do + test "add/subtract/multiply — f32" do + a = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + b = Nx.tensor([4.0, 5.0, 6.0], backend: EMLX.Backend) + + for fun <- [&Nx.add/2, &Nx.subtract/2, &Nx.multiply/2] do + check_equiv(fun, [a, b]) + end + end + + test "divide/pow/atan2 — f32" do + a = Nx.tensor([4.0, 8.0, 1.0], backend: EMLX.Backend) + b = Nx.tensor([2.0, 4.0, 3.0], backend: EMLX.Backend) + + for fun <- [&Nx.divide/2, &Nx.pow/2, &Nx.atan2/2] do + check_equiv(fun, [a, b]) + end + end + + test "min/max — f32" do + a = Nx.tensor([1.0, 5.0, 3.0], backend: EMLX.Backend) + b = Nx.tensor([4.0, 2.0, 3.0], backend: EMLX.Backend) + check_equiv(&Nx.min/2, [a, b]) + check_equiv(&Nx.max/2, [a, b]) + end + + test "quotient/remainder — s32 positive" do + a = Nx.tensor([7, 9, 15], type: :s32, backend: EMLX.Backend) + b = Nx.tensor([3, 4, 7], type: :s32, backend: EMLX.Backend) + check_equiv(&Nx.quotient/2, [a, b]) + check_equiv(&Nx.remainder/2, [a, b]) + end + + test "remainder — s32 negative dividend" do + a = Nx.tensor([-7, -9, 7], type: :s32, backend: EMLX.Backend) + b = Nx.tensor([3, 4, -3], type: :s32, backend: EMLX.Backend) + check_equiv(&Nx.remainder/2, [a, b]) + end + + test "bitwise_and/or/xor/left_shift/right_shift — s32" do + a = Nx.tensor([0b1010, 0xFF, 5], type: :s32, backend: EMLX.Backend) + b = Nx.tensor([0b1100, 0x0F, 2], type: :s32, backend: EMLX.Backend) + + for fun <- [ + &Nx.bitwise_and/2, + &Nx.bitwise_or/2, + &Nx.bitwise_xor/2, + &Nx.left_shift/2, + &Nx.right_shift/2 + ] do + check_equiv(fun, [a, b]) + end + end + + test "mixed dtypes: add(s32, f32) → f32 with implicit upcast" do + a = Nx.tensor([1, 2, 3], type: :s32, backend: EMLX.Backend) + b = Nx.tensor([0.5, 1.5, 2.5], type: :f32, backend: EMLX.Backend) + check_equiv(&mixed_add/2, [a, b]) + end + + test "binary ops — bf16" do + a = Nx.tensor([1.0, 2.0, 4.0], type: :bf16, backend: EMLX.Backend) + b = Nx.tensor([2.0, 1.0, 2.0], type: :bf16, backend: EMLX.Backend) + + for fun <- [&Nx.add/2, &Nx.subtract/2, &Nx.multiply/2] do + check_equiv(fun, [a, b], tol: 1.0e-2) + end + end + end + + describe "compare and logical" do + test "equal/not_equal/greater/less/greater_equal/less_equal — f32" do + a = Nx.tensor([1.0, 2.0, 2.0, 3.0], backend: EMLX.Backend) + b = Nx.tensor([2.0, 2.0, 1.0, 1.0], backend: EMLX.Backend) + + for fun <- [ + &Nx.equal/2, + &Nx.not_equal/2, + &Nx.greater/2, + &Nx.less/2, + &Nx.greater_equal/2, + &Nx.less_equal/2 + ] do + check_equiv(fun, [a, b]) + end + end + + test "compare — s32" do + a = Nx.tensor([-1, 0, 1, 2], type: :s32, backend: EMLX.Backend) + b = Nx.tensor([0, 0, 0, 1], type: :s32, backend: EMLX.Backend) + + for fun <- [&Nx.equal/2, &Nx.less/2, &Nx.greater/2] do + check_equiv(fun, [a, b]) + end + end + + test "logical_and/or/xor — u8" do + a = Nx.tensor([0, 0, 1, 1], type: :u8, backend: EMLX.Backend) + b = Nx.tensor([0, 1, 0, 1], type: :u8, backend: EMLX.Backend) + + for fun <- [&Nx.logical_and/2, &Nx.logical_or/2, &Nx.logical_xor/2] do + check_equiv(fun, [a, b]) + end + end + + test "compare output dtype is u8 (bool)" do + a = Nx.tensor([1.0, 3.0], backend: EMLX.Backend) + b = Nx.tensor([2.0, 1.0], backend: EMLX.Backend) + + compiled = + Nx.Defn.compile(>_f32/2, [Nx.template({2}, :f32), Nx.template({2}, :f32)], + compiler: EMLX + ) + + result = compiled.(a, b) + assert result.type == {:u, 8} + end + + test "compare with mixed dtypes s32/f32 — output u8" do + a = Nx.tensor([1, 3], type: :s32, backend: EMLX.Backend) + b = Nx.tensor([2.0, 1.0], type: :f32, backend: EMLX.Backend) + check_equiv(&cmp_mixed/2, [a, b]) + end + end + + describe "reshape" do + test "1D → 2D, 2D → 1D" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], backend: EMLX.Backend) + check_equiv(fn t -> Nx.reshape(t, {2, 3}) end, [x]) + check_equiv(fn t -> Nx.reshape(t, {3, 2}) end, [x]) + end + + test "2D → 1D (flatten)" do + x = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.reshape(t, {4}) end, [x]) + end + + test "rank-changing — s32" do + x = Nx.tensor([[[1, 2], [3, 4]]], type: :s32, backend: EMLX.Backend) + check_equiv(fn t -> Nx.reshape(t, {2, 2}) end, [x]) + check_equiv(fn t -> Nx.reshape(t, {4}) end, [x]) + end + end + + describe "squeeze" do + test "remove singleton dimension" do + x = Nx.tensor([[1.0, 2.0, 3.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.squeeze(t, axes: [0]) end, [x]) + end + + test "multiple axes, negative axis" do + x = Nx.tensor([[[1.0], [2.0], [3.0]]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.squeeze(t, axes: [0, 2]) end, [x]) + check_equiv(fn t -> Nx.squeeze(t, axes: [-1]) end, [x]) + end + end + + describe "transpose" do + test "2D matrix transpose" do + x = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.transpose(t) end, [x]) + check_equiv(fn t -> Nx.transpose(t, axes: [1, 0]) end, [x]) + end + + test "3D permutation, negative perm" do + x = Nx.tensor([[[1.0, 2.0], [3.0, 4.0]]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.transpose(t, axes: [2, 0, 1]) end, [x]) + check_equiv(fn t -> Nx.transpose(t, axes: [-1, -3, -2]) end, [x]) + end + end + + describe "as_type" do + test "f32 → s32 and back" do + x = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + check_equiv(fn t -> Nx.as_type(t, {:s, 32}) end, [x]) + xi = Nx.tensor([1, 2, 3], type: :s32, backend: EMLX.Backend) + check_equiv(fn t -> Nx.as_type(t, {:f, 32}) end, [xi]) + end + + test "f32 → bf16 and back" do + x = Nx.tensor([0.5, 1.5, 2.5], backend: EMLX.Backend) + check_equiv(fn t -> Nx.as_type(t, {:bf, 16}) end, [x], tol: 1.0e-2) + xb = Nx.tensor([0.5, 1.5, 2.5], type: {:bf, 16}, backend: EMLX.Backend) + check_equiv(fn t -> Nx.as_type(t, {:f, 32}) end, [xb]) + end + end + + describe "bitcast" do + test "u8 → s8 same bit pattern" do + # Values chosen so the bit pattern is unambiguous for both u8 and s8. + x = Nx.tensor([1, 2, 3, 127], type: :u8, backend: EMLX.Backend) + check_equiv(fn t -> Nx.bitcast(t, {:s, 8}) end, [x]) + end + + test "f32 → u32 round-trip" do + x = Nx.tensor([1.0, 2.0, 0.0], type: :f32, backend: EMLX.Backend) + check_equiv(fn t -> Nx.bitcast(t, {:u, 32}) end, [x]) + end + end + + describe "broadcast" do + test "scalar → 1D" do + x = Nx.tensor(1.0, backend: EMLX.Backend) + check_equiv(fn t -> Nx.broadcast(t, {4}) end, [x]) + end + + test "1D → 2D row broadcast" do + x = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + check_equiv(fn t -> Nx.broadcast(t, {2, 3}) end, [x]) + end + + test "1D column broadcast (axes: [0])" do + x = Nx.tensor([1.0, 2.0], backend: EMLX.Backend) + check_equiv(fn t -> Nx.broadcast(t, {2, 3}, axes: [0]) end, [x]) + end + + test "2D → 3D, broadcast_in_dim style" do + x = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.broadcast(t, {3, 2, 2}) end, [x]) + end + end + + describe "pad" do + test "zero-padding on 1D" do + x = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + check_equiv(fn t -> Nx.pad(t, 0.0, [{1, 1, 0}]) end, [x]) + end + + test "zero-padding on 2D, asymmetric" do + x = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.pad(t, 0.0, [{0, 1, 0}, {1, 0, 0}]) end, [x]) + end + + test "scalar pad value" do + x = Nx.tensor([1.0, 2.0], backend: EMLX.Backend) + check_equiv(fn t -> Nx.pad(t, -1.0, [{2, 2, 0}]) end, [x]) + end + + test "interior padding on 1D" do + x = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + check_equiv(fn t -> Nx.pad(t, 0.0, [{0, 0, 2}]) end, [x]) + end + + test "interior padding on 2D, both axes, non-scalar-friendly pad value" do + x = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.pad(t, -9.0, [{0, 0, 1}, {0, 0, 2}]) end, [x]) + end + + test "negative lo/hi crops on 2D" do + x = Nx.iota({4, 5}, type: :f32, backend: EMLX.Backend) + check_equiv(fn t -> Nx.pad(t, 0.0, [{-1, -1, 0}, {0, -2, 0}]) end, [x]) + end + + test "mixed positive/negative/interior padding on 2D" do + x = Nx.iota({3, 4}, type: :f32, backend: EMLX.Backend) + check_equiv(fn t -> Nx.pad(t, 0.0, [{2, -1, 1}, {-1, 3, 2}]) end, [x]) + end + + test "interior padding on 3D" do + x = Nx.iota({2, 3, 4}, type: :f32, backend: EMLX.Backend) + check_equiv(fn t -> Nx.pad(t, 0.0, [{1, 1, 1}, {-1, 0, 0}, {0, 2, 2}]) end, [x]) + end + end + + describe "reverse" do + test "1D reverse" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0], backend: EMLX.Backend) + check_equiv(fn t -> Nx.reverse(t) end, [x]) + end + + test "2D reverse single axis, both axes" do + x = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.reverse(t, axes: [0]) end, [x]) + check_equiv(fn t -> Nx.reverse(t, axes: [1]) end, [x]) + check_equiv(fn t -> Nx.reverse(t, axes: [0, 1]) end, [x]) + end + + test "negative axis" do + x = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.reverse(t, axes: [-1]) end, [x]) + end + end + + describe "concatenate" do + test "concat 1D along axis 0" do + a = Nx.tensor([1.0, 2.0], backend: EMLX.Backend) + b = Nx.tensor([3.0, 4.0, 5.0], backend: EMLX.Backend) + check_equiv(fn x, y -> Nx.concatenate([x, y], axis: 0) end, [a, b]) + end + + test "concat 2D along axis 0 and axis 1" do + a = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + b = Nx.tensor([[5.0, 6.0]], backend: EMLX.Backend) + c = Nx.tensor([[7.0, 8.0], [9.0, 10.0]], backend: EMLX.Backend) + check_equiv(fn x, y -> Nx.concatenate([x, y], axis: 0) end, [a, b]) + check_equiv(fn x, y -> Nx.concatenate([x, y], axis: 1) end, [a, c]) + end + + test "three tensors" do + a = Nx.tensor([1.0], backend: EMLX.Backend) + b = Nx.tensor([2.0, 3.0], backend: EMLX.Backend) + c = Nx.tensor([4.0], backend: EMLX.Backend) + check_equiv(fn x, y, z -> Nx.concatenate([x, y, z]) end, [a, b, c]) + end + end + + describe "stack" do + test "stack 1D tensors → 2D along axis 0" do + a = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + b = Nx.tensor([4.0, 5.0, 6.0], backend: EMLX.Backend) + check_equiv(fn x, y -> Nx.stack([x, y]) end, [a, b]) + end + + test "stack along axis 1" do + a = Nx.tensor([1.0, 2.0], backend: EMLX.Backend) + b = Nx.tensor([3.0, 4.0], backend: EMLX.Backend) + check_equiv(fn x, y -> Nx.stack([x, y], axis: 1) end, [a, b]) + end + + test "negative axis" do + a = Nx.tensor([1.0, 2.0], backend: EMLX.Backend) + b = Nx.tensor([3.0, 4.0], backend: EMLX.Backend) + check_equiv(fn x, y -> Nx.stack([x, y], axis: -1) end, [a, b]) + end + end + + describe "squeeze without explicit axes" do + test "squeeze all singleton dims" do + x = Nx.tensor([[[1.0]]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.squeeze(t) end, [x]) + end + end + + describe "sum" do + test "sum all axes f32" do + x = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.sum(t) end, [x]) + end + + test "sum along axis 0 with keep_axes" do + x = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.sum(t, axes: [0], keep_axes: true) end, [x]) + end + + test "sum along axis 1 f32" do + x = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.sum(t, axes: [1]) end, [x]) + end + + test "sum 3D along multiple axes" do + x = Nx.tensor([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.sum(t, axes: [0, 2]) end, [x]) + end + end + + describe "product" do + test "product all axes f32" do + x = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.product(t) end, [x]) + end + + test "product along axis 0" do + x = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.product(t, axes: [0]) end, [x]) + end + end + + describe "all / any" do + test "all on boolean-like tensor" do + x = Nx.tensor([[1, 1], [1, 0]], type: :u8, backend: EMLX.Backend) + check_equiv(fn t -> Nx.all(t) end, [x]) + end + + test "all along axis 0" do + x = Nx.tensor([[1, 0], [1, 1]], type: :u8, backend: EMLX.Backend) + check_equiv(fn t -> Nx.all(t, axes: [0]) end, [x]) + end + + test "any all axes" do + x = Nx.tensor([[0, 0], [0, 1]], type: :u8, backend: EMLX.Backend) + check_equiv(fn t -> Nx.any(t) end, [x]) + end + + test "any along axis 1 with keep_axes" do + x = Nx.tensor([[0, 1], [0, 0]], type: :u8, backend: EMLX.Backend) + check_equiv(fn t -> Nx.any(t, axes: [1], keep_axes: true) end, [x]) + end + end + + describe "custom-fun reduce (static unroll)" do + test "1d sum reducer" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0], backend: EMLX.Backend) + check_reduce_equiv(fn t -> Nx.reduce(t, 0.0, fn a, b -> Nx.add(a, b) end) end, [x]) + end + + test "1d product reducer" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0], backend: EMLX.Backend) + check_reduce_equiv(fn t -> Nx.reduce(t, 1.0, fn a, b -> Nx.multiply(a, b) end) end, [x]) + end + + test "non-commutative affine reducer (validates fold order)" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0], backend: EMLX.Backend) + + check_reduce_equiv( + fn t -> Nx.reduce(t, 0.0, fn a, b -> Nx.add(Nx.multiply(a, 2.0), b) end) end, + [x] + ) + end + + test "2d reduce along single axis" do + x = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], backend: EMLX.Backend) + + check_reduce_equiv( + fn t -> Nx.reduce(t, 0.0, [axes: [1]], fn a, b -> Nx.add(a, b) end) end, + [x] + ) + end + + test "2d reduce along single axis keep_axes" do + x = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], backend: EMLX.Backend) + + check_reduce_equiv( + fn t -> Nx.reduce(t, 0.0, [axes: [1], keep_axes: true], fn a, b -> Nx.add(a, b) end) end, + [x] + ) + end + + test "3d reduce over multiple axes" do + x = Nx.iota({2, 3, 4}, type: :f32, backend: EMLX.Backend) + + check_reduce_equiv( + fn t -> Nx.reduce(t, 0.0, [axes: [0, 2]], fn a, b -> Nx.add(a, b) end) end, + [x] + ) + end + + test "integer max reducer" do + x = Nx.tensor([3, 1, 4, 1, 5, 9, 2, 6], backend: EMLX.Backend) + check_reduce_equiv(fn t -> Nx.reduce(t, 0, fn a, b -> Nx.max(a, b) end) end, [x]) + end + + test "runtime accumulator input" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0], backend: EMLX.Backend) + acc = Nx.tensor(10.0, backend: EMLX.Backend) + check_reduce_equiv(fn t, a -> Nx.reduce(t, a, fn x, y -> Nx.add(x, y) end) end, [x, acc]) + end + + test "dtype-changing reduce (s32 input, f32 acc → f32)" do + x = Nx.tensor([1, 2, 3, 4], type: :s32, backend: EMLX.Backend) + check_reduce_equiv(fn t -> Nx.reduce(t, 0.0, fn a, b -> Nx.add(a, b) end) end, [x]) + end + end + + describe "custom-fun window_reduce (static unroll)" do + test "1d window sum reducer (valid, default strides)" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0], backend: EMLX.Backend) + + check_reduce_equiv(fn t -> Nx.window_reduce(t, 0.0, {2}, fn a, b -> Nx.add(a, b) end) end, [ + x + ]) + end + + test "1d window max reducer with :same padding" do + x = Nx.tensor([1.0, 5.0, 2.0, 4.0, 3.0], backend: EMLX.Backend) + + check_reduce_equiv( + fn t -> + Nx.window_reduce(t, -100.0, {3}, [padding: :same], fn a, b -> Nx.max(a, b) end) + end, + [x] + ) + end + + test "2d window sum with strides" do + x = Nx.iota({4, 4}, type: :f32, backend: EMLX.Backend) + + check_reduce_equiv( + fn t -> + Nx.window_reduce(t, 0.0, {2, 2}, [strides: [2, 2]], fn a, b -> Nx.add(a, b) end) + end, + [x] + ) + end + + test "1d window with dilations" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0], backend: EMLX.Backend) + + check_reduce_equiv( + fn t -> + Nx.window_reduce(t, 0.0, {2}, [window_dilations: [2]], fn a, b -> Nx.add(a, b) end) + end, + [x] + ) + end + + test "non-commutative affine reducer (validates window fold order)" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0], backend: EMLX.Backend) + + check_reduce_equiv( + fn t -> Nx.window_reduce(t, 0.0, {3}, fn a, b -> Nx.add(Nx.multiply(a, 2.0), b) end) end, + [x] + ) + end + + test "explicit padding config (asymmetric lo/hi)" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0], backend: EMLX.Backend) + + check_reduce_equiv( + fn t -> + Nx.window_reduce(t, 0.0, {2}, [padding: [{1, 0}]], fn a, b -> Nx.add(a, b) end) + end, + [x] + ) + end + + test "integer window_reduce (s32 in/out)" do + # window_reduce output dtype tracks the input tensor (not the acc), so the + # dtype coverage here is the integer path; acc casts to the s32 out type. + x = Nx.tensor([1, 2, 3, 4, 5], type: :s32, backend: EMLX.Backend) + + check_reduce_equiv(fn t -> Nx.window_reduce(t, 0, {2}, fn a, b -> Nx.add(a, b) end) end, [x]) + end + + test "runtime accumulator input" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0], backend: EMLX.Backend) + acc = Nx.tensor(0.5, backend: EMLX.Backend) + + check_reduce_equiv( + fn t, a -> Nx.window_reduce(t, a, {2}, fn x, y -> Nx.add(x, y) end) end, + [x, acc] + ) + end + end + + describe "reduce_max / reduce_min" do + test "reduce_max all axes f32" do + x = Nx.tensor([[1.0, 5.0], [3.0, 2.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.reduce_max(t) end, [x]) + end + + test "reduce_max along axis 1 keep_axes" do + x = Nx.tensor([[1.0, 5.0], [3.0, 2.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.reduce_max(t, axes: [1], keep_axes: true) end, [x]) + end + + test "reduce_min all axes" do + x = Nx.tensor([[4.0, 2.0], [1.0, 3.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.reduce_min(t) end, [x]) + end + + test "reduce_min along axis 0" do + x = Nx.tensor([[4.0, 2.0], [1.0, 3.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.reduce_min(t, axes: [0]) end, [x]) + end + end + + describe "argmax / argmin" do + test "argmax along axis 1" do + x = Nx.tensor([[1.0, 3.0, 2.0], [4.0, 0.0, 5.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.argmax(t, axis: 1) end, [x]) + end + + test "argmax along axis 0 keep_axis" do + x = Nx.tensor([[1.0, 3.0], [4.0, 2.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.argmax(t, axis: 0, keep_axis: true) end, [x]) + end + + test "argmin along axis 0" do + x = Nx.tensor([[4.0, 2.0], [1.0, 5.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.argmin(t, axis: 0) end, [x]) + end + + test "argmax global (no axis)" do + x = Nx.tensor([[1.0, 7.0], [3.0, 2.0]], backend: EMLX.Backend) + check_equiv(fn t -> Nx.argmax(t) end, [x]) + end + end + + describe "dot (non-batched)" do + test "matmul {2,3} × {3,4} → {2,4}" do + a = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], backend: EMLX.Backend) + + b = + Nx.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]], + backend: EMLX.Backend + ) + + check_equiv(fn x, y -> Nx.dot(x, y) end, [a, b]) + end + + test "inner product {4} · {4} → scalar" do + a = Nx.tensor([1.0, 2.0, 3.0, 4.0], backend: EMLX.Backend) + b = Nx.tensor([5.0, 6.0, 7.0, 8.0], backend: EMLX.Backend) + check_equiv(fn x, y -> Nx.dot(x, y) end, [a, b]) + end + + test "tensordot explicit contraction axes" do + a = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + b = Nx.tensor([[1.0, 0.0], [0.0, 1.0]], backend: EMLX.Backend) + check_equiv(fn x, y -> Nx.dot(x, [1], [], y, [0], []) end, [a, b]) + end + end + + describe "dot (batched)" do + test "batched matmul {2,3,4} · {2,4,5} → {2,3,5}" do + a = Nx.iota({2, 3, 4}, type: :f32, backend: EMLX.Backend) + b = Nx.iota({2, 4, 5}, type: :f32, backend: EMLX.Backend) + check_equiv(fn x, y -> Nx.dot(x, [2], [0], y, [1], [0]) end, [a, b], tol: 1.0e-3) + end + end + + describe "conv (1D)" do + test "1D conv {1,1,5} input, {1,1,3} kernel" do + # 3D: {batch=1, in_channels=1, length=5}; kernel {out=1, in=1, size=3} + input = Nx.tensor([[[1.0, 2.0, 3.0, 4.0, 5.0]]], backend: EMLX.Backend) + kernel = Nx.tensor([[[1.0, 0.0, -1.0]]], backend: EMLX.Backend) + check_equiv(fn x, k -> Nx.conv(x, k) end, [input, kernel], tol: 1.0e-4) + end + + test "1D conv with stride 2 and padding" do + input = Nx.tensor([[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]], backend: EMLX.Backend) + kernel = Nx.tensor([[[1.0, 1.0]]], backend: EMLX.Backend) + + check_equiv(fn x, k -> Nx.conv(x, k, strides: [2], padding: :same) end, [input, kernel], + tol: 1.0e-4 + ) + end + end + + describe "conv (2D)" do + test "2D conv identity kernel {1,1,3,3} × 1-ch input" do + input = Nx.iota({1, 1, 4, 4}, type: :f32, backend: EMLX.Backend) + + kernel = + Nx.tensor([[[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]]], backend: EMLX.Backend) + + check_equiv(fn x, k -> Nx.conv(x, k, padding: :same) end, [input, kernel], tol: 1.0e-4) + end + + test "2D conv with strides and dilations" do + input = Nx.iota({1, 1, 6, 6}, type: :f32, backend: EMLX.Backend) + kernel = Nx.tensor([[[[1.0, 1.0], [1.0, 1.0]]]], backend: EMLX.Backend) + check_equiv(fn x, k -> Nx.conv(x, k, strides: [2, 2]) end, [input, kernel], tol: 1.0e-4) + end + + test "2D conv multi-channel" do + input = Nx.iota({1, 3, 4, 4}, type: :f32, backend: EMLX.Backend) + kernel = Nx.iota({2, 3, 2, 2}, type: :f32, backend: EMLX.Backend) + check_equiv(fn x, k -> Nx.conv(x, k) end, [input, kernel], tol: 1.0e-3) + end + end + + defn sum_axis1_keep(x), do: Nx.sum(x, axes: [1], keep_axes: true) + defn argmax_axis0(x), do: Nx.argmax(x, axis: 0) + defn matmul_defn(a, b), do: Nx.dot(a, b) + defn reduce_max_defn(x), do: Nx.reduce_max(x) + + describe "E2E jit smoke (reductions/dot/conv)" do + test "sum with keep_axes=true via jit" do + x = Nx.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], backend: EMLX.Backend) + result = Nx.Defn.jit(&sum_axis1_keep/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&sum_axis1_keep/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(result, eager, 1.0e-4) + end + + test "argmax via jit" do + x = Nx.tensor([[1.0, 3.0], [4.0, 2.0]], backend: EMLX.Backend) + result = Nx.Defn.jit(&argmax_axis0/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&argmax_axis0/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(result, eager, 1.0e-4) + end + + test "matmul via jit" do + a = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], backend: EMLX.Backend) + b = Nx.tensor([[5.0, 6.0], [7.0, 8.0]], backend: EMLX.Backend) + result = Nx.Defn.jit(&matmul_defn/2, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(&matmul_defn/2, compiler: Nx.Defn.Evaluator).(a, b) + assert_close(result, eager, 1.0e-3) + end + + test "reduce_max via jit" do + x = Nx.tensor([[4.0, 1.0], [2.0, 3.0]], backend: EMLX.Backend) + result = Nx.Defn.jit(&reduce_max_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&reduce_max_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(result, eager, 1.0e-4) + end + end + + defn select_defn(pred, a, b), do: Nx.select(pred, a, b) + defn clip_defn(x, lo, hi), do: Nx.clip(x, lo, hi) + defn slice_static_defn(x), do: Nx.slice(x, [1, 0], [2, 3]) + defn slice_strided_defn(x), do: Nx.slice(x, [0, 0], [2, 3], strides: [1, 2]) + defn put_slice_static_defn(x, patch), do: Nx.put_slice(x, [1, 1], patch) + defn gather_defn(x, idx), do: Nx.gather(x, idx) + defn take_defn(x, idx), do: Nx.take(x, idx, axis: 1) + defn take_along_axis_defn(x, idx), do: Nx.take_along_axis(x, idx, axis: 0) + defn indexed_put_defn(x, idx, updates), do: Nx.indexed_put(x, idx, updates) + defn indexed_add_defn(x, idx, updates), do: Nx.indexed_add(x, idx, updates) + + describe "select" do + test "equivalence vs EMLX.Backend" do + pred = Nx.tensor([1, 0, 1], type: :u8, backend: EMLX.Backend) + a = Nx.tensor([10.0, 20.0, 30.0], backend: EMLX.Backend) + b = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + + native = Nx.Defn.jit(&select_defn/3, compiler: EMLX).(pred, a, b) + eager = Nx.Defn.jit(&select_defn/3, compiler: Nx.Defn.Evaluator).(pred, a, b) + assert_close(native, eager) + end + + test "select with mixed-type true/false is cast to out_type" do + pred = Nx.tensor([1, 0], type: :u8, backend: EMLX.Backend) + a = Nx.tensor([1, 2], type: :s32, backend: EMLX.Backend) + b = Nx.tensor([3, 4], type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(&select_defn/3, compiler: EMLX).(pred, a, b) + eager = Nx.Defn.jit(&select_defn/3, compiler: Nx.Defn.Evaluator).(pred, a, b) + assert_close(native, eager) + end + end + + describe "clip" do + test "equivalence vs EMLX.Backend — f32" do + x = Nx.iota({5}, type: :f32, backend: EMLX.Backend) |> Nx.subtract(2.0) + lo = Nx.tensor(-1.0, backend: EMLX.Backend) + hi = Nx.tensor(1.5, backend: EMLX.Backend) + native = Nx.Defn.jit(&clip_defn/3, compiler: EMLX).(x, lo, hi) + eager = Nx.Defn.jit(&clip_defn/3, compiler: Nx.Defn.Evaluator).(x, lo, hi) + assert_close(native, eager) + end + + test "equivalence vs EMLX.Backend — s32" do + x = Nx.tensor([-5, -1, 0, 3, 10], type: :s32, backend: EMLX.Backend) + lo = Nx.tensor(0, type: :s32, backend: EMLX.Backend) + hi = Nx.tensor(5, type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(&clip_defn/3, compiler: EMLX).(x, lo, hi) + eager = Nx.Defn.jit(&clip_defn/3, compiler: Nx.Defn.Evaluator).(x, lo, hi) + assert_close(native, eager) + end + end + + describe "slice (static indices)" do + test "equivalence vs EMLX.Backend — static 2D slice" do + x = Nx.iota({4, 5}, type: :f32, backend: EMLX.Backend) + native = Nx.Defn.jit(&slice_static_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&slice_static_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "equivalence vs EMLX.Backend — strided slice" do + x = Nx.iota({4, 6}, type: :f32, backend: EMLX.Backend) + native = Nx.Defn.jit(&slice_strided_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&slice_strided_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + end + + describe "slice (dynamic index)" do + test "equivalence vs EMLX.Backend — dynamic start along axis 0" do + # Slices rows [start, start+2) of a 5×4 tensor, start is a runtime scalar. + x = Nx.iota({5, 4}, type: :f32, backend: EMLX.Backend) + + dynamic_slice = fn x, start -> + Nx.slice(x, [start, 0], [2, 4]) + end + + start_val = Nx.tensor(2, type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(dynamic_slice, compiler: EMLX).(x, start_val) + eager = Nx.Defn.jit(dynamic_slice, compiler: Nx.Defn.Evaluator).(x, start_val) + assert_close(native, eager) + end + + test "equivalence vs EMLX.Backend — dynamic start clamped at boundary" do + x = Nx.iota({4, 4}, type: :f32, backend: EMLX.Backend) + + dynamic_slice = fn x, start -> + Nx.slice(x, [start, 0], [2, 4]) + end + + # start=10 should be clamped to 2 (4-2) + start_val = Nx.tensor(10, type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(dynamic_slice, compiler: EMLX).(x, start_val) + eager = Nx.Defn.jit(dynamic_slice, compiler: Nx.Defn.Evaluator).(x, start_val) + assert_close(native, eager) + end + end + + describe "put_slice (static indices)" do + test "equivalence vs EMLX.Backend — static 2D put_slice" do + x = Nx.iota({4, 4}, type: :f32, backend: EMLX.Backend) + patch = Nx.broadcast(Nx.tensor(99.0, backend: EMLX.Backend), {2, 2}) + native = Nx.Defn.jit(&put_slice_static_defn/2, compiler: EMLX).(x, patch) + eager = Nx.Defn.jit(&put_slice_static_defn/2, compiler: Nx.Defn.Evaluator).(x, patch) + assert_close(native, eager) + end + + test "equivalence vs EMLX.Backend — dynamic put_slice (KV-cache pattern)" do + cache = Nx.broadcast(Nx.tensor(0.0, backend: EMLX.Backend), {8, 4}) + + kv_update = fn cache, pos, new_row -> + Nx.put_slice(cache, [pos, 0], Nx.new_axis(new_row, 0)) + end + + pos = Nx.tensor(3, type: :s32, backend: EMLX.Backend) + new_row = Nx.tensor([1.0, 2.0, 3.0, 4.0], backend: EMLX.Backend) + + native = Nx.Defn.jit(kv_update, compiler: EMLX).(cache, pos, new_row) + eager = Nx.Defn.jit(kv_update, compiler: Nx.Defn.Evaluator).(cache, pos, new_row) + assert_close(native, eager) + end + end + + describe "gather" do + test "equivalence vs EMLX.Backend — 2D gather on axis 0" do + x = Nx.iota({4, 3}, type: :f32, backend: EMLX.Backend) + idx = Nx.tensor([[0], [2], [1]], type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(&gather_defn/2, compiler: EMLX).(x, idx) + eager = Nx.Defn.jit(&gather_defn/2, compiler: Nx.Defn.Evaluator).(x, idx) + assert_close(native, eager) + end + + test "equivalence vs EMLX.Backend — multi-axis gather" do + x = Nx.iota({3, 4}, type: :f32, backend: EMLX.Backend) + idx = Nx.tensor([[0, 1], [2, 3]], type: :s32, backend: EMLX.Backend) + + gather_multi = fn x, idx -> Nx.gather(x, idx, axes: [0, 1]) end + + native = Nx.Defn.jit(gather_multi, compiler: EMLX).(x, idx) + eager = Nx.Defn.jit(gather_multi, compiler: Nx.Defn.Evaluator).(x, idx) + assert_close(native, eager) + end + end + + describe "take" do + test "equivalence vs EMLX.Backend" do + x = Nx.iota({3, 4}, type: :f32, backend: EMLX.Backend) + idx = Nx.tensor([2, 0, 3, 1], type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(&take_defn/2, compiler: EMLX).(x, idx) + eager = Nx.Defn.jit(&take_defn/2, compiler: Nx.Defn.Evaluator).(x, idx) + assert_close(native, eager) + end + end + + describe "take_along_axis" do + test "equivalence vs EMLX.Backend" do + x = Nx.iota({3, 4}, type: :f32, backend: EMLX.Backend) + idx = Nx.tensor([[2, 0, 1, 2]], type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(&take_along_axis_defn/2, compiler: EMLX).(x, idx) + eager = Nx.Defn.jit(&take_along_axis_defn/2, compiler: Nx.Defn.Evaluator).(x, idx) + assert_close(native, eager) + end + end + + describe "indexed_put" do + test "equivalence vs EMLX.Backend" do + x = Nx.iota({4, 3}, type: :f32, backend: EMLX.Backend) + idx = Nx.tensor([[0], [2]], type: :s32, backend: EMLX.Backend) + updates = Nx.tensor([[99.0, 99.0, 99.0], [88.0, 88.0, 88.0]], backend: EMLX.Backend) + native = Nx.Defn.jit(&indexed_put_defn/3, compiler: EMLX).(x, idx, updates) + eager = Nx.Defn.jit(&indexed_put_defn/3, compiler: Nx.Defn.Evaluator).(x, idx, updates) + assert_close(native, eager) + end + end + + describe "indexed_add" do + test "equivalence vs EMLX.Backend" do + x = Nx.broadcast(Nx.tensor(0.0, backend: EMLX.Backend), {4, 3}) + idx = Nx.tensor([[0], [0], [2]], type: :s32, backend: EMLX.Backend) + + updates = + Nx.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0], [5.0, 5.0, 5.0]], backend: EMLX.Backend) + + native = Nx.Defn.jit(&indexed_add_defn/3, compiler: EMLX).(x, idx, updates) + eager = Nx.Defn.jit(&indexed_add_defn/3, compiler: Nx.Defn.Evaluator).(x, idx, updates) + assert_close(native, eager) + end + end + + describe "E2E jit smoke (select/slice/gather/indexed)" do + test "select via jit" do + pred = Nx.tensor([1, 0, 1, 0], type: :u8, backend: EMLX.Backend) + a = Nx.iota({4}, type: :f32, backend: EMLX.Backend) + b = Nx.multiply(Nx.iota({4}, type: :f32, backend: EMLX.Backend), -1.0) + native = Nx.Defn.jit(&select_defn/3, compiler: EMLX).(pred, a, b) + eager = Nx.Defn.jit(&select_defn/3, compiler: Nx.Defn.Evaluator).(pred, a, b) + assert_close(native, eager) + end + + test "static slice + add via jit" do + x = Nx.iota({4, 5}, type: :f32, backend: EMLX.Backend) + add_slice = fn x -> Nx.add(Nx.slice(x, [0, 0], [2, 3]), 1.0) end + native = Nx.Defn.jit(add_slice, compiler: EMLX).(x) + eager = Nx.Defn.jit(add_slice, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + end + + # ── private helpers ─────────────────────────────────────────────────────── + + defp compile_nif!(worker, %EMLX.Native.Program{} = wire) do + EMLX.NIF.compile_program(worker, wire) + |> unwrap!() + |> await_worker!() + end + + defp eval_nif!(worker, prog_ref, input_refs) do + EMLX.NIF.eval_program(worker, prog_ref, input_refs) + |> unwrap!() + |> await_worker!() + end + + defn sort_asc_defn(x), do: Nx.sort(x, axis: 0, direction: :asc) + defn sort_desc_defn(x), do: Nx.sort(x, axis: 1, direction: :desc) + defn argsort_asc_defn(x), do: Nx.argsort(x, axis: 0, direction: :asc) + defn argsort_desc_defn(x), do: Nx.argsort(x, axis: 1, direction: :desc) + + defn window_sum_defn(x), do: Nx.window_sum(x, {2, 2}) + defn window_max_defn(x), do: Nx.window_max(x, {2, 2}) + defn window_min_defn(x), do: Nx.window_min(x, {2, 2}) + defn window_product_defn(x), do: Nx.window_product(x, {2, 2}) + + defn cumulative_sum_defn(x), do: Nx.cumulative_sum(x, axis: 1) + defn cumulative_product_defn(x), do: Nx.cumulative_product(x, axis: 0) + defn cumulative_min_defn(x), do: Nx.cumulative_min(x, axis: 0) + defn cumulative_max_defn(x), do: Nx.cumulative_max(x, axis: 0, reverse: true) + + defn fft_defn(x), do: Nx.fft(x) + defn ifft_defn(x), do: Nx.ifft(x) + defn fft2_defn(x), do: Nx.fft2(x) + defn rfft_defn(x), do: Nx.rfft(x) + + defn iota_flat_defn(), do: Nx.iota({3, 4}) + defn iota_axis1_defn(), do: Nx.iota({3, 4}, axis: 1) + defn iota_f32_defn(), do: Nx.iota({5}, type: :f32) + defn eye_3x3_defn(), do: Nx.eye({3, 3}) + defn eye_2x4_defn(), do: Nx.eye({2, 4}) + + defn rng_uniform_defn(key) do + {samples, _key} = Nx.Random.uniform(key, shape: {8}) + samples + end + + defn rng_normal_defn(key) do + {samples, _key} = Nx.Random.normal(key, shape: {8}) + samples + end + + describe "sort" do + test "sort ascending, axis 0 vs EMLX.Backend — f32" do + x = + Nx.tensor([[3.0, 1.0, 2.0], [9.0, 4.0, 7.0]], backend: EMLX.Backend) + + native = Nx.Defn.jit(&sort_asc_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&sort_asc_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "sort descending, axis 1 vs EMLX.Backend — f32" do + x = + Nx.tensor([[3.0, 1.0, 2.0], [9.0, 4.0, 7.0]], backend: EMLX.Backend) + + native = Nx.Defn.jit(&sort_desc_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&sort_desc_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "sort with NaN values — NaN goes to end in asc" do + x = Nx.tensor([1.0, :nan, 2.0, :nan, 0.0], backend: EMLX.Backend) + + native = + Nx.Defn.jit(fn t -> Nx.sort(t, axis: 0, direction: :asc) end, compiler: EMLX).(x) + + eager = + Nx.Defn.jit(fn t -> Nx.sort(t, axis: 0, direction: :asc) end, + compiler: Nx.Defn.Evaluator + ).(x) + + # Both should have NaN at the end; compare non-NaN prefix. + native_list = Nx.to_flat_list(native) + eager_list = Nx.to_flat_list(eager) + + Enum.zip(native_list, eager_list) + |> Enum.each(fn {n, e} -> + if e == :nan, do: assert(n == :nan), else: assert_in_delta(n, e, 1.0e-4) + end) + end + + test "sort s32 ascending" do + x = Nx.tensor([5, 3, 1, 4, 2], type: :s32, backend: EMLX.Backend) + + native = + Nx.Defn.jit(fn t -> Nx.sort(t, axis: 0) end, compiler: EMLX).(x) + + eager = + Nx.Defn.jit(fn t -> Nx.sort(t, axis: 0) end, compiler: Nx.Defn.Evaluator).(x) + + assert_close(native, eager) + end + end + + describe "argsort" do + test "argsort ascending, axis 0 vs EMLX.Backend" do + x = + Nx.tensor([[3.0, 1.0, 2.0], [9.0, 4.0, 7.0]], backend: EMLX.Backend) + + native = Nx.Defn.jit(&argsort_asc_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&argsort_asc_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "argsort descending, axis 1 vs EMLX.Backend" do + x = + Nx.tensor([[3.0, 1.0, 2.0], [9.0, 4.0, 7.0]], backend: EMLX.Backend) + + native = Nx.Defn.jit(&argsort_desc_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&argsort_desc_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "argsort output type matches out tensor type" do + x = Nx.tensor([3.0, 1.0, 2.0], backend: EMLX.Backend) + prog = Expr.lower(Nx.Defn.debug_expr_apply(fn t -> Nx.argsort(t, axis: 0) end, [x])) + # argsort output should be :u64 (Nx default for argsort) + assert Enum.any?(prog.instructions, fn {_, op, _, _} -> op == :argsort end) + end + end + + describe "window reductions" do + test "window_sum 2x2, no padding vs EMLX.Backend — f32" do + x = Nx.iota({4, 4}, type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&window_sum_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&window_sum_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "window_max 2x2, no padding vs EMLX.Backend — f32" do + x = Nx.iota({4, 4}, type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&window_max_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&window_max_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "window_min 2x2, no padding vs EMLX.Backend — f32" do + x = Nx.iota({4, 4}, type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&window_min_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&window_min_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "window_product 2x2, no padding vs EMLX.Backend — f32" do + x = + Nx.iota({4, 4}, type: :f32, backend: EMLX.Backend) + |> Nx.add(1.0) + + native = Nx.Defn.jit(&window_product_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&window_product_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "window_sum with padding vs EMLX.Backend" do + x = Nx.iota({3, 3}, type: :f32, backend: EMLX.Backend) + + f = fn t -> Nx.window_sum(t, {2, 2}, padding: :same) end + native = Nx.Defn.jit(f, compiler: EMLX).(x) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "window_max with strides vs EMLX.Backend" do + x = Nx.iota({6, 6}, type: :f32, backend: EMLX.Backend) + + f = fn t -> Nx.window_max(t, {2, 2}, strides: [2, 2]) end + native = Nx.Defn.jit(f, compiler: EMLX).(x) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "window_sum 1D vs EMLX.Backend" do + x = Nx.iota({6}, type: :f32, backend: EMLX.Backend) + + f = fn t -> Nx.window_sum(t, {3}) end + native = Nx.Defn.jit(f, compiler: EMLX).(x) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + end + + describe "cumulative reductions" do + test "cumulative_sum axis 1 vs EMLX.Backend — f32" do + x = Nx.iota({2, 4}, type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&cumulative_sum_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&cumulative_sum_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "cumulative_product axis 0 vs EMLX.Backend — f32" do + x = Nx.iota({3, 3}, type: :f32, backend: EMLX.Backend) |> Nx.add(1.0) + + native = Nx.Defn.jit(&cumulative_product_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&cumulative_product_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "cumulative_min axis 0 vs EMLX.Backend — s32" do + x = Nx.tensor([[3, 1, 4], [1, 5, 9], [2, 6, 5]], type: :s32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&cumulative_min_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&cumulative_min_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "cumulative_max axis 0 reverse vs EMLX.Backend — f32" do + x = Nx.iota({4, 3}, type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&cumulative_max_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&cumulative_max_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "cumulative_sum s32 axis 0 vs EMLX.Backend" do + x = Nx.tensor([1, 2, 3, 4], type: :s32, backend: EMLX.Backend) + + f = fn t -> Nx.cumulative_sum(t, axis: 0) end + native = Nx.Defn.jit(f, compiler: EMLX).(x) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + end + + describe "fft / ifft" do + test "fft 1D vs EMLX.Backend" do + x = Nx.tensor([1.0, 1.0, 0.0, 0.0], type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&fft_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&fft_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_complex_close(native, eager) + end + + test "ifft 1D vs EMLX.Backend" do + x = Nx.tensor([1.0, 1.0, 0.0, 0.0], type: :f32, backend: EMLX.Backend) + x_fft = Nx.Defn.jit(&fft_defn/1, compiler: EMLX).(x) + + native = Nx.Defn.jit(&ifft_defn/1, compiler: EMLX).(x_fft) + eager = Nx.Defn.jit(&ifft_defn/1, compiler: Nx.Defn.Evaluator).(x_fft) + assert_complex_close(native, eager) + end + + test "fft with explicit length vs EMLX.Backend" do + x = Nx.tensor([1.0, 1.0], type: :f32, backend: EMLX.Backend) + + f = fn t -> Nx.fft(t, length: 4) end + native = Nx.Defn.jit(f, compiler: EMLX).(x) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(x) + assert_complex_close(native, eager) + end + + test "fft2 2D vs EMLX.Backend" do + x = + Nx.tensor([[1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0]], + type: :f32, + backend: EMLX.Backend + ) + + native = Nx.Defn.jit(&fft2_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&fft2_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_complex_close(native, eager) + end + + test "rfft via default_expr descent vs EMLX.Backend" do + x = Nx.tensor([1.0, 1.0, 0.0, 0.0], type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&rfft_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&rfft_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_complex_close(native, eager) + end + end + + describe "iota" do + test "iota flat: IR has :iota instruction, no operands" do + prog = Expr.lower(Nx.Defn.debug_expr_apply(&iota_flat_defn/0, [])) + + assert Enum.any?(prog.instructions, fn {_, op, operands, _} -> + op == :iota and operands == [] + end) + end + + test "iota flat lowering: iattrs encode shape and axis=-1" do + prog = Expr.lower(Nx.Defn.debug_expr_apply(&iota_flat_defn/0, [])) + + {_, :iota, [], [_dtype, n_dims, axis_int | shape]} = + Enum.find(prog.instructions, fn {_, op, _, _} -> op == :iota end) + + assert n_dims == 2 + assert axis_int == -1 + assert shape == [3, 4] + end + + test "iota flat: C++ replay (to_native) matches Nx.iota" do + prog = Expr.lower(Nx.Defn.debug_expr_apply(&iota_flat_defn/0, [])) + [nif_out] = run_nif(prog, []) + expected = Nx.iota({3, 4}, backend: EMLX.Backend) + assert_close(nif_out, expected) + end + + test "iota flat: E2E jit vs Nx.Defn.Evaluator" do + native = Nx.Defn.jit(&iota_flat_defn/0, compiler: EMLX).() + eager = Nx.Defn.jit(&iota_flat_defn/0, compiler: Nx.Defn.Evaluator).() + assert_close(native, eager) + end + + test "iota with axis=1: E2E jit vs Nx.Defn.Evaluator" do + native = Nx.Defn.jit(&iota_axis1_defn/0, compiler: EMLX).() + eager = Nx.Defn.jit(&iota_axis1_defn/0, compiler: Nx.Defn.Evaluator).() + assert_close(native, eager) + end + + test "iota f32 flat: E2E jit vs Nx.Defn.Evaluator" do + native = Nx.Defn.jit(&iota_f32_defn/0, compiler: EMLX).() + eager = Nx.Defn.jit(&iota_f32_defn/0, compiler: Nx.Defn.Evaluator).() + assert_close(native, eager) + end + end + + describe "eye" do + test "eye: IR has :eye instruction, no operands" do + prog = Expr.lower(Nx.Defn.debug_expr_apply(&eye_3x3_defn/0, [])) + + assert Enum.any?(prog.instructions, fn {_, op, operands, _} -> + op == :eye and operands == [] + end) + end + + test "eye 3x3: iattrs encode [dtype, 3, 3]" do + prog = Expr.lower(Nx.Defn.debug_expr_apply(&eye_3x3_defn/0, [])) + + {_, :eye, [], [_dtype, m, n]} = + Enum.find(prog.instructions, fn {_, op, _, _} -> op == :eye end) + + assert m == 3 + assert n == 3 + end + + test "eye 3x3: C++ replay (to_native) matches Nx.eye" do + prog = Expr.lower(Nx.Defn.debug_expr_apply(&eye_3x3_defn/0, [])) + [nif_out] = run_nif(prog, []) + expected = Nx.eye({3, 3}, backend: EMLX.Backend) + assert_close(nif_out, expected) + end + + test "eye 3x3: E2E jit vs Nx.Defn.Evaluator" do + native = Nx.Defn.jit(&eye_3x3_defn/0, compiler: EMLX).() + eager = Nx.Defn.jit(&eye_3x3_defn/0, compiler: Nx.Defn.Evaluator).() + assert_close(native, eager) + end + + test "eye 2x4 rectangular: E2E jit vs Nx.Defn.Evaluator" do + native = Nx.Defn.jit(&eye_2x4_defn/0, compiler: EMLX).() + eager = Nx.Defn.jit(&eye_2x4_defn/0, compiler: Nx.Defn.Evaluator).() + assert_close(native, eager) + end + end + + describe "Nx.Random" do + test "Nx.Random.uniform: native matches evaluator for fixed key" do + key = Nx.Random.key(42) |> Nx.backend_transfer(EMLX.Backend) + + native = Nx.Defn.jit(&rng_uniform_defn/1, compiler: EMLX).(key) + eager = Nx.Defn.jit(&rng_uniform_defn/1, compiler: Nx.Defn.Evaluator).(key) + + assert_close(native, eager, 1.0e-5) + # Samples should be in [0, 1) + assert Nx.all(Nx.greater_equal(native, 0.0)) |> Nx.to_number() == 1 + assert Nx.all(Nx.less(native, 1.0)) |> Nx.to_number() == 1 + end + + test "Nx.Random.uniform: same key → same samples (deterministic)" do + key = Nx.Random.key(99) |> Nx.backend_transfer(EMLX.Backend) + + out1 = Nx.Defn.jit(&rng_uniform_defn/1, compiler: EMLX).(key) + out2 = Nx.Defn.jit(&rng_uniform_defn/1, compiler: EMLX).(key) + + assert_close(out1, out2) + end + + test "Nx.Random.normal: native matches evaluator for fixed key" do + key = Nx.Random.key(7) |> Nx.backend_transfer(EMLX.Backend) + + native = Nx.Defn.jit(&rng_normal_defn/1, compiler: EMLX).(key) + eager = Nx.Defn.jit(&rng_normal_defn/1, compiler: Nx.Defn.Evaluator).(key) + + assert_close(native, eager, 1.0e-4) + end + end + + # cond: simple two-branch predicate on a scalar. + defn cond_two_branch(x) do + cond do + Nx.less(x, 0) -> Nx.negate(x) + true -> x + end + end + + # cond: three branches, each returning a different transformed value. + defn cond_three_branch(x) do + cond do + Nx.less(x, -1) -> Nx.multiply(x, -2) + Nx.less(x, 1) -> Nx.multiply(x, 3) + true -> Nx.multiply(x, 4) + end + end + + # while: count up to 10 from a given start. + defn count_to_10(x), do: while(x, Nx.less(x, 10), do: Nx.add(x, 1)) + + # while: carried state with two tensors. + defn while_two_carry(a, b) do + while {a, b}, Nx.less(a, 5) do + {Nx.add(a, 1), Nx.multiply(b, 2)} + end + end + + defn while_counter_only_cond(x, i, n) do + while {x, i, n}, Nx.less(i, n) do + {Nx.add(x, 1), Nx.add(i, 1), n} + end + end + + # A) A `while` followed by a deeply-shared post-DAG. `add(acc, acc)` references + # `acc` twice, so after N levels the graph is N nodes but 2^N tree paths. The + # split's rewrite pass must be id-memoized or this never terminates. + deftransformp deep_shared(v) do + Enum.reduce(1..28, v, fn _, acc -> Nx.add(acc, acc) end) + end + + defn while_then_shared(x) do + {acc, _} = + while {a = x, k = x}, Nx.less(Nx.sum(a), 5.0) do + {Nx.add(a, k), k} + end + + deep_shared(acc) + end + + # B) A `while` whose initial carry is computed through a runtime_call + # (EMLX.Fast.rms_norm packs its operands in a `{x, weight}` tuple). The before + # stage must collect the runtime_call's operand parameters, or it under-counts + # its args and the remapped param indices overflow the stage input list. + defn while_after_runtime_call(x, w, k) do + s = Nx.add(EMLX.Fast.rms_norm(x, w, 1.0e-6), k) + + {acc, _} = + while {a = s, kk = k}, Nx.less(Nx.sum(a), 10.0) do + {Nx.add(a, kk), kk} + end + + acc + end + + # C) A while-chain whose output is a MAP container; Graph.run/3 must return the + # non-tuple container from the final stage instead of trying to Tuple.to_list it. + defn while_then_map(x) do + {acc, _} = + while {a = x, k = x}, Nx.less(Nx.sum(a), 10.0) do + {Nx.add(a, k), k} + end + + %{value: Nx.add(acc, 1.0), doubled: Nx.multiply(acc, 2.0)} + end + + describe "cond" do + test "cond two-branch: native matches evaluator (negative input)" do + x = Nx.tensor(-3.0, backend: EMLX.Backend) + native = Nx.Defn.jit(&cond_two_branch/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&cond_two_branch/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "cond two-branch: native matches evaluator (positive input)" do + x = Nx.tensor(3.0, backend: EMLX.Backend) + native = Nx.Defn.jit(&cond_two_branch/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&cond_two_branch/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "cond three-branch: all three branches produce correct results" do + for {input, _branch} <- [{-2.0, :first}, {0.0, :second}, {5.0, :third}] do + x = Nx.tensor(input, backend: EMLX.Backend) + native = Nx.Defn.jit(&cond_three_branch/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&cond_three_branch/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + end + end + + describe "while" do + test "while: count from 0 to 10" do + x = Nx.tensor(0, type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(&count_to_10/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&count_to_10/1, compiler: Nx.Defn.Evaluator).(x) + assert Nx.to_number(native) == Nx.to_number(eager) + end + + test "while: trip count depends on runtime value (count from 5)" do + x = Nx.tensor(5, type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(&count_to_10/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&count_to_10/1, compiler: Nx.Defn.Evaluator).(x) + assert Nx.to_number(native) == Nx.to_number(eager) + end + + test "while: already past condition — zero iterations" do + x = Nx.tensor(15, type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(&count_to_10/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&count_to_10/1, compiler: Nx.Defn.Evaluator).(x) + assert Nx.to_number(native) == Nx.to_number(eager) + end + + test "while: two-element tuple carry" do + a = Nx.tensor(0, type: :s32, backend: EMLX.Backend) + b = Nx.tensor(1, type: :s32, backend: EMLX.Backend) + native = Nx.Defn.jit(&while_two_carry/2, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(&while_two_carry/2, compiler: Nx.Defn.Evaluator).(a, b) + {na, nb} = native + {ea, eb} = eager + assert Nx.to_number(na) == Nx.to_number(ea) + assert Nx.to_number(nb) == Nx.to_number(eb) + end + + test "while: counter-only cond ignores a scalar carry slot" do + x = Nx.tensor(0, type: :s32, backend: EMLX.Backend) + i0 = Nx.tensor(0, type: :s32, backend: EMLX.Backend) + n = Nx.tensor(5, type: :s32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&while_counter_only_cond/3, compiler: EMLX).(x, i0, n) + eager = Nx.Defn.jit(&while_counter_only_cond/3, compiler: Nx.Defn.Evaluator).(x, i0, n) + {nx, ni, _} = native + {ex, ei, _} = eager + assert Nx.to_number(nx) == Nx.to_number(ex) + assert Nx.to_number(ni) == Nx.to_number(ei) + assert Nx.to_number(ni) == 5 + end + + test "while: counter-only cond ignores a non-scalar carry slot" do + x = Nx.tensor([1.0, 2.0, 3.0], type: :f32, backend: EMLX.Backend) + i0 = Nx.tensor(0, type: :s32, backend: EMLX.Backend) + n = Nx.tensor(4, type: :s32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&while_counter_only_cond/3, compiler: EMLX).(x, i0, n) + eager = Nx.Defn.jit(&while_counter_only_cond/3, compiler: Nx.Defn.Evaluator).(x, i0, n) + {nx, ni, _} = native + {ex, ei, _} = eager + assert_close(nx, ex) + assert Nx.to_number(ni) == Nx.to_number(ei) + assert Nx.to_number(ni) == 4 + end + end + + describe "splitter regressions" do + # A) Without id-memoization in rewrite_subtree this hangs (2^28 tree walk), + # so a generous-but-bounded timeout turns the hang into a test failure. + @tag timeout: 30_000 + test "while + deeply-shared post-DAG compiles without exponential blowup" do + x = Nx.tensor([1.0, 1.0], type: :f32, backend: EMLX.Backend) + native = Nx.Defn.jit(&while_then_shared/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&while_then_shared/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + # B) runtime_call (rms_norm) feeding a while carry: the before stage must + # collect the runtime_call's tuple operands or param indices overflow. + test "runtime_call operands feeding a while carry are fully collected" do + x = Nx.tensor([[1.0, 2.0, 3.0, 4.0]], type: :f32, backend: EMLX.Backend) + w = Nx.tensor([1.0, 1.0, 1.0, 1.0], type: :f32, backend: EMLX.Backend) + k = Nx.tensor([[0.1, 0.1, 0.1, 0.1]], type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&while_after_runtime_call/3, compiler: EMLX).(x, w, k) + eager = Nx.Defn.jit(&while_after_runtime_call/3, compiler: Nx.Defn.Evaluator).(x, w, k) + assert_close(native, eager) + end + + # C) Map container as the final stage output of a while-chain. + test "while-chain returning a map container runs end-to-end" do + x = Nx.tensor([1.0, 1.0], type: :f32, backend: EMLX.Backend) + native = Nx.Defn.jit(&while_then_map/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&while_then_map/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native.value, eager.value) + assert_close(native.doubled, eager.doubled) + end + end + + defn cholesky_defn(x), do: Nx.LinAlg.cholesky(x) + defn det_defn(x), do: Nx.LinAlg.determinant(x) + + defn all_close_defn(a, b), do: Nx.all_close(a, b) + defn phase_defn(x), do: Nx.phase(x) + defn top_k_defn(x), do: Nx.top_k(x, k: 3) + + # A small helper: a well-conditioned SPD matrix t·tᵀ + n·I. + defp spd_matrix(n) do + t = Nx.iota({n, n}, type: :f32, backend: EMLX.Backend) |> Nx.add(1.0) + Nx.add(Nx.dot(t, Nx.transpose(t)), Nx.multiply(Nx.eye(n, backend: EMLX.Backend), n * 1.0)) + end + + describe "LinAlg.cholesky (native)" do + test "cholesky of an SPD matrix vs EMLX.Backend — f32" do + x = Nx.tensor([[4.0, 2.0], [2.0, 3.0]], type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&cholesky_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&cholesky_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "cholesky composed with surrounding ops (in-graph) vs EMLX.Backend" do + f = fn t -> + spd = Nx.add(Nx.dot(t, Nx.transpose(t)), Nx.multiply(Nx.eye(3), 1.0)) + Nx.LinAlg.cholesky(spd) |> Nx.add(1.0) + end + + t = Nx.iota({3, 3}, type: :f32, backend: EMLX.Backend) |> Nx.add(1.0) + native = Nx.Defn.jit(f, compiler: EMLX).(t) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(t) + assert_close(native, eager) + end + end + + describe "LinAlg.solve / triangular_solve (native)" do + test "solve A x = b vs EMLX.Backend" do + a = Nx.tensor([[3.0, 1.0], [1.0, 2.0]], type: :f32, backend: EMLX.Backend) + b = Nx.tensor([9.0, 8.0], type: :f32, backend: EMLX.Backend) + + f = fn a, b -> Nx.LinAlg.solve(a, b) end + native = Nx.Defn.jit(f, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(a, b) + assert_close(native, eager) + end + + test "triangular_solve (lower, left side) vs EMLX.Backend" do + a = Nx.tensor([[2.0, 0.0], [3.0, 1.0]], type: :f32, backend: EMLX.Backend) + b = Nx.tensor([4.0, 5.0], type: :f32, backend: EMLX.Backend) + + f = fn a, b -> Nx.LinAlg.triangular_solve(a, b, lower: true) end + native = Nx.Defn.jit(f, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(a, b) + assert_close(native, eager) + end + + test "chained cholesky -> solve composes natively (contiguous across linalg ops)" do + spd = Nx.tensor([[4.0, 2.0], [2.0, 3.0]], type: :f32, backend: EMLX.Backend) + b = Nx.tensor([1.0, 2.0], type: :f32, backend: EMLX.Backend) + + f = fn s, b -> Nx.LinAlg.solve(Nx.LinAlg.cholesky(s), b) end + native = Nx.Defn.jit(f, compiler: EMLX).(spd, b) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(spd, b) + assert_close(native, eager) + end + end + + describe "LinAlg batched" do + test "batched cholesky (rank-3) vs EMLX.Backend" do + a = + Nx.tensor( + [[[4.0, 2.0], [2.0, 3.0]], [[9.0, 3.0], [3.0, 5.0]]], + type: :f32, + backend: EMLX.Backend + ) + + native = Nx.Defn.jit(&Nx.LinAlg.cholesky/1, compiler: EMLX).(a) + eager = Nx.Defn.jit(&Nx.LinAlg.cholesky/1, compiler: Nx.Defn.Evaluator).(a) + assert_close(native, eager) + end + end + + describe "LinAlg.qr / eigh / svd (native, reconstruction)" do + test "qr: Q·R reconstructs A and Q is orthonormal" do + a = + Nx.iota({3, 3}, type: :f32, backend: EMLX.Backend) + |> Nx.add(Nx.eye(3, backend: EMLX.Backend)) + + {q, r} = Nx.Defn.jit(&Nx.LinAlg.qr/1, compiler: EMLX).(a) + assert_close(Nx.dot(q, r), a) + assert_close(Nx.dot(Nx.transpose(q), q), Nx.eye(3, type: :f32, backend: EMLX.Backend)) + end + + test "eigh: V·diag(W)·Vᵀ reconstructs a symmetric A" do + a = spd_matrix(3) + n = 3 + + {w, v} = Nx.Defn.jit(&Nx.LinAlg.eigh/1, compiler: EMLX).(a) + recon = Nx.dot(Nx.multiply(v, Nx.reshape(w, {1, n})), Nx.transpose(v)) + assert_close(recon, a) + end + + test "svd: U·diag(S)·Vᵀ reconstructs A" do + a = + Nx.iota({3, 3}, type: :f32, backend: EMLX.Backend) + |> Nx.add(Nx.eye(3, backend: EMLX.Backend)) + + n = 3 + + {u, s, vt} = Nx.Defn.jit(&Nx.LinAlg.svd/1, compiler: EMLX).(a) + recon = Nx.dot(Nx.multiply(u, Nx.reshape(s, {1, n})), vt) + assert_close(recon, a) + end + end + + describe "LinAlg.lu (native)" do + test "lu factors vs EMLX.Backend" do + a = Nx.tensor([[4.0, 3.0], [6.0, 3.0]], type: :f32, backend: EMLX.Backend) + + {pn, ln, un} = Nx.Defn.jit(&Nx.LinAlg.lu/1, compiler: EMLX).(a) + {pe, le, ue} = Nx.Defn.jit(&Nx.LinAlg.lu/1, compiler: Nx.Defn.Evaluator).(a) + assert_close(pn, pe) + assert_close(ln, le) + assert_close(un, ue) + end + + test "lu: P·L·U reconstructs A" do + a = Nx.tensor([[4.0, 3.0], [6.0, 3.0]], type: :f32, backend: EMLX.Backend) + + {p, l, u} = Nx.Defn.jit(&Nx.LinAlg.lu/1, compiler: EMLX).(a) + assert_close(Nx.dot(p, Nx.dot(l, u)), a) + end + end + + describe "LinAlg.determinant (default_expr descent)" do + test "determinant 2x2 (pure primitives) vs EMLX.Backend" do + x = Nx.tensor([[1.0, 2.0], [3.0, 4.0]], type: :f32, backend: EMLX.Backend) + native = Nx.Defn.jit(&det_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&det_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "determinant 3x3 (pure primitives) vs EMLX.Backend" do + x = + Nx.tensor([[1.0, 2.0, 3.0], [1.0, -2.0, 3.0], [7.0, 8.0, 9.0]], + type: :f32, + backend: EMLX.Backend + ) + + native = Nx.Defn.jit(&det_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&det_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "determinant 3x3 sign (row-swap permutation) = -1" do + # A single row swap permutation has det = -1; exercises the cofactor sign. + x = + Nx.tensor([[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + type: :f32, + backend: EMLX.Backend + ) + + native = Nx.Defn.jit(&det_defn/1, compiler: EMLX).(x) + assert_in_delta(Nx.to_number(native), -1.0, 1.0e-5) + end + + test "determinant 4x4 (descends through native LU) vs BinaryBackend reference" do + x = spd_matrix(4) + native = Nx.Defn.jit(&det_defn/1, compiler: EMLX).(x) + + # EMLX.Backend's eager N>3 determinant has a pre-existing type bug, so + # compute the reference entirely on the BinaryBackend (Nx.LinAlg.determinant + # is a defn whose intermediate ops otherwise use the global default backend). + prev = Nx.default_backend(Nx.BinaryBackend) + + ref = + try do + x |> Nx.backend_transfer(Nx.BinaryBackend) |> Nx.LinAlg.determinant() + after + Nx.default_backend(prev) + end + + # f32 LU accumulates rounding error proportional to the magnitude of + # the determinant (~2e5 here), so an absolute delta of 1.0 is too + # tight and flakes under normal f32 rounding — use a relative delta. + ref_num = Nx.to_number(ref) + assert_in_delta(Nx.to_number(native), ref_num, abs(ref_num) * 1.0e-4) + end + end + + describe "block-descent completeness (AllClose/Phase/TopK)" do + test "all_close: close tensors vs EMLX.Backend" do + a = Nx.tensor([1.0, 2.0, 3.0], type: :f32, backend: EMLX.Backend) + b = Nx.tensor([1.0, 2.0, 3.0 + 1.0e-7], type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&all_close_defn/2, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(&all_close_defn/2, compiler: Nx.Defn.Evaluator).(a, b) + assert_close(native, eager) + assert Nx.to_number(native) == 1 + end + + test "all_close: not-close tensors vs EMLX.Backend" do + a = Nx.tensor([1.0, 2.0, 3.0], type: :f32, backend: EMLX.Backend) + b = Nx.tensor([1.0, 2.0, 4.0], type: :f32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&all_close_defn/2, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(&all_close_defn/2, compiler: Nx.Defn.Evaluator).(a, b) + assert_close(native, eager) + assert Nx.to_number(native) == 0 + end + + test "phase: complex tensor vs EMLX.Backend" do + x = + Nx.complex( + Nx.tensor([1.0, -2.0, 0.0], type: :f32, backend: EMLX.Backend), + Nx.tensor([2.0, 1.0, -3.0], type: :f32, backend: EMLX.Backend) + ) + + native = Nx.Defn.jit(&phase_defn/1, compiler: EMLX).(x) + eager = Nx.Defn.jit(&phase_defn/1, compiler: Nx.Defn.Evaluator).(x) + assert_close(native, eager) + end + + test "top_k (tuple-output block): values + indices vs EMLX.Backend" do + x = Nx.tensor([3.0, 1.0, 4.0, 1.0, 5.0, 9.0, 2.0], type: :f32, backend: EMLX.Backend) + + {native_values, native_indices} = Nx.Defn.jit(&top_k_defn/1, compiler: EMLX).(x) + {eager_values, eager_indices} = Nx.Defn.jit(&top_k_defn/1, compiler: Nx.Defn.Evaluator).(x) + + assert_close(native_values, eager_values) + assert Nx.to_flat_list(native_indices) == Nx.to_flat_list(eager_indices) + end + + test "top_k on a batched input (rank 2) vs EMLX.Backend" do + x = + Nx.tensor([[3.0, 1.0, 4.0, 1.0, 5.0], [9.0, 2.0, 6.0, 5.0, 3.0]], + type: :f32, + backend: EMLX.Backend + ) + + {native_values, native_indices} = Nx.Defn.jit(&top_k_defn/1, compiler: EMLX).(x) + {eager_values, eager_indices} = Nx.Defn.jit(&top_k_defn/1, compiler: Nx.Defn.Evaluator).(x) + + assert_close(native_values, eager_values) + assert Nx.to_flat_list(native_indices) == Nx.to_flat_list(eager_indices) + end + end + + # Routing/IR-shape assertions are pure Elixir (no NIF) — run without GPU. + describe "fused-kernel recognition (lowering)" do + test "rms_norm runtime_call lowers to a single :fast_rms_norm instruction" do + fun = fn x, w -> EMLX.Fast.rms_norm(x, w, 1.0e-5) end + expr = Nx.Defn.debug_expr_apply(fun, [Nx.template({2, 16}, :f32), Nx.template({16}, :f32)]) + prog = Expr.lower(expr) + + assert [{_, :fast_rms_norm, [_x, _w], [eps_bits]}] = prog.instructions + assert_in_delta(Expr.bits_to_f64(eps_bits), 1.0e-5, 1.0e-12) + end + + test "f64_bits/1 ↔ bits_to_f64/1 round-trips through the int64 attr channel" do + for v <- [1.0e-6, 1.0e-5, 0.08838834764831845, 10_000.0, 1_000_000.0] do + assert Expr.bits_to_f64(Expr.f64_bits(v)) == v + end + end + + test "swiglu / layer_norm / sdpa each lower to a single fused opcode" do + cases = [ + {fn g, u -> EMLX.Fast.swiglu(g, u) end, + [Nx.template({2, 8}, :f32), Nx.template({2, 8}, :f32)], :fast_swiglu}, + {fn x, w, b -> EMLX.Fast.layer_norm(x, w, b, 1.0e-5) end, + [Nx.template({2, 16}, :f32), Nx.template({16}, :f32), Nx.template({16}, :f32)], + :fast_layer_norm}, + {fn x, w -> EMLX.Fast.layer_norm(x, w, 1.0e-5) end, + [Nx.template({2, 16}, :f32), Nx.template({16}, :f32)], :fast_layer_norm_no_bias}, + {fn q, k, v -> EMLX.Fast.scaled_dot_product_attention_causal(q, k, v, 0.125) end, + [ + Nx.template({1, 2, 4, 8}, :f32), + Nx.template({1, 2, 4, 8}, :f32), + Nx.template({1, 2, 4, 8}, :f32) + ], :fast_sdpa_causal}, + {fn q, k, v, s -> EMLX.Fast.scaled_dot_product_attention(q, k, v, 0.125, sinks: s) end, + [ + Nx.template({1, 2, 4, 8}, :f32), + Nx.template({1, 2, 4, 8}, :f32), + Nx.template({1, 2, 4, 8}, :f32), + Nx.template({2}, :f32) + ], :fast_sdpa_sinks}, + {fn q, k, v, s -> + EMLX.Fast.scaled_dot_product_attention_causal(q, k, v, 0.125, sinks: s) + end, + [ + Nx.template({1, 2, 4, 8}, :f32), + Nx.template({1, 2, 4, 8}, :f32), + Nx.template({1, 2, 4, 8}, :f32), + Nx.template({2}, :f32) + ], :fast_sdpa_causal_sinks} + ] + + for {fun, templates, opcode} <- cases do + prog = Expr.lower(Nx.Defn.debug_expr_apply(fun, templates)) + assert [{_, ^opcode, _operands, _attrs}] = prog.instructions + end + end + + test "prefill RoPE with positions (T>1) lowers to a single :fast_rope_positions instruction" do + fun = fn a, pos -> EMLX.Fast.rope_with_positions(a, pos, 64, false, 10_000.0, 1.0) end + + expr = + Nx.Defn.debug_expr_apply(fun, [ + Nx.template({2, 4, 2, 64}, :f32), + Nx.template({2, 4}, :s32) + ]) + + prog = Expr.lower(expr) + assert [{_, :fast_rope_positions, [_a, _pos], _attrs}] = prog.instructions + end + + test "prefill RoPE with freqs (T>1) lowers to a single :fast_rope_with_freqs_positions instruction" do + fun = fn a, pos, freqs -> EMLX.Fast.rope_with_freqs(a, pos, 64, false, 1.0, freqs) end + + expr = + Nx.Defn.debug_expr_apply(fun, [ + Nx.template({2, 4, 2, 64}, :f32), + Nx.template({2, 4}, :s32), + Nx.template({32}, :f32) + ]) + + prog = Expr.lower(expr) + + assert [{_, :fast_rope_with_freqs_positions, [_a, _pos, _freqs], _attrs}] = + prog.instructions + end + end + + describe "fused kernels vs eager + primitive (Metal)" do + @describetag :metal + + test "rms_norm: fused replay matches eager and hand-written primitive" do + fun = fn x, w -> EMLX.Fast.rms_norm(x, w, 1.0e-5) end + x = Nx.iota({2, 4, 16}, type: :f32) |> Nx.divide(50) |> gpu_t() + w = Nx.broadcast(Nx.tensor(1.0, type: :f32), {16}) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(x, w) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(x, w) + assert_all_close(native, eager, tol: 1.0e-3) + + prim_fun = fn x, w -> + rms = Nx.sqrt(Nx.add(Nx.mean(Nx.pow(x, 2), axes: [-1], keep_axes: true), 1.0e-5)) + Nx.divide(x, rms) |> Nx.multiply(w) + end + + prim = Nx.Defn.jit(prim_fun, compiler: EMLX, device: :gpu).(x, w) + assert_all_close(native, prim, tol: 1.0e-3) + end + + test "layer_norm (with bias) and no-bias: fused vs eager and primitive" do + x = Nx.iota({2, 16}, type: :f32) |> Nx.divide(30) |> gpu_t() + w = Nx.broadcast(Nx.tensor(1.2, type: :f32), {16}) |> gpu_t() + b = Nx.broadcast(Nx.tensor(0.3, type: :f32), {16}) |> gpu_t() + + with_bias = fn x, w, b -> EMLX.Fast.layer_norm(x, w, b, 1.0e-5) end + native = Nx.Defn.jit(with_bias, compiler: EMLX, device: :gpu).(x, w, b) + eager = Nx.Defn.jit(with_bias, compiler: Nx.Defn.Evaluator).(x, w, b) + assert_all_close(native, eager, tol: 1.0e-3) + + ln_prim = fn x, w, b -> + mean = Nx.mean(x, axes: [-1], keep_axes: true) + var = Nx.mean(Nx.pow(Nx.subtract(x, mean), 2), axes: [-1], keep_axes: true) + normed = Nx.divide(Nx.subtract(x, mean), Nx.sqrt(Nx.add(var, 1.0e-5))) + Nx.add(Nx.multiply(normed, w), b) + end + + prim = Nx.Defn.jit(ln_prim, compiler: EMLX, device: :gpu).(x, w, b) + assert_all_close(native, prim, tol: 1.0e-3) + + no_bias = fn x, w -> EMLX.Fast.layer_norm(x, w, 1.0e-5) end + nb_native = Nx.Defn.jit(no_bias, compiler: EMLX, device: :gpu).(x, w) + nb_eager = Nx.Defn.jit(no_bias, compiler: Nx.Defn.Evaluator).(x, w) + assert_all_close(nb_native, nb_eager, tol: 1.0e-3) + end + + test "swiglu: fused replay matches hand-written silu(gate)*up" do + fun = fn g, u -> EMLX.Fast.swiglu(g, u) end + gate = Nx.iota({2, 8}, type: :f32) |> Nx.divide(10) |> gpu_t() + up = Nx.iota({2, 8}, type: :f32) |> Nx.divide(7) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(gate, up) + prim = Nx.multiply(Nx.multiply(gate, Nx.sigmoid(gate)), up) + assert_all_close(native, prim, tol: 1.0e-4) + end + + test "sdpa (no mask): fused replay matches eager and softmax(QKᵀ)·V primitive" do + scale = 0.125 + fun = fn q, k, v -> EMLX.Fast.scaled_dot_product_attention(q, k, v, scale) end + q = Nx.iota({1, 2, 4, 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() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v) + assert_all_close(native, eager, tol: 1.0e-3) + + prim_fun = fn q, k, v -> + scores = Nx.dot(q, [3], [0, 1], k, [3], [0, 1]) |> Nx.multiply(scale) + Nx.dot(Nx.exp(scores) |> normalize_rows(), [3], [0, 1], v, [2], [0, 1]) + end + + prim = Nx.Defn.jit(prim_fun, compiler: EMLX, device: :gpu).(q, k, v) + assert_all_close(native, prim, tol: 1.0e-3) + end + + test "sdpa (causal): fused replay matches eager and masked-softmax primitive" do + scale = 0.125 + fun = fn q, k, v -> EMLX.Fast.scaled_dot_product_attention_causal(q, k, v, scale) end + q = Nx.iota({1, 2, 4, 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() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v) + assert_all_close(native, eager, tol: 1.0e-3) + + # Primitive causal attention (T_q == T_kv): row i attends keys 0..i. + prim_fun = fn q, k, v -> + scores = Nx.dot(q, [3], [0, 1], k, [3], [0, 1]) |> Nx.multiply(scale) + rows = Nx.iota({4, 4}, axis: 0) + cols = Nx.iota({4, 4}, axis: 1) + bias = Nx.select(Nx.greater_equal(rows, cols), 0.0, -1.0e9) |> Nx.reshape({1, 1, 4, 4}) + attn = normalize_rows(Nx.exp(Nx.add(scores, bias))) + Nx.dot(attn, [3], [0, 1], v, [2], [0, 1]) + end + + prim = Nx.Defn.jit(prim_fun, compiler: EMLX, device: :gpu).(q, k, v) + assert_all_close(native, prim, tol: 1.0e-3) + end + + test "sdpa (additive mask): fused replay matches eager and masked-softmax primitive" do + scale = 0.125 + fun = fn q, k, v, m -> EMLX.Fast.scaled_dot_product_attention(q, k, v, scale, m) end + q = Nx.iota({1, 2, 4, 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() + # Mask out the last key position for every query (non-trivial additive mask). + mask = + Nx.tensor([[[[0.0, 0.0, 0.0, -1.0e9]]]], type: :f32) + |> Nx.broadcast({1, 1, 4, 4}) + |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v, mask) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v, mask) + assert_all_close(native, eager, tol: 1.0e-3) + + prim_fun = fn q, k, v, m -> + scores = Nx.dot(q, [3], [0, 1], k, [3], [0, 1]) |> Nx.multiply(scale) + attn = normalize_rows(Nx.exp(Nx.add(scores, m))) + Nx.dot(attn, [3], [0, 1], v, [2], [0, 1]) + end + + prim = Nx.Defn.jit(prim_fun, compiler: EMLX, device: :gpu).(q, k, v, mask) + assert_all_close(native, prim, tol: 1.0e-3) + end + + test "sdpa (causal + key_mask, decode): fused replay matches eager (padded and all-present)" do + fun = fn q, k, v, km -> + EMLX.Fast.scaled_dot_product_attention_causal_key_masked(q, k, v, 0.125, km) + end + + q = Nx.iota({2, 2, 1, 8}, type: :f32) |> Nx.divide(100) |> gpu_t() + k = Nx.iota({2, 2, 4, 8}, type: :f32) |> Nx.divide(90) |> gpu_t() + v = Nx.iota({2, 2, 4, 8}, type: :f32) |> Nx.divide(80) |> gpu_t() + + # Padded batch: the eager NIF builds an additive mask; so does the opcode. + padded = Nx.tensor([[1, 1, 1, 0], [1, 1, 0, 0]], type: :s32) |> gpu_t() + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v, padded) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v, padded) + assert_all_close(native, eager, tol: 1.0e-3) + + # All keys present: the eager NIF host-fast-paths to pure "causal" (no mask + # alloc) while the compiled opcode always builds the additive mask. They + # must still agree — this exercises that branch divergence. + all_present = Nx.broadcast(Nx.tensor(1, type: :s32), {2, 4}) |> gpu_t() + native_ap = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v, all_present) + eager_ap = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v, all_present) + assert_all_close(native_ap, eager_ap, tol: 1.0e-3) + end + + test "sdpa (no mask, + sinks): fused replay matches eager" do + scale = 0.125 + fun = fn q, k, v, s -> EMLX.Fast.scaled_dot_product_attention(q, k, v, scale, sinks: s) end + q = Nx.iota({1, 2, 4, 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() + sinks = Nx.tensor([0.1, -0.2], type: :f32) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v, sinks) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v, sinks) + assert_all_close(native, eager, tol: 1.0e-3) + end + + test "sdpa (additive mask, + sinks): fused replay matches eager" do + scale = 0.125 + + fun = fn q, k, v, m, s -> + EMLX.Fast.scaled_dot_product_attention(q, k, v, scale, m, sinks: s) + end + + q = Nx.iota({1, 2, 4, 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() + + mask = + Nx.tensor([[[[0.0, 0.0, 0.0, -1.0e9]]]], type: :f32) + |> Nx.broadcast({1, 1, 4, 4}) + |> gpu_t() + + sinks = Nx.tensor([0.1, -0.2], type: :f32) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v, mask, sinks) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v, mask, sinks) + assert_all_close(native, eager, tol: 1.0e-3) + end + + test "sdpa (causal, + sinks): fused replay matches eager" do + scale = 0.125 + + fun = fn q, k, v, s -> + EMLX.Fast.scaled_dot_product_attention_causal(q, k, v, scale, sinks: s) + end + + q = Nx.iota({1, 2, 4, 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() + sinks = Nx.tensor([0.1, -0.2], type: :f32) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v, sinks) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v, sinks) + assert_all_close(native, eager, tol: 1.0e-3) + end + + test "sdpa (causal + key_mask, decode, + sinks): fused replay matches eager (padded and all-present)" do + fun = fn q, k, v, km, s -> + EMLX.Fast.scaled_dot_product_attention_causal_key_masked(q, k, v, 0.125, km, sinks: s) + end + + q = Nx.iota({2, 2, 1, 8}, type: :f32) |> Nx.divide(100) |> gpu_t() + k = Nx.iota({2, 2, 4, 8}, type: :f32) |> Nx.divide(90) |> gpu_t() + v = Nx.iota({2, 2, 4, 8}, type: :f32) |> Nx.divide(80) |> gpu_t() + sinks = Nx.tensor([0.1, -0.2], type: :f32) |> gpu_t() + + padded = Nx.tensor([[1, 1, 1, 0], [1, 1, 0, 0]], type: :s32) |> gpu_t() + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v, padded, sinks) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v, padded, sinks) + assert_all_close(native, eager, tol: 1.0e-3) + + all_present = Nx.broadcast(Nx.tensor(1, type: :s32), {2, 4}) |> gpu_t() + native_ap = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(q, k, v, all_present, sinks) + eager_ap = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(q, k, v, all_present, sinks) + assert_all_close(native_ap, eager_ap, tol: 1.0e-3) + end + + test "rope (scalar offset): fused replay matches eager" do + fun = fn a -> EMLX.Fast.rope(a, 64, false, 10_000.0, 1.0, 3) end + a = Nx.iota({1, 4, 1, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(a) + assert_all_close(native, eager, tol: 1.0e-3) + end + + test "rope_with_positions (decode T=1): fused replay matches eager" do + fun = fn a, pos -> EMLX.Fast.rope_with_positions(a, pos, 64, false, 10_000.0, 1.0) end + a = Nx.iota({2, 1, 2, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[3], [5]], type: :s32) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(a, pos) + assert_all_close(native, eager, tol: 1.0e-3) + end + + test "rope_with_freqs (decode T=1): fused replay matches eager" do + fun = fn a, pos, freqs -> EMLX.Fast.rope_with_freqs(a, pos, 64, false, 1.0, freqs) end + a = Nx.iota({2, 1, 2, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[3], [5]], type: :s32) |> gpu_t() + freqs = Nx.iota({32}, type: :f32) |> Nx.add(1) |> Nx.divide(1000) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos, freqs) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(a, pos, freqs) + assert_all_close(native, eager, tol: 1.0e-3) + 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 + @describetag :metal + + test "rope_with_positions (T>1, sequential positions): native matches eager" do + fun = fn a, pos -> EMLX.Fast.rope_with_positions(a, pos, 64, false, 10_000.0, 1.0) end + a = Nx.iota({2, 4, 2, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[3, 4, 5, 6], [10, 11, 12, 13]], type: :s32) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(a, pos) + assert_all_close(native, eager, tol: 1.0e-3) + end + + test "rope_with_positions (T>1, left-padded/non-sequential positions): native matches eager" do + fun = fn a, pos -> EMLX.Fast.rope_with_positions(a, pos, 64, false, 10_000.0, 1.0) end + a = Nx.iota({2, 5, 2, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + # Left-padded batch: row 0 has 2 pad tokens (position 0 repeated), row 1 has none. + pos = Nx.tensor([[0, 0, 1, 2, 3], [8, 9, 10, 11, 12]], type: :s32) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(a, pos) + assert_all_close(native, eager, tol: 1.0e-3) + end + + test "rope_with_positions (T>1, dims < D): native matches eager on the pass-through tail" do + fun = fn a, pos -> EMLX.Fast.rope_with_positions(a, pos, 32, false, 10_000.0, 1.0) end + a = Nx.iota({2, 3, 2, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[0, 1, 2], [4, 5, 6]], type: :s32) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(a, pos) + assert_all_close(native, eager, tol: 1.0e-3) + end + + # `mlx::fast::rope`'s freqs overload reciprocates internally (both the CPU + # fallback and the Metal kernel compute `inv_freq = 1/freqs[i]` — see + # mlx/fast.cpp's default_inv_freqs-vs-freqs branch and + # backend/metal/kernels/rope.metal's `1.0 / freqs[...]`), so `freqs` here is + # a *raw frequency* tensor, not an inv_freq tensor: realistic magnitudes + # are ~1..base (reciprocal lands back in the usual ~1e-4..1 inv_freq + # range). Using inv_freq-scale values directly as `freqs` (as one might + # naively expect from the name) reciprocates into huge angles and makes + # native-vs-eager float32 cos/sin agreement numerically meaningless — a + # test-data pitfall, not a lowering bug. + test "rope_with_freqs (T>1, H=1, sequential positions): native matches eager" do + fun = fn a, pos, freqs -> EMLX.Fast.rope_with_freqs(a, pos, 64, false, 1.0, freqs) end + a = Nx.iota({2, 4, 1, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[3, 4, 5, 6], [10, 11, 12, 13]], type: :s32) |> gpu_t() + freqs = Nx.pow(10_000.0, Nx.divide(Nx.iota({32}, type: :f32), 32)) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos, freqs) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(a, pos, freqs) + assert_all_close(native, eager, tol: 1.0e-3) + end + + test "rope_with_freqs (T>1, H=2, sequential positions): native matches pure-Nx primitive" do + fun = fn a, pos, freqs -> EMLX.Fast.rope_with_freqs(a, pos, 64, false, 1.0, freqs) end + a = Nx.iota({2, 4, 2, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[3, 4, 5, 6], [10, 11, 12, 13]], type: :s32) |> gpu_t() + freqs = Nx.pow(10_000.0, Nx.divide(Nx.iota({32}, type: :f32), 32)) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos, freqs) + prim_fun = fn a, pos, freqs -> rope_freqs_prim(a, pos, freqs, 64, 1.0) end + prim = Nx.Defn.jit(prim_fun, compiler: EMLX, device: :gpu).(a, pos, freqs) + assert_all_close(native, prim, tol: 1.0e-3) + end + + test "rope_with_freqs (T>1, H=2, left-padded/non-sequential positions): native matches pure-Nx primitive" do + fun = fn a, pos, freqs -> EMLX.Fast.rope_with_freqs(a, pos, 64, false, 1.0, freqs) end + a = Nx.iota({2, 5, 2, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[0, 0, 1, 2, 3], [8, 9, 10, 11, 12]], type: :s32) |> gpu_t() + freqs = Nx.pow(10_000.0, Nx.divide(Nx.iota({32}, type: :f32), 32)) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos, freqs) + prim_fun = fn a, pos, freqs -> rope_freqs_prim(a, pos, freqs, 64, 1.0) end + prim = Nx.Defn.jit(prim_fun, compiler: EMLX, device: :gpu).(a, pos, freqs) + assert_all_close(native, prim, tol: 1.0e-3) + end + + test "rope_with_freqs (T>1, H=2, dims < D): native matches pure-Nx primitive on the pass-through tail" do + fun = fn a, pos, freqs -> EMLX.Fast.rope_with_freqs(a, pos, 32, false, 1.0, freqs) end + a = Nx.iota({2, 6, 2, 64}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[0, 0, 0, 1, 2, 3], [20, 21, 22, 23, 24, 25]], type: :s32) |> gpu_t() + # dims=32 → freqs shape {dims/2} = {16}. + freqs = Nx.pow(10_000.0, Nx.divide(Nx.iota({16}, type: :f32), 16)) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos, freqs) + prim_fun = fn a, pos, freqs -> rope_freqs_prim(a, pos, freqs, 32, 1.0) end + prim = Nx.Defn.jit(prim_fun, compiler: EMLX, device: :gpu).(a, pos, freqs) + assert_all_close(native, prim, tol: 1.0e-3) + end + end + + describe ":fun node unreachability (doc audit)" do + test "custom-fun reduce: :fun never becomes an instruction" do + templates = [Nx.template({5}, :f32)] + + expr = + Nx.Defn.debug_expr_apply( + fn t -> Nx.reduce(t, 0.0, fn a, b -> Nx.add(a, b) end) end, + templates + ) + + prog = Expr.lower(expr) + + assert Enum.all?(prog.instructions, fn {_id, op, _operands, _attrs} -> op != :fun end) + end + + test "custom-fun window_reduce: :fun never becomes an instruction" do + templates = [Nx.template({6}, :f32)] + + expr = + Nx.Defn.debug_expr_apply( + fn t -> Nx.window_reduce(t, 0.0, {2}, fn a, b -> Nx.add(a, b) end) end, + templates + ) + + prog = Expr.lower(expr) + + assert Enum.all?(prog.instructions, fn {_id, op, _operands, _attrs} -> op != :fun end) + end + end + + describe "while-in-default_expr descent (static unroll)" do + # QR `mode: :complete` and SVD `full_matrices?: false` both fall through + # to `expand_block_via_default`; QR's Householder decomposition carries a + # statically-counted `while` (trip count fixed by the input shape at + # trace time) that previously raised "does not yet lower op :while" and + # fired the Evaluator fallback. `expand_node`'s new `:while` clause + # detects that shape and unrolls it in place. + test "qr :complete lowers natively — Q orthonormal, Q·R reconstructs A (square)" do + a = + Nx.iota({3, 3}, type: :f32, backend: EMLX.Backend) + |> Nx.add(Nx.eye(3, backend: EMLX.Backend)) + + {q, r} = Nx.Defn.jit(fn t -> Nx.LinAlg.qr(t, mode: :complete) end, compiler: EMLX).(a) + + assert q.shape == {3, 3} + assert r.shape == {3, 3} + assert_close(Nx.dot(q, r), a) + assert_close(Nx.dot(Nx.transpose(q), q), Nx.eye(3, type: :f32, backend: EMLX.Backend)) + end + + test "qr :complete lowers natively — tall input, Q is m×m (not m×n)" do + a = Nx.iota({4, 3}, type: :f32, backend: EMLX.Backend) |> Nx.add(1) + + {q, r} = Nx.Defn.jit(fn t -> Nx.LinAlg.qr(t, mode: :complete) end, compiler: EMLX).(a) + + assert q.shape == {4, 4} + assert r.shape == {4, 3} + assert_close(Nx.dot(q, r), a) + assert_close(Nx.dot(Nx.transpose(q), q), Nx.eye(4, type: :f32, backend: EMLX.Backend)) + end + + test "qr :complete matches eager EMLX.Backend (Evaluator) on Q·R and orthonormality" do + a = Nx.iota({5, 2}, type: :f32, backend: EMLX.Backend) |> Nx.add(1) + + {q_native, r_native} = + Nx.Defn.jit(fn t -> Nx.LinAlg.qr(t, mode: :complete) end, compiler: EMLX).(a) + + {q_eager, r_eager} = + Nx.Defn.jit(fn t -> Nx.LinAlg.qr(t, mode: :complete) end, compiler: Nx.Defn.Evaluator).(a) + + assert_close(Nx.dot(q_native, r_native), Nx.dot(q_eager, r_eager)) + + assert_close( + Nx.dot(Nx.transpose(q_native), q_native), + Nx.eye(5, type: :f32, backend: EMLX.Backend) + ) + end + + test "svd full_matrices?: false lowers natively — U·diag(S)·Vᵗ reconstructs A (tall)" do + a = Nx.iota({4, 3}, type: :f32, backend: EMLX.Backend) |> Nx.add(1) + + {u, s, vt} = + Nx.Defn.jit(fn t -> Nx.LinAlg.svd(t, full_matrices?: false) end, compiler: EMLX).(a) + + assert u.shape == {4, 3} + assert s.shape == {3} + assert vt.shape == {3, 3} + recon = Nx.dot(Nx.multiply(u, Nx.reshape(s, {1, 3})), vt) + assert_close(recon, a) + end + + test "svd full_matrices?: false lowers natively — wide and square inputs" do + wide = Nx.iota({3, 4}, type: :f32, backend: EMLX.Backend) |> Nx.add(1) + + {u, s, vt} = + Nx.Defn.jit(fn t -> Nx.LinAlg.svd(t, full_matrices?: false) end, compiler: EMLX).(wide) + + assert u.shape == {3, 3} + assert s.shape == {3} + assert vt.shape == {3, 4} + assert_close(Nx.dot(Nx.multiply(u, Nx.reshape(s, {1, 3})), vt), wide) + + square = + Nx.iota({3, 3}, type: :f32, backend: EMLX.Backend) + |> Nx.add(Nx.eye(3, backend: EMLX.Backend)) + + {u2, s2, vt2} = + Nx.Defn.jit(fn t -> Nx.LinAlg.svd(t, full_matrices?: false) end, compiler: EMLX).(square) + + assert_close(Nx.dot(Nx.multiply(u2, Nx.reshape(s2, {1, 3})), vt2), square) + end + + test "svd full_matrices?: false matches eager EMLX.Backend singular values" do + a = Nx.iota({4, 3}, type: :f32, backend: EMLX.Backend) |> Nx.add(1) + + {_u, s_native, _vt} = + Nx.Defn.jit(fn t -> Nx.LinAlg.svd(t, full_matrices?: false) end, compiler: EMLX).(a) + + {_u, s_eager, _vt} = + Nx.Defn.jit(fn t -> Nx.LinAlg.svd(t, full_matrices?: false) end, + compiler: Nx.Defn.Evaluator + ).(a) + + assert_close(s_native, s_eager) + end + + test "qr :complete lowers with no :while instruction in the compiled program" do + templates = [Nx.template({4, 3}, :f32)] + + expr = + Nx.Defn.debug_expr_apply( + fn t -> Nx.LinAlg.qr(t, mode: :complete) end, + templates + ) + + prog = Expr.lower(expr, 1) + + assert Enum.all?(prog.instructions, fn {_id, op, _operands, _attrs} -> op != :while end) + end + + test "triangular_solve left_side: false (2D b) vs EMLX.Backend" do + a = Nx.tensor([[3.0, 0.0, 0.0], [2.0, 1.0, 0.0], [1.0, 1.0, 1.0]], backend: EMLX.Backend) + b = Nx.tensor([[1.0, 2.0, 3.0]], backend: EMLX.Backend) + + f = fn a, b -> Nx.LinAlg.triangular_solve(a, b, left_side: false) end + native = Nx.Defn.jit(f, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(a, b) + assert_close(native, eager) + end + + test "triangular_solve left_side: false (1D b) vs EMLX.Backend" do + a = Nx.tensor([[3.0, 0.0, 0.0], [2.0, 1.0, 0.0], [1.0, 1.0, 1.0]], backend: EMLX.Backend) + b = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + + f = fn a, b -> Nx.LinAlg.triangular_solve(a, b, left_side: false) end + native = Nx.Defn.jit(f, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(a, b) + assert_close(native, eager) + end + + test "triangular_solve transform_a: :transpose vs EMLX.Backend" do + a = Nx.tensor([[3.0, 1.0, 2.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]], backend: EMLX.Backend) + b = Nx.tensor([4.0, 5.0, 6.0], backend: EMLX.Backend) + + f = fn a, b -> Nx.LinAlg.triangular_solve(a, b, lower: false, transform_a: :transpose) end + native = Nx.Defn.jit(f, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(a, b) + assert_close(native, eager) + end + + test "triangular_solve left_side: false + transform_a: :transpose vs EMLX.Backend" do + a = Nx.tensor([[3.0, 1.0, 2.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]], backend: EMLX.Backend) + b = Nx.tensor([[4.0, 5.0, 6.0]], backend: EMLX.Backend) + + f = fn a, b -> + Nx.LinAlg.triangular_solve(a, b, left_side: false, lower: false, transform_a: :transpose) + end + + native = Nx.Defn.jit(f, compiler: EMLX).(a, b) + eager = Nx.Defn.jit(f, compiler: Nx.Defn.Evaluator).(a, b) + assert_close(native, eager) + end + end + + describe "hooks (token/attach_token, extra-output design)" do + # `Nx.Defn.Kernel.hook/2,3` is fire-and-forget, not control flow (see + # `EMLX.Native.Expr`'s moduledoc "Hooks" section): `attach_token`'s value + # is its wrapped expr unchanged, so the hook is lowered in the *same* + # single NIF-call program as everything else -- the hooked value(s) ride + # along as extra outputs, and the callback fires host-side right after + # `eval_program` returns. No `Nx.Defn.Graph.split` host round-trip needed. + test "top-level hook fires once with the correct value; result unaffected" do + a = Nx.tensor(1, backend: EMLX.Backend) + b = Nx.tensor(2, backend: EMLX.Backend) + + result = Nx.Defn.jit(&hook_top_level/2, compiler: EMLX).(a, b) + assert Nx.to_number(result) == 5 + + assert_receive {:mid, hooked_value} + assert Nx.to_number(hooked_value) == 6 + refute_receive {:mid, _} + end + + test "a hook whose value is unreferenced by the output never fires (matches Evaluator's dead-code elimination)" do + a = Nx.tensor(4, backend: EMLX.Backend) + b = Nx.tensor(3, backend: EMLX.Backend) + + native = Nx.Defn.jit(&hook_unused_value/2, compiler: EMLX).(a, b) + assert Nx.to_number(native) == 7 + refute_receive {:dbg, _} + + eager = Nx.Defn.jit(&hook_unused_value/2, compiler: Nx.Defn.Evaluator).(a, b) + assert Nx.to_number(eager) == 7 + refute_receive {:dbg, _} + end + + test "a name-only hook (no trace-time callback, no runtime override) is a silent no-op" do + a = Nx.tensor(4, backend: EMLX.Backend) + b = Nx.tensor(3, backend: EMLX.Backend) + + result = Nx.Defn.jit(&hook_name_only/2, compiler: EMLX).(a, b) + assert Nx.to_number(result) == 7 + end + + test "a hook on a tuple payload fires with the matching container shape" do + a = Nx.tensor(4, backend: EMLX.Backend) + b = Nx.tensor(3, backend: EMLX.Backend) + + {s, d} = Nx.Defn.jit(&hook_tuple_payload/2, compiler: EMLX).(a, b) + assert Nx.to_number(s) == 7 + assert Nx.to_number(d) == 1 + + assert_receive {:pair, sum, diff} + assert Nx.to_number(sum) == 7 + assert Nx.to_number(diff) == 1 + end + + test "a hook nested inside a cond branch raises instead of silently double-firing" do + a = Nx.tensor(1, backend: EMLX.Backend) + b = Nx.tensor(2, backend: EMLX.Backend) + templates = [Nx.template(a.shape, a.type), Nx.template(b.shape, b.type)] + + expr = Nx.Defn.debug_expr_apply(&hook_in_cond_branch/2, templates) + + assert_raise ArgumentError, ~r/cannot lower a hook nested inside a cond branch/, fn -> + Expr.lower(expr, 2) + end + end + + test "a hook inside a while body fires once per iteration, matching Evaluator" do + a = Nx.tensor(3, backend: EMLX.Backend) + + native = Nx.Defn.jit(&hook_in_while_body/1, compiler: EMLX).(a) + assert Nx.to_number(native) == 6 + assert_receive {:iter, v1} + assert_receive {:iter, v2} + assert_receive {:iter, v3} + native_values = Enum.map([v1, v2, v3], &Nx.to_number/1) + refute_receive {:iter, _} + + eager = Nx.Defn.jit(&hook_in_while_body/1, compiler: Nx.Defn.Evaluator).(a) + assert Nx.to_number(eager) == 6 + assert_receive {:iter, e1} + assert_receive {:iter, e2} + assert_receive {:iter, e3} + eager_values = Enum.map([e1, e2, e3], &Nx.to_number/1) + refute_receive {:iter, _} + + assert native_values == eager_values + end + + # Regression: hooks straddling a non-bare `while` (surrounding work on + # both sides) route through `Nx.Defn.Graph.split`'s multi-stage chain + # (`EMLX.build_while_chain_eval_fn`). This previously crashed the NIF + # with a wire-arity mismatch ("vector in NIF.eval_program/2") because + # `Nx.Defn.Graph`'s `do_rewrite_subtree/3` had no `:token` clause, so a + # hook payload depending on a stage-boundary-hoisted value kept its + # stale, pre-remap parameter position. Fixed upstream in the `nx` fork + # (see Results); this test pins the fix from the EMLX side. + test "a hook before AND after a while (Graph.split chain) matches Evaluator" do + a = Nx.tensor(2, backend: EMLX.Backend) + + native = Nx.Defn.jit(&hook_around_while/1, compiler: EMLX).(a) + assert Nx.to_number(native) == 11 + assert_receive {:seed, seed_v} + assert_receive {:final, final_v} + assert Nx.to_number(seed_v) == 4 + assert Nx.to_number(final_v) == 11 + + eager = Nx.Defn.jit(&hook_around_while/1, compiler: Nx.Defn.Evaluator).(a) + assert Nx.to_number(eager) == 11 + assert_receive {:seed, seed_v2} + assert_receive {:final, final_v2} + assert Nx.to_number(seed_v2) == 4 + assert Nx.to_number(final_v2) == 11 + end + + # Regression for a bug the reviewer subagent caught: a hook inside a + # custom-fun `reduce` body was wrongly rejected by the cond-branch guard + # (false positive -- `scope_ids/1` never walks into a `:fun` body, so the + # hook's id looked "not top-scope" even though it isn't cond-branch-local + # either), AND separately, `lower_fun_body/3` was dropping any hooks + # registered while lowering the body (only `instructions`/`captures`/ + # `constants`/`inputs` were carried out of its local `state`, not `hooks`). + test "a hook inside a custom-fun reduce body fires once per fold step, matching Evaluator" do + t = Nx.tensor([1, 2, 3], backend: EMLX.Backend) + + native = Nx.Defn.jit(&hook_in_reduce_body/1, compiler: EMLX).(t) + assert Nx.to_number(native) == 6 + assert_receive {:step, v1} + assert_receive {:step, v2} + assert_receive {:step, v3} + native_values = Enum.map([v1, v2, v3], &Nx.to_number/1) + refute_receive {:step, _} + assert native_values == [1, 3, 6] + + bin_t = Nx.backend_copy(Nx.tensor([1, 2, 3]), Nx.BinaryBackend) + prev = Nx.default_backend(Nx.BinaryBackend) + + eager = + try do + Nx.Defn.jit(&hook_in_reduce_body/1, compiler: Nx.Defn.Evaluator).(bin_t) + after + Nx.default_backend(prev) + end + + assert Nx.to_number(eager) == 6 + assert_receive {:step, e1} + assert_receive {:step, e2} + assert_receive {:step, e3} + eager_values = Enum.map([e1, e2, e3], &Nx.to_number/1) + refute_receive {:step, _} + + assert native_values == eager_values + end + + test "a hook inside a cond that's nested inside a reduce body still raises" do + t = Nx.tensor(1, backend: EMLX.Backend) + template = Nx.template({3}, t.type) + + expr = Nx.Defn.debug_expr_apply(&hook_in_cond_in_reduce_body/1, [template]) + + assert_raise ArgumentError, ~r/cannot lower a hook nested inside a cond branch/, fn -> + Expr.lower(expr, 1) + end + end + end + + describe "quantized Nx.dot input (root-caused)" do + test "the same defn runs correctly under Nx.Defn.Evaluator (not a model/graph bug)" do + weight = + Nx.iota({128, 64}, type: :f32) |> Nx.divide(100) |> Nx.backend_transfer(EMLX.Backend) + + qw = EMLX.quantize(weight, []) + x = Nx.iota({4, 128}, type: :f32) |> Nx.backend_transfer(EMLX.Backend) + + result = Nx.Defn.jit(&Nx.dot/2, compiler: Nx.Defn.Evaluator).(x, qw) + assert Nx.shape(result) == {4, 64} + end + end + + describe "full fix: quantized Nx.dot via call-time program specialization" do + # See workdir/native-compiler/25-quantized-dot-full-fix.md. A quantized + # right-operand `Nx.dot` now specializes to a `:quantized_matmul` opcode + # once real (call-time) tensors reveal the quantization signature — + # equivalence-tested against eager EMLX.Backend.dot/7 and + # Nx.Defn.Evaluator, both used as independent references. + test "a quantized weight bound to a native-compiled defn now runs end-to-end" do + weight = + Nx.iota({128, 64}, type: :f32) |> Nx.divide(100) |> Nx.backend_transfer(EMLX.Backend) + + qw = EMLX.quantize(weight, []) + x = Nx.iota({4, 128}, type: :f32) |> Nx.divide(37) |> Nx.backend_transfer(EMLX.Backend) + + native = Nx.Defn.jit(&Nx.dot/2, compiler: EMLX).(x, qw) + eager = EMLX.Backend.dot(Nx.template({4, 64}, native.type), x, [1], [], qw, [0], []) + + assert Nx.shape(native) == {4, 64} + assert_all_close(native, eager) + end + + test "repeated calls with the same quantized weight reuse the cached specialized program" do + weight = + Nx.iota({128, 64}, type: :f32) |> Nx.divide(100) |> Nx.backend_transfer(EMLX.Backend) + + qw = EMLX.quantize(weight, []) + x1 = Nx.iota({4, 128}, type: :f32) |> Nx.divide(37) |> Nx.backend_transfer(EMLX.Backend) + x2 = Nx.iota({4, 128}, type: :f32) |> Nx.divide(11) |> Nx.backend_transfer(EMLX.Backend) + + jitted = Nx.Defn.jit(&Nx.dot/2, compiler: EMLX) + r1 = jitted.(x1, qw) + r2 = jitted.(x2, qw) + + eager1 = EMLX.Backend.dot(Nx.template({4, 64}, r1.type), x1, [1], [], qw, [0], []) + eager2 = EMLX.Backend.dot(Nx.template({4, 64}, r2.type), x2, [1], [], qw, [0], []) + + assert_all_close(r1, eager1) + assert_all_close(r2, eager2) + end + + test "a defn with two independently-quantized weights specializes both dots" do + w1 = + Nx.iota({128, 64}, type: :f32) |> Nx.divide(100) |> Nx.backend_transfer(EMLX.Backend) + + w2 = + Nx.iota({128, 64}, type: :f32) |> Nx.divide(50) |> Nx.backend_transfer(EMLX.Backend) + + qw1 = EMLX.quantize(w1, []) + qw2 = EMLX.quantize(w2, []) + x = Nx.iota({4, 128}, type: :f32) |> Nx.divide(37) |> Nx.backend_transfer(EMLX.Backend) + + native = Nx.Defn.jit(&two_quantized_dots/3, compiler: EMLX).(x, qw1, qw2) + + eager1 = EMLX.Backend.dot(Nx.template({4, 64}, native.type), x, [1], [], qw1, [0], []) + eager2 = EMLX.Backend.dot(Nx.template({4, 64}, native.type), x, [1], [], qw2, [0], []) + expected = Nx.concatenate([eager1, eager2], axis: 1) + + assert_all_close(native, expected) + end + + test "a microscaled-quantized weight (no biases) specializes correctly" do + weight = + Nx.iota({128, 64}, type: :f32) |> Nx.divide(100) |> Nx.backend_transfer(EMLX.Backend) + + qw = EMLX.quantize(weight, mode: "mxfp4", group_size: 32) + refute qw.data.quantization_config.biases + + x = Nx.iota({4, 128}, type: :f32) |> Nx.divide(37) |> Nx.backend_transfer(EMLX.Backend) + + native = Nx.Defn.jit(&Nx.dot/2, compiler: EMLX).(x, qw) + eager = EMLX.Backend.dot(Nx.template({4, 64}, native.type), x, [1], [], qw, [0], []) + + assert_all_close(native, eager) + end + + test "a quantized weight threaded through unchanged (pass-through output) round-trips" do + # Mirrors the shape a `while`-split stage boundary produces: a + # loop-invariant quantized weight carried through without being + # consumed by any op in *this* stage — its output leaf is a bare + # :parameter node. See build_native_eval_fn/5's output_param_positions. + weight = + Nx.iota({128, 64}, type: :f32) |> Nx.divide(100) |> Nx.backend_transfer(EMLX.Backend) + + qw = EMLX.quantize(weight, []) + x = Nx.iota({4, 128}, type: :f32) |> Nx.divide(37) |> Nx.backend_transfer(EMLX.Backend) + + passthrough_and_dot = fn x, qw -> {qw, Nx.dot(x, qw)} end + + {qw_out, native} = Nx.Defn.jit(passthrough_and_dot, compiler: EMLX).(x, qw) + eager = EMLX.Backend.dot(Nx.template({4, 64}, native.type), x, [1], [], qw, [0], []) + + assert Nx.shape(qw_out) == Nx.shape(qw) + assert qw_out.data.quantization_config + assert_all_close(EMLX.dequantize(qw_out), EMLX.dequantize(qw)) + assert_all_close(native, eager) + end + + test "a quantized left operand still raises a clear ArgumentError (matches EMLX.Backend.dot/7)" do + weight = + Nx.iota({128, 64}, type: :f32) |> Nx.divide(100) |> Nx.backend_transfer(EMLX.Backend) + + qw = EMLX.quantize(weight, []) + y = Nx.iota({64, 4}, type: :f32) |> Nx.backend_transfer(EMLX.Backend) + + assert_raise ArgumentError, + ~r/does not yet lower op :dot with a quantized left operand/, + fn -> + Nx.Defn.jit(&Nx.dot/2, compiler: EMLX).(qw, y) + end + end + end + + describe "unrecognized :runtime_call as a graph-split point (like while)" do + # See workdir/native-compiler/31-runtime-call-graph-split.md. An + # unrecognized :runtime_call (e.g. EMLX.Quantization.dequantize's + # callback, not an EMLX.Fast.* fused kernel) previously hard-raised + # "does not yet lower op :runtime_call". It now becomes a + # Nx.Defn.Graph.split point exactly like `while`: the callback runs once, + # directly, with real materialised tensors, so ordinary surrounding + # `compiler: EMLX` computation still compiles to flat native programs. + test "a bare tail runtime_call (no surrounding work) runs directly" do + weight = + Nx.iota({4, 64}, type: :f32) |> Nx.divide(10) |> Nx.backend_transfer(EMLX.Backend) + + qw = EMLX.quantize(weight, []) + + native = Nx.Defn.jit(&dequant_only/1, compiler: EMLX).(qw) + eager = EMLX.dequantize(qw) + + assert_all_close(native, eager) + end + + test "a runtime_call surrounded by ordinary ops splits into a native/host/native chain" do + weight = + Nx.iota({4, 64}, type: :f32) |> Nx.divide(10) |> Nx.backend_transfer(EMLX.Backend) + + qw = EMLX.quantize(weight, []) + x = Nx.iota({4, 64}, type: :f32) |> Nx.divide(3) |> Nx.backend_transfer(EMLX.Backend) + + native = Nx.Defn.jit(&dequant_surrounded/2, compiler: EMLX).(x, qw) + eager = Nx.Defn.jit(&dequant_surrounded/2, compiler: Nx.Defn.Evaluator).(x, qw) + + assert_all_close(native, eager) + end + + test "two independent runtime_calls in one defn both split correctly" do + w1 = Nx.iota({4, 64}, type: :f32) |> Nx.divide(10) |> Nx.backend_transfer(EMLX.Backend) + w2 = Nx.iota({4, 64}, type: :f32) |> Nx.divide(5) |> Nx.backend_transfer(EMLX.Backend) + qw1 = EMLX.quantize(w1, []) + qw2 = EMLX.quantize(w2, []) + x = Nx.iota({4, 64}, type: :f32) |> Nx.divide(3) |> Nx.backend_transfer(EMLX.Backend) + + native = Nx.Defn.jit(&two_runtime_calls/3, compiler: EMLX).(x, qw1, qw2) + eager = Nx.Defn.jit(&two_runtime_calls/3, compiler: Nx.Defn.Evaluator).(x, qw1, qw2) + + assert_all_close(native, eager) + end + + test "a runtime_call inside a while body re-enters the compiler correctly" do + weight = + Nx.iota({2, 64}, type: :f32) |> Nx.divide(10) |> Nx.backend_transfer(EMLX.Backend) + + qw = EMLX.quantize(weight, []) + x0 = Nx.broadcast(Nx.tensor(0.0, type: :f32, backend: EMLX.Backend), {2, 64}) + n = Nx.tensor(3, type: :s32, backend: EMLX.Backend) + + native = Nx.Defn.jit(&runtime_call_inside_while/3, compiler: EMLX).(x0, n, qw) + eager = Nx.Defn.jit(&runtime_call_inside_while/3, compiler: Nx.Defn.Evaluator).(x0, n, qw) + + assert_all_close(native, eager) + end + + test "a runtime_call with a tuple (multi-tensor) operand container splits correctly" do + # Unlike dequantize's single bare-tensor operand, quantized_matmul's + # runtime_call operand is a tuple {activation, qw} -- the same + # multi-tensor-container shape as the real target use case + # (EMLXAxon.native_kv_attn_callback/2), which no other test here + # exercises. + w = Nx.iota({4, 64}, type: :f32) |> Nx.divide(10) |> Nx.backend_transfer(EMLX.Backend) + qw = EMLX.quantize(w, []) + x = Nx.iota({2, 64}, type: :f32) |> Nx.divide(3) |> Nx.backend_transfer(EMLX.Backend) + + native = Nx.Defn.jit(&quantized_matmul_surrounded/2, compiler: EMLX).(x, qw) + eager = Nx.Defn.jit(&quantized_matmul_surrounded/2, compiler: Nx.Defn.Evaluator).(x, qw) + + assert_all_close(native, eager) + end + + test "a recognized EMLX.Fast.* runtime_call is still lowered in-graph, not split" do + fun = fn x, w -> EMLX.Fast.rms_norm(x, w, 1.0e-6) end + expr = Nx.Defn.debug_expr_apply(fun, [Nx.template({2, 4}, :f32), Nx.template({4}, :f32)]) + prog = Expr.lower(expr) + + assert [{_, :fast_rms_norm, [_x, _w], [_eps_bits]}] = prog.instructions + + x = Nx.iota({2, 4}, type: :f32, backend: EMLX.Backend) + w = Nx.broadcast(Nx.tensor(1.0, type: :f32, backend: EMLX.Backend), {4}) + native = Nx.Defn.jit(fun, compiler: EMLX).(x, w) + eager = fun.(x, w) + + assert_all_close(native, eager) + end + end + + defp dispatch_cache_entries_mentioning(shape) do + :emlx_native_dispatch_cache + |> :ets.tab2list() + |> Enum.filter(fn {key, _resource, _hooks, _runtime_calls} -> term_mentions?(key, shape) end) + |> Enum.map(&elem(&1, 0)) + |> Enum.uniq() + end + + defp term_mentions?(term, needle) when term == needle, do: true + + defp term_mentions?(term, needle) when is_tuple(term) do + term |> Tuple.to_list() |> term_mentions?(needle) + end + + defp term_mentions?(term, needle) when is_list(term) do + Enum.any?(term, &term_mentions?(&1, needle)) + end + + defp term_mentions?(term, needle) when is_map(term) do + term |> Map.to_list() |> term_mentions?(needle) + end + + defp term_mentions?(_term, _needle), do: false + + describe "runtime_call split-point dispatch cache (compile once, reuse)" do + test "calling the same runtime_call-split defn twice compiles its flat stages once" do + shape = {19, 128} + w1 = Nx.iota(shape, type: :f32) |> Nx.divide(10) |> Nx.backend_transfer(EMLX.Backend) + w2 = Nx.iota(shape, type: :f32) |> Nx.divide(7) |> Nx.backend_transfer(EMLX.Backend) + qw1 = EMLX.quantize(w1, []) + qw2 = EMLX.quantize(w2, []) + x = Nx.iota(shape, type: :f32) |> Nx.divide(3) |> Nx.backend_transfer(EMLX.Backend) + + native1 = Nx.Defn.jit(&dequant_surrounded/2, compiler: EMLX).(x, qw1) + entries_after_first = dispatch_cache_entries_mentioning(shape) + assert entries_after_first != [] + + native2 = Nx.Defn.jit(&dequant_surrounded/2, compiler: EMLX).(x, qw2) + entries_after_second = dispatch_cache_entries_mentioning(shape) + + assert entries_after_second == entries_after_first + + eager1 = Nx.Defn.jit(&dequant_surrounded/2, compiler: Nx.Defn.Evaluator).(x, qw1) + eager2 = Nx.Defn.jit(&dequant_surrounded/2, compiler: Nx.Defn.Evaluator).(x, qw2) + assert_all_close(native1, eager1) + assert_all_close(native2, eager2) + end + + test "two structurally-identical but distinct call sites share one cache entry" do + shape = {23, 192} + w1 = Nx.iota(shape, type: :f32) |> Nx.divide(10) |> Nx.backend_transfer(EMLX.Backend) + w2 = Nx.iota(shape, type: :f32) |> Nx.divide(7) |> Nx.backend_transfer(EMLX.Backend) + qw1 = EMLX.quantize(w1, []) + qw2 = EMLX.quantize(w2, []) + x = Nx.iota(shape, type: :f32) |> Nx.divide(3) |> Nx.backend_transfer(EMLX.Backend) + + # Two separately-defined defns with the identical op sequence stand in + # for "two of Qwen3's 28 structurally-identical attention layers": + # each traces to a fresh `Expr` (different ids), but the same shape of + # computation. + native1 = Nx.Defn.jit(&dequant_surrounded/2, compiler: EMLX).(x, qw1) + native2 = Nx.Defn.jit(&dequant_surrounded_other_site/2, compiler: EMLX).(x, qw2) + + entries = dispatch_cache_entries_mentioning(shape) + assert length(entries) == 1 + + eager1 = Nx.Defn.jit(&dequant_surrounded/2, compiler: Nx.Defn.Evaluator).(x, qw1) + eager2 = Nx.Defn.jit(&dequant_surrounded_other_site/2, compiler: Nx.Defn.Evaluator).(x, qw2) + assert_all_close(native1, eager1) + assert_all_close(native2, eager2) + end + end + + describe "native block dispatch-key caching (skips default_expr for recognized blocks)" do + test "svd full_matrices?: true: distinct retraces of the same shape share one dispatch-key entry" do + shape = {17, 17} + a1 = Nx.iota(shape, type: :f32) |> Nx.divide(3) |> Nx.backend_transfer(EMLX.Backend) + + a2 = + Nx.iota(shape, type: :f32) + |> Nx.divide(7) + |> Nx.add(1) + |> Nx.backend_transfer(EMLX.Backend) + + fun = fn t -> Nx.LinAlg.svd(t, full_matrices?: true) end + + # `Nx.Defn.jit_apply/3` retraces `fun` from scratch on every call (fresh + # `Expr` ids each time), so `dispatch_key/3`'s id-based fast path always + # misses and the structural signature walk always runs -- this is the + # retrace-per-call pattern the PRD targets. Both calls must still land + # in the *same* dispatch-cache entry: `default_expr`'s content (data- + # independent, since tracing is symbolic) was never what discriminated + # these before this change either, so this also guards against a + # regression that would split them. + {u1, s1, vt1} = Nx.Defn.jit_apply(fun, [a1], compiler: EMLX) + {u2, s2, vt2} = Nx.Defn.jit_apply(fun, [a2], compiler: EMLX) + + entries = dispatch_cache_entries_mentioning(shape) + assert length(entries) == 1 + + assert_close(Nx.dot(Nx.multiply(u1, Nx.reshape(s1, {1, 17})), vt1), a1) + assert_close(Nx.dot(Nx.multiply(u2, Nx.reshape(s2, {1, 17})), vt2), a2) + end + + test "svd full_matrices?: false: default_expr still discriminates the dispatch key" do + # full_matrices?: false lowers via `expand_block_via_default`'s + # while-unroll (see "while-in-default_expr descent" above) -- it + # genuinely consults `default_expr`, so two traces whose `default_expr` + # differs (here, via `max_iter`, which changes the unrolled Jacobi + # iteration count) must land in *different* dispatch-cache entries. + # G3 regression guard: `native_lowerable_block?/2` must not (and does + # not) match `full_matrices?: false`, so this path is untouched by P1. + shape = {13, 13} + a = Nx.iota(shape, type: :f32) |> Nx.add(1) |> Nx.backend_transfer(EMLX.Backend) + + fun_50 = fn t -> Nx.LinAlg.svd(t, full_matrices?: false, max_iter: 50) end + fun_100 = fn t -> Nx.LinAlg.svd(t, full_matrices?: false, max_iter: 100) end + + {u1, s1, vt1} = Nx.Defn.jit_apply(fun_50, [a], compiler: EMLX) + entries_after_50 = dispatch_cache_entries_mentioning(shape) + + {u2, s2, vt2} = Nx.Defn.jit_apply(fun_100, [a], compiler: EMLX) + entries_after_both = dispatch_cache_entries_mentioning(shape) + + assert length(entries_after_both) > length(entries_after_50) + + # Loose tolerance: the Jacobi rotation algorithm's reconstruction + # accuracy at these iteration counts isn't what this test is about -- + # only that both traces still compute a valid SVD. + assert_close(Nx.dot(Nx.multiply(u1, Nx.reshape(s1, {1, 13})), vt1), a, 1.0e-2) + assert_close(Nx.dot(Nx.multiply(u2, Nx.reshape(s2, {1, 13})), vt2), a, 1.0e-2) + end + end + + # Softmax normalisation over the last axis (primitive SDPA reference helper). + defp normalize_rows(t) do + Nx.divide(t, Nx.sum(t, axes: [-1], keep_axes: true)) + end + + # Pure-Nx primitive reference for prefill RoPE against a precomputed `freqs` + # tensor (mirrors emlx_compiler.cpp's fast_rope_with_freqs_positions lambda: + # inv_freq = reciprocal(freqs), half-rotate cos/sin blend). Needed for H>1 + # cases because the eager EMLX.Fast.rope_with_freqs reference calls + # mlx::core::fast::rope directly, which miscomputes non-head-0 rotations for + # multi-head input — see mlx-fast-rope-multihead-bugreport.md. `a` is + # {B, T, H, D}; `pos` is {B, T}; `freqs` is {dims/2}. + defp rope_freqs_prim(a, pos, freqs, dims, scale) do + {b, t, _h, d} = Nx.shape(a) + half = div(dims, 2) + + inv_freq = Nx.divide(1.0, freqs) |> Nx.reshape({1, 1, half}) + pos_bt1 = Nx.as_type(pos, :f32) |> Nx.reshape({b, t, 1}) + angles = Nx.multiply(Nx.multiply(pos_bt1, inv_freq), scale) + + cos_full = + Nx.cos(angles) |> Nx.reshape({b, t, 1, half}) |> then(&Nx.concatenate([&1, &1], axis: 3)) + + sin_full = + Nx.sin(angles) |> Nx.reshape({b, t, 1, half}) |> then(&Nx.concatenate([&1, &1], axis: 3)) + + x1 = a[[.., .., .., 0..(half - 1)]] + x2 = a[[.., .., .., half..(dims - 1)]] + rotated = Nx.concatenate([Nx.negate(x2), x1], axis: 3) + + a_head = a[[.., .., .., 0..(dims - 1)]] + rope_head = Nx.add(Nx.multiply(a_head, cos_full), Nx.multiply(rotated, sin_full)) + + if dims == d do + rope_head + else + tail = a[[.., .., .., dims..(d - 1)]] + Nx.concatenate([rope_head, tail], axis: 3) + end + end + + defp gpu_t(t), do: Nx.backend_transfer(t, {EMLX.Backend, device: :gpu}) + + defp unwrap!({:ok, v}), do: v + defp unwrap!({:error, e}), do: raise(EMLX.NIFError, List.to_string(e)) + + defp await_worker!(job_ref) do + receive do + {^job_ref, :ok} -> :ok + {^job_ref, {:ok, result}} -> result + {^job_ref, {:error, reason}} -> raise(EMLX.NIFError, List.to_string(reason)) + end + end + + # Compile and eval a lowered Expr program via the C++ NIF. + # Returns a list of output Nx.Tensor{} values. + defp run_nif(%Expr{} = prog, inputs) do + device = EMLX.default_device() + {worker, _} = EMLX.resolve_worker(device) + wire = Expr.to_native(prog) + + input_refs = + Enum.map(inputs, fn %Nx.Tensor{data: %EMLX.Backend{ref: {_, r}}} -> r end) + + prog_ref = compile_nif!(worker, wire) + out_refs = eval_nif!(worker, prog_ref, input_refs) + + Enum.map(out_refs, fn ref -> EMLX.Backend.to_nx({device, ref}) end) + end + + # Compare complex tensors by comparing real and imaginary parts separately. + defp assert_complex_close(a, b, tol \\ 1.0e-4) do + assert_all_close(Nx.real(a), Nx.real(b), tol: tol) + assert_all_close(Nx.imag(a), Nx.imag(b), tol: tol) + end + + defp assert_close(a, b), do: assert_close(a, b, 1.0e-4) + + defp assert_all_close(a, b, opts \\ []) do + tol = Keyword.get(opts, :tol, 1.0e-5) + + 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 + + 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 diff --git a/emlx/test/emlx/nx_doctest_test.exs b/emlx/test/emlx/nx_doctest_test.exs index 24bed31..810aa4b 100644 --- a/emlx/test/emlx/nx_doctest_test.exs +++ b/emlx/test/emlx/nx_doctest_test.exs @@ -66,7 +66,9 @@ defmodule EMLX.Nx.DoctestTest do population_count: 1, count_leading_zeros: 1, # f8_e4m3fn is not supported by MLX; skip all tensor/2 doctests - tensor: 2 + tensor: 2, + # reflect deprecated for pad_outer + reflect: 2 ] doctest Nx, diff --git a/emlx/test/emlx/nx_linalg_test.exs b/emlx/test/emlx/nx_linalg_test.exs index 4bd4705..886e589 100644 --- a/emlx/test/emlx/nx_linalg_test.exs +++ b/emlx/test/emlx/nx_linalg_test.exs @@ -9,7 +9,9 @@ defmodule EMLX.Nx.LinalgTest do # Ops that use defn fallbacks and may differ from LAPACK doctests: # - matrix_power: pure defn # - least_squares: no MLX native - # - triangular_solve: small float differences from solve_triangular vs LAPACK reference + # - triangular_solve: small float differences vs LAPACK reference, plus a + # handful of doctests using `Nx.vectorize/2`-batched inputs (unsupported — + # separate pre-existing gap, not related to left_side/transform_a) # - cholesky: small float differences # - lu: MLX raises on singular matrices (doctest uses [[1,2,3],[4,5,6],[7,8,9]]) # - qr: sign convention differences (-0.0 vs 0.0) and float precision @@ -90,6 +92,41 @@ defmodule EMLX.Nx.LinalgTest do end end + describe "native linalg: svd/1 (GPU eager fallback via defn_evaluator)" do + @describetag :metal + + test "full_matrices?: true on a GPU-resident tensor doesn't crash under compiler: Nx.Defn.Evaluator" do + # `EMLX.Backend.block/4` has no native `:gpu` kernel for `full_matrices?: + # true` SVD, so it falls back to eagerly running the plain-Nx composite + # algorithm (`Nx.LinAlg.SVD.svd/2`) op-by-op. That composite creates a + # literal tensor internally without an explicit `:backend` (`qdwh/2`'s + # `while`-loop condition seed, `Nx.iota({}, type: :u8, ...) + 1`), which + # used to land on `Nx.default_backend()` instead of matching the real + # (GPU) operand's backend/device -- crashing with a `FunctionClauseError` + # in `Nx.BinaryBackend.put_slice/5` the first time the composite's + # `while` loop tried to merge state across the mismatched backends. + # `EMLX.Backend.block/4`'s `eager_fallback/3,4` helper pins the default + # backend to the operand's own backend/device for the duration of the + # fallback to prevent this. + a = + Nx.iota({6, 6}, type: :f32) + |> Nx.add(Nx.eye(6)) + |> Nx.backend_transfer({EMLX.Backend, device: :gpu}) + + {u, s, vt} = + Nx.Defn.jit(fn t -> Nx.LinAlg.svd(t, full_matrices?: true) end, + compiler: Nx.Defn.Evaluator + ).(a) + + assert u.shape == {6, 6} + assert s.shape == {6} + assert vt.shape == {6, 6} + + recon = Nx.dot(Nx.multiply(u, Nx.reshape(s, {1, 6})), vt) + assert_all_close(recon, Nx.backend_transfer(a), rtol: 1.0e-3) + end + end + describe "native linalg: cholesky/1" do test "L @ L^T reconstruction" do a = diff --git a/emlx/test/emlx/nx_test.exs b/emlx/test/emlx/nx_test.exs index dbaa188..a4042d5 100644 --- a/emlx/test/emlx/nx_test.exs +++ b/emlx/test/emlx/nx_test.exs @@ -1052,19 +1052,21 @@ defmodule EMLX.NxTest do result, Nx.tensor([ [ - [15.0, 15.0], - [51.0, 51.0], - [87.0, 87.0] - ], - [ - [123.0, 123.0], - [159.0, 159.0], - [195.0, 195.0] - ], - [ - [231.0, 231.0], - [267.0, 267.0], - [303.0, 303.0] + [ + [15.0, 15.0], + [51.0, 51.0], + [87.0, 87.0] + ], + [ + [123.0, 123.0], + [159.0, 159.0], + [195.0, 195.0] + ], + [ + [231.0, 231.0], + [267.0, 267.0], + [303.0, 303.0] + ] ] ]) ) @@ -1080,7 +1082,7 @@ defmodule EMLX.NxTest do assert_equal( result, - Nx.tensor([[0, 0, -1, 1, 0, -2], [-3, 0, 4, -4, 0, 5]]) + Nx.tensor([[[[[0, 0, -1, 1, 0, -2], [-3, 0, 4, -4, 0, 5]]]]]) ) end end diff --git a/emlx/test/emlx/sdpa_sinks_test.exs b/emlx/test/emlx/sdpa_sinks_test.exs new file mode 100644 index 0000000..f3a32e3 --- /dev/null +++ b/emlx/test/emlx/sdpa_sinks_test.exs @@ -0,0 +1,287 @@ +defmodule EMLX.SdpaSinksTest do + @moduledoc """ + SDPA attention sinks. Equivalence-tests + `EMLX.Fast.scaled_dot_product_attention*`'s `:sinks` opt against the + fallback softmax-with-sinks math: + + row_max = max(reduce_max(logits), sinks_broadcast) + probs = exp(logits - row_max) / (sum(exp(logits - row_max)) + exp(sinks - row_max)) + """ + use EMLX.Case, async: true + + @moduletag :metal + + alias EMLX.Fast + + # `EMLX.Case`'s setup defaults to `EMLX.Backend`; override it to + # `Nx.BinaryBackend` here so all the plain (non-`gpu/1`) tensor construction + # and reference math below runs on plain Nx instead of round-tripping + # through MLX. `gpu/1` still explicitly transfers to `EMLX.Backend` + # regardless of the default. + setup do + Nx.default_backend(Nx.BinaryBackend) + :ok + end + + defp gpu(tensor), do: Nx.backend_transfer(tensor, {EMLX.Backend, device: :gpu}) + + # Reference (non-fused) softmax-with-sinks, operating on plain (BinaryBackend) Nx + # tensors. `logits_mask` is an optional {t_q, t_kv} boolean keep-mask (broadcast + # to {b, n, t_q, t_kv}); nil means no masking (full attention). + defp reference_sdpa_sinks(q, k, v, scale, sinks, logits_mask \\ nil) do + logits = Nx.dot(q, [3], [0, 1], k, [3], [0, 1]) |> Nx.multiply(scale) + {b, n, t_q, t_kv} = Nx.shape(logits) + + logits = + if logits_mask do + keep = Nx.reshape(logits_mask, {1, 1, t_q, t_kv}) |> Nx.broadcast({b, n, t_q, t_kv}) + neg_inf = Nx.Constants.neg_infinity(Nx.type(logits)) |> Nx.broadcast(Nx.shape(logits)) + Nx.select(keep, logits, neg_inf) + else + logits + end + + sinks_b = Nx.reshape(sinks, {1, n, 1, 1}) |> Nx.broadcast({b, n, t_q, 1}) + row_max = Nx.max(Nx.reduce_max(logits, axes: [3], keep_axes: true), sinks_b) + exp_logits = Nx.exp(Nx.subtract(logits, row_max)) + exp_sinks = Nx.exp(Nx.subtract(sinks_b, row_max)) + denom = Nx.add(Nx.sum(exp_logits, axes: [3], keep_axes: true), exp_sinks) + probs = Nx.divide(exp_logits, denom) + + Nx.dot(probs, [3], [0, 1], v, [2], [0, 1]) + end + + defp causal_keep_mask(t_q, t_kv) do + row = Nx.iota({t_q, t_kv}, axis: 0) + col = Nx.iota({t_q, t_kv}, axis: 1) + Nx.less_equal(col, row) + end + + # f32 tolerance for this exact fallback-vs-fused-kernel comparison + # (max-abs-diff observed ~2e-7). + @tol 1.0e-4 + + describe "EMLX.Fast.scaled_dot_product_attention/5 (opts, no mask) with :sinks" do + test "matches the softmax-with-sinks fallback math" do + b = 1 + n = 4 + t_q = 5 + t_kv = 5 + d = 8 + scale = 1.0 / :math.sqrt(d) + + q_cpu = Nx.iota({b, n, t_q, d}, type: :f32) |> Nx.divide(37) + k_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(41) + v_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(53) + sinks_cpu = Nx.tensor([0.1, -0.2, 0.5, 0.0], type: :f32) + + out_fast = + Fast.scaled_dot_product_attention(gpu(q_cpu), gpu(k_cpu), gpu(v_cpu), scale, + sinks: gpu(sinks_cpu) + ) + |> Nx.backend_transfer() + + out_ref = reference_sdpa_sinks(q_cpu, k_cpu, v_cpu, scale, sinks_cpu) + + assert_all_close(out_fast, out_ref, atol: @tol) + end + + test "GQA (N_q != N_kv): sinks shape follows N_q" do + b = 1 + n_q = 8 + n_kv = 4 + t_q = 3 + t_kv = 6 + d = 16 + scale = 1.0 / :math.sqrt(d) + + q = Nx.iota({b, n_q, t_q, d}, type: :f32) |> Nx.divide(29) |> gpu() + k = Nx.iota({b, n_kv, t_kv, d}, type: :f32) |> Nx.divide(31) |> gpu() + v = Nx.iota({b, n_kv, t_kv, d}, type: :f32) |> Nx.divide(37) |> gpu() + sinks = Nx.iota({n_q}, type: :f32) |> Nx.divide(10) |> gpu() + + out = Fast.scaled_dot_product_attention(q, k, v, scale, sinks: sinks) + + assert Nx.shape(out) == {b, n_q, t_q, d} + assert Nx.type(out) == {:f, 32} + end + + test "omitting :sinks is unchanged from the pre-sinks arity" do + b = 1 + n = 4 + t_q = 4 + t_kv = 4 + d = 16 + scale = 1.0 / :math.sqrt(d) + + q = Nx.iota({b, n, t_q, d}, type: :f32) |> gpu() + k = Nx.iota({b, n, t_kv, d}, type: :f32) |> gpu() + v = Nx.iota({b, n, t_kv, d}, type: :f32) |> gpu() + + out_4arity = Fast.scaled_dot_product_attention(q, k, v, scale) + out_empty_opts = Fast.scaled_dot_product_attention(q, k, v, scale, []) + + assert Nx.all_close(out_4arity, out_empty_opts) |> Nx.to_number() == 1 + end + end + + describe "EMLX.Fast.scaled_dot_product_attention/6 (mask + :sinks)" do + test "matches the softmax-with-sinks fallback math with an additive causal mask" do + b = 1 + n = 4 + t_q = 5 + t_kv = 5 + d = 8 + scale = 1.0 / :math.sqrt(d) + + q_cpu = Nx.iota({b, n, t_q, d}, type: :f32) |> Nx.divide(37) + k_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(41) + v_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(53) + sinks_cpu = Nx.tensor([0.1, -0.2, 0.5, 0.0], type: :f32) + + mask = + Nx.tril(Nx.broadcast(0.0, {t_q, t_kv})) + |> Nx.subtract(Nx.triu(Nx.broadcast(1.0e9, {t_q, t_kv}), k: 1)) + |> Nx.reshape({1, 1, t_q, t_kv}) + + out_fast = + Fast.scaled_dot_product_attention(gpu(q_cpu), gpu(k_cpu), gpu(v_cpu), scale, gpu(mask), + sinks: gpu(sinks_cpu) + ) + |> Nx.backend_transfer() + + out_ref = + reference_sdpa_sinks(q_cpu, k_cpu, v_cpu, scale, sinks_cpu, causal_keep_mask(t_q, t_kv)) + + assert_all_close(out_fast, out_ref, atol: @tol) + end + end + + describe "EMLX.Fast.scaled_dot_product_attention_causal/5 with :sinks" do + test "matches the softmax-with-sinks fallback math with an implicit causal mask" do + b = 1 + n = 4 + t_q = 5 + t_kv = 5 + d = 8 + scale = 1.0 / :math.sqrt(d) + + q_cpu = Nx.iota({b, n, t_q, d}, type: :f32) |> Nx.divide(37) + k_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(41) + v_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(53) + sinks_cpu = Nx.tensor([0.1, -0.2, 0.5, 0.0], type: :f32) + + out_fast = + Fast.scaled_dot_product_attention_causal(gpu(q_cpu), gpu(k_cpu), gpu(v_cpu), scale, + sinks: gpu(sinks_cpu) + ) + |> Nx.backend_transfer() + + out_ref = + reference_sdpa_sinks(q_cpu, k_cpu, v_cpu, scale, sinks_cpu, causal_keep_mask(t_q, t_kv)) + + assert_all_close(out_fast, out_ref, atol: @tol) + end + + test "omitting :sinks is unchanged from the pre-sinks arity" do + b = 1 + n = 2 + t_q = 3 + t_kv = 3 + d = 8 + scale = 1.0 / :math.sqrt(d) + + q = Nx.iota({b, n, t_q, d}, type: :f32) |> gpu() + k = Nx.iota({b, n, t_kv, d}, type: :f32) |> gpu() + v = Nx.iota({b, n, t_kv, d}, type: :f32) |> gpu() + + out_4arity = Fast.scaled_dot_product_attention_causal(q, k, v, scale) + out_empty_opts = Fast.scaled_dot_product_attention_causal(q, k, v, scale, []) + + assert Nx.all_close(out_4arity, out_empty_opts) |> Nx.to_number() == 1 + end + end + + describe "EMLX.Fast.scaled_dot_product_attention_causal_key_masked/6 with :sinks" do + test "trivial (all-ones) key_mask matches causal-with-sinks fallback math" do + b = 1 + n = 4 + t_q = 1 + t_kv = 5 + d = 8 + scale = 1.0 / :math.sqrt(d) + + q_cpu = Nx.iota({b, n, t_q, d}, type: :f32) |> Nx.divide(37) + k_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(41) + v_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(53) + sinks_cpu = Nx.tensor([0.1, -0.2, 0.5, 0.0], type: :f32) + key_mask_ones = Nx.broadcast(1, {b, t_kv}) + + out_fast = + Fast.scaled_dot_product_attention_causal_key_masked( + gpu(q_cpu), + gpu(k_cpu), + gpu(v_cpu), + scale, + gpu(key_mask_ones), + sinks: gpu(sinks_cpu) + ) + |> Nx.backend_transfer() + + # decode (T_q=1): kv_offset = T_kv - 1, so every key position is causally visible. + keep_mask = Nx.broadcast(1, {t_q, t_kv}) |> Nx.equal(1) + out_ref = reference_sdpa_sinks(q_cpu, k_cpu, v_cpu, scale, sinks_cpu, keep_mask) + + assert_all_close(out_fast, out_ref, atol: @tol) + end + + test "padded key_mask (non-trivial branch) matches combined causal+key_mask fallback math" do + b = 1 + n = 4 + t_q = 5 + t_kv = 5 + d = 8 + scale = 1.0 / :math.sqrt(d) + + q_cpu = Nx.iota({b, n, t_q, d}, type: :f32) |> Nx.divide(37) + k_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(41) + v_cpu = Nx.iota({b, n, t_kv, d}, type: :f32) |> Nx.divide(53) + sinks_cpu = Nx.tensor([0.1, -0.2, 0.5, 0.0], type: :f32) + key_mask_padded = Nx.tensor([[0, 1, 1, 1, 1]]) + + out_fast = + Fast.scaled_dot_product_attention_causal_key_masked( + gpu(q_cpu), + gpu(k_cpu), + gpu(v_cpu), + scale, + gpu(key_mask_padded), + sinks: gpu(sinks_cpu) + ) + |> Nx.backend_transfer() + + causal_keep = causal_keep_mask(t_q, t_kv) + km_keep = Nx.equal(key_mask_padded, 1) |> Nx.reshape({1, t_kv}) |> Nx.broadcast({t_q, t_kv}) + keep_mask = Nx.logical_and(causal_keep, km_keep) + + out_ref = reference_sdpa_sinks(q_cpu, k_cpu, v_cpu, scale, sinks_cpu, keep_mask) + + assert_all_close(out_fast, out_ref, atol: @tol) + end + + test "omitting :sinks is unchanged from the pre-sinks arity" do + q = Nx.broadcast(0.1, {1, 16, 1, 64}) |> Nx.as_type(:f16) |> gpu() + k = Nx.broadcast(0.1, {1, 8, 10, 64}) |> Nx.as_type(:f16) |> gpu() + v = Nx.broadcast(0.1, {1, 8, 10, 64}) |> Nx.as_type(:f16) |> gpu() + scale = 1.0 / :math.sqrt(64) + key_mask = Nx.broadcast(1, {1, 10}) |> gpu() + + out_5arity = Fast.scaled_dot_product_attention_causal_key_masked(q, k, v, scale, key_mask) + + out_empty_opts = + Fast.scaled_dot_product_attention_causal_key_masked(q, k, v, scale, key_mask, []) + + assert Nx.all_close(out_5arity, out_empty_opts) |> Nx.to_number() == 1 + end + end +end diff --git a/emlx/test/emlx/telemetry_test.exs b/emlx/test/emlx/telemetry_test.exs new file mode 100644 index 0000000..7e094e3 --- /dev/null +++ b/emlx/test/emlx/telemetry_test.exs @@ -0,0 +1,70 @@ +defmodule EMLX.TelemetryTest do + use EMLX.Case, async: false + doctest EMLX.Telemetry + + alias EMLX.Telemetry + + setup do + on_exit(fn -> + :telemetry.list_handlers([]) + |> Enum.each(&:telemetry.detach(&1.id)) + end) + + :ok + end + + # Module-captured handler to avoid telemetry's local-fn perf warning. + def forward_to_test(event, measurements, metadata, %{pid: pid}) do + send(pid, {:event, event, measurements, metadata}) + end + + defp attach(event, id) do + :telemetry.attach(id, event, &__MODULE__.forward_to_test/4, %{pid: self()}) + end + + describe "[:emlx, :eval, *]" do + test "fires start + stop around EMLX.eval/1" do + attach([:emlx, :eval, :start], "eval-start-#{inspect(self())}") + attach([:emlx, :eval, :stop], "eval-stop-#{inspect(self())}") + + t = Nx.tensor([1, 2, 3], backend: EMLX.Backend) + EMLX.eval(EMLX.Backend.from_nx(t)) + + assert_receive {:event, [:emlx, :eval, :start], _, %{}} + assert_receive {:event, [:emlx, :eval, :stop], %{duration: dur}, %{}} + assert is_integer(dur) and dur >= 0 + end + end + + describe "[:emlx, :to_binary, *]" do + test "fires on EMLX.Backend.to_binary/2 (the Nx.to_binary path)" do + attach([:emlx, :to_binary, :stop], "to-bin-stop-#{inspect(self())}") + + t = Nx.tensor([1.0, 2.0, 3.0], backend: EMLX.Backend) + _ = Nx.to_binary(t) + + assert_receive {:event, [:emlx, :to_binary, :stop], %{duration: dur}, + %{shape: shape, dtype: dtype, byte_size: bs}} + + assert shape == {3} + assert dtype == {:f, 32} + assert bs == 12 + assert is_integer(dur) and dur >= 0 + end + end + + describe "[:emlx, :memory, :stats]" do + test "memory_stats/0 emits active/peak/cache measurements" do + attach([:emlx, :memory, :stats], "mem-stats-#{inspect(self())}") + + result = Telemetry.memory_stats() + + assert %{active_memory: a, peak_memory: p, cache_memory: c} = result + assert is_integer(a) and a >= 0 + assert is_integer(p) and p >= 0 + assert is_integer(c) and c >= 0 + + assert_receive {:event, [:emlx, :memory, :stats], ^result, _} + end + end +end diff --git a/emlx/test/test_helper.exs b/emlx/test/test_helper.exs index a6ed859..051682b 100644 --- a/emlx/test/test_helper.exs +++ b/emlx/test/test_helper.exs @@ -1,3 +1,10 @@ +# MLX's CPU backend JIT-compiles fused kernels via `popen`/`pclose`, which +# fails with `ECHILD` under the BEAM's default SIGCHLD=SIG_IGN disposition +# (see `EMLX.Application`'s moduledoc for the full explanation). Library +# code doesn't get to make this call for consumers, but our own test suite +# is exactly the kind of deployment where we do control this trade-off. +:os.set_signal(:sigchld, :default) + use_gpu? = String.downcase(System.get_env("EMLX_TEST_DEFAULT_GPU", "false")) in [ "1", @@ -38,6 +45,7 @@ distributed_exclude = match?({:ok, _}, Node.start(:"emlx_primary@127.0.0.1", :longnames)) -> {:ok, _pid, peer} = :peer.start(%{name: :"emlx_peer@127.0.0.1"}) true = :erpc.call(peer, :code, :set_path, [:code.get_path()]) + :ok = :erpc.call(peer, :os, :set_signal, [:sigchld, :default]) :ok = :erpc.call(peer, Application, :put_env, [:emlx, :default_device, default_device]) :ok = :erpc.call(peer, Application, :put_env, [:nx, :default_backend, backend]) {:ok, _} = :erpc.call(peer, :application, :ensure_all_started, [:nx]) @@ -55,4 +63,11 @@ gpu_exclude = {:error, _} -> [:metal] end -ExUnit.start(exclude: distributed_exclude ++ gpu_exclude) +# debug_flags_functional_test.exs needs EMLX compiled with :detect_non_finites +# on (compile_env, baked in at compile time) — excluded unless that build was +# requested, so a normal `mix test` (flags off, per production defaults) +# doesn't fail on it. +debug_flags_exclude = + if System.get_env("EMLX_DEBUG_FLAGS") == "1", do: [], else: [:debug_flags_functional] + +ExUnit.start(exclude: distributed_exclude ++ gpu_exclude ++ debug_flags_exclude) diff --git a/emlx_axon/Makefile b/emlx_axon/Makefile new file mode 100644 index 0000000..88a8681 --- /dev/null +++ b/emlx_axon/Makefile @@ -0,0 +1,58 @@ +# Helper to escape spaces so that paths work as Make targets/prerequisites +# and survive expansion into shell recipes (backslash-escaped spaces are +# treated as literal spaces by both Make and the shell). +empty := +space := $(empty) $(empty) +esc = $(subst $(space),\$(space),$(1)) + +# Builds the standalone qwen3 compute plugin — pure MLX graph-building +# code, no erl_nif dependency whatsoever (see c_src/qwen3_plugin.cpp and +# emlx's c_src/emlx_fast/qwen3_plugin_abi.hpp). `dlopen`'d at runtime by +# emlx's host NIF via `EMLX.NIF.load_plugin("qwen3", path)` — see +# lib/emlx_axon/application.ex. +# +# MLX_INCLUDE_DIR/MLX_LIB_DIR point at emlx's own already-resolved libmlx +# copy (priv/mlx/{include,lib}) so this plugin links against the exact +# same libmlx, without duplicating emlx's version/cache-dir resolution +# logic — see mix.exs's `make_env`. +PRIV_DIR = $(MIX_APP_PATH)/priv +QWEN3_SO = $(PRIV_DIR)/libemlx_qwen3.so + +CFLAGS = -fPIC -I$(call esc,$(MLX_INCLUDE_DIR)) -I$(call esc,$(QWEN3_ABI_INCLUDE_DIR)) -Wall -std=c++20 -O3 +LDFLAGS = -L$(call esc,$(MLX_LIB_DIR)) -lmlx -shared + +UNAME_S := $(shell uname -s) +ifeq ($(UNAME_S),Darwin) + LDFLAGS += -undefined dynamic_lookup -flat_namespace -rpath @loader_path/mlx/lib + LDFLAGS += -framework Metal -framework Foundation -framework Accelerate +else + LDFLAGS += -Wl,-rpath,'$$ORIGIN/mlx/lib' +endif + +SOURCES = c_src/qwen3_plugin.cpp +HEADERS = $(call esc,$(QWEN3_ABI_INCLUDE_DIR))/qwen3_plugin_abi.hpp +OBJECTS = $(patsubst c_src/%.cpp,$(call esc,$(PRIV_DIR))/objs/%.o,$(SOURCES)) + +all: $(call esc,$(QWEN3_SO)) + @ echo > /dev/null + +$(call esc,$(PRIV_DIR)): + @ mkdir -p "$(PRIV_DIR)" + +$(call esc,$(PRIV_DIR))/objs/%.o: c_src/%.cpp $(HEADERS) + @ mkdir -p "$(dir $@)" + $(CXX) $(CFLAGS) -c "$<" -o "$@" + +# Emlx's own mlx lib copy (priv/mlx/lib, see MLX_LIB_DIR in mix.exs) must +# already exist — this plugin's rpath (`@loader_path/mlx/lib`) resolves +# relative to *its own* location, so we also copy the lib alongside it +# here rather than relying on emlx's priv dir at runtime. +$(call esc,$(QWEN3_SO)): $(call esc,$(PRIV_DIR)) $(OBJECTS) + @ mkdir -p "$(PRIV_DIR)/mlx/lib" + @ cp -a "$(MLX_LIB_DIR)"/* "$(PRIV_DIR)/mlx/lib" + $(CXX) $(OBJECTS) -o "$(QWEN3_SO)" $(LDFLAGS) + +clean: + rm -rf "$(PRIV_DIR)" + +.PHONY: all clean diff --git a/emlx_axon/bench/profile_eval.exs b/emlx_axon/bench/profile_eval.exs deleted file mode 100644 index 7a77605..0000000 --- a/emlx_axon/bench/profile_eval.exs +++ /dev/null @@ -1,131 +0,0 @@ -# bench/profile_eval.exs -# -# C01 measurement script — EMLX eval-dispatch profiling. -# -# Uses bumblebee-testing/tiny-random-LlamaForCausalLM for the KV-cache -# decode loop. The tiny model (hidden=32, layers=2, heads=2) has the same -# compute graph structure as a production Llama/Qwen3 — all the same NIF -# call patterns, just smaller. The counter measurements transfer directly. -# -# Run (from the emlx_axon/ directory): -# cd emlx_axon -# EMLX_PROFILE_EVAL=1 VQ_MAX_NEW=20 mix run bench/profile_eval.exs -# -# Scriptable Metal capture (no Xcode GUI): -# MTL_CAPTURE_ENABLED=1 METAL_DEVICE_WRAPPER_TYPE=1 \ -# xctrace record --output trace.xctrace \ -# --template "Metal System Trace" \ -# --time-limit 60s \ -# -- iex -S mix run bench/profile_eval.exs -# -# NIF call tracing (in a separate iex session on the same node): -# :recon_trace.calls({EMLX, :_, :_}, 10_000, scope: :local) -# :recon_trace.calls({EMLX.NIF, :_, :_}, 10_000, scope: :local) - -Nx.default_backend({EMLX.Backend, device: :gpu}) - -max_new = System.get_env("VQ_MAX_NEW", "20") |> String.to_integer() -seq_length = System.get_env("VQ_SEQ_LENGTH", "64") |> String.to_integer() - -unless EMLX.Profiling.enabled?() do - IO.puts(""" - [profile_eval] WARNING: EMLX_PROFILE_EVAL not set — counters are inactive. - Re-run with EMLX_PROFILE_EVAL=1 to enable atomic counters. - """) -end - -IO.puts(""" -==> C01 profiling max_new=#{max_new} seq_length=#{seq_length} - model: bumblebee-testing/tiny-random-LlamaForCausalLM -""") - -# ── Load model (no tokenizer needed — we supply input IDs directly) ──────────── - -{:ok, %{model: model, params: params, spec: spec}} = - Bumblebee.load_model({:hf, "bumblebee-testing/tiny-random-LlamaForCausalLM"}) - -IO.puts("==> Model loaded architecture=#{inspect(spec.architecture)} vocab_size=#{spec.vocab_size}") - -# ── JIT-compile the decode step ────────────────────────────────────────────── -# -# Axon.predict already goes through the Nx.Defn JIT path when a non-default -# compiler is configured via defn_options. Since we set the default backend -# to EMLX, we call Axon.predict directly and it dispatches through EMLX. - -# Ensure params live on the GPU backend (they should already, after load_model -# with default EMLX backend set, but make it explicit). -params = Nx.backend_transfer(params, {EMLX.Backend, device: :gpu}) - -predict = fn inputs -> - Axon.predict(model, params, inputs, compiler: EMLX) -end - -# Build initial inputs: batch=1, seq=seq_length (padded) -prefill_len = 5 -input_ids = Nx.broadcast(0, {1, seq_length}) |> Nx.put_slice([0, 0], Nx.tensor([[10, 20, 30, 40, 50]])) -ones = Nx.broadcast(1, {1, prefill_len}) -zeros = Nx.broadcast(0, {1, seq_length - prefill_len}) -attention_mask = Nx.concatenate([ones, zeros], axis: 1) - -IO.puts("==> Running prefill warmup (compilation) ...") -inputs = %{"input_ids" => input_ids, "attention_mask" => attention_mask} -_warmup = predict.(inputs) - -# ── Measure decode-like loop ─────────────────────────────────────────────────── -# -# Each "step" re-runs the full forward pass with the same input shape — -# this is what Nx.Defn.Evaluator does on every decode step (no caching of -# the compiled function's output). The NIF dispatch pattern per step is -# what we want to measure. - -IO.puts("==> Measuring #{max_new} decode steps ...") -EMLX.Profiling.reset() - -step_times = - for step <- 1..max_new do - # Shift the sequence by one token per step (simulates the decode position) - pos = prefill_len + step - 1 - shifted_ids = Nx.put_slice(input_ids, [0, pos], Nx.tensor([[100 + step]])) - shifted_mask = Nx.put_slice(attention_mask, [0, pos], Nx.tensor([[1]])) - inputs = %{"input_ids" => shifted_ids, "attention_mask" => shifted_mask} - - {step_us, outputs} = :timer.tc(fn -> predict.(inputs) end) - - # Extracting the next token — this calls to_blob/item, forcing eval - _next_token = outputs.logits[[0, pos, ..]] |> Nx.argmax() |> Nx.to_number() - - step_us / 1_000 - end - -{:ok, counts} = EMLX.Profiling.read() - -total_ms = Enum.sum(step_times) -median_ms = Enum.sort(step_times) |> Enum.at(div(max_new, 2)) - -IO.puts(""" - -==> Results (#{max_new} decode steps): - total wall time: #{round(total_ms)} ms - median step time: #{Float.round(median_ms, 2)} ms → #{Float.round(1_000 / median_ms, 1)} tok/s (approx) - EMLX.eval calls: #{counts.eval} (#{Float.round(counts.eval / max_new, 1)} / step) - EMLX.item calls: #{counts.item} (#{Float.round(counts.item / max_new, 1)} / step) - EMLX.to_blob calls: #{counts.to_blob} (#{Float.round(counts.to_blob / max_new, 1)} / step) - -==> Acceptance table (C01) — default :sdpa path - -| metric | tiny-Llama (proxy) | -|---------------------------------|--------------------| -| EMLX.eval calls / decode step | #{Float.round(counts.eval / max_new, 1)} | -| EMLX.item calls / decode step | #{Float.round(counts.item / max_new, 1)} | -| EMLX.to_blob calls / decode step| #{Float.round(counts.to_blob / max_new, 1)} | -| median step wall time (ms) | #{Float.round(median_ms, 2)} | - -NOTE: Metal cmd buffers / step requires Metal frame capture. - xctrace record --output trace.xctrace \\ - --template "Metal System Trace" ... - Then open in Instruments → Metal GPU Trace Viewer. - -INTERPRETATION: - eval/step ≈ 1 → MLX lazy eval is batching correctly → bottleneck is NIF overhead - eval/step ≫ 1 → MLX lazy eval interrupted mid-step → fix forced evals first -""") diff --git a/emlx_axon/bench/qwen3_e2e_bench.exs b/emlx_axon/bench/qwen3_e2e_bench.exs index 3910dac..c641beb 100644 --- a/emlx_axon/bench/qwen3_e2e_bench.exs +++ b/emlx_axon/bench/qwen3_e2e_bench.exs @@ -36,16 +36,21 @@ max_new = 100 IO.puts("==> Warming up (20 tokens greedy) ...") {_, _} = Generate.generate(input_ids, state, max_new_tokens: 20, sampler: :greedy) + IO.puts("==> Benchmarking #{max_new} tokens per sampler ...\n") bench_sampler = fn sampler, label -> {tokens, %{timing: t}} = - Generate.generate(input_ids, state, max_new_tokens: max_new, sampler: sampler) + Generate.generate(input_ids, state, + max_new_tokens: max_new, + sampler: sampler, + profile_timing: true + ) sorted_ms = Enum.sort(t.per_token_ms) n = length(sorted_ms) - median_ms = Enum.at(sorted_ms, div(n, 2)) - p95_ms = Enum.at(sorted_ms, floor(n * 0.95)) + median_ms = if n > 0, do: Enum.at(sorted_ms, div(n, 2)), else: 0.0 + p95_ms = if n > 0, do: Enum.at(sorted_ms, floor(n * 0.95)), else: 0.0 e2e_tok_s = if t.total_ms > 0, do: Float.round(length(tokens) / t.total_ms * 1000, 1), else: 0.0 kernel_tok_s = if median_ms > 0, do: Float.round(1000.0 / median_ms, 1), else: 0.0 diff --git a/emlx_axon/bench/validate_qwen3.exs b/emlx_axon/bench/validate_qwen3.exs deleted file mode 100644 index d70dcfe..0000000 --- a/emlx_axon/bench/validate_qwen3.exs +++ /dev/null @@ -1,237 +0,0 @@ -# emlx_axon/bench/validate_qwen3.exs -# -# Throughput benchmark for EMLXAxon.rewrite/2 on Qwen3-0.6B-MLX-4bit. -# Measures e2e tok/s for: -# - Bumblebee + stock Axon graph (no EMLXAxon.rewrite) -# - Bumblebee + full rewrite (default: all rewrites, including :native_attention) -# - EMLXAxon.TextGeneration (native bypass) -# -# Run from the emlx_axon directory: -# -# EMLX_QWEN3_MODEL_PATH=~/models/Qwen3-0.6B-MLX-4bit \ -# mix run bench/validate_qwen3.exs -# -# Environment variables (all optional): -# EMLX_QWEN3_MODEL_PATH — local model dir or HF repo id (default: ~/models/Qwen3-0.6B-MLX-4bit) -# EMLX_QWEN3_MAX_NEW — max new tokens per run (default: 60) -# EMLX_QWEN3_BENCH_RUNS — number of timed runs (default: 5) -# EMLX_QWEN3_WARMUP_RUNS — number of warmup runs before timing (default: 2) -# EMLX_QWEN3_SEQUENCE_LENGTH — sequence_length for Bumblebee compile (default: 1024) -# EMLX_QWEN3_NATIVE_PROFILE_TIMING — set to "0" to disable per-token monotonic_time in native decode (default: on) - -Nx.default_backend({EMLX.Backend, device: :gpu}) - -# ── Config ──────────────────────────────────────────────────────────────────── - -model_path_raw = - System.get_env("EMLX_QWEN3_MODEL_PATH") || "~/models/Qwen3-0.6B-MLX-4bit" - -model_source = - if File.exists?(Path.expand(model_path_raw)) do - {:local, Path.expand(model_path_raw)} - else - {:hf, model_path_raw} - end - -max_new = String.to_integer(System.get_env("EMLX_QWEN3_MAX_NEW", "60")) -bench_runs = String.to_integer(System.get_env("EMLX_QWEN3_BENCH_RUNS", "5")) -warmup_runs = String.to_integer(System.get_env("EMLX_QWEN3_WARMUP_RUNS", "2")) -seq_len = String.to_integer(System.get_env("EMLX_QWEN3_SEQUENCE_LENGTH", "1024")) - -native_profile_timing? = System.get_env("EMLX_QWEN3_NATIVE_PROFILE_TIMING") != "0" - -# Qwen3 instruct chat template — long enough that EOS won't hit within max_new tokens. -prompt = - "<|im_start|>user\nList twenty programming languages with a one-line description each.<|im_end|>\n<|im_start|>assistant\n" - -IO.puts(""" -=== emlx_axon/bench/validate_qwen3.exs === -model: #{inspect(model_source)} -max_new: #{max_new} bench_runs: #{bench_runs} warmup_runs: #{warmup_runs} -seq_len: #{seq_len} -""") - -# ── Load ────────────────────────────────────────────────────────────────────── - -IO.puts("==> Loading model structure ...") -t0 = System.monotonic_time(:millisecond) -# Load model spec/structure only — params are loaded separately below because -# the MLX-4bit checkpoint stores weights in a packed uint32 format that -# Bumblebee's safetensors loader cannot dequantize on its own. -{:ok, model_info} = Bumblebee.load_model(model_source, - backend: {EMLX.Backend, device: :gpu}, - type: :bf16) -t1 = System.monotonic_time(:millisecond) -IO.puts(" model structure loaded in #{t1 - t0} ms") - -IO.puts("==> Loading MLX-4bit params (dequantize → Bumblebee layout → re-quantize) ...") -t2a = System.monotonic_time(:millisecond) -params = EMLXAxon.MLX4BitParams.load(Path.expand(model_path_raw)) -params = EMLXAxon.QuantizeParams.quantize(params) -model_info = %{model_info | params: params} -t2b = System.monotonic_time(:millisecond) -IO.puts(" params loaded in #{t2b - t2a} ms") - -IO.puts("==> Loading tokenizer ...") -{:ok, tokenizer} = Bumblebee.load_tokenizer(model_source) - -IO.puts("==> Loading generation config ...") -{:ok, generation_config} = Bumblebee.load_generation_config(model_source) -generation_config = Bumblebee.configure(generation_config, - max_new_tokens: max_new, - strategy: %{type: :greedy_search} -) - -# ── Rewrite ─────────────────────────────────────────────────────────────────── - -model_base = model_info.model - -IO.puts("==> Applying EMLXAxon.rewrite/2 (all rewrites — best performing path) ...") -t3 = System.monotonic_time(:millisecond) -model_rewritten = EMLXAxon.rewrite(model_base) -t4 = System.monotonic_time(:millisecond) -IO.puts(" rewrite done in #{t4 - t3} ms") - -# ── Build serving ───────────────────────────────────────────────────────────── - -IO.puts("==> Building Bumblebee.Text.generation serving (base — no rewrite) ...") -serving_base = - Bumblebee.Text.generation( - %{model_info | model: model_base}, - tokenizer, - generation_config, - compile: [batch_size: 1, sequence_length: seq_len], - defn_options: [compiler: EMLX] - ) - -IO.puts("==> Building Bumblebee.Text.generation serving (rewritten) ...") -serving_rewrite = - Bumblebee.Text.generation( - %{model_info | model: model_rewritten}, - tokenizer, - generation_config, - compile: [batch_size: 1, sequence_length: seq_len], - defn_options: [compiler: EMLX] - ) - -# ── Benchmark helpers ───────────────────────────────────────────────────────── - -defmodule Bench do - def run_serving(serving, prompt, extract_fn) do - t_start = System.monotonic_time(:millisecond) - result = Nx.Serving.run(serving, prompt) - t_end = System.monotonic_time(:millisecond) - {text, n_tokens} = extract_fn.(result) - elapsed_ms = t_end - t_start - tok_s = if elapsed_ms > 0, do: Float.round(n_tokens / elapsed_ms * 1000.0, 1), else: 0.0 - {elapsed_ms, n_tokens, tok_s, text} - end - - def warmup(label, serving, prompt, extract_fn, runs) do - IO.puts("\n==> [#{label}] Warmup (#{runs} run(s), not timed) ...") - for i <- 1..runs do - {ms, n, _, text} = run_serving(serving, prompt, extract_fn) - preview = text |> String.slice(0, 60) |> String.replace("\n", " ") - IO.puts(" warmup #{i}: #{n} tokens / #{ms} ms \"#{preview}...\"") - end - end - - def bench(label, serving, prompt, extract_fn, runs) do - IO.puts("\n==> [#{label}] Benchmark (#{runs} run(s)) ...") - for i <- 1..runs do - {ms, n, tok_s, _} = run_serving(serving, prompt, extract_fn) - IO.puts(" run #{i}: #{n} tokens / #{ms} ms = #{tok_s} tok/s") - %{elapsed_ms: ms, n_tokens: n, tok_s: tok_s} - end - end - - def stats(label, results) do - toks_list = Enum.map(results, & &1.tok_s) - sorted = Enum.sort(toks_list) - n = length(sorted) - median = Enum.at(sorted, div(n, 2)) - mean = Float.round(Enum.sum(toks_list) / n, 1) - min_v = List.first(sorted) - max_v = List.last(sorted) - variance = Enum.reduce(toks_list, 0.0, fn x, acc -> acc + (x - mean) * (x - mean) end) / n - stddev = Float.round(:math.sqrt(variance), 1) - IO.puts(" #{label}: median=#{median} mean=#{mean}±#{stddev} min/max=#{min_v}/#{max_v} tok/s") - %{label: label, median: median, mean: mean, stddev: stddev, min: min_v, max: max_v} - end -end - -# ── Bumblebee paths (base vs rewrite) ─────────────────────────────────────── - -# Bumblebee.Text.generation returns %{results: [%{text: ..., token_summary: ...}]} -bb_extract = fn %{results: [%{text: text, token_summary: summary}]} -> - {text, summary.output} -end - -Bench.warmup("bb base", serving_base, prompt, bb_extract, warmup_runs) -base_bb_results = Bench.bench("bb base", serving_base, prompt, bb_extract, bench_runs) - -Bench.warmup("bb+rewrite", serving_rewrite, prompt, bb_extract, warmup_runs) -bb_results = Bench.bench("bb+rewrite", serving_rewrite, prompt, bb_extract, bench_runs) - -# ── Native TextGeneration path ──────────────────────────────────────────────── - -IO.puts("\n==> Loading EMLXAxon.TextGeneration (native 28-layer bypass) ...") -t_native = System.monotonic_time(:millisecond) -native_serving = - EMLXAxon.TextGeneration.from_mlx4bit( - Path.expand(model_path_raw), - tokenizer, - max_new_tokens: max_new, - sampler: :greedy, - profile_timing: native_profile_timing? - ) -IO.puts(" loaded in #{System.monotonic_time(:millisecond) - t_native} ms") - -# Native serving includes :num_tokens (tensor seq dim); avoids tokenizer round-trip for bench throughput. -native_extract = fn - %{results: [%{generated_text: text, num_tokens: n}]} -> - {text, n} - - %{results: [%{generated_text: text}]} -> - %{"input_ids" => ids} = - Nx.with_default_backend(Nx.BinaryBackend, fn -> - Bumblebee.apply_tokenizer(tokenizer, text) - end) - - {text, elem(Nx.shape(ids), 1)} -end - -Bench.warmup("native", native_serving, prompt, native_extract, warmup_runs) -native_results = Bench.bench("native", native_serving, prompt, native_extract, bench_runs) - -# ── Summary ─────────────────────────────────────────────────────────────────── - -IO.puts(""" - -=== Summary (#{bench_runs} runs, #{max_new} max_new_tokens, Qwen3-0.6B-MLX-4bit) === -""") - -base_stats = Bench.stats("bb base (Bumblebee, no rewrite) ", base_bb_results) -bb_stats = Bench.stats("bb+rewrite (Bumblebee + EMLXAxon.rewrite) ", bb_results) -native_stats = Bench.stats("native (EMLXAxon.TextGeneration) ", native_results) - -rewrite_vs_base = - if base_stats.median > 0, - do: Float.round(bb_stats.median / base_stats.median, 2), - else: :n_a - -native_vs_base = - if base_stats.median > 0, - do: Float.round(native_stats.median / base_stats.median, 2), - else: :n_a - -native_vs_rewrite = - if bb_stats.median > 0, - do: Float.round(native_stats.median / bb_stats.median, 2), - else: :n_a - -IO.puts(""" - bb+rewrite / bb base: #{rewrite_vs_base}× - native / bb base: #{native_vs_base}× - native / bb+rewrite: #{native_vs_rewrite}× -""") diff --git a/emlx_axon/bench/validate_qwen3_standalone.livemd b/emlx_axon/bench/validate_qwen3_standalone.livemd new file mode 100644 index 0000000..35998f0 --- /dev/null +++ b/emlx_axon/bench/validate_qwen3_standalone.livemd @@ -0,0 +1,1514 @@ + + +# Qwen3 EMLX / Bumblebee / Emily Benchmark + +```elixir +Mix.install( + [ + {:emlx, path: Path.expand("../../emlx", __DIR__)}, + {:emlx_axon, path: Path.expand("..", __DIR__)}, + # {:emlx, github: "elixir-nx/emlx", branch: "main", sparse: "emlx", override: true}, + # {:emlx_axon, github: "elixir-nx/emlx", branch: "main", sparse: "emlx_axon", override: true}, + + {:nx, github: "elixir-nx/nx", branch: "main", sparse: "nx", override: true}, + {:axon, "~> 0.7"}, + {:bumblebee, "~> 0.7"}, + {:emily, "~> 0.7"}, + {:kino, "~> 0.14"} + ], + config: [nx: [default_backend: {EMLX.Backend, device: :gpu}]] +) +``` + +## Overview + +Dense (bf16) and MLX-4bit Qwen3-0.6B checkpoints; lanes per checkpoint: + +* **bb base** — Bumblebee, stock Axon graph +* **bb+rewrite** — Bumblebee + `EMLXAxon.rewrite/1` +* **native** — `EMLXAxon.TextGeneration` +* **emily-eager** / **emily-native** — dense only; [Emily](https://github.com/ausimian/emily) backend/compiler, unmodified Bumblebee graph +* **emily-quantized** — mlx4bit entry only; Emily int4-quantizes the *dense* checkpoint (not the lmstudio mlx4bit file on disk) + +Each lane runs in a fresh `:peer` node (~150 ms overhead) so MLX allocator state does not bleed across lanes. + +## Peer helper + +```elixir +defmodule PeerBench.Registry do + @moduledoc false + use Agent + + def start_link(_), do: Agent.start_link(fn -> %{} end, name: __MODULE__) + + def put(module, binary) do + ensure_started() + Agent.update(__MODULE__, &Map.put(&1, module, binary)) + end + + def get(module) do + ensure_started() + Agent.get(__MODULE__, &Map.fetch!(&1, module)) + end + + defp ensure_started do + if is_nil(Process.whereis(__MODULE__)), do: start_link([]) + end +end + +defmodule PeerBench do + @moduledoc false + + defmacro defshippable(name, do: block) do + quote do + {:module, unquote(name), bin, _} = + defmodule unquote(name) do + unquote(block) + end + + PeerBench.Registry.put(unquote(name), bin) + end + end + + @apps [:nx, :emlx, :emlx_axon, :axon, :bumblebee, :emily] + + def run(modules, mod, fun, args) do + peer_args = Enum.flat_map(:code.get_path(), fn p -> [~c"-pa", p] end) + + {:ok, pid, _node} = + :peer.start_link(%{name: :peer.random_name(), args: peer_args, connection: :standard_io}) + + try do + for m <- modules do + bin = PeerBench.Registry.get(m) + {:module, ^m} = :peer.call(pid, :code, :load_binary, [m, ~c"nofile", bin], :infinity) + end + + for app <- @apps, do: :peer.call(pid, Application, :ensure_all_started, [app], :infinity) + + :peer.call(pid, :erlang, :apply, [mod, fun, args], :infinity) + after + :peer.stop(pid) + end + end +end + +:ok +``` + + + +``` +:ok +``` + +## Benchmark helpers + +```elixir +require PeerBench + +PeerBench.defshippable Bench do + def run_serving(serving, prompt, extract_fn) do + t_start = System.monotonic_time(:millisecond) + result = Nx.Serving.run(serving, prompt) + t_end = System.monotonic_time(:millisecond) + {text, n_tokens} = extract_fn.(result) + elapsed_ms = t_end - t_start + tok_s = if elapsed_ms > 0, do: Float.round(n_tokens / elapsed_ms * 1000.0, 1), else: 0.0 + {elapsed_ms, n_tokens, tok_s, text} + end + + def warmup(label, serving, prompt, extract_fn, runs) do + IO.puts("\n==> [#{label}] Warmup (#{runs} run(s), not timed) ...") + + for i <- 1..runs do + {ms, n, _, text} = run_serving(serving, prompt, extract_fn) + preview = text |> String.slice(0, 60) |> String.replace("\n", " ") + IO.puts(" warmup #{i}: #{n} tokens / #{ms} ms \"#{preview}...\"") + end + end + + def bench(label, serving, prompt, extract_fn, runs) do + IO.puts("\n==> [#{label}] Benchmark (#{runs} run(s)) ...") + + for i <- 1..runs do + {ms, n, tok_s, _} = run_serving(serving, prompt, extract_fn) + IO.puts(" run #{i}: #{n} tokens / #{ms} ms = #{tok_s} tok/s") + %{elapsed_ms: ms, n_tokens: n, tok_s: tok_s} + end + end + + def stats(label, results) do + toks_list = Enum.map(results, & &1.tok_s) + sorted = Enum.sort(toks_list) + n = length(sorted) + median = Enum.at(sorted, div(n, 2)) + mean = Float.round(Enum.sum(toks_list) / n, 1) + min_v = List.first(sorted) + max_v = List.last(sorted) + variance = Enum.reduce(toks_list, 0.0, fn x, acc -> acc + (x - mean) * (x - mean) end) / n + stddev = Float.round(:math.sqrt(variance), 1) + IO.puts(" #{label}: median=#{median} mean=#{mean}±#{stddev} min/max=#{min_v}/#{max_v} tok/s") + %{label: label, median: median, mean: mean, stddev: stddev, min: min_v, max: max_v} + end +end + +PeerBench.defshippable Extract do + def bb(%{results: [%{text: text, token_summary: summary}]}) do + {text, summary.output} + end + + def native(%{results: [%{generated_text: text, num_tokens: n}]}, _tokenizer) do + {text, n} + end + + def native(%{results: [%{generated_text: text}]}, tokenizer) do + %{"input_ids" => ids} = Bumblebee.apply_tokenizer(tokenizer, text) + {text, elem(Nx.shape(ids), 1)} + end +end + +PeerBench.defshippable ModelSource do + def resolve(path_raw) do + if File.exists?(Path.expand(path_raw)) do + {:local, Path.expand(path_raw)} + else + {:hf, path_raw} + end + end +end + +:ok +``` + + + +``` +:ok +``` + +## Emily lanes + +Quantize rewrite vendored from Emily's `livebooks/qwen3_quantized.livemd`. + +```elixir +require PeerBench + +PeerBench.defshippable EmilyQuantize do + alias Emily.Quantization.Layers + alias Emily.QuantizedWeight + + def quantize(model, model_state, opts \\ []) do + {rewrite_graph(model), quantize_state(model, model_state, opts)} + end + + defp rewrite_graph(model) do + Axon.rewrite_nodes(model, fn + %Axon.Node{op: :dense, meta: meta, name: name_fn} -> + fn [x], _output -> + quantized_dense_layer(x, meta[:units], use_bias: meta[:use_bias], name: name_fn) + end + + _ -> + :skip + end) + end + + defp quantized_dense_layer(x, units, opts) do + kernel_shape = fn input_shape -> + in_features = elem(input_shape, tuple_size(input_shape) - 1) + {in_features, units} + end + + bias_shape = fn _input_shape -> {units} end + + kernel = Axon.param("kernel", kernel_shape, initializer: :glorot_uniform) + + {inputs, op} = + if opts[:use_bias] do + bias = Axon.param("bias", bias_shape, initializer: :zeros) + {[x, kernel, bias], &Layers.quantized_dense/4} + else + {[x, kernel], &Layers.quantized_dense/3} + end + + Axon.layer(op, inputs, + name: opts[:name], + meta: %{units: units, use_bias: opts[:use_bias]}, + op_name: :quantized_dense + ) + end + + defp quantize_state(model, state, opts) do + transpose = Keyword.fetch!(opts, :transpose) + + dense_names = + model + |> Axon.properties() + |> Enum.filter(fn {_name, op} -> op == :dense end) + |> Enum.map(fn {name, _} -> name end) + + Enum.reduce(dense_names, state, fn name, acc -> + update_in(acc, [Access.key!(:data), name, "kernel"], fn kernel -> + source = if transpose, do: Nx.transpose(kernel), else: kernel + + QuantizedWeight.from_dense(source, + group_size: Keyword.fetch!(opts, :group_size), + bits: Keyword.fetch!(opts, :bits), + transpose: transpose + ) + end) + end) + end +end + +PeerBench.defshippable EmilyLane do + def run(kind, source, checkpoint_label, prompt, cfg) do + if cfg.skip_emily? do + nil + else + lane_label = "#{checkpoint_label} / emily-#{kind}" + Nx.default_backend(Emily.Backend) + + try do + IO.puts("\n==> [#{lane_label}] Loading model on Emily.Backend ...") + {:ok, model_info} = Bumblebee.load_model(source, backend: Emily.Backend, type: :bf16) + {:ok, tokenizer} = Bumblebee.load_tokenizer(source) + {:ok, generation_config} = Bumblebee.load_generation_config(source) + + generation_config = + Bumblebee.configure(generation_config, + max_new_tokens: cfg.max_new, + strategy: %{type: :greedy_search} + ) + + {defn_options, model_info} = build_defn_options(kind, model_info, cfg) + + serving = + Bumblebee.Text.generation(model_info, tokenizer, generation_config, defn_options: defn_options) + + Bench.warmup(lane_label, serving, prompt, &Extract.bb/1, cfg.warmup_runs) + Bench.bench(lane_label, serving, prompt, &Extract.bb/1, cfg.bench_runs) + rescue + exception -> + IO.puts("\n==> [#{lane_label}] FAILED: #{Exception.message(exception)}") + nil + end + end + end + + defp build_defn_options(:eager, model_info, _cfg) do + {[compiler: Nx.Defn.Evaluator], model_info} + end + + defp build_defn_options(:native, model_info, _cfg) do + {[compiler: Emily.Compiler, native: true, native_fallback: :raise], model_info} + end + + defp build_defn_options(:quantized, model_info, cfg) do + IO.puts( + " quantizing every Axon.dense node " <> + "(bits=#{cfg.emily_quant_bits}, group_size=#{cfg.emily_quant_group_size}) ..." + ) + + {qmodel, qparams} = + EmilyQuantize.quantize(model_info.model, model_info.params, + bits: cfg.emily_quant_bits, + group_size: cfg.emily_quant_group_size, + transpose: true + ) + + {[compiler: Emily.Compiler, native: true, native_fallback: :raise], + %{model_info | model: qmodel, params: qparams}} + end +end + +:ok +``` + + + +``` +:ok +``` + +## Checkpoint runner + +```elixir +require PeerBench + +PeerBench.defshippable Runner do + @shipped_modules [Bench, Extract, ModelSource, EmilyQuantize, EmilyLane, Runner] + + def run_checkpoint(%{kind: kind, label: label, path_raw: path_raw} = checkpoint, cfg) do + IO.puts(""" + + ################################################################################ + === #{label} — #{inspect(ModelSource.resolve(path_raw))} === + ################################################################################ + """) + + stats = + for lane <- applicable_lanes(kind, cfg), into: %{} do + IO.puts("\n==> Spawning isolated peer for lane #{inspect(lane)} ...") + {lane, PeerBench.run(@shipped_modules, Runner, :run_lane, [lane, checkpoint, cfg])} + end + + %{ + label: label, + base: stats[:bb_base], + rewrite: stats[:bb_rewrite], + native: stats[:native], + emily_eager: stats[:emily_eager], + emily_native: stats[:emily_native], + emily_quantized: stats[:emily_quantized] + } + end + + defp applicable_lanes(kind, cfg) do + [ + bb_base: not cfg.skip_bb_base?, + bb_rewrite: not cfg.skip_bb_rewrite?, + native: true, + emily_eager: kind == :dense and not cfg.skip_emily?, + emily_native: kind == :dense and not cfg.skip_emily?, + emily_quantized: kind == :mlx4bit and not cfg.skip_emily? + ] + |> Enum.filter(fn {_lane, enabled?} -> enabled? end) + |> Enum.map(fn {lane, _} -> lane end) + end + + def run_lane(lane, %{kind: kind, label: label, path_raw: path_raw}, cfg) do + Nx.default_backend({EMLX.Backend, device: :gpu}) + source = ModelSource.resolve(path_raw) + + {model_info, tokenizer, generation_config} = + if lane in [:bb_base, :bb_rewrite] do + {:ok, model_info} = Bumblebee.load_model(source, backend: {EMLX.Backend, device: :gpu}, type: :bf16) + model_info = load_params(model_info, kind, path_raw) + {:ok, tokenizer} = Bumblebee.load_tokenizer(source) + {:ok, generation_config} = Bumblebee.load_generation_config(source) + + generation_config = + Bumblebee.configure(generation_config, + max_new_tokens: cfg.max_new, + strategy: %{type: :greedy_search} + ) + + {model_info, tokenizer, generation_config} + else + tokenizer = + if lane == :native do + {:ok, tokenizer} = Bumblebee.load_tokenizer(source) + tokenizer + end + + {nil, tokenizer, nil} + end + + dispatch_lane(lane, kind, path_raw, label, source, model_info, tokenizer, generation_config, cfg) + end + + defp dispatch_lane(:bb_base, _kind, _path_raw, label, _source, model_info, tokenizer, generation_config, cfg) do + IO.puts("==> Building Bumblebee.Text.generation serving (base — no rewrite) ...") + + serving = + Bumblebee.Text.generation(model_info, tokenizer, generation_config, + compile: [batch_size: 1, sequence_length: cfg.seq_len], + defn_options: [compiler: EMLX] + ) + + lane_label = "#{label} / bb base" + Bench.warmup(lane_label, serving, cfg.prompt, &Extract.bb/1, cfg.warmup_runs) + results = Bench.bench(lane_label, serving, cfg.prompt, &Extract.bb/1, cfg.bench_runs) + Bench.stats("bb base (Bumblebee, no rewrite) ", results) + end + + defp dispatch_lane(:bb_rewrite, _kind, _path_raw, label, _source, model_info, tokenizer, generation_config, cfg) do + IO.puts("==> Applying EMLXAxon.rewrite/1 (all rewrites — best performing path) ...") + model_rewritten = EMLXAxon.rewrite(model_info.model) + + serving = + Bumblebee.Text.generation(%{model_info | model: model_rewritten}, tokenizer, generation_config, + compile: [batch_size: 1, sequence_length: cfg.seq_len], + defn_options: [compiler: EMLX] + ) + + lane_label = "#{label} / bb+rewrite" + Bench.warmup(lane_label, serving, cfg.prompt, &Extract.bb/1, cfg.warmup_runs) + results = Bench.bench(lane_label, serving, cfg.prompt, &Extract.bb/1, cfg.bench_runs) + Bench.stats("bb+rewrite (Bumblebee + EMLXAxon.rewrite) ", results) + end + + defp dispatch_lane(:native, kind, path_raw, label, _source, _model_info, tokenizer, _generation_config, cfg) do + serving = build_native_serving(kind, path_raw, tokenizer, cfg) + extract = fn result -> Extract.native(result, tokenizer) end + + lane_label = "#{label} / native" + Bench.warmup(lane_label, serving, cfg.prompt, extract, cfg.warmup_runs) + results = Bench.bench(lane_label, serving, cfg.prompt, extract, cfg.bench_runs) + Bench.stats("native (EMLXAxon.TextGeneration) ", results) + end + + defp dispatch_lane(:emily_eager, _kind, _path_raw, label, source, _model_info, _tokenizer, _generation_config, cfg) do + EmilyLane.run(:eager, source, label, cfg.prompt, cfg) + |> wrap_stats("emily-eager (Evaluator + Emily.Backend) ") + end + + defp dispatch_lane(:emily_native, _kind, _path_raw, label, source, _model_info, _tokenizer, _generation_config, cfg) do + EmilyLane.run(:native, source, label, cfg.prompt, cfg) + |> wrap_stats("emily-native (Emily.Compiler, native: true) ") + end + + defp dispatch_lane(:emily_quantized, _kind, _path_raw, label, _source, _model_info, _tokenizer, _generation_config, cfg) do + dense_source = ModelSource.resolve(cfg.dense_path_raw) + + EmilyLane.run(:quantized, dense_source, label, cfg.prompt, cfg) + |> wrap_stats("emily-quantized (Emily int4, native: true) ") + end + + defp wrap_stats(nil, _label), do: nil + defp wrap_stats(results, label), do: Bench.stats(label, results) + + defp load_params(model_info, :dense, _path_raw), do: model_info + + defp load_params(model_info, :mlx4bit, path_raw) do + IO.puts("==> Loading MLX-4bit params (dequantize → Bumblebee layout → re-quantize) ...") + t0 = System.monotonic_time(:millisecond) + params = EMLXAxon.MLX4BitParams.load(Path.expand(path_raw)) + params = EMLXAxon.QuantizeParams.quantize(params) + t1 = System.monotonic_time(:millisecond) + IO.puts(" params loaded in #{t1 - t0} ms") + %{model_info | params: params} + end + + defp build_native_serving(kind, path_raw, tokenizer, cfg) do + IO.puts("\n==> Loading EMLXAxon.TextGeneration (native 28-layer bypass) ...") + t0 = System.monotonic_time(:millisecond) + + serving = + case kind do + :mlx4bit -> + EMLXAxon.TextGeneration.from_mlx4bit( + Path.expand(path_raw), + tokenizer, + defn_options: [compiler: Nx.Defn.Evaluator], + max_new_tokens: cfg.max_new, + sampler: :greedy, + profile_timing: cfg.native_profile_timing? + ) + + :dense -> + {:ok, state} = EMLXAxon.Qwen3.DenseLoader.from_safetensors_dir(Path.expand(path_raw), type: :bf16) + + EMLXAxon.TextGeneration.serving( + tokenizer, + state, + defn_options: [compiler: Nx.Defn.Evaluator], + max_new_tokens: cfg.max_new, + sampler: :greedy, + profile_timing: cfg.native_profile_timing? + ) + end + + IO.puts(" loaded in #{System.monotonic_time(:millisecond) - t0} ms") + serving + end +end + +:ok +``` + + + +``` +:ok +``` + +## Summary rendering + +```elixir +defmodule Summary do + def ratio(nil, _), do: nil + def ratio(_, nil), do: nil + def ratio(row, col) when col.median > 0, do: Float.round(row.median / col.median, 2) + def ratio(_, _), do: nil + + def to_markdown({:ok, result}) do + lanes = lanes(result) + + matrix = + if length(lanes) < 2 do + "_fewer than two lanes ran — nothing to compare._" + else + matrix_markdown(lanes) + end + + bullets = + Enum.map_join(lanes, "\n", fn {name, stats} -> + "- **#{name}**: median=#{stats.median} mean=#{stats.mean}±#{stats.stddev} min/max=#{stats.min}/#{stats.max} tok/s" + end) + + "### #{result.label}\n\n#{matrix}\n\n#{bullets}" + end + + def to_markdown({:error, label, formatted}) do + fence = String.duplicate("`", 3) + "### #{label} — FAILED\n\n#{fence}\n#{formatted}\n#{fence}" + end + + def stats_rows(results) do + for {:ok, result} <- results, + {name, stats} <- lanes(result) do + %{ + checkpoint: result.label, + lane: name, + median_tok_s: stats.median, + mean_tok_s: stats.mean, + stddev: stats.stddev, + min_tok_s: stats.min, + max_tok_s: stats.max + } + end + end + + defp lanes(result) do + [ + {"bb base", result.base}, + {"bb+rewrite", result.rewrite}, + {"native", result.native}, + {"emily-eager", result.emily_eager}, + {"emily-native", result.emily_native}, + {"emily-quantized", result.emily_quantized} + ] + |> Enum.reject(fn {_name, stats} -> is_nil(stats) end) + end + + defp matrix_markdown(lanes) do + header = Enum.map_join(lanes, "", fn {name, _} -> " #{name} |" end) + sep = Enum.map_join(lanes, "", fn _ -> " --- |" end) + + rows = + for {row_name, row_stats} <- lanes do + cells = + Enum.map_join(lanes, "", fn {col_name, col_stats} -> + " #{format_cell(row_name, col_name, row_stats, col_stats)} |" + end) + + "| **#{row_name}** |#{cells}" + end + + Enum.join(["| row ∖ col (row / col) |#{header}", "| --- |#{sep}" | rows], "\n") + end + + defp format_cell(row_name, col_name, _row_stats, _col_stats) when row_name == col_name, do: "–" + + defp format_cell(_row_name, _col_name, row_stats, col_stats) do + case ratio(row_stats, col_stats) do + nil -> "n/a" + value -> "#{value}×" + end + end +end + +:ok +``` + + + +``` +:ok +``` + +## Configuration + +```elixir +dense_path_input = + Kino.Input.text("Dense model path", + default: System.get_env("EMLX_QWEN3_DENSE_MODEL_PATH") || "~/models/Qwen3-0.6B" + ) + +mlx4bit_path_input = + Kino.Input.text("MLX-4bit model path", + default: System.get_env("EMLX_QWEN3_4BIT_MODEL_PATH") || "~/models/Qwen3-0.6B-MLX-4bit" + ) + +skip_dense_input = Kino.Input.checkbox("Skip dense checkpoint", default: false) +skip_4bit_input = Kino.Input.checkbox("Skip MLX-4bit checkpoint", default: false) +skip_emily_input = Kino.Input.checkbox("Skip Emily lanes", default: false) +skip_bb_base_input = Kino.Input.checkbox("Skip bb base lane", default: false) +skip_bb_rewrite_input = Kino.Input.checkbox("Skip bb+rewrite lane", default: false) +native_profile_timing_input = Kino.Input.checkbox("Native per-token profile timing", default: true) + +max_new_input = Kino.Input.number("Max new tokens", default: 60) +bench_runs_input = Kino.Input.number("Bench runs", default: 5) +warmup_runs_input = Kino.Input.number("Warmup runs", default: 2) +seq_len_input = Kino.Input.number("Sequence length", default: 1024) +emily_quant_bits_input = Kino.Input.number("Emily quant bits", default: 4) +emily_quant_group_size_input = Kino.Input.number("Emily quant group size", default: 64) + +Kino.Layout.grid( + [ + dense_path_input, + mlx4bit_path_input, + skip_dense_input, + skip_4bit_input, + skip_emily_input, + skip_bb_base_input, + skip_bb_rewrite_input, + native_profile_timing_input, + max_new_input, + bench_runs_input, + warmup_runs_input, + seq_len_input, + emily_quant_bits_input, + emily_quant_group_size_input + ], + columns: 2 +) +``` + +```elixir +dense_path_raw = Kino.Input.read(dense_path_input) + +cfg = %{ + max_new: Kino.Input.read(max_new_input), + bench_runs: Kino.Input.read(bench_runs_input), + warmup_runs: Kino.Input.read(warmup_runs_input), + seq_len: Kino.Input.read(seq_len_input), + native_profile_timing?: Kino.Input.read(native_profile_timing_input), + skip_bb_rewrite?: Kino.Input.read(skip_bb_rewrite_input), + skip_bb_base?: Kino.Input.read(skip_bb_base_input), + skip_emily?: Kino.Input.read(skip_emily_input), + emily_quant_bits: Kino.Input.read(emily_quant_bits_input), + emily_quant_group_size: Kino.Input.read(emily_quant_group_size_input), + dense_path_raw: dense_path_raw, + prompt: + "<|im_start|>user\nList twenty programming languages with a one-line description each.<|im_end|>\n<|im_start|>assistant\n" +} + +checkpoints = [ + %{ + kind: :dense, + label: "Qwen3-0.6B (dense bf16)", + skip?: Kino.Input.read(skip_dense_input), + path_raw: dense_path_raw + }, + %{ + kind: :mlx4bit, + label: "Qwen3-0.6B-MLX-4bit", + skip?: Kino.Input.read(skip_4bit_input), + path_raw: Kino.Input.read(mlx4bit_path_input) + } +] + +IO.puts(""" +max_new: #{cfg.max_new} bench_runs: #{cfg.bench_runs} warmup_runs: #{cfg.warmup_runs} +seq_len: #{cfg.seq_len} +""") + +:ok +``` + + + +``` +max_new: 60 bench_runs: 5 warmup_runs: 2 +seq_len: 1024 + +``` + + + +``` +:ok +``` + +## Run benchmarks + +```elixir +results = + checkpoints + |> Enum.reject(& &1.skip?) + |> Enum.shuffle() + |> Enum.map(fn checkpoint -> + try do + {:ok, Runner.run_checkpoint(checkpoint, cfg)} + rescue + exception -> {:error, checkpoint.label, Exception.format(:error, exception, __STACKTRACE__)} + end + end) + +:ok +``` + + + +``` + +################################################################################ +=== Qwen3-0.6B-MLX-4bit — {:local, "/Users/valente/models/Qwen3-0.6B-MLX-4bit"} === +################################################################################ + + +==> Spawning isolated peer for lane :bb_base ... + +09:15:23.657 [debug] the following PyTorch parameters were unused: + + * model.embed_tokens.biases + * model.embed_tokens.scales + * model.layers.0.mlp.down_proj.biases + * model.layers.0.mlp.down_proj.scales + * model.layers.0.mlp.gate_proj.biases + * model.layers.0.mlp.gate_proj.scales + * model.layers.0.mlp.up_proj.biases + * model.layers.0.mlp.up_proj.scales + * model.layers.0.self_attn.k_proj.biases + * model.layers.0.self_attn.k_proj.scales + * model.layers.0.self_attn.o_proj.biases + * model.layers.0.self_attn.o_proj.scales + * model.layers.0.self_attn.q_proj.biases + * model.layers.0.self_attn.q_proj.scales + * model.layers.0.self_attn.v_proj.biases + * model.layers.0.self_attn.v_proj.scales + * model.layers.1.mlp.down_proj.biases + * model.layers.1.mlp.down_proj.scales + * model.layers.1.mlp.gate_proj.biases + * model.layers.1.mlp.gate_proj.scales + * model.layers.1.mlp.up_proj.biases + * model.layers.1.mlp.up_proj.scales + * model.layers.1.self_attn.k_proj.biases + * model.layers.1.self_attn.k_proj.scales + * model.layers.1.self_attn.o_proj.biases + * model.layers.1.self_attn.o_proj.scales + * model.layers.1.self_attn.q_proj.biases + * model.layers.1.self_attn.q_proj.scales + * model.layers.1.self_attn.v_proj.biases + * model.layers.1.self_attn.v_proj.scales + * model.layers.10.mlp.down_proj.biases + * model.layers.10.mlp.down_proj.scales + * model.layers.10.mlp.gate_proj.biases + * model.layers.10.mlp.gate_proj.scales + * model.layers.10.mlp.up_proj.biases + * model.layers.10.mlp.up_proj.scales + * model.layers.10.self_attn.k_proj.biases + * model.layers.10.self_attn.k_proj.scales + * model.layers.10.self_attn.o_proj.biases + * model.layers.10.self_attn.o_proj.scales + * model.layers.10.self_attn.q_proj.biases + * model.layers.10.self_attn.q_proj.scales + * model.layers.10.self_attn.v_proj.biases + * model.layers.10.self_attn.v_proj.scales + * model.layers.11.mlp.down_proj.biases + * model.layers.11.mlp.down_proj.scales + * model.layers.11.mlp.gate_proj.biases + * model.layers.11.mlp.gate_proj.scales + * model.layers.11.mlp.up_proj.biases + * model.layers.11.mlp.up_proj.scales + * model.layers.11.self_attn.k_proj.biases + * model.layers.11.self_attn.k_proj.scales + * model.layers.11.self_attn.o_proj.biases + * model.layers.11.self_attn.o_proj.scales + * model.layers.11.self_attn.q_proj.biases + * model.layers.11.self_attn.q_proj.scales + * model.layers.11.self_attn.v_proj.biases + * model.layers.11.self_attn.v_proj.scales + * model.layers.12.mlp.down_proj.biases + * model.layers.12.mlp.down_proj.scales + * model.layers.12.mlp.gate_proj.biases + * model.layers.12.mlp.gate_proj.scales + * model.layers.12.mlp.up_proj.biases + * model.layers.12.mlp.up_proj.scales + * model.layers.12.self_attn.k_proj.biases + * model.layers.12.self_attn.k_proj.scales + * model.layers.12.self_attn.o_proj.biases + * model.layers.12.self_attn.o_proj.scales + * model.layers.12.self_attn.q_proj.biases + * model.layers.12.self_attn.q_proj.scales + * model.layers.12.self_attn.v_proj.biases + * model.layers.12.self_attn.v_proj.scales + * model.layers.13.mlp.down_proj.biases + * model.layers.13.mlp.down_proj.scales + * model.layers.13.mlp.gate_proj.biases + * model.layers.13.mlp.gate_proj.scales + * model.layers.13.mlp.up_proj.biases + * model.layers.13.mlp.up_proj.scales + * model.layers.13.self_attn.k_proj.biases + * model.layers.13.self_attn.k_proj.scales + * model.layers.13.self_attn.o_proj.biases + * model.layers.13.self_attn.o_proj.scales + * model.layers.13.self_attn.q_proj.biases + * model.layers.13.self_attn.q_proj.scales + * model.layers.13.self_attn.v_proj.biases + * model.layers.13.self_attn.v_proj.scales + * model.layers.14.mlp.down_proj.biases + * model.layers.14.mlp.down_proj.scales + * model.layers.14.mlp.gate_proj.biases + * model.layers.14.mlp.gate_proj.scales + * model.layers.14.mlp.up_proj.biases + * model.layers.14.mlp.up_proj.scales + * model.layers.14.self_attn.k_proj.biases + * model.layers.14.self_attn.k_proj.scales + * model.layers.14.self_attn.o_proj.biases + * model.layers.14.self_attn.o_proj.scales + * model.layers.14.self_attn.q_proj.biases + * model.layers.14.self_attn.q_proj.scales + * model.layers.14.self_attn.v_proj.biases + * model.layers.14.self_attn.v_proj.scales + * model.layers.15.mlp.down_proj.biases + * model.layers.15.mlp.down_proj.scales + * model.layers.15.mlp.gate_proj.biases + * model.layers.15.mlp.gate_proj.scales + * model.layers.15.mlp.up_proj.biases + * model.layers.15.mlp.up_proj.scales + * model.layers.15.self_attn.k_proj.biases + * model.layers.15.self_attn.k_proj.scales + * model.layers.15.self_attn.o_proj.biases + * model.layers.15.self_attn.o_proj.scales + * model.layers.15.self_attn.q_proj.biases + * model.layers.15.self_attn.q_proj.scales + * model.layers.15.self_attn.v_proj.biases + * model.layers.15.self_attn.v_proj.scales + * model.layers.16.mlp.down_proj.biases + * model.layers.16.mlp.down_proj.scales + * model.layers.16.mlp.gate_proj.biases + * model.layers.16.mlp.gate_proj.scales + * model.layers.16.mlp.up_proj.biases + * model.layers.16.mlp.up_proj.scales + * model.layers.16.self_attn.k_proj.biases + * model.layers.16.self_attn.k_proj.scales + * model.layers.16.self_attn.o_proj.biases + * model.layers.16.self_attn.o_proj.scales + * model.layers.16.self_attn.q_proj.biases + * model.layers.16.self_attn.q_proj.scales + * model.layers.16.self_attn.v_proj.biases + * model.layers.16.self_attn.v_proj.scales + * model.layers.17.mlp.down_proj.biases + * model.layers.17.mlp.down_proj.scales + * model.layers.17.mlp.gate_proj.biases + * model.layers.17.mlp.gate_proj.scales + * model.layers.17.mlp.up_proj.biases + * model.layers.17.mlp.up_proj.scales + * model.layers.17.self_attn.k_proj.biases + * model.layers.17.self_attn.k_proj.scales + * model.layers.17.self_attn.o_proj.biases + * model.layers.17.self_attn.o_proj.scales + * model.layers.17.self_attn.q_proj.biases + * model.layers.17.self_attn.q_proj.scales + * model.layers.17.self_attn.v_proj.biases + * model.layers.17.self_attn.v_proj.scales + * model.layers.18.mlp.down_proj.biases + * model.layers.18.mlp.down_proj.scales + * model.layers.18.mlp.gate_proj.biases + * model.layers.18.mlp.gate_proj.scales + * model.layers.18.mlp.up_proj.biases + * model.layers.18.mlp.up_proj.scales + * model.layers.18.self_attn.k_proj.biases + * model.layers.18.self_attn.k_proj.scales + * model.layers.18.self_attn.o_proj.biases + * model.layers.18.self_attn.o_proj.scales + * model.layers.18.self_attn.q_proj.biases + * model.layers.18.self_attn.q_proj.scales + * model.layers.18.self_attn.v_proj.biases + * model.layers.18.self_attn.v_proj.scales + * model.layers.19.mlp.down_proj.biases + * model.layers.19.mlp.down_proj.scales + * model.layers.19.mlp.gate_proj.biases + * model.layers.19.mlp.gate_proj.scales + * model.layers.19.mlp.up_proj.biases + * model.layers.19.mlp.up_proj.scales + * model.layers.19.self_attn.k_proj.biases + * model.layers.19.self_attn.k_proj.scales + * model.layers.19.self_attn.o_proj.biases + * model.layers.19.self_attn.o_proj.scales + * model.layers.19.self_attn.q_proj.biases + * model.layers.19.self_attn.q_proj.scales + * model.layers.19.self_attn.v_proj.biases + * model.layers.19.self_attn.v_proj.scales + * model.layers.2.mlp.down_proj.biases + * model.layers.2.mlp.down_proj.scales + * model.layers.2.mlp.gate_proj.biases + * model.layers.2.mlp.gate_proj.scales + * model.layers.2.mlp.up_proj.biases + * model.layers.2.mlp.up_proj.scales + * model.layers.2.self_attn.k_proj.biases + * model.layers.2.self_attn.k_proj.scales + * model.layers.2.self_attn.o_proj.biases + * model.layers.2.self_attn.o_proj.scales + * model.layers.2.self_attn.q_proj.biases + * model.layers.2.self_attn.q_proj.scales + * model.layers.2.self_attn.v_proj.biases + * model.layers.2.self_attn.v_proj.scales + * model.layers.20.mlp.down_proj.biases + * model.layers.20.mlp.down_proj.scales + * model.layers.20.mlp.gate_proj.biases + * model.layers.20.mlp.gate_proj.scales + * model.layers.20.mlp.up_proj.biases + * model.layers.20.mlp.up_proj.scales + * model.layers.20.self_attn.k_proj.biases + * model.layers.20.sel (truncated) +==> Loading MLX-4bit params (dequantize → Bumblebee layout → re-quantize) ... + +09:15:23.661 [debug] the following parameters were ignored, because of non-matching shape: + + * decoder.blocks.19.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.12.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.17.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.5.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.17.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.1.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.11.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.12.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.0.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.7.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.21.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.14.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.11.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.5.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.11.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.15.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.26.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.17.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.3.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.21.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.6.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.16.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.16.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.8.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.16.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.26.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.10.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.17.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.3.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.20.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.4.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.15.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.13.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.14.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.2.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.9.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.6.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.8.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.26.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.25.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.8.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.21.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.16.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.4.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.7.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.12.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.14.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.14.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.26.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.6.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.13.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.18.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.24.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.13.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.24.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.20.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.24.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.8.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.19.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.9.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.10.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.18.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.2.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.18.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.26.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.23.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.22.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.18.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * language_modeling_head.output.kernel (expected {151936, 1024}, got: {151936, 128}) + * decoder.blocks.27.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.3.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.0.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.23.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.13.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.11.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.9.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.4.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.4.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.25.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.6.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.7.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.13.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.21.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.2.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.22.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.15.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.6.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.14.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.22.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.6.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.22.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.1.ffn.intermediate.kernel (expect (truncated) + params loaded in 97 ms +==> Building Bumblebee.Text.generation serving (base — no rewrite) ... + +==> [Qwen3-0.6B-MLX-4bit / bb base] Warmup (2 run(s), not timed) ... + warmup 1: 60 tokens / 1261 ms " Okay, the user is asking for twenty programming lang..." + warmup 2: 60 tokens / 954 ms " Okay, the user is asking for twenty programming lang..." + +==> [Qwen3-0.6B-MLX-4bit / bb base] Benchmark (5 run(s)) ... + run 1: 60 tokens / 1034 ms = 58.0 tok/s + run 2: 60 tokens / 1028 ms = 58.4 tok/s + run 3: 60 tokens / 958 ms = 62.6 tok/s + run 4: 60 tokens / 942 ms = 63.7 tok/s + run 5: 60 tokens / 969 ms = 61.9 tok/s + bb base (Bumblebee, no rewrite) : median=61.9 mean=60.9±2.3 min/max=58.0/63.7 tok/s + +==> Spawning isolated peer for lane :bb_rewrite ... + +09:15:35.284 [debug] the following PyTorch parameters were unused: + + * model.embed_tokens.biases + * model.embed_tokens.scales + * model.layers.0.mlp.down_proj.biases + * model.layers.0.mlp.down_proj.scales + * model.layers.0.mlp.gate_proj.biases + * model.layers.0.mlp.gate_proj.scales + * model.layers.0.mlp.up_proj.biases + * model.layers.0.mlp.up_proj.scales + * model.layers.0.self_attn.k_proj.biases + * model.layers.0.self_attn.k_proj.scales + * model.layers.0.self_attn.o_proj.biases + * model.layers.0.self_attn.o_proj.scales + * model.layers.0.self_attn.q_proj.biases + * model.layers.0.self_attn.q_proj.scales + * model.layers.0.self_attn.v_proj.biases + * model.layers.0.self_attn.v_proj.scales + * model.layers.1.mlp.down_proj.biases + * model.layers.1.mlp.down_proj.scales + * model.layers.1.mlp.gate_proj.biases + * model.layers.1.mlp.gate_proj.scales + * model.layers.1.mlp.up_proj.biases + * model.layers.1.mlp.up_proj.scales + * model.layers.1.self_attn.k_proj.biases + * model.layers.1.self_attn.k_proj.scales + * model.layers.1.self_attn.o_proj.biases + * model.layers.1.self_attn.o_proj.scales + * model.layers.1.self_attn.q_proj.biases + * model.layers.1.self_attn.q_proj.scales + * model.layers.1.self_attn.v_proj.biases + * model.layers.1.self_attn.v_proj.scales + * model.layers.10.mlp.down_proj.biases + * model.layers.10.mlp.down_proj.scales + * model.layers.10.mlp.gate_proj.biases + * model.layers.10.mlp.gate_proj.scales + * model.layers.10.mlp.up_proj.biases + * model.layers.10.mlp.up_proj.scales + * model.layers.10.self_attn.k_proj.biases + * model.layers.10.self_attn.k_proj.scales + * model.layers.10.self_attn.o_proj.biases + * model.layers.10.self_attn.o_proj.scales + * model.layers.10.self_attn.q_proj.biases + * model.layers.10.self_attn.q_proj.scales + * model.layers.10.self_attn.v_proj.biases + * model.layers.10.self_attn.v_proj.scales + * model.layers.11.mlp.down_proj.biases + * model.layers.11.mlp.down_proj.scales + * model.layers.11.mlp.gate_proj.biases + * model.layers.11.mlp.gate_proj.scales + * model.layers.11.mlp.up_proj.biases + * model.layers.11.mlp.up_proj.scales + * model.layers.11.self_attn.k_proj.biases + * model.layers.11.self_attn.k_proj.scales + * model.layers.11.self_attn.o_proj.biases + * model.layers.11.self_attn.o_proj.scales + * model.layers.11.self_attn.q_proj.biases + * model.layers.11.self_attn.q_proj.scales + * model.layers.11.self_attn.v_proj.biases + * model.layers.11.self_attn.v_proj.scales + * model.layers.12.mlp.down_proj.biases + * model.layers.12.mlp.down_proj.scales + * model.layers.12.mlp.gate_proj.biases + * model.layers.12.mlp.gate_proj.scales + * model.layers.12.mlp.up_proj.biases + * model.layers.12.mlp.up_proj.scales + * model.layers.12.self_attn.k_proj.biases + * model.layers.12.self_attn.k_proj.scales + * model.layers.12.self_attn.o_proj.biases + * model.layers.12.self_attn.o_proj.scales + * model.layers.12.self_attn.q_proj.biases + * model.layers.12.self_attn.q_proj.scales + * model.layers.12.self_attn.v_proj.biases + * model.layers.12.self_attn.v_proj.scales + * model.layers.13.mlp.down_proj.biases + * model.layers.13.mlp.down_proj.scales + * model.layers.13.mlp.gate_proj.biases + * model.layers.13.mlp.gate_proj.scales + * model.layers.13.mlp.up_proj.biases + * model.layers.13.mlp.up_proj.scales + * model.layers.13.self_attn.k_proj.biases + * model.layers.13.self_attn.k_proj.scales + * model.layers.13.self_attn.o_proj.biases + * model.layers.13.self_attn.o_proj.scales + * model.layers.13.self_attn.q_proj.biases + * model.layers.13.self_attn.q_proj.scales + * model.layers.13.self_attn.v_proj.biases + * model.layers.13.self_attn.v_proj.scales + * model.layers.14.mlp.down_proj.biases + * model.layers.14.mlp.down_proj.scales + * model.layers.14.mlp.gate_proj.biases + * model.layers.14.mlp.gate_proj.scales + * model.layers.14.mlp.up_proj.biases + * model.layers.14.mlp.up_proj.scales + * model.layers.14.self_attn.k_proj.biases + * model.layers.14.self_attn.k_proj.scales + * model.layers.14.self_attn.o_proj.biases + * model.layers.14.self_attn.o_proj.scales + * model.layers.14.self_attn.q_proj.biases + * model.layers.14.self_attn.q_proj.scales + * model.layers.14.self_attn.v_proj.biases + * model.layers.14.self_attn.v_proj.scales + * model.layers.15.mlp.down_proj.biases + * model.layers.15.mlp.down_proj.scales + * model.layers.15.mlp.gate_proj.biases + * model.layers.15.mlp.gate_proj.scales + * model.layers.15.mlp.up_proj.biases + * model.layers.15.mlp.up_proj.scales + * model.layers.15.self_attn.k_proj.biases + * model.layers.15.self_attn.k_proj.scales + * model.layers.15.self_attn.o_proj.biases + * model.layers.15.self_attn.o_proj.scales + * model.layers.15.self_attn.q_proj.biases + * model.layers.15.self_attn.q_proj.scales + * model.layers.15.self_attn.v_proj.biases + * model.layers.15.self_attn.v_proj.scales + * model.layers.16.mlp.down_proj.biases + * model.layers.16.mlp.down_proj.scales + * model.layers.16.mlp.gate_proj.biases + * model.layers.16.mlp.gate_proj.scales + * model.layers.16.mlp.up_proj.biases + * model.layers.16.mlp.up_proj.scales + * model.layers.16.self_attn.k_proj.biases + * model.layers.16.self_attn.k_proj.scales + * model.layers.16.self_attn.o_proj.biases + * model.layers.16.self_attn.o_proj.scales + * model.layers.16.self_attn.q_proj.biases + * model.layers.16.self_attn.q_proj.scales + * model.layers.16.self_attn.v_proj.biases + * model.layers.16.self_attn.v_proj.scales + * model.layers.17.mlp.down_proj.biases + * model.layers.17.mlp.down_proj.scales + * model.layers.17.mlp.gate_proj.biases + * model.layers.17.mlp.gate_proj.scales + * model.layers.17.mlp.up_proj.biases + * model.layers.17.mlp.up_proj.scales + * model.layers.17.self_attn.k_proj.biases + * model.layers.17.self_attn.k_proj.scales + * model.layers.17.self_attn.o_proj.biases + * model.layers.17.self_attn.o_proj.scales + * model.layers.17.self_attn.q_proj.biases + * model.layers.17.self_attn.q_proj.scales + * model.layers.17.self_attn.v_proj.biases + * model.layers.17.self_attn.v_proj.scales + * model.layers.18.mlp.down_proj.biases + * model.layers.18.mlp.down_proj.scales + * model.layers.18.mlp.gate_proj.biases + * model.layers.18.mlp.gate_proj.scales + * model.layers.18.mlp.up_proj.biases + * model.layers.18.mlp.up_proj.scales + * model.layers.18.self_attn.k_proj.biases + * model.layers.18.self_attn.k_proj.scales + * model.layers.18.self_attn.o_proj.biases + * model.layers.18.self_attn.o_proj.scales + * model.layers.18.self_attn.q_proj.biases + * model.layers.18.self_attn.q_proj.scales + * model.layers.18.self_attn.v_proj.biases + * model.layers.18.self_attn.v_proj.scales + * model.layers.19.mlp.down_proj.biases + * model.layers.19.mlp.down_proj.scales + * model.layers.19.mlp.gate_proj.biases + * model.layers.19.mlp.gate_proj.scales + * model.layers.19.mlp.up_proj.biases + * model.layers.19.mlp.up_proj.scales + * model.layers.19.self_attn.k_proj.biases + * model.layers.19.self_attn.k_proj.scales + * model.layers.19.self_attn.o_proj.biases + * model.layers.19.self_attn.o_proj.scales + * model.layers.19.self_attn.q_proj.biases + * model.layers.19.self_attn.q_proj.scales + * model.layers.19.self_attn.v_proj.biases + * model.layers.19.self_attn.v_proj.scales + * model.layers.2.mlp.down_proj.biases + * model.layers.2.mlp.down_proj.scales + * model.layers.2.mlp.gate_proj.biases + * model.layers.2.mlp.gate_proj.scales + * model.layers.2.mlp.up_proj.biases + * model.layers.2.mlp.up_proj.scales + * model.layers.2.self_attn.k_proj.biases + * model.layers.2.self_attn.k_proj.scales + * model.layers.2.self_attn.o_proj.biases + * model.layers.2.self_attn.o_proj.scales + * model.layers.2.self_attn.q_proj.biases + * model.layers.2.self_attn.q_proj.scales + * model.layers.2.self_attn.v_proj.biases + * model.layers.2.self_attn.v_proj.scales + * model.layers.20.mlp.down_proj.biases + * model.layers.20.mlp.down_proj.scales + * model.layers.20.mlp.gate_proj.biases + * model.layers.20.mlp.gate_proj.scales + * model.layers.20.mlp.up_proj.biases + * model.layers.20.mlp.up_proj.scales + * model.layers.20.self_attn.k_proj.biases + * model.layers.20.sel (truncated) +==> Loading MLX-4bit params (dequantize → Bumblebee layout → re-quantize) ... + +09:15:35.289 [debug] the following parameters were ignored, because of non-matching shape: + + * decoder.blocks.19.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.12.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.17.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.5.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.17.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.1.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.11.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.12.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.0.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.7.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.21.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.14.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.11.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.5.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.11.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.15.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.26.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.17.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.3.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.21.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.6.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.16.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.16.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.8.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.16.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.26.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.10.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.17.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.3.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.20.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.4.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.15.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.13.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.14.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.2.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.9.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.6.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.8.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.26.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.25.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.8.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.21.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.16.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.4.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.7.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.12.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.14.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.14.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.26.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.6.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.13.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.18.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.24.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.13.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.24.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.20.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.24.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.8.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.19.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.9.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.10.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.18.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.2.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.18.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.26.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.23.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.22.ffn.gate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.18.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * language_modeling_head.output.kernel (expected {151936, 1024}, got: {151936, 128}) + * decoder.blocks.27.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.3.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.0.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.23.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.13.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.11.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.9.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.4.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.4.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.25.self_attention.value.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.6.ffn.output.kernel (expected {3072, 1024}, got: {384, 1024}) + * decoder.blocks.7.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.13.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.21.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.2.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.22.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.15.self_attention.query.kernel (expected {1024, 2048}, got: {128, 2048}) + * decoder.blocks.6.ffn.intermediate.kernel (expected {1024, 3072}, got: {128, 3072}) + * decoder.blocks.14.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.22.self_attention.key.kernel (expected {1024, 1024}, got: {128, 1024}) + * decoder.blocks.6.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.22.self_attention.output.kernel (expected {2048, 1024}, got: {256, 1024}) + * decoder.blocks.1.ffn.intermediate.kernel (expect (truncated) + params loaded in 103 ms +==> Applying EMLXAxon.rewrite/1 (all rewrites — best performing path) ... + +==> [Qwen3-0.6B-MLX-4bit / bb+rewrite] Warmup (2 run(s), not timed) ... + warmup 1: 60 tokens / 897 ms " Okay, the user is asking for twenty programming lang..." + warmup 2: 60 tokens / 708 ms " Okay, the user is asking for twenty programming lang..." + +==> [Qwen3-0.6B-MLX-4bit / bb+rewrite] Benchmark (5 run(s)) ... + run 1: 60 tokens / 762 ms = 78.7 tok/s + run 2: 60 tokens / 766 ms = 78.3 tok/s + run 3: 60 tokens / 683 ms = 87.8 tok/s + run 4: 60 tokens / 718 ms = 83.6 tok/s + run 5: 60 tokens / 754 ms = 79.6 tok/s + bb+rewrite (Bumblebee + EMLXAxon.rewrite) : median=79.6 mean=81.6±3.6 min/max=78.3/87.8 tok/s + +==> Spawning isolated peer for lane :native ... + +==> Loading EMLXAxon.TextGeneration (native 28-layer bypass) ... + loaded in 120 ms + +==> [Qwen3-0.6B-MLX-4bit / native] Warmup (2 run(s), not timed) ... + warmup 1: 60 tokens / 399 ms " Okay, the user is asking for twenty programming lang..." + warmup 2: 60 tokens / 483 ms " Okay, the user is asking for twenty programming lang..." + +==> [Qwen3-0.6B-MLX-4bit / native] Benchmark (5 run(s)) ... + run 1: 60 tokens / 465 ms = 129.0 tok/s + run 2: 60 tokens / 432 ms = 138.9 tok/s + run 3: 60 tokens / 385 ms = 155.8 tok/s + run 4: 60 tokens / 452 ms = 132.7 tok/s + run 5: 60 tokens / 378 ms = 158.7 tok/s + native (EMLXAxon.TextGeneration) : median=138.9 mean=143.0±12.1 min/max=129.0/158.7 tok/s + +==> Spawning isolated peer for lane :emily_quantized ... + +==> [Qwen3-0.6B-MLX-4bit / emily-quantized] Loading model on Emily.Backend ... + quantizing every Axon.dense node (bits=4, group_size=64) ... + +==> [Qwen3-0.6B-MLX-4bit / emily-quantized] Warmup (2 run(s), not timed) ... + warmup 1: 60 tokens / 1824 ms " Okay, the user is asking for twenty programming lang..." + warmup 2: 60 tokens / 1758 ms " Okay, the user is asking for twenty programming lang..." + +==> [Qwen3-0.6B-MLX-4bit / emily-quantized] Benchmark (5 run(s)) ... + run 1: 60 tokens / 1707 ms = 35.1 tok/s + run 2: 60 tokens / 1742 ms = 34.4 tok/s + run 3: 60 tokens / 1727 ms = 34.7 tok/s + run 4: 60 tokens / 1702 ms = 35.3 tok/s + run 5: 60 tokens / 1775 ms = 33.8 tok/s + emily-quantized (Emily int4, native: true) : median=34.7 mean=34.7±0.5 min/max=33.8/35.3 tok/s + +################################################################################ +=== Qwen3-0.6B (dense bf16) — {:local, "/Users/valente/models/Qwen3-0.6B"} === +################################################################################ + + +==> Spawning isolated peer for lane :bb_base ... +==> Building Bumblebee.Text.generation serving (base — no rewrite) ... + +==> [Qwen3-0.6B (dense bf16) / bb base] Warmup (2 run(s), not timed) ... + warmup 1: 60 tokens / 1162 ms " Okay, the user wants a list of twenty programming la..." + warmup 2: 60 tokens / 990 ms " Okay, the user wants a list of twenty programming la..." + +==> [Qwen3-0.6B (dense bf16) / bb base] Benchmark (5 run(s)) ... + run 1: 60 tokens / 1001 ms = 59.9 tok/s + run 2: 60 tokens / 1018 ms = 58.9 tok/s + run 3: 60 tokens / 1025 ms = 58.5 tok/s + run 4: 60 tokens / 986 ms = 60.9 tok/s + run 5: 60 tokens / 989 ms = 60.7 tok/s + bb base (Bumblebee, no rewrite) : median=59.9 mean=59.8±1.0 min/max=58.5/60.9 tok/s + +==> Spawning isolated peer for lane :bb_rewrite ... +==> Applying EMLXAxon.rewrite/1 (all rewrites — best performing path) ... + +==> [Qwen3-0.6B (dense bf16) / bb+rewrite] Warmup (2 run(s), not timed) ... + warmup 1: 60 tokens / 848 ms " Okay, the user wants a list of twenty programming la..." + warmup 2: 60 tokens / 739 ms " Okay, the user wants a list of twenty programming la..." + +==> [Qwen3-0.6B (dense bf16) / bb+rewrite] Benchmark (5 run(s)) ... + run 1: 60 tokens / 770 ms = 77.9 tok/s + run 2: 60 tokens / 730 ms = 82.2 tok/s + run 3: 60 tokens / 764 ms = 78.5 tok/s + run 4: 60 tokens / 728 ms = 82.4 tok/s + run 5: 60 tokens / 730 ms = 82.2 tok/s + bb+rewrite (Bumblebee + EMLXAxon.rewrite) : median=82.2 mean=80.6±2.0 min/max=77.9/82.4 tok/s + +==> Spawning isolated peer for lane :native ... + +==> Loading EMLXAxon.TextGeneration (native 28-layer bypass) ... + loaded in 207 ms + +==> [Qwen3-0.6B (dense bf16) / native] Warmup (2 run(s), not timed) ... + warmup 1: 60 tokens / 790 ms " Okay, the user wants a list of twenty programming la..." + warmup 2: 60 tokens / 708 ms " Okay, the user wants a list of twenty programming la..." + +==> [Qwen3-0.6B (dense bf16) / native] Benchmark (5 run(s)) ... + run 1: 60 tokens / 707 ms = 84.9 tok/s + run 2: 60 tokens / 716 ms = 83.8 tok/s + run 3: 60 tokens / 702 ms = 85.5 tok/s + run 4: 60 tokens / 701 ms = 85.6 tok/s + run 5: 60 tokens / 709 ms = 84.6 tok/s + native (EMLXAxon.TextGeneration) : median=84.9 mean=84.9±0.7 min/max=83.8/85.6 tok/s + +==> Spawning isolated peer for lane :emily_eager ... + +==> [Qwen3-0.6B (dense bf16) / emily-eager] Loading model on Emily.Backend ... + +==> [Qwen3-0.6B (dense bf16) / emily-eager] Warmup (2 run(s), not timed) ... + warmup 1: 60 tokens / 6632 ms " Okay, the user wants a list of twenty programming la..." + warmup 2: 60 tokens / 6448 ms " Okay, the user wants a list of twenty programming la..." + +==> [Qwen3-0.6B (dense bf16) / emily-eager] Benchmark (5 run(s)) ... + run 1: 60 tokens / 6391 ms = 9.4 tok/s + run 2: 60 tokens / 6139 ms = 9.8 tok/s + run 3: 60 tokens / 6240 ms = 9.6 tok/s + run 4: 60 tokens / 6828 ms = 8.8 tok/s + run 5: 60 tokens / 6318 ms = 9.5 tok/s + emily-eager (Evaluator + Emily.Backend) : median=9.5 mean=9.4±0.3 min/max=8.8/9.8 tok/s + +==> Spawning isolated peer for lane :emily_native ... + +==> [Qwen3-0.6B (dense bf16) / emily-native] Loading model on Emily.Backend ... + +==> [Qwen3-0.6B (dense bf16) / emily-native] Warmup (2 run(s), not timed) ... + warmup 1: 60 tokens / 793 ms " Okay, the user wants a list of twenty programming la..." + warmup 2: 60 tokens / 888 ms " Okay, the user wants a list of twenty programming la..." + +==> [Qwen3-0.6B (dense bf16) / emily-native] Benchmark (5 run(s)) ... + run 1: 60 tokens / 776 ms = 77.3 tok/s + run 2: 60 tokens / 838 ms = 71.6 tok/s + run 3: 60 tokens / 722 ms = 83.1 tok/s + run 4: 60 tokens / 851 ms = 70.5 tok/s + run 5: 60 tokens / 923 ms = 65.0 tok/s + emily-native (Emily.Compiler, native: true) : median=71.6 mean=73.5±6.2 min/max=65.0/83.1 tok/s +``` + + + +``` +:ok +``` + +## Results + +> Matrix cells: row median tok/s ÷ col median tok/s (> 1 ⇒ row faster). + +```elixir +results +|> Enum.map(&Summary.to_markdown/1) +|> Enum.join("\n\n---\n\n") +|> Kino.Markdown.new() +``` + +```elixir +results +|> Summary.stats_rows() +|> Kino.DataTable.new(name: "Per-lane tok/s stats") +``` + + + +```text +[%{stddev: 2.3, checkpoint: "Qwen3-0.6B-MLX-4bit", lane: "bb base", median_tok_s: 61.9, mean_tok_s: 60.9, min_tok_s: 58.0, max_tok_s: 63.7}, %{stddev: 3.6, checkpoint: "Qwen3-0.6B-MLX-4bit", lane: "bb+rewrite", median_tok_s: 79.6, mean_tok_s: 81.6, min_tok_s: 78.3, max_tok_s: 87.8}, %{stddev: 12.1, checkpoint: "Qwen3-0.6B-MLX-4bit", lane: "native", median_tok_s: 138.9, mean_tok_s: 143.0, min_tok_s: 129.0, max_tok_s: 158.7}, %{stddev: 0.5, checkpoint: "Qwen3-0.6B-MLX-4bit", lane: "emily-quantized", median_tok_s: 34.7, mean_tok_s: 34.7, min_tok_s: 33.8, max_tok_s: 35.3}, %{stddev: 1.0, checkpoint: "Qwen3-0.6B (dense bf16)", lane: "bb base", median_tok_s: 59.9, mean_tok_s: 59.8, min_tok_s: 58.5, max_tok_s: 60.9}, %{stddev: 2.0, checkpoint: "Qwen3-0.6B (dense bf16)", lane: "bb+rewrite", median_tok_s: 82.2, mean_tok_s: 80.6, min_tok_s: 77.9, max_tok_s: 82.4}, %{stddev: 0.7, checkpoint: "Qwen3-0.6B (dense bf16)", lane: "native", median_tok_s: 84.9, mean_tok_s: 84.9, min_tok_s: 83.8, max_tok_s: 85.6}, %{stddev: 0.3, checkpoint: "Qwen3-0.6B (dense bf16)", lane: "emily-eager", median_tok_s: 9.5, mean_tok_s: 9.4, min_tok_s: 8.8, max_tok_s: 9.8}, %{stddev: 6.2, checkpoint: "Qwen3-0.6B (dense bf16)", lane: "emily-native", median_tok_s: 71.6, mean_tok_s: 73.5, min_tok_s: 65.0, max_tok_s: 83.1}] +``` diff --git a/emlx_axon/c_src/qwen3_plugin.cpp b/emlx_axon/c_src/qwen3_plugin.cpp new file mode 100644 index 0000000..fe335ff --- /dev/null +++ b/emlx_axon/c_src/qwen3_plugin.cpp @@ -0,0 +1,1210 @@ +#include "qwen3_plugin_abi.hpp" + +#include +#include +#include + +// Qwen3 compute plugin — pure MLX graph-building code, no Erlang/erl_nif +// dependency whatsoever. Built as its own shared library (libemlx_qwen3.so, +// see this project's Makefile) and `dlopen`'d by emlx's host NIF library +// (emlx's c_src/emlx_fast/qwen3.cpp) via +// `EMLX.NIF.load_plugin("qwen3", path)`. See qwen3_plugin_abi.hpp (in +// emlx's c_src/emlx_fast, pulled in via this project's include path) for +// the ABI and the rationale for this split. + +using namespace emlx_qwen3_plugin; + +namespace { + +mlx::core::Shape to_shape(std::initializer_list v) { + return mlx::core::Shape(v.begin(), v.end()); +} + +// ── Shape/argument validation ──────────────────────────────────────────── + +bool check_rank(const mlx::core::array &tensor, int expected, const char *name, + std::string &error) { + if (tensor.ndim() != expected) { + std::ostringstream msg; + msg << name << " expects rank " << expected << ", got rank " << tensor.ndim(); + error = msg.str(); + return false; + } + return true; +} + +bool check_positive(int value, const char *name, std::string &error) { + if (value <= 0) { + error = std::string(name) + " must be positive"; + return false; + } + return true; +} + +bool check_non_negative(int value, const char *name, std::string &error) { + if (value < 0) { + error = std::string(name) + " must be non-negative"; + return false; + } + return true; +} + +bool check_dim(const mlx::core::array &tensor, int axis, int expected, + const char *name, const char *dim_name, std::string &error) { + if (tensor.shape(axis) != expected) { + std::ostringstream msg; + msg << name << " " << dim_name << " must be " << expected << ", got " + << tensor.shape(axis); + error = msg.str(); + return false; + } + return true; +} + +bool check_rank_positive(const mlx::core::array &tensor, int rank, + const char *name, std::string &error) { + if (!check_rank(tensor, rank, name, error)) { + return false; + } + for (int axis = 0; axis < rank; ++axis) { + if (tensor.shape(axis) <= 0) { + error = std::string(name) + " dimensions must be positive"; + return false; + } + } + return true; +} + +bool check_rank4_positive(const mlx::core::array &t, const char *n, std::string &e) { + return check_rank_positive(t, 4, n, e); +} +bool check_rank3_positive(const mlx::core::array &t, const char *n, std::string &e) { + return check_rank_positive(t, 3, n, e); +} +bool check_rank2_positive(const mlx::core::array &t, const char *n, std::string &e) { + return check_rank_positive(t, 2, n, e); +} + +bool check_rank1_dim(const mlx::core::array &tensor, int expected, + const char *name, std::string &error) { + if (!check_rank(tensor, 1, name, error)) { + return false; + } + return check_dim(tensor, 0, expected, name, "size", error); +} + +bool validate_projection_width(const mlx::core::array &projection, int input_width, + int head_dim, const char *name, std::string &error) { + if (!check_rank2_positive(projection, name, error)) { + return false; + } + if (!check_dim(projection, 0, input_width, name, "input width", error)) { + return false; + } + if ((projection.shape(1) % head_dim) != 0) { + error = std::string(name) + " output width must be divisible by head_dim"; + return false; + } + return true; +} + +bool validate_kv_cache_bn(const mlx::core::array &k_cache, + const mlx::core::array &v_cache, int batch, + int num_kv_heads, int offset, int token_count, + int head_dim, std::string &error) { + if (!check_rank4_positive(k_cache, "k_cache", error) || + !check_rank4_positive(v_cache, "v_cache", error)) { + return false; + } + if (!check_dim(k_cache, 0, batch, "k_cache", "batch", error) || + !check_dim(v_cache, 0, batch, "v_cache", "batch", error) || + !check_dim(k_cache, 1, num_kv_heads, "k_cache", "heads", error) || + !check_dim(v_cache, 1, num_kv_heads, "v_cache", "heads", error) || + !check_dim(k_cache, 3, head_dim, "k_cache", "head_dim", error) || + !check_dim(v_cache, 3, head_dim, "v_cache", "head_dim", error)) { + return false; + } + if (v_cache.shape(2) != k_cache.shape(2)) { + error = "k_cache and v_cache capacity must match"; + return false; + } + int64_t required_len = static_cast(offset) + static_cast(token_count); + int capacity = k_cache.shape(2); + if (required_len > capacity) { + std::ostringstream msg; + msg << "KV cache capacity " << capacity << " is smaller than required length " + << required_len; + error = msg.str(); + return false; + } + return true; +} + +bool validate_qkv_cache_attention(const mlx::core::array &q, + const mlx::core::array &new_k, + const mlx::core::array &new_v, + const mlx::core::array &k_cache, + const mlx::core::array &v_cache, int offset, + int head_dim, std::string &error) { + if (!check_rank4_positive(q, "q", error) || + !check_rank4_positive(new_k, "new_k", error) || + !check_rank4_positive(new_v, "new_v", error) || + !check_non_negative(offset, "offset", error) || + !check_positive(head_dim, "head_dim", error)) { + return false; + } + + int B = q.shape(0); + int T_new = q.shape(1); + int N_q = q.shape(2); + int D = q.shape(3); + int N_kv = new_k.shape(2); + + if (D != head_dim) { + error = "q last dimension must match head_dim"; + return false; + } + if ((N_q % N_kv) != 0) { + error = "query heads must be divisible by key/value heads"; + return false; + } + if (!check_dim(new_k, 0, B, "new_k", "batch", error) || + !check_dim(new_v, 0, B, "new_v", "batch", error) || + !check_dim(new_k, 1, T_new, "new_k", "sequence length", error) || + !check_dim(new_v, 1, T_new, "new_v", "sequence length", error) || + !check_dim(new_v, 2, N_kv, "new_v", "heads", error) || + !check_dim(new_k, 3, D, "new_k", "head_dim", error) || + !check_dim(new_v, 3, D, "new_v", "head_dim", error)) { + return false; + } + + return validate_kv_cache_bn(k_cache, v_cache, B, N_kv, offset, T_new, D, error); +} + +// ── 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) { + 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); + } + return mlx::core::matmul(x, weight, device); +} + +mlx::core::array linear_out_in(const mlx::core::array &x, const mlx::core::array &weight, + const mlx::core::Device &device) { + return mlx::core::tensordot(x, weight, std::vector{static_cast(x.ndim()) - 1}, + std::vector{1}, device); +} + +mlx::core::array apply_linear(const mlx::core::array &x, const LinearWeight &w, + const mlx::core::Device &device) { + 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); + } + return w.transpose ? linear_out_in(x, *w.weight, device) + : linear_in_out(x, *w.weight, device); +} + +int linear_weight_out_features(const LinearWeight &w) { + if (w.quantized) { + return static_cast(w.scales->shape(0)); + } + return w.transpose ? static_cast(w.weight->shape(0)) + : static_cast(w.weight->shape(1)); +} + +int64_t token_to_int64(mlx::core::array &token) { + mlx::core::eval(token); + auto dtype = token.dtype(); + if (dtype == mlx::core::uint8) return static_cast(token.item()); + if (dtype == mlx::core::uint16) return static_cast(token.item()); + if (dtype == mlx::core::uint32) return static_cast(token.item()); + if (dtype == mlx::core::uint64) return static_cast(token.item()); + if (dtype == mlx::core::int8) return static_cast(token.item()); + if (dtype == mlx::core::int16) return static_cast(token.item()); + if (dtype == mlx::core::int32) return static_cast(token.item()); + return token.item(); +} + +bool check_linear_weight_in(const LinearWeight &w, int expected_in, const char *name, + std::string &error) { + if (w.quantized) { + if (!check_rank2_positive(*w.weight, name, error) || + !check_rank2_positive(*w.scales, name, error)) { + return false; + } + if (w.bits <= 0 || w.bits > 32 || (32 % w.bits) != 0) { + error = std::string(name) + " has an invalid bits value"; + return false; + } + int actual_in = static_cast(w.weight->shape(1)) * (32 / w.bits); + if (actual_in != expected_in) { + std::ostringstream msg; + msg << name << " input width must be " << expected_in << ", got " << actual_in; + error = msg.str(); + return false; + } + return true; + } + + if (!check_rank2_positive(*w.weight, name, error)) { + return false; + } + int actual_in = + w.transpose ? static_cast(w.weight->shape(1)) : static_cast(w.weight->shape(0)); + if (actual_in != expected_in) { + std::ostringstream msg; + msg << name << " input width must be " << expected_in << ", got " << actual_in; + error = msg.str(); + return false; + } + return true; +} + +bool validate_generalized_layer(const mlx::core::array &hidden, const LayerParamsQ &layer, + const KVCache &kv, int offset, int head_dim, + std::string &error) { + if (!check_rank3_positive(hidden, "hidden", error) || + !check_non_negative(offset, "offset", error) || + !check_positive(head_dim, "head_dim", error)) { + return false; + } + + int B = hidden.shape(0); + int T_new = hidden.shape(1); + int H = hidden.shape(2); + int D = head_dim; + + if (!check_rank1_dim(*layer.norm1, H, "norm1", error) || + !check_rank1_dim(*layer.norm2, H, "norm2", error) || + !check_rank1_dim(*layer.q_norm, D, "q_norm", error) || + !check_rank1_dim(*layer.k_norm, D, "k_norm", error) || + !check_linear_weight_in(layer.q_proj, H, "q_proj", error) || + !check_linear_weight_in(layer.k_proj, H, "k_proj", error) || + !check_linear_weight_in(layer.v_proj, H, "v_proj", error) || + !check_linear_weight_in(layer.gate_proj, H, "gate_proj", error) || + !check_linear_weight_in(layer.up_proj, H, "up_proj", error)) { + return false; + } + + int q_out = linear_weight_out_features(layer.q_proj); + int k_out = linear_weight_out_features(layer.k_proj); + int v_out = linear_weight_out_features(layer.v_proj); + + if ((q_out % D) != 0 || (k_out % D) != 0) { + error = "q_proj/k_proj output width must be divisible by head_dim"; + return false; + } + if (v_out != k_out) { + error = "v_proj output width must match k_proj output width"; + return false; + } + + int N_q = q_out / D; + int N_kv = k_out / D; + if ((N_q % N_kv) != 0) { + error = "query heads must be divisible by key/value heads"; + return false; + } + + int attn_width = N_q * D; + if (!check_linear_weight_in(layer.o_proj, attn_width, "o_proj", error)) { + return false; + } + if (linear_weight_out_features(layer.o_proj) != H) { + error = "o_proj output width must match hidden width"; + return false; + } + + int gate_out = linear_weight_out_features(layer.gate_proj); + int up_out = linear_weight_out_features(layer.up_proj); + if (up_out != gate_out) { + error = "up_proj output width must match gate_proj output width"; + return false; + } + + if (!check_linear_weight_in(layer.down_proj, gate_out, "down_proj", error)) { + return false; + } + if (linear_weight_out_features(layer.down_proj) != H) { + error = "down_proj output width must match hidden width"; + return false; + } + + return validate_kv_cache_bn(*kv.k, *kv.v, B, N_kv, offset, T_new, D, error); +} + +bool validate_dense_layer(const mlx::core::array &hidden, const LayerParams &layer, + const KVCache &kv, int offset, int head_dim, std::string &error) { + if (!check_rank3_positive(hidden, "hidden", error) || + !check_non_negative(offset, "offset", error) || + !check_positive(head_dim, "head_dim", error)) { + return false; + } + + int B = hidden.shape(0); + int T_new = hidden.shape(1); + int H = hidden.shape(2); + int D = head_dim; + + if (!check_rank1_dim(*layer.norm1, H, "norm1", error) || + !check_rank1_dim(*layer.norm2, H, "norm2", error) || + !validate_projection_width(*layer.q_proj, H, D, "q_proj", error) || + !validate_projection_width(*layer.k_proj, H, D, "k_proj", error) || + !validate_projection_width(*layer.v_proj, H, D, "v_proj", error) || + !check_dim(*layer.v_proj, 1, layer.k_proj->shape(1), "v_proj", "output width", error) || + !check_rank1_dim(*layer.q_norm, D, "q_norm", error) || + !check_rank1_dim(*layer.k_norm, D, "k_norm", error)) { + return false; + } + + int N_q = layer.q_proj->shape(1) / D; + int N_kv = layer.k_proj->shape(1) / D; + int attn_width = N_q * D; + + if ((N_q % N_kv) != 0) { + error = "query heads must be divisible by key/value heads"; + return false; + } + if (!check_rank2_positive(*layer.o_proj, "o_proj", error) || + !check_dim(*layer.o_proj, 0, attn_width, "o_proj", "input width", error) || + !check_dim(*layer.o_proj, 1, H, "o_proj", "output width", error) || + !validate_kv_cache_bn(*kv.k, *kv.v, B, N_kv, offset, T_new, D, error) || + !check_rank2_positive(*layer.gate_proj, "gate_proj", error) || + !check_rank2_positive(*layer.up_proj, "up_proj", error) || + !check_rank2_positive(*layer.down_proj, "down_proj", error) || + !check_dim(*layer.gate_proj, 0, H, "gate_proj", "input width", error) || + !check_dim(*layer.up_proj, 0, H, "up_proj", "input width", error) || + !check_dim(*layer.up_proj, 1, layer.gate_proj->shape(1), "up_proj", "output width", error) || + !check_dim(*layer.down_proj, 0, layer.gate_proj->shape(1), "down_proj", "input width", + error) || + !check_dim(*layer.down_proj, 1, H, "down_proj", "output width", error)) { + return false; + } + + return true; +} + +// 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) { + 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); + 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 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); +} + +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) { + return (T_new == 1) + ? mlx::core::fast::scaled_dot_product_attention( + q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, device) + : 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); +} + +// ── Generalized (dense-or-quantized) per-layer compute ─────────────────── +// Mirrors `layer_dense_impl` exactly, but threads every projection through +// `apply_linear` instead of assuming a dense weight. +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::array *k_out, mlx::core::array *v_out) { + int B = hidden.shape(0); + int T_new = hidden.shape(1); + int D = head_dim; + int N_q = linear_weight_out_features(layer.q_proj) / D; + int N_kv = linear_weight_out_features(layer.k_proj) / D; + 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 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); + + q = mlx::core::fast::rms_norm(q, *layer.q_norm, eps, device); + k = mlx::core::fast::rms_norm(k, *layer.k_norm, eps, device); + + 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_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 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); + 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); + + 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 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 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); + + 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); +} + +// ── 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, + mlx::core::array *k_out, mlx::core::array *v_out) { + int B = hidden.shape(0); + int T_new = hidden.shape(1); + int D = head_dim; + int N_q = layer.q_proj->shape(1) / D; + int N_kv = layer.k_proj->shape(1) / D; + 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 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); + + q = mlx::core::fast::rms_norm(q, *layer.q_norm, eps, device); + k = mlx::core::fast::rms_norm(k, *layer.k_norm, eps, device); + + 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_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 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); + 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); + + 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 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 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); + + 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); +} + +// ── VTable entrypoints ──────────────────────────────────────────────────── + +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, + mlx::core::array &k_upd, mlx::core::array &v_upd, + std::string &error) { + try { + if (!validate_qkv_cache_attention(q, new_k, new_v, k_cache, v_cache, offset, head_dim, + error)) { + return false; + } + + int B = q.shape(0); + int T_new = q.shape(1); + int N_q = q.shape(2); + int D = q.shape(3); + 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_rope = + mlx::core::fast::rope(q_bn, head_dim, false, (float)theta, 1.0f, offset, std::nullopt, device); + auto k_rope = + mlx::core::fast::rope(k_bn, head_dim, false, (float)theta, 1.0f, offset, std::nullopt, device); + + 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); + 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); + + 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 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); + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin kv_cache_attention"; + return false; + } +} + +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, + mlx::core::array &out, std::string &error) { + try { + if (!check_rank3_positive(hidden, "hidden", error)) { + return false; + } + int H = hidden.shape(2); + if (!check_rank1_dim(norm, H, "norm", error) || + !check_rank2_positive(gate_proj, "gate_proj", error) || + !check_rank2_positive(up_proj, "up_proj", error) || + !check_rank2_positive(down_proj, "down_proj", error) || + !check_dim(gate_proj, 0, H, "gate_proj", "input width", error) || + !check_dim(up_proj, 0, H, "up_proj", "input width", error) || + !check_dim(up_proj, 1, gate_proj.shape(1), "up_proj", "output width", error) || + !check_dim(down_proj, 0, gate_proj.shape(1), "down_proj", "input width", error) || + !check_dim(down_proj, 1, H, "down_proj", "output width", error)) { + 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); + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin mlp"; + return false; + } +} + +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, + mlx::core::array &out, std::string &error) { + try { + if (!check_rank3_positive(hidden, "hidden", error) || + !check_rank3_positive(attn_out, "attn_out", error)) { + return false; + } + int B = hidden.shape(0); + int T = hidden.shape(1); + int H = hidden.shape(2); + if (!check_dim(attn_out, 0, B, "attn_out", "batch", error) || + !check_dim(attn_out, 1, T, "attn_out", "sequence length", error) || + !check_rank2_positive(o_proj, "o_proj", error) || + !check_dim(o_proj, 0, attn_out.shape(2), "o_proj", "input width", error) || + !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); + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin attention_residual"; + return false; + } +} + +bool v_attention_block(const mlx::core::array &hidden, const mlx::core::array &norm, + const mlx::core::array &q_proj, const mlx::core::array &k_proj, + const mlx::core::array &v_proj, const mlx::core::array &o_proj, + 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::array &k_upd, mlx::core::array &v_upd, std::string &error) { + try { + if (!check_rank3_positive(hidden, "hidden", error) || + !check_non_negative(offset, "offset", error) || + !check_positive(head_dim, "head_dim", error)) { + return false; + } + + int B = hidden.shape(0); + int T_new = hidden.shape(1); + int H = hidden.shape(2); + int D = head_dim; + if (!check_rank1_dim(norm, H, "norm", error) || !check_rank1_dim(q_norm, D, "q_norm", error) || + !check_rank1_dim(k_norm, D, "k_norm", error) || + !check_rank2_positive(q_proj, "q_proj", error) || + !check_rank2_positive(k_proj, "k_proj", error) || + !check_rank2_positive(v_proj, "v_proj", error) || + !check_dim(q_proj, 0, H, "q_proj", "input width", error) || + !check_dim(k_proj, 0, H, "k_proj", "input width", error) || + !check_dim(v_proj, 0, H, "v_proj", "input width", error)) { + return false; + } + + if ((q_proj.shape(1) % D) != 0 || (k_proj.shape(1) % D) != 0) { + error = "projection output widths must be divisible by head_dim"; + return false; + } + if (v_proj.shape(1) != k_proj.shape(1)) { + error = "v_proj output width must match k_proj output width"; + return false; + } + + int N_q = q_proj.shape(1) / D; + int N_kv = k_proj.shape(1) / D; + int attn_width = N_q * D; + + if ((N_q % N_kv) != 0) { + error = "query heads must be divisible by key/value heads"; + return false; + } + if (!check_rank2_positive(o_proj, "o_proj", error) || + !check_dim(o_proj, 0, attn_width, "o_proj", "input width", error) || + !check_dim(o_proj, 1, H, "o_proj", "output width", error) || + !validate_kv_cache_bn(k_cache, v_cache, B, N_kv, offset, T_new, D, error)) { + return false; + } + 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 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); + + q = mlx::core::fast::rms_norm(q, q_norm, (float)eps, device); + k = mlx::core::fast::rms_norm(k, k_norm, (float)eps, device); + + 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_rope = + mlx::core::fast::rope(q_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, device); + auto k_rope = + mlx::core::fast::rope(k_bn, D, false, (float)theta, 1.0f, offset, std::nullopt, device); + + 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); + 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); + + 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 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); + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin attention_block"; + return false; + } +} + +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::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); + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin layer_dense"; + return false; + } +} + +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::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); + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin layer_quantized"; + return false; + } +} + +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, + mlx::core::array &out, std::string &error) { + try { + if (!check_rank3_positive(hidden, "hidden", error)) { + return false; + } + int B = hidden.shape(0); + int T = hidden.shape(1); + int H = hidden.shape(2); + if (!check_rank1_dim(norm, H, "norm", error) || + !check_rank2_positive(lm_head, "lm_head", error) || + !check_dim(lm_head, 1, H, "lm_head", "hidden width", error)) { + return false; + } + + auto last = (T == 1) ? mlx::core::reshape(hidden, {B, H}, device) + : mlx::core::reshape(mlx::core::slice(hidden, to_shape({0, T - 1, 0}), + to_shape({B, T, H}), device), + {B, H}, device); + + auto normed = mlx::core::fast::rms_norm(last, norm, (float)eps, device); + auto logits = + mlx::core::tensordot(normed, lm_head, std::vector{1}, std::vector{1}, device); + out = mlx::core::argmax(logits, 1, false, device); + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin final_greedy"; + return false; + } +} + +bool v_forward_greedy_from_hidden(const mlx::core::array &hidden, std::vector &layers, + 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, + int64_t &token_id_out, std::vector &k_out, + std::vector &v_out, std::string &error) { + try { + if (layers.size() != kv.size()) { + error = "Qwen3 greedy forward layers and kv_cache length mismatch"; + return false; + } + if (!check_rank3_positive(hidden, "hidden", error) || + !check_non_negative(offset, "offset", error) || + !check_positive(head_dim, "head_dim", error)) { + return false; + } + + auto current = hidden; + k_out.clear(); + v_out.clear(); + k_out.reserve(layers.size()); + v_out.reserve(layers.size()); + std::vector eval_arrays; + eval_arrays.reserve((layers.size() * 2) + 1); + + for (size_t i = 0; i < layers.size(); ++i) { + std::string layer_error; + if (!validate_dense_layer(current, layers[i], kv[i], offset, head_dim, layer_error)) { + error = layer_error; + return false; + } + + mlx::core::array k_new = *kv[i].k; + mlx::core::array v_new = *kv[i].v; + current = layer_dense_impl(current, layers[i], kv[i], offset, (float)scale, head_dim, + (float)theta, (float)eps, device, &k_new, &v_new); + + eval_arrays.push_back(k_new); + eval_arrays.push_back(v_new); + k_out.push_back(k_new); + v_out.push_back(v_new); + } + + int B = current.shape(0); + int T = current.shape(1); + int H = current.shape(2); + if (return_token_id && B != 1) { + error = "token_id return paths require batch size 1"; + return false; + } + if (!check_rank1_dim(norm, H, "norm", error) || + !check_rank2_positive(lm_head, "lm_head", error) || + !check_dim(lm_head, 1, H, "lm_head", "hidden width", error)) { + return false; + } + + auto last = (T == 1) ? mlx::core::reshape(current, {B, H}, device) + : mlx::core::reshape(mlx::core::slice(current, to_shape({0, T - 1, 0}), + to_shape({B, T, H}), device), + {B, H}, device); + + 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); + + eval_arrays.push_back(token); + mlx::core::async_eval(eval_arrays); + + if (return_token_id) { + token_id_out = token_to_int64(token); + } else { + token_out = token; + } + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin forward_greedy_from_hidden"; + return false; + } +} + +bool v_forward_greedy_ids_chunk(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 mlx::core::array &lm_head, int offset, int count, + double scale, int head_dim, double theta, double eps, + const mlx::core::Device &device, + std::vector &token_out, + std::vector &k_out, + std::vector &v_out, std::string &error) { + try { + if (count <= 0) { + error = "forward_greedy_ids_chunk expects positive count"; + return false; + } + if (layers.size() != initial_kv.size()) { + error = "forward_greedy_ids_chunk layers and kv_cache length mismatch"; + return false; + } + + size_t layer_count = layers.size(); + + if (!check_rank2_positive(input_ids, "input_ids", error) || + !check_rank2_positive(embed_tokens, "embed_tokens", error) || + !check_non_negative(offset, "offset", error) || + !check_positive(head_dim, "head_dim", error) || + !check_rank1_dim(norm, embed_tokens.shape(1), "norm", error) || + !check_rank2_positive(lm_head, "lm_head", error) || + !check_dim(lm_head, 1, embed_tokens.shape(1), "lm_head", "hidden width", error)) { + return false; + } + if (input_ids.shape(0) != 1) { + error = "forward_greedy_ids_chunk requires batch size 1"; + return false; + } + if (input_ids.shape(1) != 1) { + error = "forward_greedy_ids_chunk requires sequence length 1"; + return false; + } + + std::vector k_cache; + std::vector v_cache; + std::vector next_k_cache; + std::vector next_v_cache; + std::vector token_arrays; + k_cache.reserve(layer_count); + v_cache.reserve(layer_count); + next_k_cache.reserve(layer_count); + next_v_cache.reserve(layer_count); + token_arrays.reserve(count); + + auto current_ids = input_ids; + int current_offset = offset; + + 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 current = mlx::core::reshape(mlx::core::take(embed_tokens, ids, 0, device), + {B, T, embed_tokens.shape(1)}, device); + + next_k_cache.clear(); + next_v_cache.clear(); + + for (size_t layer_idx = 0; layer_idx < layer_count; ++layer_idx) { + KVCache kv = (step == 0) ? initial_kv[layer_idx] + : KVCache{&k_cache[layer_idx], &v_cache[layer_idx]}; + + std::string layer_error; + if (!validate_dense_layer(current, layers[layer_idx], kv, current_offset, head_dim, + layer_error)) { + error = layer_error; + return false; + } + + 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); + + next_k_cache.push_back(k_new); + next_v_cache.push_back(v_new); + } + + int B_out = current.shape(0); + int T_out = current.shape(1); + int H_out = current.shape(2); + + auto last = + (T_out == 1) + ? mlx::core::reshape(current, {B_out, H_out}, device) + : 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); + + 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); + + token_arrays.push_back(token); + current_ids = mlx::core::reshape(token, {B_out, 1}, device); + k_cache.swap(next_k_cache); + v_cache.swap(next_v_cache); + current_offset += 1; + } + + std::vector eval_arrays; + eval_arrays.reserve(token_arrays.size() + (layer_count * 2)); + eval_arrays.insert(eval_arrays.end(), token_arrays.begin(), token_arrays.end()); + for (size_t layer_idx = 0; layer_idx < layer_count; ++layer_idx) { + eval_arrays.push_back(k_cache[layer_idx]); + eval_arrays.push_back(v_cache[layer_idx]); + } + mlx::core::async_eval(eval_arrays); + + token_out = std::move(token_arrays); + k_out = std::move(k_cache); + v_out = std::move(v_cache); + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin forward_greedy_ids_chunk"; + return false; + } +} + +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, + std::vector &token_out, std::vector &k_out, + std::vector &v_out, std::string &error) { + try { + if (count <= 0) { + error = "forward_greedy_ids_chunk_quantized expects positive count"; + return false; + } + if (layers.size() != initial_kv.size()) { + error = "forward_greedy_ids_chunk_quantized layers and kv_cache length mismatch"; + return false; + } + + size_t layer_count = layers.size(); + + if (!check_rank2_positive(input_ids, "input_ids", error) || + !check_rank2_positive(embed_tokens, "embed_tokens", error) || + !check_non_negative(offset, "offset", error) || + !check_positive(head_dim, "head_dim", error) || + !check_rank1_dim(norm, embed_tokens.shape(1), "norm", error) || + !check_linear_weight_in(lm_head, embed_tokens.shape(1), "lm_head", error)) { + return false; + } + if (input_ids.shape(0) != 1) { + error = "forward_greedy_ids_chunk_quantized requires batch size 1"; + return false; + } + if (input_ids.shape(1) != 1) { + error = "forward_greedy_ids_chunk_quantized requires sequence length 1"; + return false; + } + + std::vector k_cache; + std::vector v_cache; + std::vector next_k_cache; + std::vector next_v_cache; + std::vector token_arrays; + k_cache.reserve(layer_count); + v_cache.reserve(layer_count); + next_k_cache.reserve(layer_count); + next_v_cache.reserve(layer_count); + token_arrays.reserve(count); + + auto current_ids = input_ids; + int current_offset = offset; + + 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 current = mlx::core::reshape(mlx::core::take(embed_tokens, ids, 0, device), + {B, T, embed_tokens.shape(1)}, device); + + next_k_cache.clear(); + next_v_cache.clear(); + + for (size_t layer_idx = 0; layer_idx < layer_count; ++layer_idx) { + KVCache kv = (step == 0) ? initial_kv[layer_idx] + : KVCache{&k_cache[layer_idx], &v_cache[layer_idx]}; + + std::string layer_error; + if (!validate_generalized_layer(current, layers[layer_idx], kv, current_offset, head_dim, + layer_error)) { + error = layer_error; + return false; + } + + mlx::core::array k_new = *kv.k; + 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); + + next_k_cache.push_back(k_new); + next_v_cache.push_back(v_new); + } + + int B_out = current.shape(0); + int T_out = current.shape(1); + int H_out = current.shape(2); + + auto last = + (T_out == 1) + ? mlx::core::reshape(current, {B_out, H_out}, device) + : 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); + + 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); + + token_arrays.push_back(token); + current_ids = mlx::core::reshape(token, {B_out, 1}, device); + k_cache.swap(next_k_cache); + v_cache.swap(next_v_cache); + current_offset += 1; + } + + std::vector eval_arrays; + eval_arrays.reserve(token_arrays.size() + (layer_count * 2)); + eval_arrays.insert(eval_arrays.end(), token_arrays.begin(), token_arrays.end()); + for (size_t layer_idx = 0; layer_idx < layer_count; ++layer_idx) { + eval_arrays.push_back(k_cache[layer_idx]); + eval_arrays.push_back(v_cache[layer_idx]); + } + mlx::core::async_eval(eval_arrays); + + token_out = std::move(token_arrays); + k_out = std::move(k_cache); + v_out = std::move(v_cache); + return true; + } catch (const std::exception &e) { + error = e.what(); + return false; + } catch (...) { + error = "Unknown error in qwen3 plugin forward_greedy_ids_chunk_quantized"; + return false; + } +} + +const VTable kVTable = { + v_kv_cache_attention, + v_mlp, + v_attention_residual, + v_attention_block, + v_layer_dense, + v_layer_quantized, + v_final_greedy, + v_forward_greedy_from_hidden, + v_forward_greedy_ids_chunk, + v_forward_greedy_ids_chunk_quantized, +}; + +} // namespace + +extern "C" const VTable *emlx_plugin_vtable() { return &kVTable; } diff --git a/emlx_axon/lib/emlx_axon.ex b/emlx_axon/lib/emlx_axon.ex index ddde54d..35248c2 100644 --- a/emlx_axon/lib/emlx_axon.ex +++ b/emlx_axon/lib/emlx_axon.ex @@ -16,7 +16,6 @@ defmodule EMLXAxon do | `:sdpa` | Bumblebee `attention_output_impl/3` | `EMLX.Fast.scaled_dot_product_attention_causal/4` or unmasked | | `:dropout` | `op_name: :dropout` (inference) | identity pass-through | | `:swiglu` | `:multiply(container(up, silu(gate)))` | `EMLX.Fast.swiglu/2` | - | `:native_attention` | Bumblebee causal self-attention | `EMLX.kv_cache_attention_masked/8` | ## Usage @@ -92,11 +91,6 @@ defmodule EMLXAxon do @bumblebee_attn_mfa {Bumblebee.Layers, :attention_output_impl, 3} @bumblebee_repeat_interleave_mfa {Bumblebee.Layers, :"-repeat_interleave/3-fun-0-", 2} @bumblebee_update_attn_cache_mfa {Bumblebee.Layers.Decoder, :update_attention_cache, 5} - @bumblebee_put_block_cache_mfa {Bumblebee.Layers.Decoder, :"-put_block_cache/3-fun-0-", 3} - @kv_cache_proc_key :"$emlx_axon_native_attention_kv_cache" - # Memoizes Nx.to_number(offset_tensor) across all layers of a single forward pass. - # Stores {MapSet.t(layer_key), cached_integer_offset}. - @step_offset_proc_key :"$emlx_axon_step_offset_cache" @default_rewrites [ :rms_norm, @@ -106,9 +100,7 @@ defmodule EMLXAxon do :dropout, :swiglu, :attn_weights, - :if_present, - :native_attention, - :nullify_block_cache + :if_present ] @doc """ @@ -118,7 +110,7 @@ defmodule EMLXAxon do * `:only` — list of atoms selecting which rewrites to apply. Defaults to `[:rms_norm, :layer_norm, :rotary_embedding, :sdpa, :dropout, :swiglu, - :attn_weights, :if_present, :native_attention, :nullify_block_cache]`. + :attn_weights, :if_present]`. Pass `:gqa_cache_fix` explicitly when targeting a Bumblebee build whose `init_cache` allocates the KV cache with `num_key_value_heads` rather than `num_attention_heads` (i.e. the upstream PR branch patch). @@ -129,7 +121,6 @@ defmodule EMLXAxon do model = EMLXAxon.rewrite(model, only: [:rms_norm, :layer_norm]) """ - @spec rewrite(Axon.t(), keyword()) :: Axon.t() def rewrite(%Axon{} = model, opts \\ []) do {:ok, opts} = Keyword.validate(opts, only: @default_rewrites) cache = :ets.new(:emlx_axon_rewrite_cache, [:set, :public]) @@ -148,8 +139,6 @@ defmodule EMLXAxon do |> maybe_add(:attn_weights, attn_weights_rewriter(), enabled) |> maybe_add(:if_present, if_present_rewriter(), enabled) |> maybe_add(:gqa_cache_fix, gqa_cache_fix_rewriter(), enabled) - |> maybe_add(:native_attention, native_attention_rewriter(), enabled) - |> maybe_add(:nullify_block_cache, nullify_block_cache_rewriter(), enabled) Axon.rewrite_nodes(model, fn node -> Enum.find_value(rewriters, :skip, fn {_key, fun} -> @@ -208,7 +197,6 @@ defmodule EMLXAxon do warnings are benign — the quantized tensors are still used correctly via the EMLX backend's quantized_matmul dispatch. """ - @spec load_quantized({:local, Path.t()}, keyword()) :: {:ok, map()} | {:error, term()} def load_quantized(source, opts \\ []) def load_quantized({:local, path}, opts) do @@ -249,8 +237,6 @@ defmodule EMLXAxon do Replaces `op_name: :rms_norm` nodes with `shift: 0.0` with an Axon layer that calls `EMLX.Fast.rms_norm/3` — a single fused Metal shader. """ - @spec rms_norm_rewriter() :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) def rms_norm_rewriter do fn %Axon.Node{op_name: :rms_norm, opts: node_opts, name: name_fn} -> @@ -300,8 +286,6 @@ defmodule EMLXAxon do Metal shader. Skips nodes where `channel_index` is not `-1` (last axis), as the kernel only normalises over the last axis. """ - @spec layer_norm_rewriter() :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) def layer_norm_rewriter do fn %Axon.Node{op_name: :layer_norm, opts: node_opts, name: name_fn} -> @@ -357,8 +341,6 @@ defmodule EMLXAxon do **Assumes sequential positions** — see `EMLXAxon` moduledoc for the limitation. """ - @spec rotary_embedding_rewriter(reference() | nil) :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) def rotary_embedding_rewriter(cache \\ nil) do fn %Axon.Node{op: op, opts: node_opts, name: name_fn} -> if function_info(op) == @bumblebee_rope_mfa do @@ -416,8 +398,6 @@ defmodule EMLXAxon do **Not appropriate for training graphs** — only enable this rewriter when the model will be used for inference only. """ - @spec dropout_rewriter() :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) def dropout_rewriter do fn %Axon.Node{op_name: :dropout, name: name_fn} -> @@ -442,8 +422,6 @@ defmodule EMLXAxon do consumes) becomes unreachable. Inference-only: attention weight tensors are never used for token generation. """ - @spec attn_weights_rewriter() :: - (Axon.Node.t() -> :skip | ([Axon.t(), ...], Axon.t() -> Axon.t())) def attn_weights_rewriter do fn %Axon.Node{op_name: :bb_attn_weights} -> @@ -478,8 +456,6 @@ defmodule EMLXAxon do **Do not enable for training graphs** — training models typically run without a KV cache and rely on the `else` branch. """ - @spec if_present_rewriter() :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) def if_present_rewriter do fn # 3-parent :if_present: [optional(condition), optional(on_true), optional(on_false)] @@ -524,8 +500,6 @@ defmodule EMLXAxon do Only applies when the key or value parent is a Bumblebee `repeat_interleave` node. Models without GQA (or where repeat_interleave is already absent) are unaffected. """ - @spec gqa_cache_fix_rewriter() :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) def gqa_cache_fix_rewriter do fn %Axon.Node{op: op, parent: [_key_id, _value_id | _]} -> @@ -561,8 +535,6 @@ defmodule EMLXAxon do Generic `:multiply` nodes (no `:container` parent, or container without a `:silu` child) are reconstructed identically. """ - @spec swiglu_rewriter() :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) def swiglu_rewriter do fn # Bumblebee's SwiGLU: multiply(container(up_proj, silu(gate))). @@ -646,8 +618,6 @@ defmodule EMLXAxon do **Inference-only**: attention dropout is elided (a no-op at inference time). Nodes with `dropout_rate > 0` are skipped to preserve training-time stochastic behaviour. """ - @spec sdpa_rewriter() :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) def sdpa_rewriter do fn %Axon.Node{op: op, opts: node_opts} -> dropout_rate = Keyword.get(node_opts, :dropout_rate, 0.0) @@ -774,317 +744,6 @@ defmodule EMLXAxon do ) end - # ── Native KV attention ───────────────────────────────────────────────────── - - @doc """ - Returns the rewriter function for Bumblebee causal self-attention nodes. - - This rewrite replaces the attention output with a single `Nx.runtime_call` - callback that updates a process-local ETS K/V cache and calls - `EMLX.kv_cache_attention_masked/8`. It intentionally only matches causal - attention without sliding-window masking; cross-attention and local attention - fall back to the original graph. - """ - @spec native_attention_rewriter() :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) - def native_attention_rewriter do - fn %Axon.Node{op: op, opts: node_opts} -> - dropout_rate = Keyword.get(node_opts, :dropout_rate, 0.0) - - if function_info(op) == @bumblebee_attn_mfa and dropout_rate == 0.0 do - original_op = op - - fn [weights_dropped_axon, v_axon], _placeholder -> - nodes = weights_dropped_axon.nodes - weights_dropped_id = weights_dropped_axon.output - - dropout_node = nodes[weights_dropped_id] - [attn_weights_id] = dropout_node.parent - attn_weights_node = nodes[attn_weights_id] - - causal = Keyword.get(attn_weights_node.opts, :causal, false) - window_size = Keyword.get(attn_weights_node.opts, :window_size) - scale_opt = Keyword.get(attn_weights_node.opts, :scale) - - # parents: [q_id, k_id, key_mask_id, head_mask_id, bias_id, offset_id] - [q_id, k_id, key_mask_id, _head_mask_id, _bias_id, _offset_id] = - attn_weights_node.parent - - with true <- causal, - true <- is_nil(window_size), - {:ok, k_key_id, _k_value_id, _k_cache_id, k_offset_id} <- - find_update_attention_cache(nodes, k_id), - {:ok, _v_key_id, v_value_id, _v_cache_id, v_offset_id} <- - find_update_attention_cache(nodes, v_axon.output), - true <- k_offset_id == v_offset_id do - q_axon = %Axon{output: q_id, nodes: nodes} - new_k_axon = maybe_strip_repeat_interleave(nodes, k_key_id) - new_v_axon = maybe_strip_repeat_interleave(nodes, v_value_id) - offset_axon = %Axon{output: k_offset_id, nodes: nodes} - - key_mask_axon = %Axon{output: key_mask_id, nodes: nodes} - - build_native_attention_layer( - q_axon, - new_k_axon, - new_v_axon, - offset_axon, - key_mask_axon, - scale_opt - ) - else - reason -> - IO.puts( - "[native_attention_rewriter] FALLTHROUGH causal=#{causal} window=#{inspect(window_size)} reason=#{inspect(reason)}" - ) - - Axon.layer(original_op, [weights_dropped_axon, v_axon]) - end - end - else - :skip - end - end - end - - defp build_native_attention_layer( - q_axon, - new_k_axon, - new_v_axon, - offset_axon, - key_mask_axon, - scale_opt - ) do - layer_key = make_ref() - - Axon.layer( - fn q, new_k, new_v, offset, key_mask, op_opts -> - out = Nx.template(Nx.shape(q), Nx.type(q)) - head_dim = elem(Nx.shape(q), 3) - scale = op_opts[:scale] || 1.0 / :math.sqrt(head_dim) - - # Both prefill and decode are handled entirely inside the callback: - # - Prefill (t_new > 1): callback computes SDPA eagerly, stores K/V in ETS. - # - Decode (t_new == 1): callback reads ETS cache, calls kv_cache_attention_masked. - Nx.runtime_call( - out, - {q, new_k, new_v, offset, key_mask}, - [layer_key: op_opts[:layer_key], scale: scale], - &__MODULE__.native_kv_attn_callback/2 - ) - end, - [q_axon, new_k_axon, new_v_axon, offset_axon, key_mask_axon], - op_name: :native_kv_attention, - layer_key: layer_key, - scale: if(is_number(scale_opt), do: scale_opt, else: nil) - ) - end - - @doc false - def native_kv_attn_callback( - {query, new_k, new_v, offset_tensor, key_mask}, - opts - ) do - layer_key = Keyword.fetch!(opts, :layer_key) - t_new = elem(Nx.shape(new_k), 1) - - if t_new > 1 do - # Prefill path: always read the actual offset tensor. - # - # The step-offset cache may hold a stale value from a previous serving's - # last decode step when this layer_key is brand new (never seen before). - # For prefill the optimization is irrelevant (only 1 GPU→CPU sync per - # layer anyway since it's a single step), so bypass get_step_offset. - # - # Also accumulate this layer_key into the seen set so decode step 1 finds - # it already "seen" and issues a fresh read instead of using the (now - # correct) prefill cached_offset at the wrong decode position. - offset = Nx.to_number(offset_tensor) - register_prefill_layer(layer_key, offset) - - if offset == 0 do - native_kv_prefill(query, new_k, new_v, key_mask, layer_key, opts) - else - native_kv_decode(query, new_k, new_v, offset, key_mask, layer_key, opts) - end - else - offset = get_step_offset(offset_tensor, layer_key) - native_kv_decode(query, new_k, new_v, offset, key_mask, layer_key, opts) - end - end - - # Registers a layer_key seen during the prefill step. - # Accumulates all prefill keys into the seen set so that decode step 1 - # will find each layer_key already "seen" and issue a fresh step-boundary read - # at the correct decode offset rather than reusing the prefill offset (0). - defp register_prefill_layer(layer_key, offset) do - new_state = - case Process.get(@step_offset_proc_key) do - nil -> {MapSet.new([layer_key]), offset} - {seen, _prev} -> {MapSet.put(seen, layer_key), offset} - end - - Process.put(@step_offset_proc_key, new_state) - end - - defp native_kv_prefill(query, new_k, new_v, key_mask, layer_key, opts) do - t_new = elem(Nx.shape(new_k), 1) - scale = Keyword.fetch!(opts, :scale) - - # Bumblebee compiles with max_length = seq_length + max_new_tokens, so - # key_mask is {B, max_length} while new_k is {B, seq_length, N_kv, D}. - # Pad new_k/new_v with zeros to max_length before storing so the - # decode path can retrieve a buffer of the expected size. - max_len = - case Nx.shape(key_mask) do - {_, max} -> max - {_, _, _, max} -> max - end - - pad_len = max_len - t_new - {b, _, nkv, d} = Nx.shape(new_k) - type = Nx.type(new_k) - - {k_full, v_full} = - if pad_len > 0 do - zeros_k = Nx.broadcast(Nx.tensor(0, type: type), {b, pad_len, nkv, d}) - zeros_v = Nx.broadcast(Nx.tensor(0, type: type), {b, pad_len, nkv, d}) - {Nx.concatenate([new_k, zeros_k], axis: 1), Nx.concatenate([new_v, zeros_v], axis: 1)} - else - {new_k, new_v} - end - - # Store raw {dev, ref} tuples — no ETS, no to_nx overhead for k/v. - cache_map = Process.get(@kv_cache_proc_key, %{}) - - Process.put( - @kv_cache_proc_key, - Map.put(cache_map, layer_key, {EMLX.Backend.from_nx(k_full), EMLX.Backend.from_nx(v_full)}) - ) - - # Compute prefill SDPA eagerly here in the callback (avoids splitting the - # computation across the Axon layer boundary and the Nx.add(sdpa, zeros) trick). - # Use k_full/v_full (padded to max_length) and the FULL key_mask rather than - # slicing both to t_new. This exactly matches the default sdpa rewriter path - # (T_kv = max_length), ensuring identical NaN-propagation behavior on Metal - # for left-padded input sequences. - - # Q/K/V: {B, T, N, D} → transpose to {B, N, T, D} for the NIF. - q_t = Nx.transpose(query, axes: [0, 2, 1, 3]) - k_t = Nx.transpose(k_full, axes: [0, 2, 1, 3]) - v_t = Nx.transpose(v_full, axes: [0, 2, 1, 3]) - - # kv_offset = 0 for prefill (lower-triangular causal mask from position 0). - sdpa_t = - EMLX.fast_sdpa_causal_key_masked( - EMLX.Backend.from_nx(q_t), - EMLX.Backend.from_nx(k_t), - EMLX.Backend.from_nx(v_t), - scale, - EMLX.Backend.from_nx(key_mask), - 0 - ) - |> EMLX.Backend.to_nx() - - # Transpose back to {B, T, N, D} and match query dtype. - out = Nx.transpose(sdpa_t, axes: [0, 2, 1, 3]) - if Nx.type(out) == Nx.type(query), do: out, else: Nx.as_type(out, Nx.type(query)) - end - - defp native_kv_decode(query, new_k, new_v, offset, key_mask, layer_key, opts) do - t_new = elem(Nx.shape(new_k), 1) - valid_len = offset + t_new - - {batch_size, max_length, mask_axis} = - case Nx.shape(key_mask) do - {batch_size, max_length} -> {batch_size, max_length, 1} - {batch_size, _heads, _query_len, max_length} -> {batch_size, max_length, 3} - end - - {_, _, kv_heads, head_dim} = Nx.shape(new_k) - full_shape = {batch_size, max_length, kv_heads, head_dim} - type = Nx.type(new_k) - scale = Keyword.fetch!(opts, :scale) - - # Get cached refs directly from process dict — no ETS lookup, no term copying. - cache_map = Process.get(@kv_cache_proc_key, %{}) - - {k_cache_ref, v_cache_ref} = - case Map.get(cache_map, layer_key) do - nil -> - # No prefill ran — initialize zero-filled cache buffer. - zeros = Nx.broadcast(Nx.tensor(0, type: type), full_shape) - {EMLX.Backend.from_nx(zeros), EMLX.Backend.from_nx(zeros)} - - refs -> - refs - end - - key_mask_sliced = Nx.slice_along_axis(key_mask, 0, valid_len, axis: mask_axis) - - {attn_ref, k_upd_ref, v_upd_ref} = - EMLX.kv_cache_attention_masked( - EMLX.Backend.from_nx(query), - EMLX.Backend.from_nx(new_k), - EMLX.Backend.from_nx(new_v), - k_cache_ref, - v_cache_ref, - offset, - scale, - EMLX.Backend.from_nx(key_mask_sliced) - ) - - # Store raw refs directly — no to_nx for k/v, no ETS insert. - Process.put(@kv_cache_proc_key, Map.put(cache_map, layer_key, {k_upd_ref, v_upd_ref})) - - attn_out = EMLX.Backend.to_nx(attn_ref) - - if Nx.type(attn_out) == Nx.type(query) do - attn_out - else - Nx.as_type(attn_out, Nx.type(query)) - end - end - - defp find_update_attention_cache(nodes, node_id, seen \\ MapSet.new()) do - cond do - MapSet.member?(seen, node_id) -> - :error - - node = nodes[node_id] -> - if function_info(node.op) == @bumblebee_update_attn_cache_mfa do - [key_id, value_id, cache_id, offset_id] = node.parent - {:ok, key_id, value_id, cache_id, offset_id} - else - seen = MapSet.put(seen, node_id) - - # For :if_present nodes, only follow parent[1] (the on_true / cache-present branch). - # parent[0] is optional(condition) which chains through put_block_cache to OTHER - # layers' UAC nodes. parent[2] is the fallback (no cache). Only parent[1] leads - # to the current layer's own UAC. - parents_to_search = - if node.op_name == :if_present do - case node.parent do - [_cond, on_true, _on_false] -> [on_true] - _ -> node.parent - end - else - node.parent - end - - Enum.find_value(parents_to_search, :error, fn parent_id -> - case find_update_attention_cache(nodes, parent_id, seen) do - {:ok, _key_id, _value_id, _cache_id, _offset_id} = found -> found - :error -> nil - end - end) - end - - true -> - :error - end - end - defp maybe_unwrap_optional(nodes, node_id) do case nodes[node_id] do %Axon.Node{op_name: :optional, parent: [inner_id]} -> inner_id @@ -1092,53 +751,6 @@ defmodule EMLXAxon do end end - @doc """ - Returns the rewriter function for Bumblebee block-cache update nodes. - - When native attention owns K/V state in an ETS table, the Axon block-cache - update chain is dead. Replacing `put_block_cache` with an identity lets DCE - prune `get_block_cache`, `update_attention_cache`, and container plumbing. - """ - @spec nullify_block_cache_rewriter() :: - (Axon.Node.t() -> ([Axon.t()], Axon.t() -> Axon.t()) | :skip) - def nullify_block_cache_rewriter do - fn %Axon.Node{op: op} -> - if function_info(op) == @bumblebee_put_block_cache_mfa do - fn [cache_axon, _block_cache_axon], _placeholder -> cache_axon end - else - :skip - end - end - end - - # Memoizes Nx.to_number(offset_tensor) within a single forward pass (decode step). - # - # All layer callbacks within one forward pass share the same offset value. Calling - # Nx.to_number once per layer forces one GPU→CPU sync per layer. Instead, call it - # once per step: detect the step boundary by tracking which layer_keys have been - # called in the current step. When a layer_key appears AGAIN (it was already seen in - # the previous step), we know a new step has begun and issue one fresh Nx.to_number; - # all other layers reuse the cache. - defp get_step_offset(offset_tensor, layer_key) do - case Process.get(@step_offset_proc_key) do - nil -> - offset = Nx.to_number(offset_tensor) - Process.put(@step_offset_proc_key, {MapSet.new([layer_key]), offset}) - offset - - {seen, cached_offset} -> - if MapSet.member?(seen, layer_key) do - # layer_key already seen in the prior step cycle → this is the start of a new step. - offset = Nx.to_number(offset_tensor) - Process.put(@step_offset_proc_key, {MapSet.new([layer_key]), offset}) - offset - else - Process.put(@step_offset_proc_key, {MapSet.put(seen, layer_key), cached_offset}) - cached_offset - end - end - end - # ── RoPE frequency precomputation ──────────────────────────────────────────── # Returns a precomputed {dims/2} inv-frequency tensor for strategies where @@ -1194,7 +806,6 @@ defmodule EMLXAxon do Note: Nx's `defnp` may compile to a closure rather than a named function, so this helper intentionally does not filter by `:erlang.fun_info(:type)`. """ - @spec function_info(term()) :: {module(), atom(), non_neg_integer()} | nil def function_info(fun) when is_function(fun) do {:module, m} = Function.info(fun, :module) {:name, n} = Function.info(fun, :name) @@ -1218,13 +829,7 @@ defmodule EMLXAxon do end end - defp expand_enabled(enabled) do - if :native_attention in enabled do - Enum.uniq([:if_present | enabled]) - else - enabled - end - end + defp expand_enabled(enabled), do: enabled defp maybe_add(acc, key, fun, enabled) do if key in enabled, do: [{key, fun} | acc], else: acc @@ -1263,7 +868,6 @@ defmodule EMLXAxon.QuantizeParams do * `:group_size` — quantization group size, must evenly divide in_features (default 64). * `:skip_vocab_threshold` — skip tensors whose first dim exceeds this (default 100_000). """ - @spec quantize(map(), keyword()) :: map() def quantize(params, opts \\ []) do bits = Keyword.get(opts, :bits, 4) group_size = Keyword.get(opts, :group_size, 64) diff --git a/emlx_axon/lib/emlx_axon/application.ex b/emlx_axon/lib/emlx_axon/application.ex new file mode 100644 index 0000000..8a2f905 --- /dev/null +++ b/emlx_axon/lib/emlx_axon/application.ex @@ -0,0 +1,34 @@ +defmodule EMLXAxon.Application do + @moduledoc """ + OTP Application for EMLXAxon. + + Eagerly loads the standalone qwen3 compute plugin (`c_src/qwen3_plugin.cpp`, + built as `priv/libemlx_qwen3.so`) into emlx's generic native plugin + registry via `EMLX.NIF.load_plugin/2`. Every `EMLX.Native.Qwen3.*` NIF + errors with `{:error, _}` until this has run — since `:emlx` is a + dependency of `:emlx_axon`, OTP starts it first, so `EMLX.NIF.load_plugin/2` + is always available by the time this runs. + """ + + use Application + + @doc false + def start(_type, _args) do + ensure_qwen3_plugin_loaded!() + Supervisor.start_link([], strategy: :one_for_one, name: __MODULE__) + end + + @doc false + def ensure_qwen3_plugin_loaded! do + path = :filename.join(:code.priv_dir(:emlx_axon), ~c"libemlx_qwen3.so") + + case EMLX.NIF.load_plugin("qwen3", List.to_string(path)) do + :ok -> + :ok + + {:error, reason} -> + raise EMLX.NIFError, + "EMLXAxon.Application could not load the qwen3 plugin: " <> List.to_string(reason) + end + end +end diff --git a/emlx_axon/lib/emlx_axon/mlx4bit_params.ex b/emlx_axon/lib/emlx_axon/mlx4bit_params.ex index b44acd6..fce9c02 100644 --- a/emlx_axon/lib/emlx_axon/mlx4bit_params.ex +++ b/emlx_axon/lib/emlx_axon/mlx4bit_params.ex @@ -38,7 +38,6 @@ defmodule EMLXAxon.MLX4BitParams do Returns `%Axon.ModelState{}` with BF16 tensors in Bumblebee key layout. All linear kernels are transposed from MLX `{out, in}` to Bumblebee `{in, out}`. """ - @spec load(Path.t()) :: Axon.ModelState.t() def load(mlx_path) do mlx_path = Path.expand(mlx_path) config = read_config(mlx_path) diff --git a/emlx_axon/lib/emlx_axon/qwen3/attention.ex b/emlx_axon/lib/emlx_axon/qwen3/attention.ex index aba6222..9b07f85 100644 --- a/emlx_axon/lib/emlx_axon/qwen3/attention.ex +++ b/emlx_axon/lib/emlx_axon/qwen3/attention.ex @@ -123,7 +123,7 @@ defmodule EMLXAxon.Qwen3.Attention do # The slice and SDPA are fused in the same lazy graph, so they land in one # Metal command buffer submission. {attn_ref, k_cache_ref, v_cache_ref} = - EMLX.qwen3_kv_cache_attention( + EMLX.Native.Qwen3.kv_cache_attention( EMLX.Backend.from_nx(q), EMLX.Backend.from_nx(k), EMLX.Backend.from_nx(v), @@ -161,7 +161,7 @@ defmodule EMLXAxon.Qwen3.Attention do theta = cfg.rope_theta {out_ref, k_cache_ref, v_cache_ref} = - EMLX.qwen3_attention_block( + EMLX.Native.Qwen3.attention_block( EMLX.Backend.from_nx(hidden), EMLX.Backend.from_nx(norm), EMLX.Backend.from_nx(q_proj), @@ -193,7 +193,7 @@ defmodule EMLXAxon.Qwen3.Attention do else residual_hidden |> EMLX.Backend.from_nx() - |> EMLX.qwen3_attention_residual(attn_ref, EMLX.Backend.from_nx(o_proj)) + |> EMLX.Native.Qwen3.attention_residual(attn_ref, EMLX.Backend.from_nx(o_proj)) |> EMLX.Backend.to_nx() end end diff --git a/emlx_axon/lib/emlx_axon/qwen3/dense_loader.ex b/emlx_axon/lib/emlx_axon/qwen3/dense_loader.ex index eea2378..75113c9 100644 --- a/emlx_axon/lib/emlx_axon/qwen3/dense_loader.ex +++ b/emlx_axon/lib/emlx_axon/qwen3/dense_loader.ex @@ -17,7 +17,6 @@ defmodule EMLXAxon.Qwen3.DenseLoader do @doc """ Converts a Bumblebee Qwen3 `model_info` map into a native dense state. """ - @spec from_model_info(map()) :: {:ok, State.t()} | {:error, term()} def from_model_info(%{params: %{data: params}, spec: %Bumblebee.Text.Qwen3{} = spec}) do config = config_from_spec(spec) @@ -71,7 +70,6 @@ defmodule EMLXAxon.Qwen3.DenseLoader do weights are transposed once at load time from safetensors `{out, in}` layout into the native dense `{in, out}` layout. """ - @spec from_safetensors_dir(Path.t(), keyword()) :: {:ok, State.t()} | {:error, term()} def from_safetensors_dir(path, opts \\ []) do path = Path.expand(path) @@ -88,8 +86,6 @@ defmodule EMLXAxon.Qwen3.DenseLoader do This is useful when a repository cache stores files by content addressed names rather than as a model directory. """ - @spec from_safetensors_files(Path.t(), [Path.t()], keyword()) :: - {:ok, State.t()} | {:error, term()} def from_safetensors_files(config_path, safetensors_paths, opts \\ []) do device = Keyword.get(opts, :device, :gpu) type = Keyword.get(opts, :type) diff --git a/emlx_axon/lib/emlx_axon/qwen3/generate.ex b/emlx_axon/lib/emlx_axon/qwen3/generate.ex index 69c38d3..30f14c3 100644 --- a/emlx_axon/lib/emlx_axon/qwen3/generate.ex +++ b/emlx_axon/lib/emlx_axon/qwen3/generate.ex @@ -76,8 +76,6 @@ defmodule EMLXAxon.Qwen3.Generate do (median ≈ steady-state) - `:total_ms` — wall time for the whole call, with microsecond resolution """ - @spec generate(Nx.Tensor.t(), Model.State.t(), keyword()) :: - {[non_neg_integer()], map()} def generate(input_ids, %Model.State{} = state, opts \\ []) do ensure_single_batch!(input_ids) diff --git a/emlx_axon/lib/emlx_axon/qwen3/loader.ex b/emlx_axon/lib/emlx_axon/qwen3/loader.ex index 92ee3fe..1a0a839 100644 --- a/emlx_axon/lib/emlx_axon/qwen3/loader.ex +++ b/emlx_axon/lib/emlx_axon/qwen3/loader.ex @@ -41,7 +41,6 @@ defmodule EMLXAxon.Qwen3.Loader do Returns `{:ok, %State{}}` or `{:error, reason}`. """ - @spec load(Path.t()) :: {:ok, State.t()} | {:error, term()} def load(path) do path = Path.expand(path) diff --git a/emlx_axon/lib/emlx_axon/qwen3/model.ex b/emlx_axon/lib/emlx_axon/qwen3/model.ex index d2384ed..79a7c4d 100644 --- a/emlx_axon/lib/emlx_axon/qwen3/model.ex +++ b/emlx_axon/lib/emlx_axon/qwen3/model.ex @@ -4,22 +4,22 @@ defmodule EMLXAxon.Qwen3.Model do ## Defn / JIT strategy - Every hot-path computation that consists purely of tensor arithmetic is - wrapped in a `defnp` kernel (declared in `Layers` and `Attention`). - `defnp` uses `Nx.Defn.Compiler.__jit__` — the same mechanism as - `Nx.Defn.jit/1` — to compile the function once per unique input shape and - cache the result. Subsequent calls skip Nx-side type/shape inference and - dispatch directly to the compiled kernel. - - Functions that still run eagerly: - - The quantized `Nx.dot` projections — the `Nx.Defn.Evaluator` (EMLX's - default compiler) cannot mix jit-argument `Nx.Defn.Expr` nodes with - captured `EMLX.Backend` tensors in closures. Wrapping individual ops - with `Nx.Defn.jit` would require a backend that implements - `Nx.Defn.Compiler` (e.g. EXLA). Dense native generation uses dedicated - EMLX Qwen3 primitives instead. - - `Nx.put_slice` KV-cache update (dynamic start index) and the valid-slice - read (dynamic end index). + For the quantized (MLX-4bit) and dense native-greedy paths actually + exercised by `bench/validate_qwen3.exs`, **every** hot-path call is a plain + eager function call against concrete `EMLX.Backend` tensors — + `EMLX.Native.Qwen3.*` NIFs, `EMLX.Fast.*` fused kernels, `Nx.dot` (quantized + dispatch) — with no `Nx.Defn.Expr` tracing or `Nx.Defn.Compiler` involved + anywhere in the loop. `Layers.swiglu/2` (the sole remaining `defn` in + `Layers`/`Attention`) is dead code (`mlp/5` below calls + `EMLX.Fast.swiglu/2` directly, not `Layers.swiglu/2`), and + `Sampler.top_p_gpu/3` is only reached by the `:top_p_gpu` sampler, not + `:greedy`. This is why a `:compiler` option threaded into + `EMLXAxon.TextGeneration.serving/3` would be a no-op regardless of whether + it's read out of `opts`: there is no `Nx.Defn` call site downstream for it + to select a compiler for. + + `Nx.put_slice` KV-cache update (dynamic start index) and the valid-slice + read (dynamic end index) also always ran eagerly, independent of the above. GPU sync: `EMLX.eval` is called once per token at the sampler boundary so the full lazy MLX graph spans all 28 layers before any CPU sync. @@ -82,7 +82,6 @@ defmodule EMLXAxon.Qwen3.Model do Returns a list of `{k_cache, v_cache}` pairs, one per transformer layer, where each cache is pre-allocated to `max_len` positions. """ - @spec init_kv_cache(State.t(), pos_integer()) :: [{Nx.Tensor.t(), Nx.Tensor.t()}] def init_kv_cache(%State{config: cfg, layers: layers}, max_len) do num_kv_heads = cfg.num_key_value_heads head_dim = cfg.head_dim @@ -111,7 +110,6 @@ defmodule EMLXAxon.Qwen3.Model do `%Nx.Tensor{}`. The native greedy Qwen3 NIFs accept this shape directly. Unsupported states fall back to `init_kv_cache/2`. """ - @spec init_native_kv_cache(State.t(), pos_integer()) :: kv_cache() def init_native_kv_cache(%State{} = state, max_len) do if native_forward_greedy?(state) do %State{config: cfg, layers: layers} = state @@ -139,8 +137,6 @@ defmodule EMLXAxon.Qwen3.Model do The cache is updated in-place via `Nx.put_slice`; `kv_cache_updated` is the same list with updated slices. """ - @spec forward(Nx.Tensor.t(), [{Nx.Tensor.t(), Nx.Tensor.t()}], non_neg_integer(), State.t()) :: - {Nx.Tensor.t(), [{Nx.Tensor.t(), Nx.Tensor.t()}]} def forward(input_ids, kv_cache, current_len, %State{} = state) do {hidden, kv_cache_updated} = forward_hidden(input_ids, kv_cache, current_len, state) @@ -159,8 +155,6 @@ defmodule EMLXAxon.Qwen3.Model do """ @type kv_cache :: [{Nx.Tensor.t() | EMLX.tensor_ref(), Nx.Tensor.t() | EMLX.tensor_ref()}] - @spec forward_greedy(Nx.Tensor.t(), kv_cache(), non_neg_integer(), State.t()) :: - {Nx.Tensor.t(), kv_cache()} def forward_greedy(input_ids, kv_cache, current_len, %State{} = state) do if native_forward_greedy?(state) do forward_native_greedy(input_ids, kv_cache, current_len, state) @@ -190,13 +184,6 @@ defmodule EMLXAxon.Qwen3.Model do a CPU `Nx.Tensor` plus backend transfer for every decode step. Other states fall back to the regular tensor based path. """ - @spec forward_greedy_decode_token_id( - non_neg_integer(), - kv_cache(), - non_neg_integer(), - State.t() - ) :: - {non_neg_integer(), kv_cache()} def forward_greedy_decode_token_id(token_id, kv_cache, current_len, %State{} = state) do if native_forward_greedy?(state) do forward_native_greedy_decode_token_id(token_id, kv_cache, current_len, state) @@ -213,17 +200,9 @@ defmodule EMLXAxon.Qwen3.Model do generated token tensors on the EMLX backend and returns the final KV cache, avoiding one Elixir/NIF boundary crossing per decoded token. """ - @spec forward_greedy_chunk( - Nx.Tensor.t(), - kv_cache(), - non_neg_integer(), - pos_integer(), - State.t() - ) :: - {[Nx.Tensor.t()], kv_cache()} | :fallback def forward_greedy_chunk(input_ids, kv_cache, current_len, count, %State{} = state) when is_integer(count) and count > 0 do - if native_forward_greedy?(state) do + if native_forward_greedy_chunk?(state) do forward_native_greedy_chunk(input_ids, kv_cache, current_len, count, state) else :fallback @@ -264,6 +243,16 @@ defmodule EMLXAxon.Qwen3.Model do defp native_forward_greedy?(%State{config: cfg, lm_head: lm_head}), do: cfg[:dense_layers?] == true and not EMLX.Quantization.quantized?(lm_head) + # Relaxed gate for `forward_greedy_chunk/5` only: `qwen3_layer_quantized`/ + # `qwen3_forward_greedy_ids_chunk_quantized` handle dense, quantized, and + # mixed per-projection weights (including a quantized `lm_head`) + # uniformly, so every `State` qualifies for the native chunked decode path + # — unlike `native_forward_greedy?/1` above, which still gates the + # single-step native paths (`forward_greedy`/`_token_id`/`_decode_token_id`) + # that call the dense-only `qwen3_forward_greedy_ids*` NIFs and are + # untouched by this change. + defp native_forward_greedy_chunk?(%State{}), do: true + defp forward_native_greedy(input_ids, kv_cache, current_len, %State{} = state) do %State{embed_tokens: embed_tokens, layers: layers, norm: norm, lm_head: lm_head, config: cfg} = state @@ -272,7 +261,7 @@ defmodule EMLXAxon.Qwen3.Model do embed_ref = EMLX.Backend.from_nx(embed_tokens) {token_ref, kv_cache_refs} = - EMLX.qwen3_forward_greedy_ids( + EMLX.Native.Qwen3.forward_greedy_ids( input_ids_ref(input_ids, embed_ref), embed_ref, layers, @@ -302,7 +291,7 @@ defmodule EMLXAxon.Qwen3.Model do embed_ref = EMLX.Backend.from_nx(embed_tokens) {token_id, kv_cache_refs} = - EMLX.qwen3_forward_greedy_ids_token_id( + EMLX.Native.Qwen3.forward_greedy_ids_token_id( input_ids_ref(input_ids, embed_ref), embed_ref, layers, @@ -327,7 +316,7 @@ defmodule EMLXAxon.Qwen3.Model do embed_ref = EMLX.Backend.from_nx(embed_tokens) {next_token_id, kv_cache_refs} = - EMLX.qwen3_forward_greedy_token_id( + EMLX.Native.Qwen3.forward_greedy_token_id( token_id, embed_ref, layers, @@ -352,20 +341,42 @@ defmodule EMLXAxon.Qwen3.Model do embed_ref = EMLX.Backend.from_nx(embed_tokens) {token_refs, kv_cache_refs} = - EMLX.qwen3_forward_greedy_ids_chunk( - input_ids_ref(input_ids, embed_ref), - embed_ref, - layers, - kv_cache, - EMLX.Backend.from_nx(norm), - EMLX.Backend.from_nx(lm_head), - current_len, - count, - scale, - head_dim, - theta, - eps - ) + if native_forward_greedy?(state) do + # All-dense layers + dense lm_head: keep using the proven, dense-only + # fused chunk NIF — zero regression risk to the working fast path. + EMLX.Native.Qwen3.forward_greedy_ids_chunk( + input_ids_ref(input_ids, embed_ref), + embed_ref, + layers, + kv_cache, + EMLX.Backend.from_nx(norm), + EMLX.Backend.from_nx(lm_head), + current_len, + count, + scale, + head_dim, + theta, + eps + ) + else + # Any quantized layer projection and/or quantized lm_head: the + # generalized chunk NIF fuses the whole chunk into 1 NIF call instead + # of falling back to per-token/per-op decoding. + EMLX.Native.Qwen3.forward_greedy_ids_chunk_quantized( + input_ids_ref(input_ids, embed_ref), + embed_ref, + layers, + kv_cache, + EMLX.Backend.from_nx(norm), + lm_head, + current_len, + count, + scale, + head_dim, + theta, + eps + ) + end tokens = Enum.map(token_refs, fn token_ref -> @@ -410,7 +421,7 @@ defmodule EMLXAxon.Qwen3.Model do else hidden |> EMLX.Backend.from_nx() - |> EMLX.qwen3_final_greedy( + |> EMLX.Native.Qwen3.final_greedy( EMLX.Backend.from_nx(norm), EMLX.Backend.from_nx(lm_head), cfg.rms_norm_eps @@ -471,25 +482,29 @@ defmodule EMLXAxon.Qwen3.Model do [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj], &EMLX.Quantization.quantized?/1 ) do - {attn_out, k_new, v_new} = - Attention.forward( - hidden, - norm1, - k_cache, - v_cache, - current_len, - q_proj, - k_proj, - v_proj, - o_proj, - q_norm, - k_norm, - cfg - ) - - hidden = mlp(attn_out, norm2, gate_proj, up_proj, down_proj, cfg.rms_norm_eps) - - {hidden, k_new, v_new} + # `qwen3_layer_quantized` is a strict superset of `qwen3_layer`: it + # accepts any dense-or-quantized mix of the 7 projections, fusing the + # ~13 per-op NIF calls the old `Attention.forward` + `mlp/5` path made + # for a quantized layer down to 1. `Attention.forward`/quantized `mlp` + # stay defined below as an unreferenced rollback option. + layer_generalized( + hidden, + norm1, + q_norm, + k_norm, + q_proj, + k_proj, + v_proj, + o_proj, + gate_proj, + up_proj, + down_proj, + norm2, + k_cache, + v_cache, + current_len, + cfg + ) else layer( hidden, @@ -535,7 +550,7 @@ defmodule EMLXAxon.Qwen3.Model do theta = cfg.rope_theta {hidden_ref, k_cache_ref, v_cache_ref} = - EMLX.qwen3_layer( + EMLX.Native.Qwen3.layer( EMLX.Backend.from_nx(hidden), EMLX.Backend.from_nx(norm1), EMLX.Backend.from_nx(q_proj), @@ -564,6 +579,61 @@ defmodule EMLXAxon.Qwen3.Model do } end + # Generalized variant of `layer/16`: q/k/v/o/gate/up/down each independently + # accept a dense or quantized (`EMLX.quantize/2`) `Nx.Tensor`, via + # `EMLX.Native.Qwen3.layer_quantized/19`. + defp layer_generalized( + hidden, + norm1, + q_norm, + k_norm, + q_proj, + k_proj, + v_proj, + o_proj, + gate_proj, + up_proj, + down_proj, + norm2, + k_cache, + v_cache, + current_len, + cfg + ) do + head_dim = cfg.head_dim + scale = 1.0 / :math.sqrt(head_dim) + theta = cfg.rope_theta + + {hidden_ref, k_cache_ref, v_cache_ref} = + EMLX.Native.Qwen3.layer_quantized( + EMLX.Backend.from_nx(hidden), + norm1, + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + EMLX.Backend.from_nx(k_cache), + EMLX.Backend.from_nx(v_cache), + norm2, + gate_proj, + up_proj, + down_proj, + current_len, + scale, + head_dim, + theta, + cfg.rms_norm_eps + ) + + { + EMLX.Backend.to_nx(hidden_ref), + EMLX.Backend.to_nx(k_cache_ref), + EMLX.Backend.to_nx(v_cache_ref) + } + end + defp dense_transformer_layer( hidden, k_cache, @@ -574,7 +644,7 @@ defmodule EMLXAxon.Qwen3.Model do {head_dim, scale, theta, eps} ) do {hidden_ref, k_cache_ref, v_cache_ref} = - EMLX.qwen3_layer( + EMLX.Native.Qwen3.layer( EMLX.Backend.from_nx(hidden), EMLX.Backend.from_nx(norm1), EMLX.Backend.from_nx(q_proj), @@ -613,7 +683,7 @@ defmodule EMLXAxon.Qwen3.Model do else hidden |> EMLX.Backend.from_nx() - |> EMLX.qwen3_mlp( + |> EMLX.Native.Qwen3.mlp( EMLX.Backend.from_nx(norm2), EMLX.Backend.from_nx(gate_proj), EMLX.Backend.from_nx(up_proj), diff --git a/emlx_axon/lib/emlx_axon/text_generation.ex b/emlx_axon/lib/emlx_axon/text_generation.ex index 2905751..820e766 100644 --- a/emlx_axon/lib/emlx_axon/text_generation.ex +++ b/emlx_axon/lib/emlx_axon/text_generation.ex @@ -56,13 +56,6 @@ defmodule EMLXAxon.TextGeneration do `EMLXAxon.Qwen3.Model.State` and want to avoid serving preprocessing overhead. The return shape matches `serving/3` for the same `:output_format`. """ - @spec run( - Bumblebee.Tokenizer.t(), - Model.State.t(), - binary() | %{text: binary()}, - keyword() - ) :: - map() def run(tokenizer, %Model.State{} = state, input, opts \\ []) do max_new = max_new_tokens!(opts, 100) configured_max_len = Keyword.get(opts, :max_len, 2048) @@ -211,13 +204,6 @@ defmodule EMLXAxon.TextGeneration do the response. Samplers that cannot use deferred host sync fall back to emitting each token. """ - @spec stream( - Bumblebee.Tokenizer.t(), - Model.State.t(), - binary() | map(), - (binary() -> term()), - keyword() - ) :: %{token_summary: map()} def stream(tokenizer, %Model.State{} = state, input, emit_fun, opts \\ []) when is_function(emit_fun, 1) do max_new = max_new_tokens!(opts, 100) @@ -323,7 +309,6 @@ defmodule EMLXAxon.TextGeneration do `:bumblebee` returns `%{text: text, token_summary: summary}` compatible with `Bumblebee.Text.generation/4` (default `:native`) """ - @spec serving(Bumblebee.Tokenizer.t(), Model.State.t(), keyword()) :: Nx.Serving.t() def serving(tokenizer, state, opts \\ []) do max_new = max_new_tokens!(opts, 100) configured_max_len = Keyword.get(opts, :max_len, 2048) @@ -338,7 +323,9 @@ defmodule EMLXAxon.TextGeneration do fn batch -> # `Nx.Batch` stays lazy until a defn/jit entry; `jit_apply(identity)` is the # supported way to concatenate stacked entries into concrete tensors. - %{"input_ids" => input_ids} = Nx.Defn.jit_apply(&Function.identity/1, [batch]) + %{"input_ids" => input_ids} = + Nx.Defn.jit_apply(&Function.identity/1, [batch], opts[:defn_options]) + ensure_single_batch!(input_ids) # Pre-alloc KV cache per call. The cache is safe to reuse because @@ -398,7 +385,6 @@ defmodule EMLXAxon.TextGeneration do The tokenizer is expected to come from the same directory (same `tokenizer.json`). Loading both from the same directory avoids chat-template / BOS-token divergence. """ - @spec from_mlx4bit(Path.t(), Bumblebee.Tokenizer.t(), keyword()) :: Nx.Serving.t() def from_mlx4bit(checkpoint_path, tokenizer, opts \\ []) do {:ok, state} = Loader.load(checkpoint_path) serving(tokenizer, state, opts) diff --git a/emlx_axon/mix.exs b/emlx_axon/mix.exs index cbe9ffa..d008d52 100644 --- a/emlx_axon/mix.exs +++ b/emlx_axon/mix.exs @@ -14,17 +14,36 @@ defmodule EMLXAxon.MixProject do docs: docs(), aliases: aliases(), description: "Axon model rewrites to swap supported nodes for EMLX.Fast Metal shaders", - package: package() + package: package(), + # elixir_make — builds the standalone qwen3 compute plugin + # (c_src/qwen3_plugin.cpp), `dlopen`'d at runtime by emlx's + # `EMLX.NIF.load_plugin/2` (see lib/emlx_axon/application.ex). A + # function ref so `Application.app_dir(:emlx, ...)` (needs the + # :emlx dep already compiled) isn't called while this project/0 + # map is being built — mirrors emlx's own `Fine.include_dir()` + # deferral in its mix.exs. + make_env: fn -> + emlx_priv_dir = Application.app_dir(:emlx, "priv") + + %{ + "MLX_INCLUDE_DIR" => Path.join(emlx_priv_dir, "mlx/include"), + "MLX_LIB_DIR" => Path.join(emlx_priv_dir, "mlx/lib"), + "QWEN3_ABI_INCLUDE_DIR" => Path.join(Mix.Project.deps_paths()[:emlx], "c_src/emlx_fast") + } + end, + compilers: [:elixir_make] ++ Mix.compilers() ] end def application do - [extra_applications: [:logger]] + [extra_applications: [:logger], mod: {EMLXAxon.Application, []}] end defp deps do [ - {:emlx, path: "../emlx", override: true}, + {:elixir_make, "~> 0.6"}, + {:emlx, path: "../emlx"}, + {:nx, github: "elixir-nx/nx", branch: "main", sparse: "nx", override: true}, {:axon, "~> 0.7"}, {:bumblebee, "~> 0.7"}, {:ex_doc, "~> 0.34", only: :docs} diff --git a/emlx_axon/mix.lock b/emlx_axon/mix.lock index 22b6d5f..2804a68 100644 --- a/emlx_axon/mix.lock +++ b/emlx_axon/mix.lock @@ -8,12 +8,13 @@ "elixir_make": {:hex, :elixir_make, "0.10.0", "16577e2583a79bb79237bbff349619ef5d80afffc07eac6e4faf0d00e2ddaf7d", [:mix], [], "hexpm", "dc1f09fb7fa68866b886abd5f0f3c83553b1a19a52359a899e92af1bb3b31982"}, "emlx": {:hex, :emlx, "0.3.0", "71fa22717d4a768ed2353f052845b2cbd6840c9e50dc67b1c9ccd5731d54ecf8", [:make, :mix], [{:elixir_make, "~> 0.6", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:nx, "~> 0.12", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "2778038adb05c4a2813a24eb5a89112c41a0b5b00350f059b348af0e01891c58"}, "ex_doc": {:hex, :ex_doc, "0.40.2", "f50edec428c4b0a457a167de42414c461122a3585a99515a69d09fff19e5597e", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "4fa426e2beb47854a162e2c488727fdec51cd4692e319b23810c2804cb1a40fe"}, + "fine": {:hex, :fine, "0.1.6", "4bf7151493443c454aac9f2fa2f34f5fefd0346a83fb5586a016c4a135c63247", [:mix], [], "hexpm", "5638eb4495488e885ebec167fa57973e5c35e1a50c344eb7666c90ec1c4e3b12"}, "jason": {:hex, :jason, "1.4.5", "2e3a008590b0b8d7388c20293e9dcc9cf3e5d642fd2a114e4cbbb52e595d940a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0 or ~> 3.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "b0c823996102bcd0239b3c2444eb00409b72f6a140c1950bc8b457d836b30684"}, "makeup": {:hex, :makeup, "1.2.1", "e90ac1c65589ef354378def3ba19d401e739ee7ee06fb47f94c687016e3713d1", [:mix], [{:nimble_parsec, "~> 1.4", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "d36484867b0bae0fea568d10131197a4c2e47056a6fbe84922bf6ba71c8d17ce"}, "makeup_elixir": {:hex, :makeup_elixir, "1.0.1", "e928a4f984e795e41e3abd27bfc09f51db16ab8ba1aebdba2b3a575437efafc2", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "7284900d412a3e5cfd97fdaed4f5ed389b8f2b4cb49efc0eb3bd10e2febf9507"}, "makeup_erlang": {:hex, :makeup_erlang, "1.1.0", "835f7e60792e08824cda445639555d7bf1bbbddb1b60b306e33cb6f6db24dc74", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "1cd6780fb1dd1a03979abaed0fe82712b0625118fd5257d3ebbf73f960c73c3c"}, "nimble_parsec": {:hex, :nimble_parsec, "1.4.2", "8efba0122db06df95bfaa78f791344a89352ba04baedd3849593bfce4d0dc1c6", [:mix], [], "hexpm", "4b21398942dda052b403bbe1da991ccd03a053668d147d53fb8c4e0efe09c973"}, - "nx": {:hex, :nx, "0.12.1", "6e9fee43a77646d04faad2ba4e449b9e270c6b23e413f203cb4d81e71c0a617f", [:mix], [{:complex, "~> 0.7", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "ee80f6ae898f68bbfe7f30216b06aab10096231ec99ded1208a827256b23d0fb"}, + "nx": {:git, "https://github.com/elixir-nx/nx.git", "da2c15c9beb03ed3d7edd79ea8ce6d9377775631", [branch: "main", sparse: "nx"]}, "nx_image": {:hex, :nx_image, "0.1.2", "0c6e3453c1dc30fc80c723a54861204304cebc8a89ed3b806b972c73ee5d119d", [:mix], [{:nx, "~> 0.4", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "9161863c42405ddccb6dbbbeae078ad23e30201509cc804b3b3a7c9e98764b81"}, "nx_signal": {:hex, :nx_signal, "0.2.0", "e1ca0318877b17c81ce8906329f5125f1e2361e4c4235a5baac8a95ee88ea98e", [:mix], [{:nx, "~> 0.6", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "7247e5e18a177a59c4cb5355952900c62fdeadeb2bad02a9a34237b68744e2bb"}, "polaris": {:hex, :polaris, "0.1.0", "dca61b18e3e801ecdae6ac9f0eca5f19792b44a5cb4b8d63db50fc40fc038d22", [:mix], [{:nx, "~> 0.5", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "13ef2b166650e533cb24b10e2f3b8ab4f2f449ba4d63156e8c569527f206e2c2"}, diff --git a/emlx_axon/test/emlx/qwen3_generate_test.exs b/emlx_axon/test/emlx/qwen3_generate_test.exs index e51f585..4f8a221 100644 --- a/emlx_axon/test/emlx/qwen3_generate_test.exs +++ b/emlx_axon/test/emlx/qwen3_generate_test.exs @@ -625,9 +625,9 @@ defmodule EMLXAxon.Qwen3GenerateTest do |> Nx.backend_transfer({EMLX.Backend, device: :gpu}) assert_raise ArgumentError, - ~r/qwen3_forward_greedy_ids_token_id requires batch size 1, got batch size 2/, + ~r/forward_greedy_ids_token_id requires batch size 1, got batch size 2/, fn -> - EMLX.qwen3_forward_greedy_ids_token_id( + EMLX.Native.Qwen3.forward_greedy_ids_token_id( EMLX.Backend.from_nx(input_ids), EMLX.Backend.from_nx(state.embed_tokens), [], @@ -656,8 +656,8 @@ defmodule EMLXAxon.Qwen3GenerateTest do embed_ref = EMLX.Backend.from_nx(state.embed_tokens) - assert_raise ArgumentError, ~r/qwen3_forward_greedy_ids_chunk requires batch size 1/, fn -> - EMLX.qwen3_forward_greedy_ids_chunk( + assert_raise ArgumentError, ~r/forward_greedy_ids_chunk requires batch size 1/, fn -> + EMLX.Native.Qwen3.forward_greedy_ids_chunk( EMLX.Backend.from_nx(input_ids), embed_ref, [], diff --git a/emlx_axon/test/emlx/qwen3_native_test.exs b/emlx_axon/test/emlx/qwen3_native_test.exs new file mode 100644 index 0000000..48f0581 --- /dev/null +++ b/emlx_axon/test/emlx/qwen3_native_test.exs @@ -0,0 +1,1511 @@ +defmodule EMLXAxon.Qwen3NativeTest do + @moduledoc """ + Unit tests for the qwen3 native accelerators — `EMLX.Native.Qwen3.*` NIFs, + backed by the standalone qwen3 compute plugin loaded from this project + (see `c_src/qwen3_plugin.cpp` and `EMLXAxon.Application`). + + Each test verifies: + 1. The NIF returns a tensor with the correct shape and dtype. + 2. Numerical output is close to the equivalent pure-Nx reference. + + All tests require Metal (tagged :metal). + """ + use ExUnit.Case, async: false + import Nx.Testing + + @moduletag :metal + + defp gpu(tensor), do: Nx.backend_transfer(tensor, {EMLX.Backend, device: :gpu}) + + describe "EMLX.Native.Qwen3.kv_cache_attention/9 validation" do + test "rejects malformed Q rank before reading shapes" do + q = Nx.broadcast(0.0, {1, 1, 4}) |> Nx.as_type(:f16) |> gpu() + k = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() + v = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() + k_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() + v_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() + + assert_raise EMLX.NIFError, ~r/q expects rank 4/, fn -> + EMLX.Native.Qwen3.kv_cache_attention( + EMLX.Backend.from_nx(q), + EMLX.Backend.from_nx(k), + EMLX.Backend.from_nx(v), + EMLX.Backend.from_nx(k_cache), + EMLX.Backend.from_nx(v_cache), + 0, + 1.0 / :math.sqrt(4), + 4, + 10_000.0 + ) + end + end + + test "rejects negative offsets" do + q = Nx.broadcast(0.0, {1, 1, 2, 4}) |> Nx.as_type(:f16) |> gpu() + k = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() + v = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() + k_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() + v_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() + + assert_raise EMLX.NIFError, ~r/offset must be non-negative/, fn -> + EMLX.Native.Qwen3.kv_cache_attention( + EMLX.Backend.from_nx(q), + EMLX.Backend.from_nx(k), + EMLX.Backend.from_nx(v), + EMLX.Backend.from_nx(k_cache), + EMLX.Backend.from_nx(v_cache), + -1, + 1.0 / :math.sqrt(4), + 4, + 10_000.0 + ) + end + end + + test "rejects cache capacity overflow" do + q = Nx.broadcast(0.0, {1, 1, 2, 4}) |> Nx.as_type(:f16) |> gpu() + k = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() + v = Nx.broadcast(0.0, {1, 1, 1, 4}) |> Nx.as_type(:f16) |> gpu() + k_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() + v_cache = Nx.broadcast(0.0, {1, 1, 4, 4}) |> Nx.as_type(:f16) |> gpu() + + assert_raise EMLX.NIFError, ~r/KV cache capacity 4 is smaller than required length 5/, fn -> + EMLX.Native.Qwen3.kv_cache_attention( + EMLX.Backend.from_nx(q), + EMLX.Backend.from_nx(k), + EMLX.Backend.from_nx(v), + EMLX.Backend.from_nx(k_cache), + EMLX.Backend.from_nx(v_cache), + 4, + 1.0 / :math.sqrt(4), + 4, + 10_000.0 + ) + end + end + + test "matches pure Nx reference for GQA prefill and updates caches" do + {q_cpu, k_cpu, v_cpu, k_cache_cpu, v_cache_cpu} = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + { + Nx.iota({1, 2, 4, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(100), + Nx.iota({1, 2, 2, 4}, type: :f32) |> Nx.add(11) |> Nx.divide(110), + Nx.iota({1, 2, 2, 4}, type: :f32) |> Nx.add(17) |> Nx.divide(120), + Nx.iota({1, 2, 5, 4}, type: :f32) |> Nx.divide(1_000), + Nx.iota({1, 2, 5, 4}, type: :f32) |> Nx.add(50) |> Nx.divide(1_000) + } + end) + + offset = 1 + head_dim = 4 + theta = 10_000.0 + scale = 1.0 / :math.sqrt(head_dim) + + {attn_ref, k_ref, v_ref} = + EMLX.Native.Qwen3.kv_cache_attention( + EMLX.Backend.from_nx(gpu(q_cpu)), + EMLX.Backend.from_nx(gpu(k_cpu)), + EMLX.Backend.from_nx(gpu(v_cpu)), + EMLX.Backend.from_nx(gpu(k_cache_cpu)), + EMLX.Backend.from_nx(gpu(v_cache_cpu)), + offset, + scale, + head_dim, + theta + ) + + {expected_attn, expected_k, expected_v} = + qwen3_kv_cache_attention_reference( + q_cpu, + k_cpu, + v_cpu, + k_cache_cpu, + v_cache_cpu, + offset, + scale, + head_dim, + theta + ) + + assert_all_close( + attn_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_attn, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_k, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_v, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + end + + test "matches pure Nx reference for batched GQA prefill with nonzero offset" do + {q_cpu, k_cpu, v_cpu, k_cache_cpu, v_cache_cpu} = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + { + Nx.iota({2, 2, 4, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(100), + Nx.iota({2, 2, 2, 4}, type: :f32) |> Nx.add(23) |> Nx.divide(130), + Nx.iota({2, 2, 2, 4}, type: :f32) |> Nx.add(37) |> Nx.divide(140), + Nx.iota({2, 2, 6, 4}, type: :f32) |> Nx.add(5) |> Nx.divide(1_000), + Nx.iota({2, 2, 6, 4}, type: :f32) |> Nx.add(95) |> Nx.divide(1_000) + } + end) + + offset = 2 + head_dim = 4 + theta = 10_000.0 + scale = 1.0 / :math.sqrt(head_dim) + + {attn_ref, k_ref, v_ref} = + EMLX.Native.Qwen3.kv_cache_attention( + EMLX.Backend.from_nx(gpu(q_cpu)), + EMLX.Backend.from_nx(gpu(k_cpu)), + EMLX.Backend.from_nx(gpu(v_cpu)), + EMLX.Backend.from_nx(gpu(k_cache_cpu)), + EMLX.Backend.from_nx(gpu(v_cache_cpu)), + offset, + scale, + head_dim, + theta + ) + + {expected_attn, expected_k, expected_v} = + qwen3_kv_cache_attention_reference( + q_cpu, + k_cpu, + v_cpu, + k_cache_cpu, + v_cache_cpu, + offset, + scale, + head_dim, + theta + ) + + assert Nx.shape(expected_attn) == {2, 2, 16} + + assert_all_close( + attn_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_attn, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_k, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_v, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + end + end + + describe "dense Qwen3 native helpers" do + test "qwen3_mlp matches pure Nx reference" do + {hidden_cpu, norm_cpu, gate_cpu, up_cpu, down_cpu} = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + { + Nx.iota({1, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(10), + Nx.tensor([1.0, 1.1, 0.9, 1.2], type: :f32), + Nx.iota({4, 6}, type: :f32) |> Nx.add(1) |> Nx.divide(50), + Nx.iota({4, 6}, type: :f32) |> Nx.add(7) |> Nx.divide(60), + Nx.iota({6, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(70) + } + end) + + eps = 1.0e-6 + + out_ref = + EMLX.Native.Qwen3.mlp( + EMLX.Backend.from_nx(gpu(hidden_cpu)), + EMLX.Backend.from_nx(gpu(norm_cpu)), + EMLX.Backend.from_nx(gpu(gate_cpu)), + EMLX.Backend.from_nx(gpu(up_cpu)), + EMLX.Backend.from_nx(gpu(down_cpu)), + eps + ) + + expected = qwen3_mlp_reference(hidden_cpu, norm_cpu, gate_cpu, up_cpu, down_cpu, eps) + + assert_all_close( + out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + end + + test "qwen3_attention_block matches pure Nx reference" do + fixtures = qwen3_dense_attention_fixtures() + scale = 1.0 / :math.sqrt(fixtures.head_dim) + + {out_ref, k_ref, v_ref} = + EMLX.Native.Qwen3.attention_block( + EMLX.Backend.from_nx(gpu(fixtures.hidden)), + EMLX.Backend.from_nx(gpu(fixtures.norm1)), + EMLX.Backend.from_nx(gpu(fixtures.q_proj)), + EMLX.Backend.from_nx(gpu(fixtures.k_proj)), + EMLX.Backend.from_nx(gpu(fixtures.v_proj)), + EMLX.Backend.from_nx(gpu(fixtures.o_proj)), + EMLX.Backend.from_nx(gpu(fixtures.q_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_cache)), + EMLX.Backend.from_nx(gpu(fixtures.v_cache)), + fixtures.offset, + scale, + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + + {expected, expected_k, expected_v} = qwen3_attention_block_reference(fixtures) + + assert_all_close( + out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_k, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_v, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + end + + test "qwen3_attention_block matches pure Nx reference for batched GQA" do + fixtures = qwen3_batched_dense_attention_fixtures() + scale = 1.0 / :math.sqrt(fixtures.head_dim) + + {out_ref, k_ref, v_ref} = + EMLX.Native.Qwen3.attention_block( + EMLX.Backend.from_nx(gpu(fixtures.hidden)), + EMLX.Backend.from_nx(gpu(fixtures.norm1)), + EMLX.Backend.from_nx(gpu(fixtures.q_proj)), + EMLX.Backend.from_nx(gpu(fixtures.k_proj)), + EMLX.Backend.from_nx(gpu(fixtures.v_proj)), + EMLX.Backend.from_nx(gpu(fixtures.o_proj)), + EMLX.Backend.from_nx(gpu(fixtures.q_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_cache)), + EMLX.Backend.from_nx(gpu(fixtures.v_cache)), + fixtures.offset, + scale, + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + + {expected, expected_k, expected_v} = qwen3_attention_block_reference(fixtures) + + assert Nx.shape(expected) == {2, 2, 4} + + assert_all_close( + out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_k, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_v, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + end + + test "qwen3_layer matches pure Nx reference" do + fixtures = qwen3_dense_attention_fixtures() + scale = 1.0 / :math.sqrt(fixtures.head_dim) + + {out_ref, k_ref, v_ref} = + EMLX.Native.Qwen3.layer( + EMLX.Backend.from_nx(gpu(fixtures.hidden)), + EMLX.Backend.from_nx(gpu(fixtures.norm1)), + EMLX.Backend.from_nx(gpu(fixtures.q_proj)), + EMLX.Backend.from_nx(gpu(fixtures.k_proj)), + EMLX.Backend.from_nx(gpu(fixtures.v_proj)), + EMLX.Backend.from_nx(gpu(fixtures.o_proj)), + EMLX.Backend.from_nx(gpu(fixtures.q_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_cache)), + EMLX.Backend.from_nx(gpu(fixtures.v_cache)), + EMLX.Backend.from_nx(gpu(fixtures.norm2)), + EMLX.Backend.from_nx(gpu(fixtures.gate_proj)), + EMLX.Backend.from_nx(gpu(fixtures.up_proj)), + EMLX.Backend.from_nx(gpu(fixtures.down_proj)), + fixtures.offset, + scale, + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + + {attn_expected, expected_k, expected_v} = qwen3_attention_block_reference(fixtures) + + expected = + qwen3_mlp_reference( + attn_expected, + fixtures.norm2, + fixtures.gate_proj, + fixtures.up_proj, + fixtures.down_proj, + fixtures.eps + ) + + assert_all_close( + out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_k, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_v, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + end + + test "qwen3_layer matches pure Nx reference for batched GQA" do + fixtures = qwen3_batched_dense_attention_fixtures() + scale = 1.0 / :math.sqrt(fixtures.head_dim) + + {out_ref, k_ref, v_ref} = + EMLX.Native.Qwen3.layer( + EMLX.Backend.from_nx(gpu(fixtures.hidden)), + EMLX.Backend.from_nx(gpu(fixtures.norm1)), + EMLX.Backend.from_nx(gpu(fixtures.q_proj)), + EMLX.Backend.from_nx(gpu(fixtures.k_proj)), + EMLX.Backend.from_nx(gpu(fixtures.v_proj)), + EMLX.Backend.from_nx(gpu(fixtures.o_proj)), + EMLX.Backend.from_nx(gpu(fixtures.q_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_cache)), + EMLX.Backend.from_nx(gpu(fixtures.v_cache)), + EMLX.Backend.from_nx(gpu(fixtures.norm2)), + EMLX.Backend.from_nx(gpu(fixtures.gate_proj)), + EMLX.Backend.from_nx(gpu(fixtures.up_proj)), + EMLX.Backend.from_nx(gpu(fixtures.down_proj)), + fixtures.offset, + scale, + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + + {attn_expected, expected_k, expected_v} = qwen3_attention_block_reference(fixtures) + + expected = + qwen3_mlp_reference( + attn_expected, + fixtures.norm2, + fixtures.gate_proj, + fixtures.up_proj, + fixtures.down_proj, + fixtures.eps + ) + + assert Nx.shape(expected) == {2, 2, 4} + + assert_all_close( + out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_k, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_v, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + end + + # ── qwen3_layer_quantized (dense-or-quantized per-projection fusion) ────── + + @qwen3_layer_quantized_modes [ + {"affine", 32, 4}, + {"mxfp4", 32, 4} + ] + + for {mode, group_size, bits} <- @qwen3_layer_quantized_modes do + test "qwen3_layer_quantized (#{mode}) matches per-op quantized reference for prefill (T_new > 1)" do + fixtures = + qwen3_quantized_attention_fixtures( + [:q_proj, :k_proj, :v_proj, :o_proj, :gate_proj, :up_proj, :down_proj], + group_size: unquote(group_size), + bits: unquote(bits), + mode: unquote(mode) + ) + + {out_ref, k_ref, v_ref} = qwen3_layer_quantized_call(fixtures) + {expected, expected_k, expected_v} = qwen3_layer_quantized_reference(fixtures) + + assert_all_close( + out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected, + atol: 5.0e-2, + rtol: 5.0e-2 + ) + + assert_all_close( + k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_k, + atol: 5.0e-2, + rtol: 5.0e-2 + ) + + assert_all_close( + v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_v, + atol: 5.0e-2, + rtol: 5.0e-2 + ) + end + end + + test "qwen3_layer_quantized matches per-op quantized reference for decode (T_new == 1)" do + fixtures = + [:q_proj, :k_proj, :v_proj, :o_proj, :gate_proj, :up_proj, :down_proj] + |> qwen3_quantized_attention_fixtures() + |> Map.update!(:hidden, &Nx.slice_along_axis(&1, 0, 1, axis: 1)) + + {out_ref, k_ref, v_ref} = qwen3_layer_quantized_call(fixtures) + {expected, expected_k, expected_v} = qwen3_layer_quantized_reference(fixtures) + + assert_all_close( + out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected, + atol: 5.0e-2, + rtol: 5.0e-2 + ) + + assert_all_close( + k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_k, + atol: 5.0e-2, + rtol: 5.0e-2 + ) + + assert_all_close( + v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_v, + atol: 5.0e-2, + rtol: 5.0e-2 + ) + end + + test "qwen3_layer_quantized matches per-op reference for a mixed dense/quantized layer" do + # Only the MLP projections are quantized; q/k/v/o stay dense — exercises + # `Qwen3LinearWeight`/`qwen3_apply_linear` picking the right branch + # independently per projection within one layer. + fixtures = qwen3_quantized_attention_fixtures([:gate_proj, :up_proj, :down_proj]) + + {out_ref, k_ref, v_ref} = qwen3_layer_quantized_call(fixtures) + {expected, expected_k, expected_v} = qwen3_layer_quantized_reference(fixtures) + + assert_all_close( + out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected, + atol: 5.0e-2, + rtol: 5.0e-2 + ) + + assert_all_close( + k_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_k, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + + assert_all_close( + v_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected_v, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + end + + # ── qwen3_forward_greedy_ids_chunk_quantized (multi-step chunk fusion) ──── + + test "qwen3_forward_greedy_ids_chunk_quantized produces identical token ids to " <> + "running the quantized single-layer NIF + manual final step N times" do + # Each fused NIF call `std::move`s the k/v caches it receives (mirroring + # `qwen3_layer`'s existing move-and-return-updated-cache contract), so + # the two independent call paths below must not share the same + # underlying cache resource — build separate fixture instances rather + # than reusing one across both. + count = 3 + + chunk_tokens = qwen3_quantized_chunk_call(qwen3_quantized_chunk_fixtures(), count) + manual_tokens = qwen3_quantized_chunk_manual_tokens(qwen3_quantized_chunk_fixtures(), count) + + assert chunk_tokens == manual_tokens + end + + test "qwen3_attention_residual matches pure Nx reference" do + {hidden, attn_out, o_proj} = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + { + Nx.iota({1, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(10), + Nx.iota({1, 2, 3}, type: :f32) |> Nx.add(5) |> Nx.divide(20), + Nx.iota({3, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(30) + } + end) + + out_ref = + EMLX.Native.Qwen3.attention_residual( + EMLX.Backend.from_nx(gpu(hidden)), + EMLX.Backend.from_nx(gpu(attn_out)), + EMLX.Backend.from_nx(gpu(o_proj)) + ) + + expected = + hidden + |> Nx.add(Nx.dot(attn_out, [2], o_proj, [0])) + |> Nx.backend_transfer(Nx.BinaryBackend) + + assert_all_close( + out_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected, + atol: 1.0e-5, + rtol: 1.0e-5 + ) + end + + test "qwen3_final_greedy matches pure Nx reference" do + {hidden, norm, lm_head} = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + { + Nx.iota({1, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(10), + Nx.tensor([1.0, 1.1, 0.9, 1.2], type: :f32), + Nx.tensor( + [ + [0.2, 0.1, 0.0, 0.3], + [0.0, 0.5, 0.2, 0.1], + [0.3, 0.1, 0.4, 0.2] + ], + type: :f32 + ) + } + end) + + eps = 1.0e-6 + + token_ref = + EMLX.Native.Qwen3.final_greedy( + EMLX.Backend.from_nx(gpu(hidden)), + EMLX.Backend.from_nx(gpu(norm)), + EMLX.Backend.from_nx(gpu(lm_head)), + eps + ) + + expected = + hidden + |> qwen3_final_logits_reference(norm, lm_head, eps) + |> Nx.argmax(axis: 1) + + assert_all_close( + token_ref |> EMLX.Backend.to_nx() |> Nx.backend_transfer(Nx.BinaryBackend), + expected + ) + end + + test "qwen3_forward_greedy_token_id defaults to embedding tensor device" do + fixtures = qwen3_dense_attention_fixtures() + cpu = fn tensor -> Nx.backend_transfer(tensor, {EMLX.Backend, device: :cpu}) end + + layer = { + cpu.(fixtures.norm1), + cpu.(fixtures.norm2), + cpu.(fixtures.q_norm), + cpu.(fixtures.k_norm), + cpu.(fixtures.q_proj), + cpu.(fixtures.k_proj), + cpu.(fixtures.v_proj), + cpu.(fixtures.o_proj), + cpu.(fixtures.gate_proj), + cpu.(fixtures.up_proj), + cpu.(fixtures.down_proj) + } + + kv_cache = + {EMLX.Backend.from_nx(cpu.(fixtures.k_cache)), + EMLX.Backend.from_nx(cpu.(fixtures.v_cache))} + + embed_tokens = + Nx.tensor(List.duplicate(List.duplicate(0.1, 4), 4), + type: :f32, + backend: {EMLX.Backend, device: :cpu} + ) + + norm = Nx.tensor(List.duplicate(1.0, 4), type: :f32, backend: {EMLX.Backend, device: :cpu}) + + lm_head = + Nx.tensor(List.duplicate(List.duplicate(0.1, 4), 4), + type: :f32, + backend: {EMLX.Backend, device: :cpu} + ) + + assert {:cpu, _ref} = EMLX.Backend.from_nx(embed_tokens) + + {_token_id, [{{k_device, _k_ref}, {v_device, _v_ref}}]} = + EMLX.Native.Qwen3.forward_greedy_token_id( + 0, + EMLX.Backend.from_nx(embed_tokens), + [layer], + [kv_cache], + EMLX.Backend.from_nx(norm), + EMLX.Backend.from_nx(lm_head), + 0, + 1.0 / :math.sqrt(fixtures.head_dim), + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + + assert k_device == :cpu + assert v_device == :cpu + end + + test "qwen3_forward_greedy_ids returns tensor token matching token id path" do + fixtures = qwen3_dense_attention_fixtures() + input_ids = Nx.tensor([[0]], type: :s64) |> gpu() + + embed_tokens = + Nx.tensor(List.duplicate(List.duplicate(0.1, 4), 4), type: :f32) + |> gpu() + + norm = Nx.tensor(List.duplicate(1.0, 4), type: :f32) |> gpu() + + lm_head = + Nx.tensor(List.duplicate(List.duplicate(0.1, 4), 4), type: :f32) + |> gpu() + + layer = { + gpu(fixtures.norm1), + gpu(fixtures.norm2), + gpu(fixtures.q_norm), + gpu(fixtures.k_norm), + gpu(fixtures.q_proj), + gpu(fixtures.k_proj), + gpu(fixtures.v_proj), + gpu(fixtures.o_proj), + gpu(fixtures.gate_proj), + gpu(fixtures.up_proj), + gpu(fixtures.down_proj) + } + + {token_ref, _kv_cache} = + EMLX.Native.Qwen3.forward_greedy_ids( + EMLX.Backend.from_nx(input_ids), + EMLX.Backend.from_nx(embed_tokens), + [layer], + [{gpu(fixtures.k_cache), gpu(fixtures.v_cache)}], + EMLX.Backend.from_nx(norm), + EMLX.Backend.from_nx(lm_head), + 0, + 1.0 / :math.sqrt(fixtures.head_dim), + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + + {token_id, _kv_cache} = + EMLX.Native.Qwen3.forward_greedy_token_id( + 0, + EMLX.Backend.from_nx(embed_tokens), + [layer], + [{gpu(fixtures.k_cache), gpu(fixtures.v_cache)}], + EMLX.Backend.from_nx(norm), + EMLX.Backend.from_nx(lm_head), + 0, + 1.0 / :math.sqrt(fixtures.head_dim), + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + + token = + token_ref + |> EMLX.Backend.to_nx() + |> Nx.backend_transfer(Nx.BinaryBackend) + |> Nx.to_flat_list() + |> hd() + + assert token == token_id + end + end + + describe "EMLX.Native.Qwen3.attention_block/15 validation" do + test "qwen3_kv_cache_attention rejects required cache length overflow before graph construction" do + q = Nx.broadcast(0.0, {1, 1, 2, 2}) |> Nx.as_type(:f32) |> gpu() + k = Nx.broadcast(0.0, {1, 1, 1, 2}) |> Nx.as_type(:f32) |> gpu() + v = Nx.broadcast(0.0, {1, 1, 1, 2}) |> Nx.as_type(:f32) |> gpu() + k_cache = Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f32) |> gpu() + v_cache = Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f32) |> gpu() + + assert_raise EMLX.NIFError, + ~r/KV cache capacity 4 is smaller than required length 2147483648/, + fn -> + EMLX.Native.Qwen3.kv_cache_attention( + EMLX.Backend.from_nx(q), + EMLX.Backend.from_nx(k), + EMLX.Backend.from_nx(v), + EMLX.Backend.from_nx(k_cache), + EMLX.Backend.from_nx(v_cache), + 2_147_483_647, + 1.0 / :math.sqrt(2), + 2, + 10_000.0 + ) + end + end + + test "qwen3_layer rejects required cache length overflow before graph construction" do + fixtures = qwen3_dense_attention_fixtures() + + assert_raise EMLX.NIFError, + ~r/KV cache capacity 4 is smaller than required length 2147483649/, + fn -> + EMLX.Native.Qwen3.layer( + EMLX.Backend.from_nx(gpu(fixtures.hidden)), + EMLX.Backend.from_nx(gpu(fixtures.norm1)), + EMLX.Backend.from_nx(gpu(fixtures.q_proj)), + EMLX.Backend.from_nx(gpu(fixtures.k_proj)), + EMLX.Backend.from_nx(gpu(fixtures.v_proj)), + EMLX.Backend.from_nx(gpu(fixtures.o_proj)), + EMLX.Backend.from_nx(gpu(fixtures.q_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_cache)), + EMLX.Backend.from_nx(gpu(fixtures.v_cache)), + EMLX.Backend.from_nx(gpu(fixtures.norm2)), + EMLX.Backend.from_nx(gpu(fixtures.gate_proj)), + EMLX.Backend.from_nx(gpu(fixtures.up_proj)), + EMLX.Backend.from_nx(gpu(fixtures.down_proj)), + 2_147_483_647, + 1.0 / :math.sqrt(fixtures.head_dim), + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + end + end + + test "qwen3_attention_block rejects required cache length overflow before graph construction" do + fixtures = qwen3_dense_attention_fixtures() + + assert_raise EMLX.NIFError, + ~r/KV cache capacity 4 is smaller than required length 2147483649/, + fn -> + EMLX.Native.Qwen3.attention_block( + EMLX.Backend.from_nx(gpu(fixtures.hidden)), + EMLX.Backend.from_nx(gpu(fixtures.norm1)), + EMLX.Backend.from_nx(gpu(fixtures.q_proj)), + EMLX.Backend.from_nx(gpu(fixtures.k_proj)), + EMLX.Backend.from_nx(gpu(fixtures.v_proj)), + EMLX.Backend.from_nx(gpu(fixtures.o_proj)), + EMLX.Backend.from_nx(gpu(fixtures.q_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_norm)), + EMLX.Backend.from_nx(gpu(fixtures.k_cache)), + EMLX.Backend.from_nx(gpu(fixtures.v_cache)), + 2_147_483_647, + 1.0 / :math.sqrt(fixtures.head_dim), + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + end + end + + test "qwen3_forward_greedy_ids_token_id rejects batches before layer/cache normalization" do + input_ids = Nx.tensor([[0], [1]], type: :s64) |> gpu() + embed_tokens = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f32) |> gpu() + norm = Nx.broadcast(1.0, {4}) |> Nx.as_type(:f32) |> gpu() + lm_head = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f32) |> gpu() + + assert_raise ArgumentError, + ~r/forward_greedy_ids_token_id requires batch size 1, got batch size 2/, + fn -> + EMLX.Native.Qwen3.forward_greedy_ids_token_id( + EMLX.Backend.from_nx(input_ids), + EMLX.Backend.from_nx(embed_tokens), + [:invalid_layer], + [:invalid_cache], + EMLX.Backend.from_nx(norm), + EMLX.Backend.from_nx(lm_head), + 0, + 1.0 / :math.sqrt(2), + 2, + 10_000.0, + 1.0e-6 + ) + end + end + + test "qwen3_forward_greedy_ids_chunk rejects input that is not a decode step before layer/cache normalization" do + input_ids = Nx.tensor([[0, 1]], type: :s64) |> gpu() + embed_tokens = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f32) |> gpu() + norm = Nx.broadcast(1.0, {4}) |> Nx.as_type(:f32) |> gpu() + lm_head = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f32) |> gpu() + + assert_raise ArgumentError, + ~r/forward_greedy_ids_chunk requires sequence length 1, got sequence length 2/, + fn -> + EMLX.Native.Qwen3.forward_greedy_ids_chunk( + EMLX.Backend.from_nx(input_ids), + EMLX.Backend.from_nx(embed_tokens), + [:invalid_layer], + [:invalid_cache], + EMLX.Backend.from_nx(norm), + EMLX.Backend.from_nx(lm_head), + 0, + 1, + 1.0 / :math.sqrt(2), + 2, + 10_000.0, + 1.0e-6 + ) + end + end + + test "rejects projection widths before deriving head counts" do + hidden = Nx.broadcast(0.0, {1, 1, 4}) |> Nx.as_type(:f16) |> gpu() + norm = Nx.broadcast(1.0, {4}) |> Nx.as_type(:f16) |> gpu() + q_proj = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f16) |> gpu() + k_proj = Nx.broadcast(0.1, {4, 1}) |> Nx.as_type(:f16) |> gpu() + v_proj = Nx.broadcast(0.1, {4, 1}) |> Nx.as_type(:f16) |> gpu() + o_proj = Nx.broadcast(0.1, {4, 4}) |> Nx.as_type(:f16) |> gpu() + q_norm = Nx.broadcast(1.0, {2}) |> Nx.as_type(:f16) |> gpu() + k_norm = Nx.broadcast(1.0, {2}) |> Nx.as_type(:f16) |> gpu() + k_cache = Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f16) |> gpu() + v_cache = Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f16) |> gpu() + + assert_raise EMLX.NIFError, + ~r/projection output widths must be divisible by head_dim/, + fn -> + EMLX.Native.Qwen3.attention_block( + EMLX.Backend.from_nx(hidden), + EMLX.Backend.from_nx(norm), + EMLX.Backend.from_nx(q_proj), + EMLX.Backend.from_nx(k_proj), + EMLX.Backend.from_nx(v_proj), + EMLX.Backend.from_nx(o_proj), + EMLX.Backend.from_nx(q_norm), + EMLX.Backend.from_nx(k_norm), + EMLX.Backend.from_nx(k_cache), + EMLX.Backend.from_nx(v_cache), + 0, + 1.0 / :math.sqrt(2), + 2, + 10_000.0, + 1.0e-6 + ) + end + end + end + + defp qwen3_dense_attention_fixtures do + Nx.with_default_backend(Nx.BinaryBackend, fn -> + %{ + hidden: Nx.iota({1, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(10), + norm1: Nx.tensor([1.0, 1.1, 0.9, 1.2], type: :f32), + norm2: Nx.tensor([0.9, 1.0, 1.1, 1.2], type: :f32), + q_norm: Nx.tensor([1.0, 1.1], type: :f32), + k_norm: Nx.tensor([0.9, 1.2], type: :f32), + q_proj: Nx.iota({4, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(50), + k_proj: Nx.iota({4, 2}, type: :f32) |> Nx.add(2) |> Nx.divide(60), + v_proj: Nx.iota({4, 2}, type: :f32) |> Nx.add(3) |> Nx.divide(70), + o_proj: Nx.iota({4, 4}, type: :f32) |> Nx.add(4) |> Nx.divide(80), + gate_proj: Nx.iota({4, 6}, type: :f32) |> Nx.add(1) |> Nx.divide(50), + up_proj: Nx.iota({4, 6}, type: :f32) |> Nx.add(7) |> Nx.divide(60), + down_proj: Nx.iota({6, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(70), + k_cache: Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f32), + v_cache: Nx.broadcast(0.0, {1, 1, 4, 2}) |> Nx.as_type(:f32), + offset: 0, + head_dim: 2, + theta: 10_000.0, + eps: 1.0e-6 + } + end) + end + + defp qwen3_batched_dense_attention_fixtures do + Nx.with_default_backend(Nx.BinaryBackend, fn -> + %{ + hidden: Nx.iota({2, 2, 4}, type: :f32) |> Nx.add(1) |> Nx.divide(20), + norm1: Nx.tensor([1.0, 1.1, 0.9, 1.2], type: :f32), + norm2: Nx.tensor([1.2, 0.8, 1.1, 0.95], type: :f32), + q_norm: Nx.tensor([1.0, 1.1], type: :f32), + k_norm: Nx.tensor([0.9, 1.2], type: :f32), + q_proj: Nx.iota({4, 8}, type: :f32) |> Nx.add(1) |> Nx.divide(80), + k_proj: Nx.iota({4, 4}, type: :f32) |> Nx.add(3) |> Nx.divide(90), + v_proj: Nx.iota({4, 4}, type: :f32) |> Nx.add(5) |> Nx.divide(100), + o_proj: Nx.iota({8, 4}, type: :f32) |> Nx.add(7) |> Nx.divide(110), + gate_proj: Nx.iota({4, 6}, type: :f32) |> Nx.add(13) |> Nx.divide(120), + up_proj: Nx.iota({4, 6}, type: :f32) |> Nx.add(29) |> Nx.divide(130), + down_proj: Nx.iota({6, 4}, type: :f32) |> Nx.add(41) |> Nx.divide(140), + k_cache: Nx.iota({2, 2, 6, 2}, type: :f32) |> Nx.add(11) |> Nx.divide(1_000), + v_cache: Nx.iota({2, 2, 6, 2}, type: :f32) |> Nx.add(101) |> Nx.divide(1_000), + offset: 2, + head_dim: 2, + theta: 10_000.0, + eps: 1.0e-6 + } + end) + end + + # Passes a tensor through unchanged if it's already GPU/EMLX-backed + # (quantized fixtures come back from `EMLX.quantize/2` that way); otherwise + # transfers it, mirroring the dense `gpu/1` helper above. + defp ensure_gpu(%Nx.Tensor{data: %EMLX.Backend{}} = tensor), do: tensor + defp ensure_gpu(tensor), do: gpu(tensor) + + # Builds a Qwen3 single-layer fixture (H=32, head_dim=16, N_q=2, N_kv=1, + # intermediate=32 — all input widths are 32, satisfying `EMLX.quantize/2`'s + # supported group sizes {32, 64, 128} and every microscaled mode's fixed + # group_size) and quantizes the projections listed in `quantized_slots` via + # `EMLX.quantize/2`. + # + # For each quantized slot, also stashes a `_reference` dequantized + # tensor in the *same* {in, out} convention as the dense fixtures, so + # `qwen3_layer_quantized_reference/1` can reuse the existing dense-only + # `qwen3_attention_block_reference/1`/`qwen3_mlp_reference/6` helpers + # unchanged — isolating fusion correctness from quantization lossiness + # (both sides of the comparison use the same quantized-then-dequantized + # numbers). + defp qwen3_quantized_attention_fixtures(quantized_slots, opts \\ []) do + group_size = Keyword.get(opts, :group_size, 32) + bits = Keyword.get(opts, :bits, 4) + mode = Keyword.get(opts, :mode, "affine") + + base = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + %{ + hidden: Nx.iota({1, 2, 32}, type: :f32) |> Nx.add(1) |> Nx.divide(400), + norm1: Nx.iota({32}, type: :f32) |> Nx.add(10) |> Nx.divide(10), + norm2: Nx.iota({32}, type: :f32) |> Nx.add(20) |> Nx.divide(10), + q_norm: Nx.iota({16}, type: :f32) |> Nx.add(30) |> Nx.divide(10), + k_norm: Nx.iota({16}, type: :f32) |> Nx.add(40) |> Nx.divide(10), + q_proj: Nx.iota({32, 32}, type: :f32) |> Nx.add(1) |> Nx.divide(900), + k_proj: Nx.iota({32, 16}, type: :f32) |> Nx.add(2) |> Nx.divide(900), + v_proj: Nx.iota({32, 16}, type: :f32) |> Nx.add(3) |> Nx.divide(900), + o_proj: Nx.iota({32, 32}, type: :f32) |> Nx.add(4) |> Nx.divide(900), + gate_proj: Nx.iota({32, 32}, type: :f32) |> Nx.add(5) |> Nx.divide(900), + up_proj: Nx.iota({32, 32}, type: :f32) |> Nx.add(6) |> Nx.divide(900), + down_proj: Nx.iota({32, 32}, type: :f32) |> Nx.add(7) |> Nx.divide(900), + k_cache: Nx.broadcast(0.0, {1, 1, 4, 16}) |> Nx.as_type(:f32), + v_cache: Nx.broadcast(0.0, {1, 1, 4, 16}) |> Nx.as_type(:f32), + offset: 0, + head_dim: 16, + theta: 10_000.0, + eps: 1.0e-6 + } + end) + + Enum.reduce(quantized_slots, base, fn slot, acc -> + dense_w = Map.fetch!(acc, slot) + + qw = + dense_w + |> gpu() + |> Nx.transpose() + |> EMLX.quantize(type: {:s, bits}, group_size: group_size, mode: mode) + + dequantized_in_out = + qw + |> EMLX.dequantize() + |> Nx.transpose() + |> Nx.backend_transfer(Nx.BinaryBackend) + + acc + |> Map.put(slot, qw) + |> Map.put(:"#{slot}_reference", dequantized_in_out) + end) + end + + defp qwen3_layer_quantized_call(fixtures) do + scale = 1.0 / :math.sqrt(fixtures.head_dim) + + EMLX.Native.Qwen3.layer_quantized( + EMLX.Backend.from_nx(gpu(fixtures.hidden)), + ensure_gpu(fixtures.norm1), + ensure_gpu(fixtures.q_proj), + ensure_gpu(fixtures.k_proj), + ensure_gpu(fixtures.v_proj), + ensure_gpu(fixtures.o_proj), + ensure_gpu(fixtures.q_norm), + ensure_gpu(fixtures.k_norm), + EMLX.Backend.from_nx(gpu(fixtures.k_cache)), + EMLX.Backend.from_nx(gpu(fixtures.v_cache)), + ensure_gpu(fixtures.norm2), + ensure_gpu(fixtures.gate_proj), + ensure_gpu(fixtures.up_proj), + ensure_gpu(fixtures.down_proj), + fixtures.offset, + scale, + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + end + + defp qwen3_layer_quantized_reference(fixtures) do + dense_fixtures = + Enum.reduce( + [:q_proj, :k_proj, :v_proj, :o_proj, :gate_proj, :up_proj, :down_proj], + fixtures, + fn slot, acc -> + case Map.fetch(fixtures, :"#{slot}_reference") do + {:ok, dequantized} -> Map.put(acc, slot, dequantized) + :error -> acc + end + end + ) + + {attn_expected, expected_k, expected_v} = qwen3_attention_block_reference(dense_fixtures) + + expected = + qwen3_mlp_reference( + attn_expected, + dense_fixtures.norm2, + dense_fixtures.gate_proj, + dense_fixtures.up_proj, + dense_fixtures.down_proj, + dense_fixtures.eps + ) + + {expected, expected_k, expected_v} + end + + # Single-layer fixture for `qwen3_forward_greedy_ids_chunk_quantized` + # determinism testing: same projection shapes as + # `qwen3_quantized_attention_fixtures/2`, plus a small embedding table and a + # quantized `lm_head` (`{vocab, H}`, already in the `{out, in}` convention + # `EMLX.quantize/2` expects — no transpose needed, unlike the projections). + defp qwen3_quantized_chunk_fixtures do + base = + qwen3_quantized_attention_fixtures([ + :q_proj, + :k_proj, + :v_proj, + :o_proj, + :gate_proj, + :up_proj, + :down_proj + ]) + + embed_tokens = + Nx.iota({6, 32}, type: :f32) |> Nx.add(1) |> Nx.divide(900) |> gpu() + + lm_head = + Nx.iota({6, 32}, type: :f32) + |> Nx.add(2) + |> Nx.divide(900) + |> gpu() + |> EMLX.quantize(type: {:s, 4}, group_size: 32, mode: "affine") + + norm = Nx.iota({32}, type: :f32) |> Nx.add(50) |> Nx.divide(10) |> gpu() + + # Enough KV cache capacity for `offset + count` across the test's decode + # steps (capacity 8 comfortably covers `count = 3`). + k_cache = Nx.broadcast(0.0, {1, 1, 8, 16}) |> Nx.as_type(:f32) |> gpu() + v_cache = Nx.broadcast(0.0, {1, 1, 8, 16}) |> Nx.as_type(:f32) |> gpu() + + %{ + base + | k_cache: k_cache, + v_cache: v_cache + } + |> Map.put(:embed_tokens, embed_tokens) + |> Map.put(:lm_head, lm_head) + |> Map.put(:norm, norm) + |> Map.put(:input_id, 0) + end + + defp qwen3_quantized_chunk_call(fixtures, count) do + scale = 1.0 / :math.sqrt(fixtures.head_dim) + input_ids = Nx.tensor([[fixtures.input_id]], type: :s64) |> gpu() + + layer = { + ensure_gpu(fixtures.norm1), + ensure_gpu(fixtures.norm2), + ensure_gpu(fixtures.q_norm), + ensure_gpu(fixtures.k_norm), + ensure_gpu(fixtures.q_proj), + ensure_gpu(fixtures.k_proj), + ensure_gpu(fixtures.v_proj), + ensure_gpu(fixtures.o_proj), + ensure_gpu(fixtures.gate_proj), + ensure_gpu(fixtures.up_proj), + ensure_gpu(fixtures.down_proj) + } + + {token_refs, _kv_cache} = + EMLX.Native.Qwen3.forward_greedy_ids_chunk_quantized( + EMLX.Backend.from_nx(input_ids), + EMLX.Backend.from_nx(fixtures.embed_tokens), + [layer], + [{fixtures.k_cache, fixtures.v_cache}], + EMLX.Backend.from_nx(fixtures.norm), + fixtures.lm_head, + fixtures.offset, + count, + scale, + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + + Enum.map(token_refs, &(&1 |> EMLX.Backend.to_nx() |> Nx.to_flat_list())) + end + + # Manually replays the same `count` decode steps by calling the + # already-verified single-layer `qwen3_layer_quantized` NIF plus a manual + # final RMSNorm + `EMLX.quantized_matmul` lm_head + argmax step, once per + # step. Both paths run the exact same MLX ops (just organized as 1 fused + # NIF call vs `count` separate calls), so token ids should be bit-identical. + defp qwen3_quantized_chunk_manual_tokens(fixtures, count) do + scale = 1.0 / :math.sqrt(fixtures.head_dim) + + {_ids, _kv, tokens_rev} = + Enum.reduce(1..count, {fixtures.input_id, {fixtures.k_cache, fixtures.v_cache}, []}, fn + step, {token_id, {k_cache, v_cache}, acc} -> + ids = Nx.tensor([token_id], type: :s64) |> gpu() + + hidden = + fixtures.embed_tokens + |> Nx.take(ids, axis: 0) + |> Nx.new_axis(0) + + {hidden_ref, k_ref, v_ref} = + EMLX.Native.Qwen3.layer_quantized( + EMLX.Backend.from_nx(hidden), + ensure_gpu(fixtures.norm1), + ensure_gpu(fixtures.q_proj), + ensure_gpu(fixtures.k_proj), + ensure_gpu(fixtures.v_proj), + ensure_gpu(fixtures.o_proj), + ensure_gpu(fixtures.q_norm), + ensure_gpu(fixtures.k_norm), + EMLX.Backend.from_nx(k_cache), + EMLX.Backend.from_nx(v_cache), + ensure_gpu(fixtures.norm2), + ensure_gpu(fixtures.gate_proj), + ensure_gpu(fixtures.up_proj), + ensure_gpu(fixtures.down_proj), + fixtures.offset + step - 1, + scale, + fixtures.head_dim, + fixtures.theta, + fixtures.eps + ) + + hidden_out = EMLX.Backend.to_nx(hidden_ref) + k_new = EMLX.Backend.to_nx(k_ref) + v_new = EMLX.Backend.to_nx(v_ref) + + normed = + hidden_out + |> Nx.reshape({1, 32}) + |> EMLX.Fast.rms_norm(fixtures.norm, fixtures.eps) + + logits = EMLX.Quantization.quantized_matmul(normed, fixtures.lm_head) + token = Nx.argmax(logits, axis: 1) + token_id = token |> Nx.to_flat_list() |> hd() + + {token_id, {k_new, v_new}, [Nx.to_flat_list(token) | acc]} + end) + + Enum.reverse(tokens_rev) + end + + defp qwen3_kv_cache_attention_reference( + q, + new_k, + new_v, + k_cache, + v_cache, + offset, + scale, + head_dim, + theta + ) do + {batch, seq_len, q_heads, _head_dim} = Nx.shape(q) + {_batch, _seq_len, kv_heads, _head_dim} = Nx.shape(new_k) + + q_bn = + q + |> Nx.transpose(axes: [0, 2, 1, 3]) + |> qwen3_rope_reference(head_dim, theta, offset) + |> Nx.backend_transfer(Nx.BinaryBackend) + + k_bn = + new_k + |> Nx.transpose(axes: [0, 2, 1, 3]) + |> qwen3_rope_reference(head_dim, theta, offset) + |> Nx.backend_transfer(Nx.BinaryBackend) + + v_bn = new_v |> Nx.transpose(axes: [0, 2, 1, 3]) |> Nx.backend_transfer(Nx.BinaryBackend) + + k_cache = Nx.put_slice(k_cache, [0, 0, offset, 0], k_bn) + v_cache = Nx.put_slice(v_cache, [0, 0, offset, 0], v_bn) + + valid_len = offset + seq_len + k_valid = Nx.slice_along_axis(k_cache, 0, valid_len, axis: 2) + v_valid = Nx.slice_along_axis(v_cache, 0, valid_len, axis: 2) + + groups = div(q_heads, kv_heads) + k_repeated = repeat_kv_heads_bn(k_valid, groups) + v_repeated = repeat_kv_heads_bn(v_valid, groups) + + scores = + q_bn + |> Nx.new_axis(3) + |> Nx.multiply(Nx.new_axis(k_repeated, 2)) + |> Nx.sum(axes: [4]) + |> Nx.multiply(scale) + |> apply_causal_mask(offset, seq_len, valid_len) + + weights = softmax_reference(scores, 3) + + attn = + weights + |> Nx.new_axis(4) + |> Nx.multiply(Nx.new_axis(v_repeated, 2)) + |> Nx.sum(axes: [3]) + + attn_out = + attn + |> Nx.transpose(axes: [0, 2, 1, 3]) + |> Nx.reshape({batch, seq_len, q_heads * head_dim}) + + {attn_out, k_cache, v_cache} + end + + defp qwen3_mlp_reference(hidden, norm, gate_proj, up_proj, down_proj, eps) do + xn = rms_norm_reference(hidden, norm, eps) + gate = Nx.dot(xn, [2], gate_proj, [0]) + up = Nx.dot(xn, [2], up_proj, [0]) + mlp = Nx.multiply(Nx.multiply(gate, Nx.sigmoid(gate)), up) + out = Nx.dot(mlp, [2], down_proj, [0]) + Nx.add(hidden, out) + end + + defp qwen3_attention_block_reference(fixtures) do + hidden = fixtures.hidden + offset = fixtures.offset + head_dim = fixtures.head_dim + scale = 1.0 / :math.sqrt(head_dim) + + {batch, seq_len, _hidden_size} = Nx.shape(hidden) + + xn = rms_norm_reference(hidden, fixtures.norm1, fixtures.eps) + + q_flat = Nx.dot(xn, [2], fixtures.q_proj, [0]) + k_flat = Nx.dot(xn, [2], fixtures.k_proj, [0]) + v_flat = Nx.dot(xn, [2], fixtures.v_proj, [0]) + + q_heads = div(elem(Nx.shape(fixtures.q_proj), 1), head_dim) + kv_heads = div(elem(Nx.shape(fixtures.k_proj), 1), head_dim) + + q = Nx.reshape(q_flat, {batch, seq_len, q_heads, head_dim}) + k = Nx.reshape(k_flat, {batch, seq_len, kv_heads, head_dim}) + v = Nx.reshape(v_flat, {batch, seq_len, kv_heads, head_dim}) + + q = rms_norm_reference(q, fixtures.q_norm, fixtures.eps) + k = rms_norm_reference(k, fixtures.k_norm, fixtures.eps) + + q_bn = + q + |> Nx.transpose(axes: [0, 2, 1, 3]) + |> qwen3_rope_reference(head_dim, fixtures.theta, offset) + + k_bn = + k + |> Nx.transpose(axes: [0, 2, 1, 3]) + |> qwen3_rope_reference(head_dim, fixtures.theta, offset) + + q_bn = Nx.backend_transfer(q_bn, Nx.BinaryBackend) + k_bn = Nx.backend_transfer(k_bn, Nx.BinaryBackend) + v_bn = v |> Nx.transpose(axes: [0, 2, 1, 3]) |> Nx.backend_transfer(Nx.BinaryBackend) + + k_cache = Nx.put_slice(fixtures.k_cache, [0, 0, offset, 0], k_bn) + v_cache = Nx.put_slice(fixtures.v_cache, [0, 0, offset, 0], v_bn) + + valid_len = offset + seq_len + k_valid = Nx.slice_along_axis(k_cache, 0, valid_len, axis: 2) + v_valid = Nx.slice_along_axis(v_cache, 0, valid_len, axis: 2) + + groups = div(q_heads, kv_heads) + k_repeated = repeat_kv_heads_bn(k_valid, groups) + v_repeated = repeat_kv_heads_bn(v_valid, groups) + + scores = + q_bn + |> Nx.new_axis(3) + |> Nx.multiply(Nx.new_axis(k_repeated, 2)) + |> Nx.sum(axes: [4]) + |> Nx.multiply(scale) + + scores = apply_causal_mask(scores, offset, seq_len, valid_len) + weights = softmax_reference(scores, 3) + + attn = + weights + |> Nx.new_axis(4) + |> Nx.multiply(Nx.new_axis(v_repeated, 2)) + |> Nx.sum(axes: [3]) + + attn_out = + attn + |> Nx.transpose(axes: [0, 2, 1, 3]) + |> Nx.reshape({batch, seq_len, q_heads * head_dim}) + + projected = Nx.dot(attn_out, [2], fixtures.o_proj, [0]) + + {Nx.add(hidden, projected), k_cache, v_cache} + end + + defp qwen3_final_logits_reference(hidden, norm, lm_head, eps) do + {_batch, seq_len, hidden_size} = Nx.shape(hidden) + + last = + hidden + |> Nx.slice([0, seq_len - 1, 0], [1, 1, hidden_size]) + |> Nx.reshape({1, hidden_size}) + + last + |> rms_norm_reference(norm, eps) + |> Nx.dot([1], lm_head, [1]) + end + + defp rms_norm_reference(tensor, weight, eps) do + tensor + |> Nx.pow(2) + |> Nx.mean(axes: [-1], keep_axes: true) + |> Nx.add(eps) + |> Nx.sqrt() + |> then(&Nx.divide(tensor, &1)) + |> Nx.multiply(weight) + end + + defp qwen3_rope_reference(tensor, dims, theta, offset) do + {_batch, _heads, seq_len, _dims} = Nx.shape(tensor) + half = div(dims, 2) + + inv_freq = + Nx.iota({half}, type: :f32) + |> Nx.multiply(2) + |> Nx.divide(dims) + |> then(&Nx.pow(theta, &1)) + |> then(&Nx.divide(1.0, &1)) + + positions = + Nx.iota({seq_len}, type: :f32) + |> Nx.add(offset) + |> Nx.reshape({seq_len, 1}) + + freqs = Nx.multiply(positions, Nx.reshape(inv_freq, {1, half})) + + cos = + Nx.concatenate([Nx.cos(freqs), Nx.cos(freqs)], axis: 1) + |> Nx.reshape({1, 1, seq_len, dims}) + + sin = + Nx.concatenate([Nx.sin(freqs), Nx.sin(freqs)], axis: 1) + |> Nx.reshape({1, 1, seq_len, dims}) + + first = tensor[[.., .., .., 0..(half - 1)//1]] + second = tensor[[.., .., .., half..(dims - 1)//1]] + rotated = Nx.concatenate([Nx.negate(second), first], axis: 3) + + Nx.add(Nx.multiply(tensor, cos), Nx.multiply(rotated, sin)) + end + + defp repeat_kv_heads_bn(tensor, 1), do: tensor + + defp repeat_kv_heads_bn(tensor, groups) do + {batch, n_kv, t_kv, head_dim} = Nx.shape(tensor) + + tensor + |> Nx.new_axis(2) + |> Nx.broadcast({batch, n_kv, groups, t_kv, head_dim}) + |> Nx.reshape({batch, n_kv * groups, t_kv, head_dim}) + end + + defp apply_causal_mask(scores, offset, seq_len, valid_len) do + query_positions = + Nx.iota({seq_len}, type: :s32, backend: Nx.BinaryBackend) + |> Nx.add(offset) + |> Nx.reshape({1, 1, seq_len, 1}) + + key_positions = + Nx.iota({valid_len}, type: :s32, backend: Nx.BinaryBackend) + |> Nx.reshape({1, 1, 1, valid_len}) + + mask = + key_positions + |> Nx.less_equal(query_positions) + |> Nx.broadcast(Nx.shape(scores)) + + Nx.select(mask, scores, Nx.broadcast(-1.0e9, Nx.shape(scores))) + end + + defp softmax_reference(tensor, axis) do + shifted = Nx.subtract(tensor, Nx.reduce_max(tensor, axes: [axis], keep_axes: true)) + exp = Nx.exp(shifted) + Nx.divide(exp, Nx.sum(exp, axes: [axis], keep_axes: true)) + end +end diff --git a/mise.toml b/mise.toml new file mode 100644 index 0000000..ff536b8 --- /dev/null +++ b/mise.toml @@ -0,0 +1,3 @@ +[tools] +elixir = "1.20.1-otp-28" +erlang = "28.5.0.1" diff --git a/mix.exs b/mix.exs index bc498fd..e7bd991 100644 --- a/mix.exs +++ b/mix.exs @@ -10,7 +10,8 @@ defmodule EMLXRoot do aliases: [ setup: cmd("deps.get"), compile: cmd("compile"), - test: cmd("test") + test: cmd("test"), + format: cmd("format") ] ] end