chore: Upgrade to MLX 0.32.4

emlx/c_src/emlx/compiler.cpp

@@ -1405,23 +1405,23 @@ static const std::unordered_map<std::string, OpFn> op_registry = {
      [](const auto &ops, const auto &attrs, mlx::core::StreamOrDevice stream) {
        return mlx::core::fast::scaled_dot_product_attention(
            ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "", std::nullopt,
-           std::nullopt, stream);
+           std::nullopt, /*force_fused=*/false, stream);
      }},
 
     // sdpa (no mask, + sinks): operands = [q, k, v, sinks]; attrs = [scale_bits].
     {"fast_sdpa_sinks",
      [](const auto &ops, const auto &attrs, mlx::core::StreamOrDevice stream) {
        return mlx::core::fast::scaled_dot_product_attention(
            ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "", std::nullopt,
-           ops[3], stream);
+           ops[3], /*force_fused=*/false, stream);
      }},
 
     // sdpa (array mask): operands = [q, k, v, mask]; attrs = [scale_bits].
     {"fast_sdpa_masked",
      [](const auto &ops, const auto &attrs, mlx::core::StreamOrDevice stream) {
        return mlx::core::fast::scaled_dot_product_attention(
            ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "array", ops[3],
-           std::nullopt, stream);
+           std::nullopt, /*force_fused=*/false, stream);
      }},
 
     // sdpa (array mask, + sinks): operands = [q, k, v, mask, sinks];
@@ -1430,23 +1430,23 @@ static const std::unordered_map<std::string, OpFn> op_registry = {
      [](const auto &ops, const auto &attrs, mlx::core::StreamOrDevice stream) {
        return mlx::core::fast::scaled_dot_product_attention(
            ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "array", ops[3],
-           ops[4], stream);
+           ops[4], /*force_fused=*/false, stream);
      }},
 
     // sdpa (causal): operands = [q, k, v]; attrs = [scale_bits].
     {"fast_sdpa_causal",
      [](const auto &ops, const auto &attrs, mlx::core::StreamOrDevice stream) {
        return mlx::core::fast::scaled_dot_product_attention(
            ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "causal",
-           std::nullopt, std::nullopt, stream);
+           std::nullopt, std::nullopt, /*force_fused=*/false, stream);
      }},
 
     // sdpa (causal, + sinks): operands = [q, k, v, sinks]; attrs = [scale_bits].
     {"fast_sdpa_causal_sinks",
      [](const auto &ops, const auto &attrs, mlx::core::StreamOrDevice stream) {
        return mlx::core::fast::scaled_dot_product_attention(
            ops[0], ops[1], ops[2], attr_to_float(attrs[0]), "causal",
-           std::nullopt, ops[3], stream);
+           std::nullopt, ops[3], /*force_fused=*/false, stream);
      }},
 
     // sdpa (causal + key_mask): operands = [q, k, v, key_mask];
@@ -1482,7 +1482,7 @@ static const std::unordered_map<std::string, OpFn> op_registry = {
        auto additive = mlx::core::where(keep, zero_val, neginf_val, stream);
 
        auto attn_t = mlx::core::fast::scaled_dot_product_attention(
-           q, k, v, scale, "array", additive, std::nullopt, stream);
+           q, k, v, scale, "array", additive, std::nullopt, /*force_fused=*/false, stream);
 
        // Replace NaN with 0 for all-masked rows (softmax(-inf,...,-inf) = NaN
        // in Flash-Attention when seq_len >= Metal tile size, but semantically
