Merge upstream MLX and preserve shared-buffer teardown safety

mlx/backend/common/gemma4_expert_qmm.h

@@ -119,11 +119,11 @@ inline Gemma4ExpertQMMRoute classify_gemma4_expert_qmm(
     return Gemma4ExpertQMMRoute::fallback_nax;
   }
   if (!input.affine || !input.transpose || !input.has_bias ||
-      !input.indices_uint32 || !input.indices_contiguous ||
-      !input.x_bfloat16 || !input.x_contiguous || !input.w_uint32 ||
-      !input.w_contiguous || !input.scales_bfloat16 ||
-      !input.scales_contiguous || !input.biases_bfloat16 ||
-      !input.biases_contiguous || input.group_size != 64 || input.bits != 4) {
+      !input.indices_uint32 || !input.indices_contiguous || !input.x_bfloat16 ||
+      !input.x_contiguous || !input.w_uint32 || !input.w_contiguous ||
+      !input.scales_bfloat16 || !input.scales_contiguous ||
+      !input.biases_bfloat16 || !input.biases_contiguous ||
+      input.group_size != 64 || input.bits != 4) {
     return Gemma4ExpertQMMRoute::fallback_quantization;
   }
   const bool gemma4 = input.expert_count == 128;
@@ -183,9 +183,8 @@ struct Gemma4ExpertQMMCounterSnapshot {
   bool armed{false};
 
   uint64_t attempts() const {
-    return hits + fallback_nax + fallback_outer_route +
-        fallback_quantization + fallback_topology +
-        fallback_assignment_count + fallback_geometry +
+    return hits + fallback_nax + fallback_outer_route + fallback_quantization +
+        fallback_topology + fallback_assignment_count + fallback_geometry +
         fallback_metallib_unavailable + fallback_sortedness_retracted;
   }
 };

mlx/backend/metal/allocator.cpp

@@ -7,8 +7,8 @@
 
 #include <mach/vm_page_size.h>
 #include <unistd.h>
-#include <cassert>
 #include <algorithm>
+#include <cassert>
 #include <cerrno>
 #include <cstdlib>
 #include <limits>
@@ -70,13 +70,13 @@ MetalAllocator::MetalAllocator(Device& d)
   // crash. The value may only LOWER the ceiling (it is clamped to the OS limit)
   // — raising it above what the hardware/OS reports would invite the very crash
   // this guards against. Strictly validated: a plain unsigned decimal that
