chore(pkg): bump to mlx 0.32.1 + nanobind 2.13.0, standalone arena residency set

.github/workflows/ci.yml

@@ -62,7 +62,7 @@ jobs:
       - name: Build extension + test deps
         run: |
           python -m pip install --upgrade pip
-          pip install "setuptools>=77" wheel "mlx==0.31.2" "nanobind==2.12.0" "cmake>=3.27" gguf numpy pytest "mlx-lm>=0.27" "transformers<5.13"
+          pip install "setuptools>=77" wheel "mlx==0.32.1" "nanobind==2.13.0" "cmake>=3.27" gguf numpy pytest "mlx-lm>=0.27" "transformers<5.13"
           pip install -e . --no-build-isolation
 
       - name: Codec lint (incl. runtime kq.codecs() leg)
@@ -155,7 +155,7 @@ jobs:
       - name: Build extension (CPU-only) + test deps
         run: |
           python -m pip install --upgrade pip
-          pip install "setuptools>=77" wheel "mlx[cpu]==0.31.2" "nanobind==2.12.0" "cmake>=3.27" gguf numpy pytest "mlx-lm>=0.27" "transformers<5.13"
+          pip install "setuptools>=77" wheel "mlx[cpu]==0.32.1" "nanobind==2.13.0" "cmake>=3.27" gguf numpy pytest "mlx-lm>=0.27" "transformers<5.13"
           pip install -e . --no-build-isolation
 
       # A non-Metal build must report no metallib (and not throw doing so).

CMakeLists.txt

@@ -3,7 +3,7 @@ cmake_minimum_required(VERSION 3.27)
 project(mlx_kquant LANGUAGES CXX C)
 
 # ----------------------------- Setup -----------------------------
-# C++20: MLX 0.31.2's public headers (device.h, stream.h) use defaulted
+# C++20: MLX 0.32.1's public headers (device.h, stream.h) use defaulted
 # comparison operators, a C++20 feature. Apple Clang accepts them as an
 # extension under -std=c++17, but GCC (the Linux toolchain) rejects them, so a
 # C++17 standard breaks the Metal-free Linux build. 20 builds cleanly on both.
@@ -37,7 +37,7 @@ find_package(
 # "incompatible function arguments". Build isolation installs the pin;
 # non-isolated dev builds (build_ext --inplace, -e --no-build-isolation) use
 # whatever the venv has - guard that path here.
-set(KQ_NANOBIND_PIN "2.12.0")
+set(KQ_NANOBIND_PIN "2.13.0")
 execute_process(
   COMMAND "${Python_EXECUTABLE}" -c "import nanobind; print(nanobind.__version__)"
   OUTPUT_STRIP_TRAILING_WHITESPACE
@@ -248,7 +248,7 @@ if(BUILD_SHARED_LIBS)
     # installed mlx wheel at runtime. _ext.so lives in site-packages/mlx_kquant/,
     # libmlx in site-packages/mlx/lib/, so ../mlx/lib reaches it. (The absolute
     # build-tree rpath MLXConfig adds works only on the build machine; this is
-    # what makes a redistributed wheel find the ABI-pinned mlx==0.31.2. delocate
+    # what makes a redistributed wheel find the ABI-pinned mlx==0.32.1. delocate
     # must therefore *exclude* libmlx from vendoring - see [tool.cibuildwheel].)
     target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path)
     target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path/../mlx/lib)

CONTRIBUTING.md

@@ -12,7 +12,7 @@ front-end. A few notes to get productive quickly.
 
   ```sh
   python -m venv .venv && source .venv/bin/activate
-  pip install "mlx==0.31.2" "nanobind==2.12.0" "cmake>=3.27"
+  pip install "mlx==0.32.1" "nanobind==2.13.0" "cmake>=3.27"
   pip install -e ".[dev,tools]"   # builds _ext + mlx_kquant.metallib
   ```
 

README.md

