build: update pmetal bindings to MLX 0.32

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.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", path = "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" }
 

mlx-rs/Cargo.toml

@@ -1,14 +1,14 @@
 [package]
-name = "mlx-rs"
-version.workspace = true
+name = "pmetal-mlx-rs"
+version = "0.25.8"
 authors.workspace = true
 edition.workspace = true
-repository.workspace = true
+repository = "https://github.com/nicholasjpaterno/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]

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;

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));
+    }
+}

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,19 +138,14 @@ 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>> {
     TASK_LOCAL_STATE.with_borrow_mut(|state| state.as_mut().map(|s| s.next()))
 }
 
 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();

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.4" # 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