From 3e78c72a3b54962d1fb5f21b82f4806ee1b17267 Mon Sep 17 00:00:00 2001
From: Nick Paterno <43416138+nicholasjpaterno@users.noreply.github.com>
Date: Sat, 28 Feb 2026 16:50:01 -0500
Subject: [PATCH 01/17] feat: zero-config metallib discovery for pmetal

- Patch MLX device.cpp via FetchContent PATCH_COMMAND to check
  PMETAL_METALLIB_PATH env var and ~/.cache/pmetal/lib/ before
  existing fallback chain
- Stage mlx-c source to OUT_DIR before cmake to avoid modifying
  the git submodule
- Auto-cache mlx.metallib to ~/.cache/pmetal/lib/ during build
  (critical for cargo install where build dir is cleaned up)
- Bump pmetal-mlx-sys to 0.2.2, pmetal-mlx-rs to 0.25.5
---
 Cargo.toml                                 |   4 +-
 mlx-rs/Cargo.toml                          |   8 +-
 mlx-sys/Cargo.toml                         |   8 +-
 mlx-sys/build.rs                           | 120 ++++++++++++++++++++-
 mlx-sys/patches/metallib-search-path.patch |  90 ++++++++++++++++
 5 files changed, 219 insertions(+), 11 deletions(-)
 create mode 100644 mlx-sys/patches/metallib-search-path.patch

diff --git a/Cargo.toml b/Cargo.toml
index 100712f0f..249188323 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -32,10 +32,10 @@ resolver = "2"
 
 [workspace.dependencies]
 # workspace local dependencies
-mlx-sys = { version = "=0.2.0", path = "mlx-sys" }
+mlx-sys = { version = "=0.2.2", path = "mlx-sys", package = "pmetal-mlx-sys" }
 mlx-macros = { version = "0.25", path = "mlx-macros" }
 mlx-internal-macros = { version = "0.25", path = "mlx-internal-macros" }
-mlx-rs = { version = "0.25", path = "mlx-rs" }
+mlx-rs = { version = "0.25.5", path = "mlx-rs", package = "pmetal-mlx-rs" }
 mlx-lm = { version = "0.0.1", path = "mlx-lm" }
 mlx-lm-utils = { version = "0.0.1", path = "mlx-lm-utils" }
 
diff --git a/mlx-rs/Cargo.toml b/mlx-rs/Cargo.toml
index c2c4efd03..b5339c16c 100644
--- a/mlx-rs/Cargo.toml
+++ b/mlx-rs/Cargo.toml
@@ -1,14 +1,14 @@
 [package]
-name = "mlx-rs"
-version.workspace = true
+name = "pmetal-mlx-rs"
+version = "0.25.5"
 authors.workspace = true
 edition.workspace = true
-repository.workspace = true
+repository = "https://github.com/oxiglade/mlx-rs"
 keywords.workspace = true
 categories.workspace = true
 license.workspace = true
 documentation.workspace = true
-description = "Unofficial rust wrapper for Apple's mlx machine learning library."
+description = "pmetal-maintained fork of mlx-rs: unofficial rust wrapper for Apple's mlx machine learning library."
 readme = "README.md"
 
 [package.metadata.docs.rs]
diff --git a/mlx-sys/Cargo.toml b/mlx-sys/Cargo.toml
index 55d5325f3..864e92d03 100644
--- a/mlx-sys/Cargo.toml
+++ b/mlx-sys/Cargo.toml
@@ -1,11 +1,11 @@
 [package]
-name = "mlx-sys"
-version = "0.2.0" # mlx-sys version should follow that of mlx-c
+name = "pmetal-mlx-sys"
+version = "0.2.2" # mlx-sys version should follow that of mlx-c
 authors.workspace = true
 edition.workspace = true
 
-description = "Low-level interface and binding generation for the mlx library"
-repository.workspace = true
+description = "pmetal-maintained fork of mlx-sys: low-level interface and binding generation for the mlx library"
+repository = "https://github.com/oxiglade/mlx-rs"
 keywords.workspace = true
 categories.workspace = true
 license.workspace = true
diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index 11f6300f0..fb2ad09dc 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -44,8 +44,85 @@ fn find_clang_rt_path() -> Option<String> {
     None
 }
 