@@ -1525,7 +1525,7 @@ static const std::unordered_map<std::string, OpFn> op_registry = {
        auto additive = mlx::core::where(keep, zero_val, neginf_val, stream);
 
        auto attn_t = mlx::core::fast::scaled_dot_product_attention(
-           q, k, v, scale, "array", additive, sinks, stream);
+           q, k, v, scale, "array", additive, sinks, /*force_fused=*/false, stream);
 
        // See "fast_sdpa_causal_key_masked" above: NaN-guard all-masked rows.
        return mlx::core::where(mlx::core::isnan(attn_t, stream),
@@ -2058,7 +2058,7 @@ static const std::unordered_map<std::string, MultiOpFn> multi_op_registry = {
        auto additive = mlx::core::where(keep, zero_val, neginf_val, stream);
 
        auto attn_out = mlx::core::fast::scaled_dot_product_attention(
-           q_t, k_t, v_t, scale, "array", additive, std::nullopt, stream);
+           q_t, k_t, v_t, scale, "array", additive, std::nullopt, /*force_fused=*/false, stream);
 
        // Replace NaN with 0 for all-masked rows (softmax(-inf,...,-inf) = NaN
        // in Flash-Attention when seq_len >= Metal tile size, but semantically
@@ -2344,17 +2344,29 @@ 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<Expr> when the BEAM resource
-// is GC'd.
+// Evicts the per-Expr entry from each MLX compile cache it is in, so stale
+// compiled graphs don't accumulate.  Called by default_dtor<Expr> when the BEAM
+// resource is GC'd.
 
 Expr::~Expr() {
   if (compile_id != 0) {
     std::lock_guard<std::mutex> lk(s_mlx_compile_mutex);
-    mlx::core::detail::compile_erase(compile_id);
+    for (const auto &cache : compile_caches)
+      mlx::core::detail::compile_erase(cache, compile_id);
   }
 }
 
+// Records the compile cache of the calling thread, the one that compiled_fn just
+// used. Called with s_mlx_compile_mutex held.
+static void remember_compile_cache(Expr &prog) {
+  auto cache = mlx::core::detail::compile_cache();
+  auto current = cache.lock();
+  for (const auto &known : prog.compile_caches)
+    if (known.lock() == current)
+      return;
+  prog.compile_caches.push_back(cache);
+}
+
 // ── NIF implementations
 // ───────────────────────────────────────────────────────
 
@@ -2496,6 +2508,7 @@ ERL_NIF_TERM eval_program(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     {
       std::lock_guard<std::mutex> lk(s_mlx_compile_mutex);
       outputs = prog->compiled_fn(inputs);
+      remember_compile_cache(*prog);
     }
 
     // A program containing an inlined `:runtime_call` node must be forced to

emlx/c_src/emlx/compiler.hpp

@@ -9,6 +9,7 @@
 
 #include "emlx_nif_shared.hpp"
 #include "emlx/plugin/registry.hpp"
+#include "mlx/compile_impl.h"
 #include <variant>
 
 namespace emlx {
@@ -117,7 +118,7 @@ struct Program {
 // 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.
+// The destructor evicts the per-ID entry from each MLX compile cache it is in.
 struct Expr {
   int num_inputs = 0;
   std::uintptr_t compile_id = 0;  // unique key for mlx::core::detail compile cache
@@ -134,6 +135,9 @@ struct Expr {
   // emlx::native::Program::num_real_outputs.
   int num_real_outputs = 0;
   emlx::function compiled_fn;
+  // MLX keeps one compile cache per thread, and compiled_fn adds its graph to
+  // the cache of the thread that calls it. These are the caches to erase from.
+  std::vector<mlx::core::detail::CompileCacheWeakPtr> compile_caches;
 
   ~Expr();
 };

emlx/c_src/emlx_fast.cpp

@@ -126,7 +126,7 @@ mlx::core::array fast_sdpa_impl(ErlNifEnv *env, TensorArg q, TensorArg k,
       sinks ? std::make_optional(**sinks) : std::nullopt;
 
   return fast::scaled_dot_product_attention(
-      *q, *k, *v, (float)scale, "", std::nullopt, sinks_opt, device);
+      *q, *k, *v, (float)scale, "", std::nullopt, sinks_opt, /*force_fused=*/false, device);
 }
 FINE_ASYNC_NIF(fast_sdpa)
 
@@ -141,7 +141,7 @@ mlx::core::array fast_sdpa_masked_impl(ErlNifEnv *env, TensorArg q,
       sinks ? std::make_optional(**sinks) : std::nullopt;
 
   return fast::scaled_dot_product_attention(
-      *q, *k, *v, (float)scale, "array", *mask, sinks_opt, device);
+      *q, *k, *v, (float)scale, "array", *mask, sinks_opt, /*force_fused=*/false, device);
 }
 FINE_ASYNC_NIF(fast_sdpa_masked)
 
@@ -316,7 +316,7 @@ mlx::core::array fast_sdpa_causal_key_masked_impl(
 
   if (trivial) {
     return fast::scaled_dot_product_attention(
-        *q, *k, *v, (float)scale, "causal", std::nullopt, sinks_opt, device);
+        *q, *k, *v, (float)scale, "causal", std::nullopt, sinks_opt, /*force_fused=*/false, 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)});
@@ -340,7 +340,7 @@ mlx::core::array fast_sdpa_causal_key_masked_impl(
     auto additive = where(keep, zero_val, neginf_val);
 
     auto attn_t = fast::scaled_dot_product_attention(
-        *q, *k, *v, (float)scale, "array", additive, sinks_opt, device);
+        *q, *k, *v, (float)scale, "array", additive, sinks_opt, /*force_fused=*/false, 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).
@@ -372,7 +372,7 @@ mlx::core::array fast_sdpa_causal_impl(ErlNifEnv *env, TensorArg q,
       sinks ? std::make_optional(**sinks) : std::nullopt;
 
   return fast::scaled_dot_product_attention(
-      *q, *k, *v, (float)scale, "causal", std::nullopt, sinks_opt, device);
+      *q, *k, *v, (float)scale, "causal", std::nullopt, sinks_opt, /*force_fused=*/false, device);
 }
 FINE_ASYNC_NIF(fast_sdpa_causal)
 
@@ -459,9 +459,9 @@ kv_cache_attention_impl(ErlNifEnv *env, TensorArg q, TensorArg new_k,
 
   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)
+          q_t, k_valid_t, v_valid_t, (float)scale, "", std::nullopt, std::nullopt, /*force_fused=*/false, 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);
+          q_t, k_valid_t, v_valid_t, (float)scale, "array", build_prefill_mask(), std::nullopt, /*force_fused=*/false, 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);
@@ -558,7 +558,7 @@ kv_cache_attention_masked_impl(ErlNifEnv *env, TensorArg q,       // {B, T_q,
       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);
