diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e70f146..24f691e 100644 --- a/.github/workflows/ci.yml +++ b/.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). diff --git a/CMakeLists.txt b/CMakeLists.txt index 40ec913..d8e87ef 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index af8ddcf..b89cb23 100644 --- a/CONTRIBUTING.md +++ b/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 ``` diff --git a/README.md b/README.md index 8832f59..567386c 100644 --- a/README.md +++ b/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 diff --git a/docs/integration.md b/docs/integration.md index 578efa0..4fc66ee 100644 --- a/docs/integration.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index 6cc8c6c..bd35153 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/kquant_arena.cpp b/src/kquant_arena.cpp index 65f13e9..a4f6885 100644 --- a/src/kquant_arena.cpp +++ b/src/kquant_arena.cpp @@ -67,13 +67,32 @@ std::pair 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(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) {