+/// Resolve the macOS deployment target.
+///
+/// Uses `MACOSX_DEPLOYMENT_TARGET` env var if set, otherwise defaults to 14.0
+/// (MLX's minimum supported version for Metal).
+#[cfg(target_os = "macos")]
+fn resolve_deployment_target() -> String {
+    env::var("MACOSX_DEPLOYMENT_TARGET").unwrap_or_else(|_| "14.0".to_string())
+}
+
+/// Copy src/mlx-c to a staging directory and inject the metallib search-path
+/// patch into the CMakeLists.txt. This avoids modifying the mlx-c git submodule
+/// while ensuring the patch is applied when MLX is fetched via FetchContent.
+fn prepare_mlx_c_source() -> PathBuf {
+    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
+    let staged = out_dir.join("mlx-c-staged");
+    let src = PathBuf::from("src/mlx-c");
+
+    // Copy the entire mlx-c source tree to the staging area
+    if staged.exists() {
+        std::fs::remove_dir_all(&staged).expect("Failed to clean staged mlx-c");
+    }
+    copy_dir_recursive(&src, &staged).expect("Failed to copy mlx-c to staging");
+
+    // Copy our patch file into the staged source
+    let patches_dir = staged.join("patches");
+    std::fs::create_dir_all(&patches_dir).expect("Failed to create patches dir");
+    std::fs::copy(
+        "patches/metallib-search-path.patch",
+        patches_dir.join("metallib-search-path.patch"),
+    )
+    .expect("Failed to copy metallib patch");
+
+    // Inject PATCH_COMMAND into the FetchContent_Declare for MLX
+    let cmake_path = staged.join("CMakeLists.txt");
+    let cmake_content =
+        std::fs::read_to_string(&cmake_path).expect("Failed to read CMakeLists.txt");
+    let patched = cmake_content.replace(
+        "GIT_TAG v0.30.6)",
+        "GIT_TAG v0.30.6\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true)",
+    );
+    std::fs::write(&cmake_path, patched).expect("Failed to write patched CMakeLists.txt");
+
+    // Tell cargo to rerun if the patch changes
+    println!("cargo:rerun-if-changed=patches/metallib-search-path.patch");
+
+    staged
+}
+
+fn copy_dir_recursive(src: &std::path::Path, dst: &std::path::Path) -> std::io::Result<()> {
+    std::fs::create_dir_all(dst)?;
+    for entry in std::fs::read_dir(src)? {
+        let entry = entry?;
+        let ty = entry.file_type()?;
+        let dest_path = dst.join(entry.file_name());
+        if ty.is_dir() {
+            copy_dir_recursive(&entry.path(), &dest_path)?;
+        } else if ty.is_symlink() {
+            // Resolve symlinks (common in git submodules)
+            let target = std::fs::read_link(entry.path())?;
+            let resolved = if target.is_absolute() {
+                target
+            } else {
+                entry.path().parent().unwrap().join(&target)
+            };
+            if resolved.is_dir() {
+                copy_dir_recursive(&resolved, &dest_path)?;
+            } else {
+                std::fs::copy(&resolved, &dest_path)?;
+            }
+        } else {
+            std::fs::copy(entry.path(), &dest_path)?;
+        }
+    }
+    Ok(())
+}
+
 fn build_and_link_mlx_c() {
-    let mut config = Config::new("src/mlx-c");
+    let mlx_c_src = prepare_mlx_c_source();
+    let mut config = Config::new(&mlx_c_src);
     config.very_verbose(true);
     config.define("CMAKE_INSTALL_PREFIX", ".");
 
@@ -104,6 +181,47 @@ fn build_and_link_mlx_c() {
         println!("cargo:rustc-link-search={}", clang_rt_path);
         println!("cargo:rustc-link-lib=static=clang_rt.osx");
     }
+
+    // Cache mlx.metallib to ~/.cache/pmetal/lib/ so the binary works regardless
+    // of where it's installed. This is critical for `cargo install` where the
+    // build directory is cleaned up after the binary is placed.
+    #[cfg(feature = "metal")]
+    {
+        let metallib = dst.join("build/lib/mlx.metallib");
+        if metallib.exists() {
+            if let Ok(home) = env::var("HOME") {
+                let cache_dir = PathBuf::from(home).join(".cache/pmetal/lib");
+                let dest = cache_dir.join("mlx.metallib");
+                let should_copy = if dest.exists() {
+                    // Replace if the build artifact is newer
+                    dest.metadata()
+                        .and_then(|d| {
+                            metallib.metadata().map(|s| {
+                                s.modified().ok().zip(d.modified().ok()).is_some_and(
+                                    |(src_t, dst_t)| src_t > dst_t,
+                                )
+                            })
+                        })
+                        .unwrap_or(false)
+                } else {
+                    true
+                };
+                if should_copy {
+                    let _ = std::fs::create_dir_all(&cache_dir);
+                    match std::fs::copy(&metallib, &dest) {
+                        Ok(_) => println!(
+                            "cargo:warning=Cached mlx.metallib to {}",
+                            dest.display()
+                        ),
+                        Err(e) => println!(
+                            "cargo:warning=Failed to cache mlx.metallib: {}",
+                            e
+                        ),
+                    }
+                }
+            }
+        }
+    }
 }
 
 fn main() {
diff --git a/mlx-sys/patches/metallib-search-path.patch b/mlx-sys/patches/metallib-search-path.patch
new file mode 100644
index 000000000..54672e0a2
--- /dev/null
+++ b/mlx-sys/patches/metallib-search-path.patch
@@ -0,0 +1,90 @@
+--- a/mlx/backend/metal/device.cpp
++++ b/mlx/backend/metal/device.cpp
+@@ -87,6 +87,28 @@
+ }
+ #endif
+
++// pmetal: Check PMETAL_METALLIB_PATH env var (set by pmetal-cli at startup)
++std::pair<MTL::Library*, NS::Error*> load_env_override_library(
++    MTL::Device* device) {
++  const char* env_path = std::getenv("PMETAL_METALLIB_PATH");
++  if (env_path && env_path[0] != '\0') {
++    return load_library_from_path(device, env_path);
++  }
++  return {nullptr, nullptr};
++}
++
++// pmetal: Check ~/.cache/pmetal/lib/mlx.metallib (XDG cache location)
++std::pair<MTL::Library*, NS::Error*> load_user_cache_library(
++    MTL::Device* device) {
++  const char* home = std::getenv("HOME");
++  if (home && home[0] != '\0') {
++    std::string cache_path =
++        std::string(home) + "/.cache/pmetal/lib/mlx.metallib";
++    return load_library_from_path(device, cache_path.c_str());
++  }
++  return {nullptr, nullptr};
++}
++
+ // Firstly, search for the metallib in the same path as this binary
+ std::pair<MTL::Library*, NS::Error*> load_colocated_library(
+     MTL::Device* device,
+@@ -134,38 +156,51 @@
+ }
+
+ MTL::Library* load_default_library(MTL::Device* device) {
+-  NS::Error* error[5];
++  NS::Error* error[7];
+   MTL::Library* lib;
++
++  // pmetal: Try env var override first (set by pmetal-cli resolver)
++  std::tie(lib, error[0]) = load_env_override_library(device);
++  if (lib) {
++    return lib;
++  }
++
++  // pmetal: Try user cache directory
++  std::tie(lib, error[1]) = load_user_cache_library(device);
++  if (lib) {
++    return lib;
++  }
++
+   // First try the colocated mlx.metallib
+-  std::tie(lib, error[0]) = load_colocated_library(device, "mlx");
++  std::tie(lib, error[2]) = load_colocated_library(device, "mlx");
+   if (lib) {
+     return lib;
+   }
+
+-  std::tie(lib, error[1]) = load_colocated_library(device, "Resources/mlx");
++  std::tie(lib, error[3]) = load_colocated_library(device, "Resources/mlx");
+   if (lib) {
+     return lib;
+   }
+
+   // Then try default.metallib in a SwiftPM bundle if we have one
+-  std::tie(lib, error[2]) = load_swiftpm_library(device, "default");
++  std::tie(lib, error[4]) = load_swiftpm_library(device, "default");
+   if (lib) {
+     return lib;
+   }
+
+   // Try lo load resources from Framework resources if SwiftPM wrapped as a
+   // dynamic framework.
+-  std::tie(lib, error[3]) = load_colocated_library(device, "Resources/default");
++  std::tie(lib, error[5]) = load_colocated_library(device, "Resources/default");
+   if (lib) {
+     return lib;
+   }
+
+   // Finally try default_mtllib_path
+-  std::tie(lib, error[4]) = load_library_from_path(device, default_mtllib_path);
++  std::tie(lib, error[6]) = load_library_from_path(device, default_mtllib_path);
+   if (!lib) {
+     std::ostringstream msg;
+     msg << "Failed to load the default metallib. ";
+-    for (int i = 0; i < 5; i++) {
++    for (int i = 0; i < 7; i++) {
+       if (error[i] != nullptr) {
+         msg << error[i]->localizedDescription()->utf8String() << " ";
+       }

From 3874ff5fc1e12a13fc2e16497f6cd599895d6524 Mon Sep 17 00:00:00 2001
From: Nick Paterno <43416138+nicholasjpaterno@users.noreply.github.com>
Date: Sat, 28 Feb 2026 21:36:20 -0500
Subject: [PATCH 02/17] fix: enforce macOS deployment target >= 14.0 in
 build.rs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Cargo defaults MACOSX_DEPLOYMENT_TARGET to 10.13 which causes the
cmake crate to inject -mmacosx-version-min=10.13 into CFLAGS. MLX's
CMakeLists.txt then rejects the build with "MLX requires macOS >= 14.0".

Set MACOSX_DEPLOYMENT_TARGET and CMAKE_OSX_DEPLOYMENT_TARGET to 14.0
(or user override) before cmake runs. Fixes builds from Tauri and
other frameworks that don't set their own deployment target.

Bumps pmetal-mlx-sys 0.2.2 → 0.2.3, pmetal-mlx-rs 0.25.5 → 0.25.6.
---
 Cargo.toml         |  4 ++--
 mlx-rs/Cargo.toml  |  2 +-
 mlx-sys/Cargo.toml |  2 +-
 mlx-sys/build.rs   | 16 ++++++++++++++++
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 249188323..d57891595 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -32,10 +32,10 @@ resolver = "2"
 
 [workspace.dependencies]
 # workspace local dependencies
-mlx-sys = { version = "=0.2.2", path = "mlx-sys", package = "pmetal-mlx-sys" }
+mlx-sys = { version = "=0.2.3", path = "mlx-sys", package = "pmetal-mlx-sys" }
 mlx-macros = { version = "0.25", path = "mlx-macros" }
 mlx-internal-macros = { version = "0.25", path = "mlx-internal-macros" }
-mlx-rs = { version = "0.25.5", path = "mlx-rs", package = "pmetal-mlx-rs" }
+mlx-rs = { version = "0.25.6", path = "mlx-rs", package = "pmetal-mlx-rs" }
 mlx-lm = { version = "0.0.1", path = "mlx-lm" }
 mlx-lm-utils = { version = "0.0.1", path = "mlx-lm-utils" }
 
diff --git a/mlx-rs/Cargo.toml b/mlx-rs/Cargo.toml
index b5339c16c..fbcae51c9 100644
--- a/mlx-rs/Cargo.toml
+++ b/mlx-rs/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pmetal-mlx-rs"
-version = "0.25.5"
+version = "0.25.6"
 authors.workspace = true
 edition.workspace = true
 repository = "https://github.com/oxiglade/mlx-rs"
diff --git a/mlx-sys/Cargo.toml b/mlx-sys/Cargo.toml
index 864e92d03..d1ba12766 100644
--- a/mlx-sys/Cargo.toml
+++ b/mlx-sys/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pmetal-mlx-sys"
-version = "0.2.2" # mlx-sys version should follow that of mlx-c
+version = "0.2.3" # mlx-sys version should follow that of mlx-c
 authors.workspace = true
 edition.workspace = true
 
diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index fb2ad09dc..7582bf3b1 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -121,11 +121,27 @@ fn copy_dir_recursive(src: &std::path::Path, dst: &std::path::Path) -> std::io::
 }
 
 fn build_and_link_mlx_c() {
+    // MLX requires macOS >= 14.0 for Metal support. Override the deployment
+    // target early so the cmake crate (and cc crate) don't inject a lower
+    // -mmacosx-version-min flag into CFLAGS/CXXFLAGS. Without this, Cargo's
+    // default target (10.13) causes MLX's CMakeLists.txt to reject the build.
+    #[cfg(target_os = "macos")]
+    {
+        let target = resolve_deployment_target();
+        env::set_var("MACOSX_DEPLOYMENT_TARGET", &target);
+    }
+
     let mlx_c_src = prepare_mlx_c_source();
     let mut config = Config::new(&mlx_c_src);
     config.very_verbose(true);
     config.define("CMAKE_INSTALL_PREFIX", ".");
 
+    #[cfg(target_os = "macos")]
+    {
+        let target = resolve_deployment_target();
+        config.define("CMAKE_OSX_DEPLOYMENT_TARGET", &target);
+    }
+
     // Use Xcode's clang to ensure compatibility with the macOS SDK
     config.define("CMAKE_C_COMPILER", "/usr/bin/cc");
     config.define("CMAKE_CXX_COMPILER", "/usr/bin/c++");

From 44db9d63685e650810b8f7ac5a93fec989ff4943 Mon Sep 17 00:00:00 2001
From: Nick Paterno <43416138+nicholasjpaterno@users.noreply.github.com>
Date: Sat, 28 Feb 2026 23:13:58 -0500
Subject: [PATCH 03/17] fix: enforce minimum macOS 14.0 deployment target for
 MLX
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

resolve_deployment_target() now enforces a floor of 14.0 instead of
blindly using MACOSX_DEPLOYMENT_TARGET. Cargo/Tauri default this to
10.13 which MLX's CMakeLists.txt rejects. Higher values (e.g. 15.0)
are preserved.

Bumps pmetal-mlx-sys 0.2.3 → 0.2.4, pmetal-mlx-rs 0.25.6 → 0.25.7.
---
 Cargo.toml         |  4 ++--
 mlx-rs/Cargo.toml  |  2 +-
 mlx-sys/Cargo.toml |  2 +-
 mlx-sys/build.rs   | 17 ++++++++++++++---
 4 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index d57891595..0481d3ab7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -32,10 +32,10 @@ resolver = "2"
 
 [workspace.dependencies]
 # workspace local dependencies
-mlx-sys = { version = "=0.2.3", path = "mlx-sys", package = "pmetal-mlx-sys" }
+mlx-sys = { version = "=0.2.4", path = "mlx-sys", package = "pmetal-mlx-sys" }
 mlx-macros = { version = "0.25", path = "mlx-macros" }
 mlx-internal-macros = { version = "0.25", path = "mlx-internal-macros" }
-mlx-rs = { version = "0.25.6", path = "mlx-rs", package = "pmetal-mlx-rs" }
+mlx-rs = { version = "0.25.7", path = "mlx-rs", package = "pmetal-mlx-rs" }
 mlx-lm = { version = "0.0.1", path = "mlx-lm" }
 mlx-lm-utils = { version = "0.0.1", path = "mlx-lm-utils" }
 
diff --git a/mlx-rs/Cargo.toml b/mlx-rs/Cargo.toml
index fbcae51c9..2ebc4ac91 100644
--- a/mlx-rs/Cargo.toml
+++ b/mlx-rs/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pmetal-mlx-rs"
-version = "0.25.6"
+version = "0.25.7"
 authors.workspace = true
 edition.workspace = true
 repository = "https://github.com/oxiglade/mlx-rs"
diff --git a/mlx-sys/Cargo.toml b/mlx-sys/Cargo.toml
index d1ba12766..f490be285 100644
--- a/mlx-sys/Cargo.toml
+++ b/mlx-sys/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pmetal-mlx-sys"
-version = "0.2.3" # mlx-sys version should follow that of mlx-c
+version = "0.2.4" # mlx-sys version should follow that of mlx-c
 authors.workspace = true
 edition.workspace = true
 
diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index 7582bf3b1..3ba98dd5c 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -46,11 +46,22 @@ fn find_clang_rt_path() -> Option<String> {
 
 /// Resolve the macOS deployment target.
 ///
-/// Uses `MACOSX_DEPLOYMENT_TARGET` env var if set, otherwise defaults to 14.0
-/// (MLX's minimum supported version for Metal).
+/// Enforces a minimum of 14.0 (MLX's requirement for Metal support).
+/// If `MACOSX_DEPLOYMENT_TARGET` is set to a higher value, that is used instead.
+/// Cargo/Tauri often default to 10.13, which MLX's CMakeLists.txt rejects.
 #[cfg(target_os = "macos")]
 fn resolve_deployment_target() -> String {
-    env::var("MACOSX_DEPLOYMENT_TARGET").unwrap_or_else(|_| "14.0".to_string())
+    const MLX_MIN_MACOS: (u32, u32) = (14, 0);
+
+    if let Ok(val) = env::var("MACOSX_DEPLOYMENT_TARGET") {
+        let parts: Vec<u32> = val.split('.').filter_map(|s| s.parse().ok()).collect();
+        let major = parts.first().copied().unwrap_or(0);
+        let minor = parts.get(1).copied().unwrap_or(0);
+        if (major, minor) >= MLX_MIN_MACOS {
+            return val;
+        }
+    }
+    format!("{}.{}", MLX_MIN_MACOS.0, MLX_MIN_MACOS.1)
 }
 
 /// Copy src/mlx-c to a staging directory and inject the metallib search-path

From 4b43ffff9bb5a614572dcffd41c1bbdeb68200e0 Mon Sep 17 00:00:00 2001
From: Nick Paterno <43416138+nicholasjpaterno@users.noreply.github.com>
Date: Sat, 21 Mar 2026 00:31:09 -0400
Subject: [PATCH 04/17] feat(memory): expose MLX Metal memory management API

Add mlx_rs::memory module wrapping mlx_sys FFI calls for:
- get_active_memory, get_peak_memory, get_cache_memory
- get_memory_limit, set_memory_limit
- set_cache_limit, set_wired_limit
- clear_cache, reset_peak_memory

Bump version to 0.25.8.
---
 mlx-rs/Cargo.toml    |   4 +-
 mlx-rs/src/lib.rs    |   1 +
 mlx-rs/src/memory.rs | 177 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 180 insertions(+), 2 deletions(-)
 create mode 100644 mlx-rs/src/memory.rs

diff --git a/mlx-rs/Cargo.toml b/mlx-rs/Cargo.toml
index 2ebc4ac91..caa0213c2 100644
--- a/mlx-rs/Cargo.toml
+++ b/mlx-rs/Cargo.toml
@@ -1,9 +1,9 @@
 [package]
 name = "pmetal-mlx-rs"
-version = "0.25.7"
+version = "0.25.8"
 authors.workspace = true
 edition.workspace = true
-repository = "https://github.com/oxiglade/mlx-rs"
+repository = "https://github.com/nicholasjpaterno/mlx-rs"
 keywords.workspace = true
 categories.workspace = true
 license.workspace = true
diff --git a/mlx-rs/src/lib.rs b/mlx-rs/src/lib.rs
index 1cc36f970..53459f34a 100644
--- a/mlx-rs/src/lib.rs
+++ b/mlx-rs/src/lib.rs
@@ -292,6 +292,7 @@ pub mod fast;
 pub mod fft;
 pub mod linalg;
 pub mod losses;
+pub mod memory;
 pub mod module;
 pub mod nested;
 pub mod nn;
diff --git a/mlx-rs/src/memory.rs b/mlx-rs/src/memory.rs
new file mode 100644
index 000000000..f07b1f969
--- /dev/null
+++ b/mlx-rs/src/memory.rs
@@ -0,0 +1,177 @@
+//! Metal memory management for MLX.
+//!
+//! MLX uses a caching allocator for Metal buffers. When arrays are freed,
+//! their underlying buffers are retained in a cache for reuse rather than
+//! being returned to the system. This module exposes controls over that cache
+//! and provides visibility into memory usage.
+//!
+//! # Memory Model
+//!
+//! - **Active memory**: Buffers currently held by live [`Array`](crate::Array) objects.
+//! - **Cache memory**: Freed buffers retained for reuse (not returned to OS).
+//! - **Peak memory**: High-water mark since process start or last [`reset_peak_memory`].
+//! - **Memory limit**: Soft cap that triggers backpressure during graph evaluation.
+//!   When active memory exceeds this limit, MLX blocks and waits for in-flight
+//!   GPU operations to complete before scheduling more work.
+//! - **Cache limit**: Maximum size of the buffer cache. Excess freed buffers are
+//!   returned to the system immediately.
+//!
+//! # Example
+//!
+//! ```rust,ignore
+//! use mlx_rs::memory;
+//!
+//! // Check current usage
+//! let active = memory::get_active_memory();
+//! let cached = memory::get_cache_memory();
+//! println!("Active: {} bytes, Cached: {} bytes", active, cached);
+//!
+//! // Clear the buffer cache to free memory
+//! memory::clear_cache();
+//!
+//! // Threshold-based clearing (like mlx-lm)
+//! if memory::get_cache_memory() > 2 * 1024 * 1024 * 1024 {
+//!     memory::clear_cache();
+//! }
+//! ```
+
+/// Get the number of bytes currently allocated by MLX's Metal allocator.
+///
+/// This is "active" memory — buffers held by live arrays. Does **not** include
+/// cached (freed but retained) buffers.
+pub fn get_active_memory() -> usize {
+    let mut res: usize = 0;
+    // SAFETY: mlx_get_active_memory writes a single size_t through a valid pointer.
+    unsafe { mlx_sys::mlx_get_active_memory(&mut res) };
+    res
+}
+
+/// Get the peak memory usage since process start or last [`reset_peak_memory`].
+pub fn get_peak_memory() -> usize {
+    let mut res: usize = 0;
+    // SAFETY: mlx_get_peak_memory writes a single size_t through a valid pointer.
+    unsafe { mlx_sys::mlx_get_peak_memory(&mut res) };
+    res
+}
+
+/// Get the number of bytes held in the buffer cache.
+///
+/// These are freed buffers retained for reuse. They count toward process RSS
+/// but are available for reallocation without a system call.
+pub fn get_cache_memory() -> usize {
+    let mut res: usize = 0;
+    // SAFETY: mlx_get_cache_memory writes a single size_t through a valid pointer.
+    unsafe { mlx_sys::mlx_get_cache_memory(&mut res) };
+    res
+}
+
+/// Get the current memory limit.
+///
+/// During graph evaluation, if active memory exceeds this limit, MLX blocks
+/// and waits for in-flight GPU operations to complete before scheduling more
+/// work. Default is 1.5× the device's recommended working set size.
+pub fn get_memory_limit() -> usize {
+    let mut res: usize = 0;
+    // SAFETY: mlx_get_memory_limit writes a single size_t through a valid pointer.
+    unsafe { mlx_sys::mlx_get_memory_limit(&mut res) };
+    res
+}
+
+/// Set the memory limit for MLX's backpressure mechanism.
+///
+/// Returns the previous limit. Setting to 0 disables the limit.
+pub fn set_memory_limit(limit: usize) -> usize {
+    let mut prev: usize = 0;
+    // SAFETY: mlx_set_memory_limit writes the old limit and sets the new one.
+    unsafe { mlx_sys::mlx_set_memory_limit(&mut prev, limit) };
+    prev
+}
+
+/// Set the maximum size of the buffer cache.
+///
+/// Freed buffers beyond this limit are returned to the system immediately.
+/// Returns the previous cache limit. Setting to 0 disables caching entirely.
+pub fn set_cache_limit(limit: usize) -> usize {
+    let mut prev: usize = 0;
+    // SAFETY: mlx_set_cache_limit writes the old limit and sets the new one.
+    unsafe { mlx_sys::mlx_set_cache_limit(&mut prev, limit) };
+    prev
+}
+
+/// Set the wired memory limit (macOS 15.0+).
+///
+/// Wired buffers are kept resident in GPU memory and not paged out.
+/// Returns the previous wired limit. Setting to 0 (default) disables
+/// residency tracking.
+pub fn set_wired_limit(limit: usize) -> usize {
+    let mut prev: usize = 0;
+    // SAFETY: mlx_set_wired_limit writes the old limit and sets the new one.
+    unsafe { mlx_sys::mlx_set_wired_limit(&mut prev, limit) };
+    prev
+}
+
+/// Clear the Metal buffer cache, returning all cached buffers to the system.
+///
+/// This frees buffers that were retained for reuse after their owning arrays
+/// were dropped. It does **not** affect buffers held by live arrays.
+///
+/// **When to call:**
+/// - After a failed initialization (e.g., ANE fallback) before loading a new model
+/// - After training completes to release memory
+/// - When cache memory exceeds a threshold (like mlx-lm's `_clear_cache`)
+///
+/// **When NOT to call:**
+/// - Between training steps (causes reallocation storms)
+/// - Between epochs (same issue — buffers are immediately re-needed)
+pub fn clear_cache() {
+    // SAFETY: mlx_clear_cache has no preconditions and is idempotent.
+    unsafe { mlx_sys::mlx_clear_cache() };
+}
+
+/// Reset the peak memory counter to zero.
+///
+/// After calling this, [`get_peak_memory`] tracks the new maximum from
+/// this point forward.
+pub fn reset_peak_memory() {
+    // SAFETY: mlx_reset_peak_memory has no preconditions and is idempotent.
+    unsafe { mlx_sys::mlx_reset_peak_memory() };
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_memory_queries_dont_crash() {
+        let _active = get_active_memory();
+        let _peak = get_peak_memory();
+        let _cache = get_cache_memory();
+        let _limit = get_memory_limit();
+    }
+
+    #[test]
+    fn test_clear_cache() {
+        clear_cache(); // idempotent, should not crash
+    }
+
+    #[test]
+    fn test_reset_peak_memory() {
+        reset_peak_memory();
+    }
+
+    #[test]
+    fn test_set_memory_limit_roundtrip() {
+        let original = get_memory_limit();
+        let prev = set_memory_limit(1024 * 1024 * 1024); // 1 GB
+        assert_eq!(prev, original);
+        set_memory_limit(original); // restore
+    }
+
+    #[test]
+    fn test_set_cache_limit_roundtrip() {
+        let original = get_memory_limit(); // cache limit defaults to memory limit
+        let prev = set_cache_limit(512 * 1024 * 1024); // 512 MB
+        // Restore (use max of prev and original to avoid going below default)
+        set_cache_limit(prev.max(original));
+    }
+}

From 2fb797e2a10a5f308c088e9d4779576e5301927d Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 11:28:17 -0700
Subject: [PATCH 05/17] build: update pmetal bindings to MLX 0.32

---
 mlx-rs/src/random.rs |  45 +++++++++++++------
 mlx-sys/build.rs     | 100 ++++++++++++++++++++++++++++++++++++-------
 2 files changed, 118 insertions(+), 27 deletions(-)

diff --git a/mlx-rs/src/random.rs b/mlx-rs/src/random.rs
index d909e9b4c..9ab43bbc5 100644
--- a/mlx-rs/src/random.rs
+++ b/mlx-rs/src/random.rs
@@ -6,14 +6,14 @@ use crate::utils::IntoOption;
 use crate::{error::Result, Array, ArrayElement, Stream};
 use mach_sys::mach_time;
 use mlx_internal_macros::{default_device, generate_macro};
-use parking_lot::Mutex;
 use std::borrow::Cow;
 use std::cell::RefCell;
-use std::sync::OnceLock;
-
-static GLOBAL_STATE: OnceLock<Mutex<RandomState>> = OnceLock::new();
 
 thread_local! {
+    // MLX 0.32 streams are thread-affine. Keeping the implicit random key in a
+    // process-global mutex lets a key array created on one Rust test/worker
+    // thread leak its stream into another thread's graph.
+    static GLOBAL_STATE: RefCell<RandomState> = RefCell::new(RandomState::new().unwrap());
     static TASK_LOCAL_STATE: RefCell<Option<RandomState>> = const { RefCell::new(None) };
 }
 
@@ -138,10 +138,6 @@ impl crate::utils::Updatable for RandomState {
     }
 }
 
-fn global_state() -> &'static Mutex<RandomState> {
-    GLOBAL_STATE.get_or_init(|| Mutex::new(RandomState::new().unwrap()))
-}
-
 /// Returns a key from the task-local state if it exists, otherwise
 /// returns `None`
 fn resolve_task_local_key() -> Option<Result<Array>> {
@@ -149,8 +145,7 @@ fn resolve_task_local_key() -> Option<Result<Array>> {
 }
 
 fn resolve_global_key() -> Result<Array> {
-    let mut state = global_state().lock();
-    state.next()
+    GLOBAL_STATE.with_borrow_mut(|state| state.next())
 }
 
 /// Use given key or generate a new one if `None`.
@@ -183,8 +178,7 @@ where
 
 /// Seed the random number generator.
 pub fn seed(seed: u64) -> Result<()> {
-    let mut state = global_state().lock();
-    state.seed(seed)
+    GLOBAL_STATE.with_borrow_mut(|state| state.seed(seed))
 }
 
 /// Get a PRNG key from a seed.
@@ -614,6 +608,33 @@ mod tests {
         assert_array_eq!(b, y, 0.01);
     }
 
+    #[test]
+    fn test_implicit_rng_state_does_not_cross_thread_streams() {
+        let first = std::thread::spawn(|| {
+            seed(3).unwrap();
+            uniform::<_, f32>(0, 1, None, None)
+                .unwrap()
+                .try_item::<f32>()
+                .unwrap()
+        })
+        .join()
+        .unwrap();
+
+        // MLX 0.32 rejects a lazy key array whose stream was created by the
+        // first thread when it is evaluated on this second thread.
+        let second = std::thread::spawn(|| {
+            uniform::<_, f32>(0, 1, None, None)
+                .unwrap()
+                .try_item::<f32>()
+                .unwrap()
+        })
+        .join()
+        .unwrap();
+
+        assert!(first.is_finite());
+        assert!(second.is_finite());
+    }
+
     #[test]
     fn test_key() {
         let k1 = key(0).unwrap();
diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index 3ba98dd5c..c13f1772e 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -64,9 +64,81 @@ fn resolve_deployment_target() -> String {
     format!("{}.{}", MLX_MIN_MACOS.0, MLX_MIN_MACOS.1)
 }
 
-/// Copy src/mlx-c to a staging directory and inject the metallib search-path
-/// patch into the CMakeLists.txt. This avoids modifying the mlx-c git submodule
-/// while ensuring the patch is applied when MLX is fetched via FetchContent.
+fn replace_once(contents: &mut String, from: &str, to: &str, context: &str) {
+    let matches = contents.matches(from).count();
+    assert_eq!(
+        matches, 1,
+        "expected exactly one {context} compatibility site, found {matches}"
+    );
+    *contents = contents.replacen(from, to, 1);
+}
+
+fn add_quant_global_scales(contents: &mut String, function: &str, arguments: &str) {
+    let start = contents
+        .find(function)
+        .unwrap_or_else(|| panic!("missing mlx-c function {function}"));
+    let end = contents[start + function.len()..]
+        .find("\nextern \"C\"")
+        .map(|offset| start + function.len() + offset)
+        .unwrap_or(contents.len());
+    let mut body = contents[start..end].to_owned();
+    replace_once(
+        &mut body,
+        "            std::string(mode),\n",
+        &format!("            std::string(mode),\n{arguments}"),
+        function,
+    );
+    contents.replace_range(start..end, &body);
+}
+
+/// Keep mlx-c 0.5's C ABI while supplying arguments added to MLX 0.32's C++
+/// API. This lets existing Rust callers move forward without an unrelated C
+/// API migration.
+fn patch_mlx_c_for_032(staged: &std::path::Path) {
+    let fft_path = staged.join("mlx/c/fft.cpp");
+    let mut fft = std::fs::read_to_string(&fft_path).expect("Failed to read mlx-c fft.cpp");
+    for op in ["fft", "ifft", "irfft", "rfft"] {
+        let from = format!("mlx::core::fft::{op}(mlx_array_get_(a), n, axis, mlx_stream_get_(s))");
+        let to = format!(
+            "mlx::core::fft::{op}(\n            mlx_array_get_(a),\n            n,\n            axis,\n            mlx::core::fft::FFTNorm::Backward,\n            mlx_stream_get_(s))"
+        );
+        replace_once(&mut fft, &from, &to, op);
+    }
+    for op in [
+        "fft2", "fftn", "ifft2", "ifftn", "irfft2", "irfftn", "rfft2", "rfftn",
+    ] {
+        let from = format!(
+            "mlx::core::fft::{op}(\n            mlx_array_get_(a),\n            mlx::core::Shape(n, n + n_num),\n            std::vector<int>(axes, axes + axes_num),\n            mlx_stream_get_(s))"
+        );
+        let to = format!(
+            "mlx::core::fft::{op}(\n            mlx_array_get_(a),\n            mlx::core::Shape(n, n + n_num),\n            std::vector<int>(axes, axes + axes_num),\n            mlx::core::fft::FFTNorm::Backward,\n            mlx_stream_get_(s))"
+        );
+        replace_once(&mut fft, &from, &to, op);
+    }
+    std::fs::write(&fft_path, fft).expect("Failed to patch mlx-c fft.cpp");
+
+    let ops_path = staged.join("mlx/c/ops.cpp");
+    let mut ops = std::fs::read_to_string(&ops_path).expect("Failed to read mlx-c ops.cpp");
+    add_quant_global_scales(
+        &mut ops,
+        "extern \"C\" int mlx_dequantize(",
+        "            std::nullopt,\n",
+    );
+    add_quant_global_scales(
+        &mut ops,
+        "extern \"C\" int mlx_qqmm(",
+        "            std::nullopt,\n            std::nullopt,\n",
+    );
+    add_quant_global_scales(
+        &mut ops,
+        "extern \"C\" int mlx_quantize(",
+        "            std::nullopt,\n",
+    );
+    std::fs::write(&ops_path, ops).expect("Failed to patch mlx-c ops.cpp");
+}
+
+/// Copy src/mlx-c to a staging directory, adapt its stable ABI to MLX 0.32,
+/// and inject the metallib search-path patch into the CMake project.
 fn prepare_mlx_c_source() -> PathBuf {
     let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
     let staged = out_dir.join("mlx-c-staged");
@@ -77,6 +149,7 @@ fn prepare_mlx_c_source() -> PathBuf {
         std::fs::remove_dir_all(&staged).expect("Failed to clean staged mlx-c");
     }
     copy_dir_recursive(&src, &staged).expect("Failed to copy mlx-c to staging");
+    patch_mlx_c_for_032(&staged);
 
     // Copy our patch file into the staged source
     let patches_dir = staged.join("patches");
@@ -93,7 +166,7 @@ fn prepare_mlx_c_source() -> PathBuf {
         std::fs::read_to_string(&cmake_path).expect("Failed to read CMakeLists.txt");
     let patched = cmake_content.replace(
         "GIT_TAG v0.30.6)",
-        "GIT_TAG v0.30.6\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true)",
+        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true)",
     );
     std::fs::write(&cmake_path, patched).expect("Failed to write patched CMakeLists.txt");
 
@@ -224,9 +297,10 @@ fn build_and_link_mlx_c() {
                     dest.metadata()
                         .and_then(|d| {
                             metallib.metadata().map(|s| {
-                                s.modified().ok().zip(d.modified().ok()).is_some_and(
-                                    |(src_t, dst_t)| src_t > dst_t,
-                                )
+                                s.modified()
+                                    .ok()
+                                    .zip(d.modified().ok())
+                                    .is_some_and(|(src_t, dst_t)| src_t > dst_t)
                             })
                         })
                         .unwrap_or(false)
@@ -236,14 +310,10 @@ fn build_and_link_mlx_c() {
                 if should_copy {
                     let _ = std::fs::create_dir_all(&cache_dir);
                     match std::fs::copy(&metallib, &dest) {
-                        Ok(_) => println!(
-                            "cargo:warning=Cached mlx.metallib to {}",
-                            dest.display()
-                        ),
-                        Err(e) => println!(
-                            "cargo:warning=Failed to cache mlx.metallib: {}",
-                            e
-                        ),
+                        Ok(_) => {
+                            println!("cargo:warning=Cached mlx.metallib to {}", dest.display())
+                        }
+                        Err(e) => println!("cargo:warning=Failed to cache mlx.metallib: {}", e),
                     }
                 }
             }

From 549f31b2f4628d71e13ba03746db30a4fe40f643 Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 11:57:44 -0700
Subject: [PATCH 06/17] perf(mlx): add expert-aligned NVFP4 gather

---
 mlx-sys/build.rs                            |   8 +-
 mlx-sys/patches/expert-aligned-gather.patch | 259 ++++++++++++++++++++
 2 files changed, 266 insertions(+), 1 deletion(-)
 create mode 100644 mlx-sys/patches/expert-aligned-gather.patch

diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index c13f1772e..c66e219a2 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -159,6 +159,11 @@ fn prepare_mlx_c_source() -> PathBuf {
         patches_dir.join("metallib-search-path.patch"),
     )
     .expect("Failed to copy metallib patch");
+    std::fs::copy(
+        "patches/expert-aligned-gather.patch",
+        patches_dir.join("expert-aligned-gather.patch"),
+    )
+    .expect("Failed to copy expert-aligned gather patch");
 
     // Inject PATCH_COMMAND into the FetchContent_Declare for MLX
     let cmake_path = staged.join("CMakeLists.txt");
@@ -166,12 +171,13 @@ fn prepare_mlx_c_source() -> PathBuf {
         std::fs::read_to_string(&cmake_path).expect("Failed to read CMakeLists.txt");
     let patched = cmake_content.replace(
         "GIT_TAG v0.30.6)",
-        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true)",
+        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true\n    COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch || true)",
     );
     std::fs::write(&cmake_path, patched).expect("Failed to write patched CMakeLists.txt");
 
     // Tell cargo to rerun if the patch changes
     println!("cargo:rerun-if-changed=patches/metallib-search-path.patch");
+    println!("cargo:rerun-if-changed=patches/expert-aligned-gather.patch");
 
     staged
 }
diff --git a/mlx-sys/patches/expert-aligned-gather.patch b/mlx-sys/patches/expert-aligned-gather.patch
new file mode 100644
index 000000000..9ab1bc2c1
--- /dev/null
+++ b/mlx-sys/patches/expert-aligned-gather.patch
@@ -0,0 +1,259 @@
+diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp
+index 9ed9ce4..f315643 100644
+--- a/mlx/backend/metal/jit_kernels.cpp
++++ b/mlx/backend/metal/jit_kernels.cpp
+@@ -1159,7 +1159,11 @@ MTL::ComputePipelineState* get_gather_qmm_nax_kernel(
+         is_affine ? metal::quantized_nax() : metal::fp_quantized_nax(),
+         get_template_definition(
+             lib_name,
+-            (is_affine ? "affine" : "fp") + std::string("_gather_qmm_rhs_nax"),
++            (is_affine ? "affine" : "fp") +
++                std::string(
++                    kernel_name.find("_expert_") != std::string::npos
++                        ? "_gather_qmm_rhs_expert_nax"
++                        : "_gather_qmm_rhs_nax"),
+             get_type_string(x.dtype()),
+             group_size,
+             bits,
+diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h
+index 0ea6af1..67c7892 100644
+--- a/mlx/backend/metal/kernels/fp_quantized_nax.h
++++ b/mlx/backend/metal/kernels/fp_quantized_nax.h
+@@ -1016,3 +1016,171 @@ template <
+     });
+   }
+ }
++
++METAL_FUNC int expert_sorted_lower_bound(
++    const device uint32_t* indices,
++    const int count,
++    const uint32_t value) {
++  int lo = 0;
++  int hi = count;
++  while (lo < hi) {
++    const int mid = lo + (hi - lo) / 2;
++    if (indices[mid] < value) {
++      lo = mid + 1;
++    } else {
++      hi = mid;
++    }
++  }
++  return lo;
++}
++
++// Expert-sorted Laguna prefill assigns four expert ids to each threadgroup.
++// This avoids reloading an expert's weight tile when its run crosses a fixed
++// row-tile boundary while preserving the stock NAX arithmetic and BF16 weight
++// staging boundary.
++template <
++    typename T,
++    int group_size,
++    const int bits,
++    int BM,
++    int BN,
++    int BK,
++    int WM,
++    int WN,
++    bool transpose,
++    typename Wtype = bfloat>
++[[kernel]] void fp_gather_qmm_rhs_expert_nax(
++    const device T* x,
++    const device uint32_t* w,
++    const device uint8_t* scales,
++    const device uint32_t* indices,
++    device T* y,
++    const constant int& M,
++    const constant int& N,
++    const constant int& K,
++    uint3 tid [[threadgroup_position_in_grid]],
++    uint lid [[thread_index_in_threadgroup]],
++    uint simd_group_id [[simdgroup_index_in_threadgroup]],
++    uint simd_lane_id [[thread_index_in_simdgroup]]) {
++  static_assert(transpose, "expert-aligned QMM requires NT weights");
++  static_assert(group_size == 16, "expert-aligned QMM requires gs16");
++  static_assert(bits == 4, "expert-aligned QMM requires NVFP4");
++
++  constexpr int pack_factor = get_pack_factor<8, bits>();
++  constexpr int bytes_per_pack = get_bytes_per_pack();
++  constexpr int BK_padded = BK + 16 / sizeof(Wtype);
++  constexpr int expert_groups = 64;
++  constexpr int experts = 256;
++
++  using loader_w_t = QuantizedBlockLoader<
++      Wtype,
++      BN,
++      BK,
++      BK_padded,
++      true,
++      WM * WN * SIMD_SIZE,
++      group_size,
++      bits>;
++
++  threadgroup Wtype Ws[BN * BK_padded];
++  threadgroup int bounds[2];
++
++  const int K_w = K * bytes_per_pack / pack_factor;
++  const int K_g = K / group_size;
++  const int K_it = K / BK;
++  const size_t stride_w = size_t(N) * K_w;
++  const size_t stride_s = size_t(N) * K_g;
++  const int y_col = tid.x * BN;
++
++  auto wl = (const device uint8_t*)w + size_t(y_col) * K_w;
++  const device uint8_t* scale_base = scales + size_t(y_col) * K_g;
++
++  constexpr short SM = BM / WM;
++  constexpr short SN = BN / WN;
++  constexpr short SK = 32;
++  constexpr short TM = SM / 16;
++  constexpr short TN = SN / 16;
++  constexpr short TK = SK / 16;
++
++  const short tm = SM * (simd_group_id / WN);
++  const short tn = SN * (simd_group_id % WN);
++
++  for (int expert_slot = 0; expert_slot < experts / expert_groups;
++       ++expert_slot) {
++    const uint32_t expert =
++        static_cast<uint32_t>(tid.y + expert_slot * expert_groups);
++
++    threadgroup_barrier(mem_flags::mem_threadgroup);
++    if (lid == 0) {
++      bounds[0] = expert_sorted_lower_bound(indices, M, expert);
++      bounds[1] = expert_sorted_lower_bound(indices, M, expert + 1);
++    }
++    threadgroup_barrier(mem_flags::mem_threadgroup);
++
++    const int run_start = bounds[0];
++    const int run_end = bounds[1];
++    for (int chunk_start = run_start; chunk_start < run_end;
++         chunk_start += BM) {
++      const short chunk_rows = short(min(BM, run_end - chunk_start));
++      const short sgp_sm = min(int(SM), max(0, int(chunk_rows) - int(tm)));
++      const bool sg_active = sgp_sm > 0;
++
++      NAXTile<float, TM, TN> Dtile;
++      Dtile.clear();
++
++      const device T* xn = x + size_t(chunk_start + tm) * K;
++      thread loader_w_t loader_w(
++          wl + size_t(expert) * stride_w,
++          scale_base + size_t(expert) * stride_s,
++          K,
++          Ws,
++          simd_group_id,
++          simd_lane_id);
++
++      for (int k = 0; k < K_it; ++k) {
++        threadgroup_barrier(mem_flags::mem_threadgroup);
++        loader_w.load_unsafe();
++        threadgroup_barrier(mem_flags::mem_threadgroup);
++
++        if (sg_active) {
++          STEEL_PRAGMA_NO_UNROLL
++          for (int kk1 = 0; kk1 < BK; kk1 += SK) {
++            NAXTile<T, TM, TK> Atile;
++            NAXTile<Wtype, TN, TK> Btile;
++            volatile int compiler_barrier;
++
++            if (sgp_sm == SM) {
++              Atile.load(xn + kk1, K);
++            } else {
++              Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm));
++            }
++            Btile.template load<Wtype, BK_padded, 1>(
++                Ws + tn * BK_padded + kk1);
++
++            tile_matmad_nax(
++                Dtile,
++                Atile,
++                metal::bool_constant<false>{},
++                Btile,
++                metal::bool_constant<true>{});
++            (void)compiler_barrier;
++          }
++        }
++
++        xn += BK;
++        loader_w.next();
++      }
++
++      threadgroup_barrier(mem_flags::mem_threadgroup);
++      if (sg_active) {
++        device T* yn = y + size_t(chunk_start + tm) * N + y_col + tn;
++        if (sgp_sm == SM) {
++          Dtile.store(yn, N);
++        } else {
++          Dtile.store_slice(
++              yn, N, short2(0, 0), short2(SN, sgp_sm));
++        }
++      }
++    }
++  }
++}
+diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp
+index 62d4871..6ba4a2e 100644
+--- a/mlx/backend/metal/quantized.cpp
++++ b/mlx/backend/metal/quantized.cpp
+@@ -1218,6 +1218,17 @@ void gather_qmm_rhs_nax(
+   int bm = 64, bn = 64, bk = 64;
+   int wm = 2, wn = 2;
+ 
++  const bool laguna_moe_shape =
++      (K == 2048 && N == 1024) || (K == 512 && N == 2048);
++  const bool expert_aligned =
++      env::get_var("MLX_EXPERT_ALIGNED_GATHER", 0) != 0 && mode != "affine" &&
++      transpose && group_size == 16 && bits == 4 && laguna_moe_shape &&
++      M >= 64;
++  if (expert_aligned) {
++    wm = 4;
++    wn = 1;
++  }
++
+   const bool align_M = (M % bm) == 0;
+   const bool align_N = (N % bn) == 0;
+   const bool align_K = (K % bk) == 0;
+@@ -1229,7 +1240,10 @@ void gather_qmm_rhs_nax(
+   concatenate(
+       kname,
+       mode +
+-          (transpose ? "_gather_qmm_rhs_nax_nt_" : "_gather_qmm_rhs_nax_nn_"),
++          (expert_aligned
++               ? "_gather_qmm_rhs_expert_nax_nt_"
++               : (transpose ? "_gather_qmm_rhs_nax_nt_"
++                            : "_gather_qmm_rhs_nax_nn_")),
+       type_string,
+       "_gs_",
+       group_size,
+@@ -1246,11 +1260,14 @@ void gather_qmm_rhs_nax(
+       "_wn_",
+       wn);
+ 
+-  metal::MTLFCList func_consts = {
+-      {&align_M, MTL::DataType::DataTypeBool, 200},
+-      {&align_N, MTL::DataType::DataTypeBool, 201},
+-      {&align_K, MTL::DataType::DataTypeBool, 202},
+-  };
++  metal::MTLFCList func_consts;
++  if (!expert_aligned) {
++    func_consts = {
++        {&align_M, MTL::DataType::DataTypeBool, 200},
++        {&align_N, MTL::DataType::DataTypeBool, 201},
++        {&align_K, MTL::DataType::DataTypeBool, 202},
++    };
++  }
+ 
+   // And the kernel hash that includes the function constants
+   std::string hash_name;
+@@ -1285,7 +1302,10 @@ void gather_qmm_rhs_nax(
+   compute_encoder.set_compute_pipeline_state(kernel);
+ 
+   MTL::Size group_dims(32, wn, wm);
+-  MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1);
++  MTL::Size grid_dims(
++      (N + bn - 1) / bn,
++      expert_aligned ? 64 : (M + bm - 1) / bm,
++      1);
+ 
+   int c = 0;
+   compute_encoder.set_input_array(x, c++);

From 1a786a1dd132bd9890910b4afda6124ee98002c3 Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 13:48:06 -0700
Subject: [PATCH 07/17] perf(mlx): fuse Laguna SwiGLU gather epilogue

---
 mlx-sys/patches/expert-aligned-gather.patch | 35 +++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/mlx-sys/patches/expert-aligned-gather.patch b/mlx-sys/patches/expert-aligned-gather.patch
index 9ab1bc2c1..0b527e541 100644
--- a/mlx-sys/patches/expert-aligned-gather.patch
+++ b/mlx-sys/patches/expert-aligned-gather.patch
@@ -19,7 +19,7 @@ diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/ke
 index 0ea6af1..67c7892 100644
 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h
 +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h
-@@ -1016,3 +1016,171 @@ template <
+@@ -1016,3 +1016,198 @@ template <
      });
    }
  }