@@ -48,14 +48,14 @@ pip install mlx-kquant            # the K-quant ops + precompiled metallib
 pip install "mlx-kquant[tools]"   # the CLI (quantize / run / chat / lora / fuse)
 ```
 
-Either pulls the ABI-matched `mlx==0.31.2` automatically.
+Either pulls the ABI-matched `mlx==0.32.1` automatically.
 
 **From source** (or to develop) needs the Metal toolchain (`xcrun metal`); the metallib compiles at
 install time, no runtime JIT:
 
 ```sh
 git clone https://github.com/asher/mlx-kquant && cd mlx-kquant
-pip install "mlx==0.31.2"         # pinned, ABI-matched stock wheel (pulls the Metal backend)
+pip install "mlx==0.32.1"         # pinned, ABI-matched stock wheel (pulls the Metal backend)
 pip install -e ".[tools]"         # builds _ext + mlx_kquant.metallib; adds mlx-lm for the CLI
 ```
 
@@ -66,7 +66,7 @@ x86_64 stays on the scalar/threaded path.) The base `mlx` wheel ships no backend
 install the CPU one explicitly:
 
 ```sh
-pip install "mlx[cpu]==0.31.2"   # base frontend + libmlx CPU backend
+pip install "mlx[cpu]==0.32.1"   # base frontend + libmlx CPU backend
 pip install -e . --no-build-isolation
 ```
 
@@ -334,7 +334,7 @@ are informed by our analysis of the mixed-precision quants that [Unsloth][unslot
 
 ## Version pinning
 
-Pinned to `mlx==0.31.2`. The kernels include MLX's steel headers and the extension links `libmlx`,
+Pinned to `mlx==0.32.1`. The kernels include MLX's steel headers and the extension links `libmlx`,
 binding it to that release's ABI and header API. To move to a newer MLX: update the bundled headers
 under `metal/mlx/backend/metal/kernels/` for that wheel, rebuild, and re-run the test suite.
 
@@ -348,10 +348,10 @@ python -m pytest tests/
 
 - **macOS 26.2 (Tahoe) or later on Apple Silicon** (M-series). Building from source needs the Metal
   toolchain (`xcrun metal`).
-- **Linux** (x86_64 or aarch64) is supported CPU-only. Build against `mlx[cpu]==0.31.2`, no Metal
+- **Linux** (x86_64 or aarch64) is supported CPU-only. Build against `mlx[cpu]==0.32.1`, no Metal
   toolchain required. See [Install](#install) and [Limitations](#limitations).
-- **Python >= 3.10** (the pinned `mlx==0.31.2` ships no cp39 wheel).
-- **`mlx==0.31.2`** exactly - the kernels include MLX's steel headers and the extension links
+- **Python >= 3.10** (the pinned `mlx==0.32.1` ships no cp39 wheel).
+- **`mlx==0.32.1`** exactly - the kernels include MLX's steel headers and the extension links
   `libmlx`, so the ABI is version-locked (see [Version pinning](#version-pinning)).
 
 ## Limitations

docs/integration.md

@@ -10,7 +10,7 @@ Everything here runs against a **stock `mlx` wheel** - importing `mlx_kquant` ad
 namespace; it does **not** register `mode="kquant"` onto `mx.*`, so integrations call `kq.*`
 explicitly.
 
-**Compatibility.** The ops require stock `mlx==0.31.2` (the C++ extension is ABI-pinned to it). The
+**Compatibility.** The ops require stock `mlx==0.32.1` (the C++ extension is ABI-pinned to it). The
 `[tools]` integration points below (loader, recipes, the mlx-lm patch) additionally need
 `mlx-lm>=0.27`. The on-disk checkpoint format is stable and mirrors what a future in-core kquant mode
 would read.

pyproject.toml

@@ -2,10 +2,10 @@
 requires = [
   "setuptools>=77",
   "cmake>=3.27",
-  "mlx==0.31.2",
+  "mlx==0.32.1",
   # Must match mlx's nanobind ABI exactly; CMakeLists.txt (KQ_NANOBIND_PIN)
   # enforces it at configure time for non-isolated builds - keep in sync.
-  "nanobind==2.12.0",
+  "nanobind==2.13.0",
 ]
 build-backend = "setuptools.build_meta"
 
@@ -41,7 +41,7 @@ classifiers = [
   "Environment :: GPU",
   "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
-dependencies = ["mlx==0.31.2"]
+dependencies = ["mlx==0.32.1"]
 
 # Model-level helpers (the create-weights toolchain, loader, and CLI) need the
 # mlx-lm model zoo and friends; the raw ``kq.*`` ops do not. Keep that weight
@@ -96,11 +96,11 @@ ignore = ["E741"]  # ambiguous var names (l/I) appear in math-y kernel glue
 # ---------------------------------------------------------------------------
 # Wheel building (cibuildwheel). macOS arm64 only - the extension is a Metal
 # module. Build isolation pulls [build-system].requires (incl. the ABI-pinned
-# mlx==0.31.2 + nanobind + the Metal toolchain from the runner's Xcode).
+# mlx==0.32.1 + nanobind + the Metal toolchain from the runner's Xcode).
 #
 # NOTE: this config is a skeleton - it has NOT yet produced a wheel validated on
 # a second machine. Before the first PyPI publish, build locally, then
-# `pip install` the wheel into a CLEAN venv with only mlx==0.31.2 and confirm
+# `pip install` the wheel into a CLEAN venv with only mlx==0.32.1 and confirm
 # `python -c "import mlx_kquant as kq; kq.metallib_loads()"` on a GPU box.
 # ---------------------------------------------------------------------------
 [tool.cibuildwheel]
@@ -120,7 +120,7 @@ archs = ["arm64"]
 # https://github.com/ml-explore/mlx/issues/3586
 environment = { MACOSX_DEPLOYMENT_TARGET = "26.2" }
 # Do NOT vendor libmlx into the wheel: the extension is ABI-pinned to the user's
-# installed mlx==0.31.2 and resolves it via the @loader_path/../mlx/lib rpath
+# installed mlx==0.32.1 and resolves it via the @loader_path/../mlx/lib rpath
 # (see CMakeLists.txt). Bundling a second libmlx copy would split MLX global
 # state and resurrect the gguflib<->libmlx symbol collision. --exclude keeps it
 # out of the wheel; --ignore-missing-dependencies stops delocate from erroring

src/kquant_arena.cpp

@@ -67,13 +67,32 @@ std::pair<mx::array, uintptr_t> arena_alloc(
 
 #ifdef _METAL_
 
-// MLX's ResidencySet wrapper methods are not exported from libmlx; its
-// inline accessor hands back the raw MTL::ResidencySet (created at device
-// init and attached to MLX's command queue), and metal-cpp is header-only,
-// so the additions run entirely in this TU.
+// mlx 0.31 exposed its queue-attached MTL::ResidencySet and these helpers
+// piggybacked on it; 0.32 split residency into size-capped sets budgeted
+// by set_wired_limit, with no raw handle out and the command queues
+// private. The arena keeps its own standalone set instead: created once
+// against MLX's MTL device with a standing requestResidency, mutations
+// applied at residency_commit(). The arena's host pages stay pinned by
+// mlock either way; this set keeps their GPU mappings resident without
+// consuming MLX's wired-limit budget, which the arena must not depend on
+// (streaming-mode servers run with that budget at zero). metal-cpp is
+// header-only, so the additions run entirely in this TU.
 static MTL::ResidencySet* kq_residency_set() {
-  auto& d = mx::metal::device(mx::Device(mx::Device::gpu));
-  return const_cast<MTL::ResidencySet*>(d.residency_set().mtl_residency_set());
+  static MTL::ResidencySet* rs = []() -> MTL::ResidencySet* {
+    if (__builtin_available(macOS 15, *)) {
+      auto& d = mx::metal::device(mx::Device(mx::Device::gpu));
+      auto* desc = MTL::ResidencySetDescriptor::alloc()->init();
+      NS::Error* error = nullptr;
+      auto* set = d.mtl_device()->newResidencySet(desc, &error);
+      desc->release();
+      if (set != nullptr) {
+        set->requestResidency();
+      }
+      return set;
+    }
+    return nullptr;
+  }();
+  return rs;
 }
 
 bool residency_insert(const mx::array& a) {