+          additive, std::nullopt, /*force_fused=*/false, 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).
@@ -579,7 +579,7 @@ kv_cache_attention_masked_impl(ErlNifEnv *env, TensorArg q,       // {B, T_q,
 
       return mlx::core::fast::scaled_dot_product_attention(
           q_t, k_valid_t, v_valid_t, (float)scale, "array",
-          additive, std::nullopt, device);
+          additive, std::nullopt, /*force_fused=*/false, device);
     }
   }();
 
@@ -683,9 +683,9 @@ kv_cache_sdpa_update_impl(ErlNifEnv *env, TensorArg q, TensorArg new_k,
 
   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)
+          *q, k_valid, v_valid, (float)scale, "", std::nullopt, std::nullopt, /*force_fused=*/false, device)
     : mlx::core::fast::scaled_dot_product_attention(
-          *q, k_valid, v_valid, (float)scale, "array", build_prefill_mask(), std::nullopt, device);
+          *q, k_valid, v_valid, (float)scale, "array", build_prefill_mask(), std::nullopt, /*force_fused=*/false, device);
   // k_cache_owned and v_cache_owned destruct here → use_count 2 → 1.
 
   return {attn_out, k_upd, v_upd};

emlx/c_src/emlx_nif.cpp

@@ -184,19 +184,10 @@ NIF(from_blob) {
                           "Binary size is too small for the requested shape");
 
   try {
-    // Allocate MLX buffer and copy data from blob
-    size_t byte_size = blob.size;
-    allocator::Buffer mlx_buf = allocator::malloc(byte_size);
-    void *buf_ptr = mlx_buf.raw_ptr();
-
-    // Copy binary data to MLX buffer
-    std::memcpy(buf_ptr, blob.data, byte_size);
-
-    // Create deleter for the buffer
-    auto deleter = [](allocator::Buffer buf) { allocator::free(buf); };
-
-    // Create MLX array from the buffer
-    TENSOR(mlx::core::array(mlx_buf, to_shape(shape), type, deleter));
+    // Allocate MLX memory, copy the binary into it, and give it to the array
+    allocator::Data data = allocator::malloc(blob.size);
+    std::memcpy(data.buffer().raw_ptr(), blob.data, blob.size);
+    TENSOR(mlx::core::array(std::move(data), to_shape(shape), type));
   } catch (const std::exception &e) {
     return nx::nif::error(env, e.what());
   } catch (...) {
@@ -1419,7 +1410,8 @@ NIF(gather_qmm) {
 
   TENSOR(mlx::core::gather_qmm(*x, *w, *scales, biases_opt, lhs_opt, rhs_opt,
                                transpose, group_size, bits, mode,
-                               sorted_indices, device));
+                               /*global_scale=*/std::nullopt, sorted_indices,
+                               device));
 }
 
 NIF(quantized_matmul) {

emlx/mix.exs

@@ -3,7 +3,7 @@ defmodule EMLX.MixProject do
 
   @app :emlx
   @version "0.5.0"
-  @mlx_version "0.32.0"
+  @mlx_version "0.32.4"
   @source_url "https://github.com/elixir-nx/emlx"
   # mlx-build Darwin archives from v0.31.0+ encode the deployment target in the
   # filename. 26.2 ships AOT NAX (Metal 4); 14.0 is the portable reduced set.

emlx_axon/c_src/llama_plugin.cpp

@@ -285,10 +285,11 @@ mlx::core::array sdpa(const mlx::core::array &q_rope, const mlx::core::array &k_
                        int valid_len, mlx::core::StreamOrDevice stream) {
   return (T_new == 1)
              ? mlx::core::fast::scaled_dot_product_attention(
-                   q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, stream)
+                   q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, /*force_fused=*/false, stream)
              : mlx::core::fast::scaled_dot_product_attention(
                    q_rope, k_valid, v_valid, scale, "array",
-                   build_prefill_mask(q_rope, T_new, valid_len, stream), std::nullopt, stream);
+                   build_prefill_mask(q_rope, T_new, valid_len, stream), std::nullopt,
+                   /*force_fused=*/false, stream);
 }
 
 mlx::core::array layer_dense_impl(const mlx::core::array &hidden, const LayerParams &layer,

emlx_axon/c_src/qwen3_plugin.cpp

@@ -574,7 +574,7 @@ bool tensor_offset_attention(
       stream);
   auto attended = mlx::core::fast::scaled_dot_product_attention(
       query_bn, k_updated, v_updated, scale, "array", mask, std::nullopt,
-      stream);
+      /*force_fused=*/false, stream);
   attention = mlx::core::reshape(
       mlx::core::transpose(attended, {0, 2, 1, 3}, stream),
       {batch, tokens, query_heads * width}, stream);
@@ -587,10 +587,11 @@ mlx::core::array sdpa(const mlx::core::array &q_rope, const mlx::core::array &k_
                        int valid_len, mlx::core::StreamOrDevice stream) {
   return (T_new == 1)
              ? mlx::core::fast::scaled_dot_product_attention(
-                   q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, stream)
+                   q_rope, k_valid, v_valid, scale, "", std::nullopt, std::nullopt, /*force_fused=*/false, stream)
              : mlx::core::fast::scaled_dot_product_attention(
                    q_rope, k_valid, v_valid, scale, "array",
-                   build_prefill_mask(q_rope, T_new, valid_len, stream), std::nullopt, stream);
+                   build_prefill_mask(q_rope, T_new, valid_len, stream), std::nullopt,
+                   /*force_fused=*/false, stream);
 }
 
 // ── Generalized (dense-or-quantized) per-layer compute ───────────────────