@@ -90,6 +90,7 @@ index 0ea6af1..67c7892 100644
 +      bits>;
 +
 +  threadgroup Wtype Ws[BN * BK_padded];
++  threadgroup bfloat* gate_up_stage = (threadgroup bfloat*)Ws;
 +  threadgroup int bounds[2];
 +
 +  const int K_w = K * bytes_per_pack / pack_factor;
@@ -179,7 +180,37 @@ index 0ea6af1..67c7892 100644
 +      }
 +
 +      threadgroup_barrier(mem_flags::mem_threadgroup);
-+      if (sg_active) {
++      if (N == 1024) {
++        // Laguna's fused bank stores each matched [gate32, up32] pair in one
++        // BN=64 column tile. Preserve the stock gather's BF16 output boundary,
++        // then apply SwiGLU in threadgroup memory and pack the 512 activated
++        // columns into the contiguous prefix of the nominal 1024-wide output.
++        if (sg_active) {
++          Dtile.template store<bfloat, BN, 1>(
++              gate_up_stage + tm * BN + tn);
++        }
++        threadgroup_barrier(mem_flags::mem_threadgroup);
++        if (sg_active && (simd_group_id % WN) == 0) {
++          constexpr int activated_cols = BN / 2;
++          for (int linear = simd_lane_id;
++               linear < int(sgp_sm) * activated_cols;
++               linear += SIMD_SIZE) {
++            const int row = linear / activated_cols;
++            const int col = linear % activated_cols;
++            const bfloat gate = gate_up_stage[(tm + row) * BN + col];
++            const bfloat up =
++                gate_up_stage[(tm + row) * BN + activated_cols + col];
++            const bfloat exp_abs = metal::exp(metal::abs(gate));
++            const bfloat denominator = bfloat(1) + exp_abs;
++            const bfloat z = bfloat(1) / denominator;
++            const bfloat sigmoid = gate < bfloat(0) ? z : bfloat(1) - z;
++            const bfloat silu = bfloat(gate * sigmoid);
++            y[size_t(chunk_start + tm + row) * (N / 2) +
++              size_t(tid.x) * activated_cols + col] = bfloat(silu * up);
++          }
++        }
++        threadgroup_barrier(mem_flags::mem_threadgroup);
++      } else if (sg_active) {
 +        device T* yn = y + size_t(chunk_start + tm) * N + y_col + tn;
 +        if (sgp_sm == SM) {
 +          Dtile.store(yn, N);

From 5024cd00b4f62369dbaa7c15c9ff64d81f4d01fd Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 14:12:02 -0700
Subject: [PATCH 08/17] fix(mlx): key fused gather pipeline separately

---
 mlx-sys/patches/expert-aligned-gather.patch | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlx-sys/patches/expert-aligned-gather.patch b/mlx-sys/patches/expert-aligned-gather.patch
index 0b527e541..6e84ae9f7 100644
--- a/mlx-sys/patches/expert-aligned-gather.patch
+++ b/mlx-sys/patches/expert-aligned-gather.patch
@@ -244,13 +244,15 @@ index 62d4871..6ba4a2e 100644
    const bool align_M = (M % bm) == 0;
    const bool align_N = (N % bn) == 0;
    const bool align_K = (K % bk) == 0;
-@@ -1229,7 +1240,10 @@ void gather_qmm_rhs_nax(
+@@ -1229,7 +1240,12 @@ void gather_qmm_rhs_nax(
    concatenate(
        kname,
        mode +
 -          (transpose ? "_gather_qmm_rhs_nax_nt_" : "_gather_qmm_rhs_nax_nn_"),
 +          (expert_aligned
-+               ? "_gather_qmm_rhs_expert_nax_nt_"
++               ? (N == 1024
++                      ? "_gather_qmm_rhs_expert_swiglu_nax_nt_"
++                      : "_gather_qmm_rhs_expert_nax_nt_")
 +               : (transpose ? "_gather_qmm_rhs_nax_nt_"
 +                            : "_gather_qmm_rhs_nax_nn_")),
        type_string,

From c61e26bc99852fc67cc9d2a420216582aac31ad6 Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 14:45:26 -0700
Subject: [PATCH 09/17] test(mlx): trace expert gather dispatch

---
 mlx-sys/patches/expert-aligned-gather.patch | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/mlx-sys/patches/expert-aligned-gather.patch b/mlx-sys/patches/expert-aligned-gather.patch
index 6e84ae9f7..0992a7967 100644
--- a/mlx-sys/patches/expert-aligned-gather.patch
+++ b/mlx-sys/patches/expert-aligned-gather.patch
@@ -226,7 +226,7 @@ diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp
 index 62d4871..6ba4a2e 100644
 --- a/mlx/backend/metal/quantized.cpp
 +++ b/mlx/backend/metal/quantized.cpp
-@@ -1218,6 +1218,17 @@ void gather_qmm_rhs_nax(
+@@ -1218,6 +1218,28 @@ void gather_qmm_rhs_nax(
    int bm = 64, bn = 64, bk = 64;
    int wm = 2, wn = 2;
  
@@ -236,6 +236,17 @@ index 62d4871..6ba4a2e 100644
 +      env::get_var("MLX_EXPERT_ALIGNED_GATHER", 0) != 0 && mode != "affine" &&
 +      transpose && group_size == 16 && bits == 4 && laguna_moe_shape &&
 +      M >= 64;
++  if (env::get_var("MLX_EXPERT_ALIGNED_TRACE", 0) != 0) {
++    fprintf(
++        stderr,
++        "mlx: expert gather dispatch=%d M=%d N=%d K=%d mode=%s sorted_rows=%zu\n",
++        int(expert_aligned),
++        M,
++        N,
++        K,
++        mode.c_str(),
++        indices.size());
++  }
 +  if (expert_aligned) {
 +    wm = 4;
 +    wn = 1;

From e33f7c2dcd36cce964f3518f3d5100e51b77e84f Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 14:52:16 -0700
Subject: [PATCH 10/17] test(mlx): trace gather eligibility

---
 mlx-sys/patches/expert-aligned-gather.patch | 26 ++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/mlx-sys/patches/expert-aligned-gather.patch b/mlx-sys/patches/expert-aligned-gather.patch
index 0992a7967..0330ed861 100644
--- a/mlx-sys/patches/expert-aligned-gather.patch
+++ b/mlx-sys/patches/expert-aligned-gather.patch
@@ -255,7 +255,7 @@ index 62d4871..6ba4a2e 100644
    const bool align_M = (M % bm) == 0;
    const bool align_N = (N % bn) == 0;
    const bool align_K = (K % bk) == 0;
-@@ -1229,7 +1240,12 @@ void gather_qmm_rhs_nax(
+@@ -1229,7 +1251,12 @@ void gather_qmm_rhs_nax(
    concatenate(
        kname,
        mode +
@@ -269,7 +269,7 @@ index 62d4871..6ba4a2e 100644
        type_string,
        "_gs_",
        group_size,
-@@ -1246,11 +1260,14 @@ void gather_qmm_rhs_nax(
+@@ -1246,11 +1273,14 @@ void gather_qmm_rhs_nax(
        "_wn_",
        wn);
  
@@ -289,7 +289,7 @@ index 62d4871..6ba4a2e 100644
  
    // And the kernel hash that includes the function constants
    std::string hash_name;
-@@ -1285,7 +1302,10 @@ void gather_qmm_rhs_nax(
+@@ -1285,7 +1315,10 @@ void gather_qmm_rhs_nax(
    compute_encoder.set_compute_pipeline_state(kernel);
  
    MTL::Size group_dims(32, wn, wm);
@@ -301,3 +301,23 @@ index 62d4871..6ba4a2e 100644
  
    int c = 0;
    compute_encoder.set_input_array(x, c++);
+@@ -1576,6 +1609,19 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+   int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4;
+   auto mode = quantization_mode_to_string(mode_);
+ 
++  if (env::get_var("MLX_EXPERT_ALIGNED_TRACE", 0) != 0) {
++    fprintf(
++        stderr,
++        "mlx: gather eval M=%d B=%d E=%d N=%d K=%d sorted=%d rhs=%zu\n",
++        M,
++        B,
++        E,
++        N,
++        K,
++        int(right_sorted_),
++        rhs_indices.size());
++  }
++
+   // We are walking x in order and w is also in order so we can batch up the
+   // matmuls and reuse reading x and w.
+   //

From ba360534a6df93785217a55f748ce2b6f6ef0207 Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 15:15:34 -0700
Subject: [PATCH 11/17] build(mlx): add opt-in NAX JIT mode

---
 mlx-sys/build.rs | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index c66e219a2..9ac6bc795 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -211,6 +211,8 @@ fn copy_dir_recursive(src: &std::path::Path, dst: &std::path::Path) -> std::io::
 }
 
 fn build_and_link_mlx_c() {
+    println!("cargo:rerun-if-env-changed=MLX_ENABLE_NAX");
+    println!("cargo:rerun-if-env-changed=MACOSX_DEPLOYMENT_TARGET");
     // MLX requires macOS >= 14.0 for Metal support. Override the deployment
     // target early so the cmake crate (and cc crate) don't inject a lower
     // -mmacosx-version-min flag into CFLAGS/CXXFLAGS. Without this, Cargo's
@@ -226,6 +228,14 @@ fn build_and_link_mlx_c() {
     config.very_verbose(true);
     config.define("CMAKE_INSTALL_PREFIX", ".");
 
+    // MLX's NAX kernels require the macOS 26.2 SDK/deployment contract. Its
+    // expert kernels are runtime-specialized, so use MLX's supported all-JIT
+    // Metal build when callers explicitly opt in instead of compiling the
+    // dynamic expert template into the ahead-of-time metallib.
+    if env::var("MLX_ENABLE_NAX").as_deref() == Ok("1") {
+        config.define("MLX_METAL_JIT", "ON");
+    }
+
     #[cfg(target_os = "macos")]
     {
         let target = resolve_deployment_target();

From 19de8c00eac6f4deea976a3577ba564eb7ffedcc Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 15:52:14 -0700
Subject: [PATCH 12/17] test(mlx): trace NAX runtime eligibility

---
 mlx-sys/patches/expert-aligned-gather.patch | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mlx-sys/patches/expert-aligned-gather.patch b/mlx-sys/patches/expert-aligned-gather.patch
index 0330ed861..f13187436 100644
--- a/mlx-sys/patches/expert-aligned-gather.patch
+++ b/mlx-sys/patches/expert-aligned-gather.patch
@@ -308,14 +308,15 @@ index 62d4871..6ba4a2e 100644
 +  if (env::get_var("MLX_EXPERT_ALIGNED_TRACE", 0) != 0) {
 +    fprintf(
 +        stderr,
-+        "mlx: gather eval M=%d B=%d E=%d N=%d K=%d sorted=%d rhs=%zu\n",
++        "mlx: gather eval M=%d B=%d E=%d N=%d K=%d sorted=%d rhs=%zu nax=%d\n",
 +        M,
 +        B,
 +        E,
 +        N,
 +        K,
 +        int(right_sorted_),
-+        rhs_indices.size());
++        rhs_indices.size(),
++        int(metal::is_nax_available()));
 +  }
 +
    // We are walking x in order and w is also in order so we can batch up the

From f78940318bdbffddce73a0be3bc59e3e8e099096 Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 16:17:30 -0700
Subject: [PATCH 13/17] build(mlx): fail when Laguna patch is invalid

---
 mlx-sys/build.rs                            | 2 +-
 mlx-sys/patches/expert-aligned-gather.patch | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index 9ac6bc795..ef26dca35 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -171,7 +171,7 @@ fn prepare_mlx_c_source() -> PathBuf {
         std::fs::read_to_string(&cmake_path).expect("Failed to read CMakeLists.txt");
     let patched = cmake_content.replace(
         "GIT_TAG v0.30.6)",
-        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true\n    COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch || true)",
+        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch\n    COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch)",
     );
     std::fs::write(&cmake_path, patched).expect("Failed to write patched CMakeLists.txt");
 
diff --git a/mlx-sys/patches/expert-aligned-gather.patch b/mlx-sys/patches/expert-aligned-gather.patch
index f13187436..4a6eab7e7 100644
--- a/mlx-sys/patches/expert-aligned-gather.patch
+++ b/mlx-sys/patches/expert-aligned-gather.patch
@@ -301,7 +301,7 @@ index 62d4871..6ba4a2e 100644
  
    int c = 0;
    compute_encoder.set_input_array(x, c++);
-@@ -1576,6 +1609,19 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+@@ -1576,6 +1609,20 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
    int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4;
    auto mode = quantization_mode_to_string(mode_);
  

From 2021a0904125d634039315588d4c55bb52c5919a Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 16:18:18 -0700
Subject: [PATCH 14/17] build(mlx): keep metallib relocation optional in JIT
 mode

---
 mlx-sys/build.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index ef26dca35..c677c3edc 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -171,7 +171,7 @@ fn prepare_mlx_c_source() -> PathBuf {
         std::fs::read_to_string(&cmake_path).expect("Failed to read CMakeLists.txt");
     let patched = cmake_content.replace(
         "GIT_TAG v0.30.6)",
-        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch\n    COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch)",
+        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true\n    COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch)",
     );
     std::fs::write(&cmake_path, patched).expect("Failed to write patched CMakeLists.txt");
 

From e5ca3b0ccf3dc66656fe726917b973441daf3c6e Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 16:49:34 -0700
Subject: [PATCH 15/17] test(mlx): trace detected Metal architecture

---
 mlx-sys/patches/expert-aligned-gather.patch | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/mlx-sys/patches/expert-aligned-gather.patch b/mlx-sys/patches/expert-aligned-gather.patch
index 4a6eab7e7..7c76af123 100644
--- a/mlx-sys/patches/expert-aligned-gather.patch
+++ b/mlx-sys/patches/expert-aligned-gather.patch
@@ -301,14 +301,15 @@ index 62d4871..6ba4a2e 100644
  
    int c = 0;
    compute_encoder.set_input_array(x, c++);
-@@ -1576,6 +1609,20 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+@@ -1576,6 +1609,23 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
    int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4;
    auto mode = quantization_mode_to_string(mode_);
  
 +  if (env::get_var("MLX_EXPERT_ALIGNED_TRACE", 0) != 0) {
 +    fprintf(
 +        stderr,
-+        "mlx: gather eval M=%d B=%d E=%d N=%d K=%d sorted=%d rhs=%zu nax=%d\n",
++        "mlx: gather eval M=%d B=%d E=%d N=%d K=%d sorted=%d rhs=%zu "
++        "nax=%d arch=%s gen=%d\n",
 +        M,
 +        B,
 +        E,
@@ -316,7 +317,9 @@ index 62d4871..6ba4a2e 100644
 +        K,
 +        int(right_sorted_),
 +        rhs_indices.size(),
-+        int(metal::is_nax_available()));
++        int(metal::is_nax_available()),
++        d.get_architecture().c_str(),
++        d.get_architecture_gen());
 +  }
 +
    // We are walking x in order and w is also in order so we can batch up the

From 6d4174bbdaf51aefe8f0d89614e014beaa627a15 Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 17:08:18 -0700
Subject: [PATCH 16/17] build(mlx): make Laguna patch application idempotent

---
 mlx-sys/build.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index c677c3edc..e32ad30a4 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -171,7 +171,7 @@ fn prepare_mlx_c_source() -> PathBuf {
         std::fs::read_to_string(&cmake_path).expect("Failed to read CMakeLists.txt");
     let patched = cmake_content.replace(
         "GIT_TAG v0.30.6)",
-        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true\n    COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch)",
+        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true\n    COMMAND sh -c \"git apply --check '$1' && git apply '$1' || git apply --reverse --check '$1'\" sh ${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch)",
     );
     std::fs::write(&cmake_path, patched).expect("Failed to write patched CMakeLists.txt");
 

From 20d791520b65f86d14decf84c157d69476379aeb Mon Sep 17 00:00:00 2001
From: Illia Polosukhin <ilblackdragon@gmail.com>
Date: Tue, 4 Aug 2026 17:09:08 -0700
Subject: [PATCH 17/17] fix(mlx): pass Laguna patch path to CMake shell

---
 mlx-sys/build.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index e32ad30a4..4a2fede91 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -171,7 +171,7 @@ fn prepare_mlx_c_source() -> PathBuf {
         std::fs::read_to_string(&cmake_path).expect("Failed to read CMakeLists.txt");
     let patched = cmake_content.replace(
         "GIT_TAG v0.30.6)",
-        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true\n    COMMAND sh -c \"git apply --check '$1' && git apply '$1' || git apply --reverse --check '$1'\" sh ${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch)",
+        "GIT_TAG v0.32.0\n    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/metallib-search-path.patch || true\n    COMMAND sh -c \"git apply --check '${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch' && git apply '${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch' || git apply --reverse --check '${CMAKE_CURRENT_SOURCE_DIR}/patches/expert-aligned-gather.patch'\")",
     );
     std::fs::write(&cmake_path, patched).expect("Failed to write patched CMakeLists.txt");