diff --git a/Cargo.toml b/Cargo.toml index 100712f0..0481d3ab 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.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" } diff --git a/mlx-rs/Cargo.toml b/mlx-rs/Cargo.toml index c2c4efd0..caa0213c 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.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] diff --git a/mlx-rs/src/lib.rs b/mlx-rs/src/lib.rs index 1cc36f97..53459f34 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 00000000..f07b1f96 --- /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)); + } +} diff --git a/mlx-rs/src/random.rs b/mlx-rs/src/random.rs index d909e9b4..9ab43bbc 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> = 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 = RefCell::new(RandomState::new().unwrap()); static TASK_LOCAL_STATE: RefCell> = const { RefCell::new(None) }; } @@ -138,10 +138,6 @@ impl crate::utils::Updatable for RandomState { } } -fn global_state() -> &'static Mutex { - 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> { @@ -149,8 +145,7 @@ fn resolve_task_local_key() -> Option> { } fn resolve_global_key() -> Result { - 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::() + .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::() + .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/Cargo.toml b/mlx-sys/Cargo.toml index 55d5325f..f490be28 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.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 diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs index 11f6300f..4a2fede9 100644 --- a/mlx-sys/build.rs +++ b/mlx-sys/build.rs @@ -44,11 +44,204 @@ fn find_clang_rt_path() -> Option { None } +/// Resolve the macOS deployment target. +/// +/// 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 { + const MLX_MIN_MACOS: (u32, u32) = (14, 0); + + if let Ok(val) = env::var("MACOSX_DEPLOYMENT_TARGET") { + let parts: Vec = 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) +} + +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(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(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"); + 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"); + patch_mlx_c_for_032(&staged); + + // 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"); + 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"); + 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.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"); + + // 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 +} + +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"); + 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 + // 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", "."); + // 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(); + 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++"); @@ -104,6 +297,44 @@ 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/expert-aligned-gather.patch b/mlx-sys/patches/expert-aligned-gather.patch new file mode 100644 index 00000000..7c76af12 --- /dev/null +++ b/mlx-sys/patches/expert-aligned-gather.patch @@ -0,0 +1,327 @@ +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,198 @@ 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 bfloat* gate_up_stage = (threadgroup bfloat*)Ws; ++ 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(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 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 Atile; ++ NAXTile 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( ++ Ws + tn * BK_padded + kk1); ++ ++ tile_matmad_nax( ++ Dtile, ++ Atile, ++ metal::bool_constant{}, ++ Btile, ++ metal::bool_constant{}); ++ (void)compiler_barrier; ++ } ++ } ++ ++ xn += BK; ++ loader_w.next(); ++ } ++ ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ 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( ++ 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); ++ } 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,28 @@ 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 (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; ++ } ++ + const bool align_M = (M % bm) == 0; + const bool align_N = (N % bn) == 0; + const bool align_K = (K % bk) == 0; +@@ -1229,7 +1251,12 @@ void gather_qmm_rhs_nax( + concatenate( + kname, + mode + +- (transpose ? "_gather_qmm_rhs_nax_nt_" : "_gather_qmm_rhs_nax_nn_"), ++ (expert_aligned ++ ? (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, + "_gs_", + group_size, +@@ -1246,11 +1273,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 +1315,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++); +@@ -1576,6 +1609,23 @@ void GatherQMM::eval_gpu(const std::vector& 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 arch=%s gen=%d\n", ++ M, ++ B, ++ E, ++ N, ++ K, ++ int(right_sorted_), ++ rhs_indices.size(), ++ 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 + // matmuls and reuse reading x and w. + // diff --git a/mlx-sys/patches/metallib-search-path.patch b/mlx-sys/patches/metallib-search-path.patch new file mode 100644 index 00000000..54672e0a --- /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 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 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 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() << " "; + }