-  // consumes the whole string, is non-zero, and does not overflow; anything else
-  // (empty, sign, junk, range error) is ignored and the OS limit stands.
+  // consumes the whole string, is non-zero, and does not overflow; anything
+  // else (empty, sign, junk, range error) is ignored and the OS limit stands.
   if (const char* rl = std::getenv("MLX_RESOURCE_LIMIT")) {
     while (*rl == ' ' || *rl == '\t') {
       ++rl;
     }
-    if (*rl >= '0' && *rl <= '9') {  // unsigned decimal only (reject sign/junk)
+    if (*rl >= '0' && *rl <= '9') { // unsigned decimal only (reject sign/junk)
       errno = 0;
       char* end = nullptr;
       unsigned long long v = std::strtoull(rl, &end, 10);
@@ -158,11 +158,13 @@ Buffer MetalAllocator::malloc(size_t size) {
   if (!buf) {
     size_t mem_required = get_active_memory() + get_cache_memory() + size;
 
+    auto pool = metal::new_scoped_memory_pool();
+
     // If we have a lot of memory pressure try to reclaim memory from the cache.
-    // NOTE: release_cached_buffers takes a BYTES-to-free target; when the buffers
-    // are tiny this frees only a few entries even though the COUNT is the binding
-    // constraint, so the byte path alone cannot bound num_resources_ (see the
-    // count-aware reclaim below).
+    // NOTE: release_cached_buffers takes a BYTES-to-free target; when the
+    // buffers are tiny this frees only a few entries even though the COUNT is
+    // the binding constraint, so the byte path alone cannot bound
+    // num_resources_ (see the count-aware reclaim below).
     if (mem_required >= gc_limit_ || num_resources_ >= resource_limit_) {
       num_resources_ -=
           buffer_cache_.release_cached_buffers(mem_required - gc_limit_);
@@ -171,18 +173,18 @@ Buffer MetalAllocator::malloc(size_t size) {
     // Count-aware reclaim (Darkbloom): the Metal resource COUNT limit
     // (resource_limit_, ~iogpu.rsrc_limit/499000) is independent of byte usage.
     // Under churn with many distinct buffer shapes (varied prompt lengths,
-    // growing KV caches, multiple co-resident models) freed buffers are recycled
-    // into the size-keyed cache and never reused at that exact size, so the cache
-    // ENTRY COUNT creeps toward the limit while byte usage stays modest — the
-    // byte-driven trim above never fires (its threshold is ~physical RAM). Once
-    // the count crosses a high-water mark, proactively clear the cache (pure
-    // reuse pool — clearing only costs re-allocation, never correctness) so the
-    // count drops back to the live working set. This makes the count limit
-    // unreachable by any request mix / batching method, while the existing byte
-    // limits keep total memory below physical RAM.
+    // growing KV caches, multiple co-resident models) freed buffers are
+    // recycled into the size-keyed cache and never reused at that exact size,
+    // so the cache ENTRY COUNT creeps toward the limit while byte usage stays
+    // modest — the byte-driven trim above never fires (its threshold is
+    // ~physical RAM). Once the count crosses a high-water mark, proactively
+    // clear the cache (pure reuse pool — clearing only costs re-allocation,
+    // never correctness) so the count drops back to the live working set. This
+    // makes the count limit unreachable by any request mix / batching method,
+    // while the existing byte limits keep total memory below physical RAM.
     if (resource_limit_ > 0 &&
         num_resources_ >= (resource_limit_ * resource_high_water_num_) /
-                              resource_high_water_den_) {
+                resource_high_water_den_) {
       num_resources_ -= buffer_cache_.clear();
     }
 

mlx/backend/metal/device.cpp

@@ -589,8 +589,7 @@ Device::Device() : device_(load_device()), residency_sets_(device_.get()) {
   auto pool = new_scoped_memory_pool();
   default_library_ = NS::TransferPtr(load_default_library(device_.get()));
 
-  std::string expert_qmm_env =
-      env::get_var("MLX_GATHER_QMM_EXPERT_SLICES", "");
+  std::string expert_qmm_env = env::get_var("MLX_GATHER_QMM_EXPERT_SLICES", "");
   std::transform(
       expert_qmm_env.begin(),
       expert_qmm_env.end(),
@@ -615,16 +614,14 @@ Device::Device() : device_(load_device()), residency_sets_(device_.get()) {
       "alN_true_bm_32_bn_32_bk_32";
   auto has_default_function = [this](const char* name) {
     auto ns_name = NS::String::string(name, NS::ASCIIStringEncoding);
-    auto function =
-        NS::TransferPtr(default_library_->newFunction(ns_name));
+    auto function = NS::TransferPtr(default_library_->newFunction(ns_name));
     return function.get() != nullptr;
   };
   // All expert-tile symbols ship from one source-matched metallib
   // (scripts/fetch-metallib.sh completeness contract), so availability is
   // all-or-nothing: a metallib missing any of them predates this revision
   // and must fail the whole route closed.
-  gemma4_expert_qmm_aot_available_ =
-      has_default_function(descriptor_kernel) &&
+  gemma4_expert_qmm_aot_available_ = has_default_function(descriptor_kernel) &&
       has_default_function(descriptor_kernel_e256) &&
       has_default_function(tile_kernel);
   if (gemma4_expert_qmm_requested_ && gemma4_expert_qmm_aot_available_) {
@@ -1043,8 +1040,7 @@ void gemma4_expert_qmm_diagnostics_snapshot(
     diagnostics->fallback_outer_route = counters.fallback_outer_route;
     diagnostics->fallback_quantization = counters.fallback_quantization;
     diagnostics->fallback_topology = counters.fallback_topology;
-    diagnostics->fallback_assignment_count =
-        counters.fallback_assignment_count;
+    diagnostics->fallback_assignment_count = counters.fallback_assignment_count;
     diagnostics->fallback_geometry = counters.fallback_geometry;
     diagnostics->fallback_metallib_unavailable =
         counters.fallback_metallib_unavailable;
@@ -1080,8 +1076,7 @@ extern "C" void mlx_metal_gemma4_expert_qmm_diagnostics_clear_and_arm(void) {
   }
 }
 
-extern "C" void
-mlx_metal_gemma4_expert_qmm_diagnostics_snapshot_and_disarm(
+extern "C" void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot_and_disarm(
     mlx_metal_gemma4_expert_qmm_diagnostics* diagnostics) {
   gemma4_expert_qmm_diagnostics_snapshot(diagnostics, true);
 }

mlx/backend/metal/device.h

@@ -10,6 +10,7 @@
 #include <unordered_set>
 
 #include "mlx/array.h"
+#include "mlx/backend/common/gemma4_expert_qmm.h"
 #include "mlx/backend/common/metal_kernel.h"
 #include "mlx/backend/common/gemma4_expert_qmm.h"
 #include "mlx/backend/metal/resident.h"
@@ -196,10 +197,6 @@ class MLX_API Device {
       const MTLFCList& func_consts = {},
       const std::vector<MTL::Function*>& linked_functions = {});
 
-  ResidencySets& residency_sets() {
-    return residency_sets_;
-  }
-
   bool gemma4_expert_qmm_requested() const {
     return gemma4_expert_qmm_requested_;
   }
@@ -240,6 +237,10 @@ class MLX_API Device {
     gemma4_expert_qmm_counters_.clear_and_arm();
   }
 
+  ResidencySets& residency_sets() {
+    return residency_sets_;
+  }
+
  private:
   NS::SharedPtr<MTL::Library> build_library_(
       const std::string& source_string,

mlx/backend/metal/kernels/quantized.h

@@ -2615,8 +2615,7 @@ template <int NE>
   for (int i = int(lid) + 1; i < M; i += int(expert_count)) {
     adjacent_ok = adjacent_ok && indices[i - 1] <= indices[i];
   }
-  const uint violation_vote =
-      simd_or((boundary_ok && adjacent_ok) ? 0u : 1u);
+  const uint violation_vote = simd_or((boundary_ok && adjacent_ok) ? 0u : 1u);
   if (simd_lid == 0) {
     violation_votes[simd_gid] = violation_vote;
   }
@@ -2650,8 +2649,7 @@ template <int NE>
     threadgroup_barrier(mem_flags::mem_threadgroup);
   }
 
-  const uint descriptor_count =
-      inclusive_tile_offsets[expert_count - 1];
+  const uint descriptor_count = inclusive_tile_offsets[expert_count - 1];
   if (lid == expert_count - 1) {
     // A retracted count keeps the tile kernel's capacity check memory-safe
     // (every threadgroup early-returns) and unambiguously signals the host:
@@ -2666,8 +2664,7 @@ template <int NE>
     uint expert_lower = 0;
     uint expert_upper = expert_count;
     while (expert_lower < expert_upper) {
-      const uint midpoint =
-          expert_lower + (expert_upper - expert_lower) / 2;
+      const uint midpoint = expert_lower + (expert_upper - expert_lower) / 2;
       if (inclusive_tile_offsets[midpoint] <= slot) {
         expert_lower = midpoint + 1;
       } else {
@@ -2677,10 +2674,8 @@ template <int NE>
     const uint expert = expert_lower;
     const uint expert_tile_begin =
         expert == 0 ? 0 : inclusive_tile_offsets[expert - 1];
-    const uint row =
-        segment_starts[expert] + (slot - expert_tile_begin) * BM;
-    const uint row_count =
-        min(BM, segment_starts[expert + 1] - row);
+    const uint row = segment_starts[expert] + (slot - expert_tile_begin) * BM;
+    const uint row_count = min(BM, segment_starts[expert + 1] - row);
     descriptors[slot] = uint4(row, row_count, expert, 0);
   }
 }
@@ -2737,8 +2732,7 @@ template <
   x += row_start * size_t(K);
   y += row_start * size_t(N);
   const device uint8_t* expert_w =
-      reinterpret_cast<const device uint8_t*>(w) +
-      expert * expert_w_stride;
+      reinterpret_cast<const device uint8_t*>(w) + expert * expert_w_stride;
   scales += expert * expert_sb_stride;
   biases += expert * expert_sb_stride;
 

mlx/backend/metal/quantized.cpp

@@ -15,9 +15,9 @@
 
 namespace mlx::core {
 
+using metal::classify_gemma4_expert_qmm;
 using metal::Gemma4ExpertQMMRoute;
 using metal::Gemma4ExpertQMMRouteInput;
-using metal::classify_gemma4_expert_qmm;
 
 namespace {
 
@@ -1691,8 +1691,7 @@ Gemma4ExpertQMMRoute try_gemma4_expert_qmm(
   compute_encoder.set_bytes(N, c++);
 
   compute_encoder.dispatch_threadgroups(
-      MTL::Size((N + bn - 1) / bn, max_tile_count, 1),
-      MTL::Size(32, wn, wm));
+      MTL::Size((N + bn - 1) / bn, max_tile_count, 1), MTL::Size(32, wn, wm));
   return Gemma4ExpertQMMRoute::hit;
 }
 
@@ -1721,8 +1720,7 @@ void gather_qmm_rhs(
       route_input.requested = true;
       route_input.outer_route = true;
       route_input.nax_available = true;
-      d.record_armed_gemma4_expert_qmm(
-          classify_gemma4_expert_qmm(route_input));
+      d.record_armed_gemma4_expert_qmm(classify_gemma4_expert_qmm(route_input));
     }
     return gather_qmm_rhs_nax(
         /* const array& x_ = */ x_,
@@ -1798,8 +1796,7 @@ void gather_qmm_rhs(
     route_input.biases_contiguous = biases_ && biases_->flags().row_contiguous;
     route_input.group_size = group_size;
     route_input.bits = bits;
-    route_input.expert_count =
-        w.size() / w.shape(-1) / w.shape(-2);
+    route_input.expert_count = w.size() / w.shape(-1) / w.shape(-2);
     route_input.assignments = M;
     route_input.index_count = indices.size();
     route_input.k = K;
@@ -2078,11 +2075,8 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
 // record below and the dispatch decision evaluate this one predicate so a
 // future tuning change cannot desynchronize them.
 // TODO: Tune 16 and 4 here a bit better.
-static constexpr bool takes_sorted_rhs_route(
-    int M,
-    int B,
-    int E,
-    bool right_sorted) {
+static constexpr bool
+takes_sorted_rhs_route(int M, int B, int E, bool right_sorted) {
   return M == 1 && B >= 16 && right_sorted && B / E >= 4;
 }
 
@@ -2116,8 +2110,7 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
     Gemma4ExpertQMMRouteInput route_input;
     route_input.requested = true;
     route_input.outer_route = false;
-    d.record_armed_gemma4_expert_qmm(
-        classify_gemma4_expert_qmm(route_input));
+    d.record_armed_gemma4_expert_qmm(classify_gemma4_expert_qmm(route_input));
   }
 
   // We are walking x in order and w is also in order so we can batch up the

tests/array_tests.cpp

@@ -601,6 +601,7 @@ TEST_CASE("test array shared buffer") {
   array b = array(buf_b, shape, float32, deleter);
 
   eval(a + b);
+  synchronize();
 }
 
 TEST_CASE("test make empty array") {