Add nanobind 3.0 split-mode support

.github/actions/test-wheel/action.yml

@@ -10,15 +10,24 @@ inputs:
 runs:
   using: 'composite'
   steps:
-    - name: Get Python version
+    - name: Get frontend ABI
       id: python
       shell: bash
-      run: python -c "import sys; print(f'version={sys.version_info.major}.{sys.version_info.minor}')" >> $GITHUB_OUTPUT
+      run: |
+        python - <<'PY' >> "$GITHUB_OUTPUT"
+        import sys
+        import sysconfig
+
+        if sysconfig.get_config_var("Py_GIL_DISABLED"):
+            print(f"abi=py{sys.version_info.major}.{sys.version_info.minor}t")
+        else:
+            print("abi=abi3")
+        PY
 
     - name: Download frontend packages
       uses: actions/download-artifact@v8
       with:
-        pattern: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ steps.python.outputs.version }}
+        pattern: frontend-${{ runner.os }}-${{ runner.arch }}-${{ steps.python.outputs.abi }}
         path: wheelhouse
 
     - name: Download backend packages

.github/workflows/release.yml

@@ -58,12 +58,18 @@ jobs:
       matrix:
         os: ['Linux', 'Windows']
         arch: ['x86_64', 'aarch64']
-        python-version: &pyver ['3.10', '3.11', '3.12', '3.13', '3.13t', '3.14', '3.14t']
-        # There is no cp310 binary for Windows on arm.
+        # One Stable ABI frontend wheel covers all supported regular CPython
+        # versions on a platform. Python 3.13t and 3.14t do not have a
+        # free-threaded Stable ABI and require version-specific linked wheels.
+        python-version: &pyver ['3.10', '3.13t', '3.14t']
         exclude:
           - os: 'Windows'
             arch: 'aarch64'
             python-version: '3.10'
+        include:
+          - os: 'Windows'
+            arch: 'aarch64'
+            python-version: '3.11'
     runs-on: |-
       ${{ case(matrix.os == 'Windows',
                case(matrix.arch == 'aarch64', 'windows-11-arm',
@@ -89,7 +95,7 @@ jobs:
       - run: unzip -l wheelhouse/*.whl
       - uses: actions/upload-artifact@v7
         with:
-          name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }}
+          name: frontend-${{ runner.os }}-${{ runner.arch }}-${{ contains(matrix.python-version, 't') && format('py{0}', matrix.python-version) || 'abi3' }}
           path: wheelhouse/mlx-*.whl
           if-no-files-found: error
 
@@ -175,7 +181,7 @@ jobs:
       - name: Upload frontend packages
         uses: actions/upload-artifact@v7
         with:
-          name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }}
+          name: frontend-${{ runner.os }}-${{ runner.arch }}-${{ contains(matrix.python-version, 't') && format('py{0}', matrix.python-version) || 'abi3' }}
           path: wheelhouse/mlx-*.whl
           if-no-files-found: error
       - name: Upload backend packages
@@ -187,30 +193,36 @@ jobs:
           if-no-files-found: error
 
   test_wheel:
-    name: Test (${{ matrix.os }}, ${{ matrix.toolkit }}, ${{ matrix.arch }})
+    name: Test (${{ matrix.os }}, python-${{ matrix.python-version }}, ${{ matrix.toolkit }}, ${{ matrix.arch }})
     if: github.repository == 'ml-explore/mlx'
     needs: [build_frontend, build_backend, build_mac_wheels]
     strategy:
       matrix:
         os: ['Linux', 'Windows']
-        arch: ['aarch64']
+        arch: ['x86_64', 'aarch64']
         toolkit: ['cpu']
+        python-version: ['3.12', '3.13t', '3.14t']
         include:
-          - os: 'Linux'
-            arch: 'x86_64'
-            toolkit: 'cpu'
           - os: 'Linux'
             arch: 'x86_64'
             toolkit: 'cuda-12.9'
+            python-version: '3.12'
           - os: 'Linux'
             arch: 'x86_64'
             toolkit: 'cuda-13.0'
-          - os: 'Windows'
-            arch: 'x86_64'
-            toolkit: 'cpu'
+            python-version: '3.12'
           - os: 'macOS'
             arch: 'aarch64'
             toolkit: 'metal'
+            python-version: '3.12'
+          - os: 'macOS'
+            arch: 'aarch64'
+            toolkit: 'metal'
+            python-version: '3.13t'
+          - os: 'macOS'
+            arch: 'aarch64'
+            toolkit: 'metal'
+            python-version: '3.14t'
     runs-on: |-
       ${{ case(matrix.os == 'Windows', case(matrix.arch == 'aarch64', 'windows-11-arm',
                                             'windows-2022'),
@@ -224,6 +236,7 @@ jobs:
       - uses: ./.github/actions/setup
         with:
           toolkit: ${{ matrix.toolkit }}
+          python-version: ${{ matrix.python-version }}
           use-ccache: false
       - uses: ./.github/actions/test-wheel
         with:

CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.25)
+cmake_minimum_required(VERSION 3.26)
 
 if(NOT MLX_VERSION)
   file(STRINGS "mlx/version.h" _mlx_h_version REGEX "^#define MLX_VERSION_.*$")
@@ -392,10 +392,26 @@ if(MLX_BUILD_PYTHON_BINDINGS)
     Python 3.10
     COMPONENTS Interpreter Development.Module
     REQUIRED)
+  execute_process(
+    COMMAND
+      "${Python_EXECUTABLE}" -c
+      "import sysconfig; print('ON' if sysconfig.get_config_var('Py_GIL_DISABLED') else 'OFF')"
+    RESULT_VARIABLE _mlx_python_gil_result
+    OUTPUT_VARIABLE MLX_PYTHON_FREE_THREADED
+    OUTPUT_STRIP_TRAILING_WHITESPACE)
+  if(NOT _mlx_python_gil_result EQUAL 0)
+    message(FATAL_ERROR "Failed to determine whether Python is free-threaded.")
+  endif()
+  if(NOT MLX_PYTHON_FREE_THREADED)
+    find_package(
+      Python 3.10
+      COMPONENTS Development.SABIModule
+      REQUIRED)
+  endif()
   FetchContent_Declare(
     nanobind
     GIT_REPOSITORY https://github.com/wjakob/nanobind.git
-    GIT_TAG v2.15.0
+    GIT_TAG v3.0.0
     GIT_SHALLOW TRUE
     EXCLUDE_FROM_ALL)
   FetchContent_MakeAvailable(nanobind)

MANIFEST.in

@@ -1,5 +1,6 @@
 include CMakeLists.txt
 include mlx.pc.in
+include _build_backend/backend.py
 recursive-include mlx *
 include cmake/*
 include python/src/*

_build_backend/backend.py

@@ -0,0 +1,23 @@
+import os
+import sysconfig
+
+from setuptools import build_meta as _setuptools_build_meta
+from setuptools.build_meta import *  # type: ignore # noqa: F403
+
+
+def _with_nanobind_backend(requires):
+    build_backend = int(os.environ.get("MLX_BUILD_BACKEND_PACKAGE", 0))
+    free_threaded = bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
+    if not build_backend and not free_threaded:
+        requires.append("nanobind-backend>=1.0.0")
+    return requires
+
+
+def get_requires_for_build_wheel(config_settings=None):
+    requires = _setuptools_build_meta.get_requires_for_build_wheel(config_settings)
+    return _with_nanobind_backend(requires)
+
+
+def get_requires_for_build_editable(config_settings=None):
+    requires = _setuptools_build_meta.get_requires_for_build_editable(config_settings)
+    return _with_nanobind_backend(requires)

docs/src/dev/extensions.rst

@@ -679,7 +679,8 @@ Finally, we build the nanobind_ bindings
 
     nanobind_add_module(
       _ext
-      NB_STATIC STABLE_ABI LTO NOMINSIZE
+      BACKEND_MODULE nanobind_backend
+      LTO NOMINSIZE
       NB_DOMAIN mlx
       ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp
     )
@@ -709,9 +710,11 @@ build utilities defined in :mod:`mlx.extension`:
             cmdclass={"build_ext": extension.CMakeBuild},
             packages=["mlx_sample_extensions"],
             package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]},
+            install_requires=["nanobind-backend>=1.0.0"],
+            options={"bdist_wheel": {"py_limited_api": "cp310"}},
             extras_require={"dev":[]},
             zip_safe=False,
-            python_requires=">=3.8",
+            python_requires=">=3.10",
         )
 
 .. note::
@@ -733,7 +736,7 @@ This results in the directory structure:
 | │   ├── __init__.py
 | │   ├── libmlx_ext.dylib # C++ extension library
 | │   ├── mlx_ext.metallib # Metal library
-| │   └── _ext.cpython-3x-darwin.so # Python Binding
+| │   └── _ext.abi3.so # Python Binding
 | ...
 
 When you try to install using the command ``python -m pip install .`` (in

examples/extensions/CMakeLists.txt

@@ -11,8 +11,8 @@ option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON)
 
 # ----------------------------- Dependencies -----------------------------
 find_package(
-  Python 3.8
-  COMPONENTS Interpreter Development.Module
+  Python 3.10
+  COMPONENTS Interpreter Development.Module Development.SABIModule
   REQUIRED)
 execute_process(
   COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
@@ -64,8 +64,8 @@ endif()
 # ----------------------------- Python Bindings -----------------------------
 nanobind_add_module(
   _ext
-  NB_STATIC
-  STABLE_ABI
+  BACKEND_MODULE
+  nanobind_backend
   LTO
   NOMINSIZE
   NB_DOMAIN

examples/extensions/pyproject.toml

@@ -1,8 +1,9 @@
 [build-system]
 requires = [
   "setuptools>=42",
-  "cmake>=3.25",
+  "cmake>=3.27",
   "mlx>=0.18.0",
-  "nanobind==2.15.0",
+  "nanobind==3.0.0",
+  "nanobind-backend>=1.0.0",
 ]
 build-backend = "setuptools.build_meta"

examples/extensions/requirements.txt

@@ -1,4 +1,5 @@
 setuptools>=42
-cmake>=3.25
+cmake>=3.27
 mlx>=0.31.2
-nanobind==2.15.0
+nanobind==3.0.0
+nanobind-backend>=1.0.0

examples/extensions/setup.py

@@ -13,6 +13,8 @@
         cmdclass={"build_ext": extension.CMakeBuild},
         packages=["mlx_sample_extensions"],
         package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]},
+        install_requires=["nanobind-backend>=1.0.0"],
+        options={"bdist_wheel": {"py_limited_api": "cp310"}},
         zip_safe=False,
         python_requires=">=3.10",
     )

pyproject.toml

@@ -1,7 +1,8 @@
 [build-system]
 requires = [
   "setuptools>=80",
-  "cmake>=3.25",
+  "cmake>=3.26",
   "typing_extensions",
 ]
-build-backend = "setuptools.build_meta"
+build-backend = "backend"
+backend-path = ["_build_backend"]

python/mlx/extension.py

@@ -17,7 +17,7 @@
 # A CMakeExtension needs a sourcedir instead of a file list.
 class CMakeExtension(Extension):
     def __init__(self, name: str, sourcedir: str = "") -> None:
-        super().__init__(name, sources=[])
+        super().__init__(name, sources=[], py_limited_api=True)
         self.sourcedir = os.fspath(Path(sourcedir).resolve())
 
 

python/src/CMakeLists.txt

@@ -1,8 +1,14 @@
+# Python 3.13t and 3.14t have no free-threaded Stable ABI, so they must use
+# linked mode instead of the regular split-mode build.
+if(MLX_PYTHON_FREE_THREADED)
+  set(MLX_NANOBIND_MODE NB_STATIC FREE_THREADED)
+else()
+  set(MLX_NANOBIND_MODE BACKEND_MODULE nanobind_backend)
+endif()
+
 nanobind_add_module(
   core
-  NB_STATIC
-  STABLE_ABI
-  FREE_THREADED
+  ${MLX_NANOBIND_MODE}
   LTO
   NOMINSIZE
   NB_DOMAIN

python/src/buffer.h

@@ -7,6 +7,28 @@
 #include "mlx/array.h"
 #include "mlx/utils.h"
 
+// Py_buffer became part of the Stable ABI in Python 3.11. CPython 3.10 uses
+// the same layout, which split-mode extensions targeting cp310 must declare.
+#if defined(Py_LIMITED_API) && Py_LIMITED_API + 0 < 0x030B0000
+typedef struct {
+  void* buf;
+  PyObject* obj;
+  Py_ssize_t len;
+  Py_ssize_t itemsize;
+  int readonly;
+  int ndim;
+  char* format;
+  Py_ssize_t* shape;
+  Py_ssize_t* strides;
+  Py_ssize_t* suboffsets;
+  void* internal;
+} Py_buffer;
+
+#define PyBUF_FORMAT 0x0004
+#define PyBUF_ND 0x0008
+#define PyBUF_STRIDES (0x0010 | PyBUF_ND)
+#endif
+
 // Only defined in >= Python 3.9
 // https://github.com/python/cpython/blob/f6cdc6b4a191b75027de342aa8b5d344fb31313e/Include/typeslots.h#L2-L3
 #ifndef Py_bf_getbuffer

python/src/convert.h

@@ -14,24 +14,6 @@ namespace nb = nanobind;
 
 namespace nanobind {
 
-template <>
-struct ndarray_traits<mx::float16_t> {
-  static constexpr bool is_complex = false;
-  static constexpr bool is_float = true;
-  static constexpr bool is_bool = false;
-  static constexpr bool is_int = false;
-  static constexpr bool is_signed = true;
-};
-
-template <>
-struct ndarray_traits<mx::bfloat16_t> {
-  static constexpr bool is_complex = false;
-  static constexpr bool is_float = true;
-  static constexpr bool is_bool = false;
-  static constexpr bool is_int = false;
-  static constexpr bool is_signed = true;
-};
-
 namespace detail {
 
 template <>

python/src/load.cpp

@@ -15,6 +15,12 @@
 #include "python/src/small_vector.h"
 #include "python/src/utils.h"
 
+// These flags became part of the Stable ABI together with Py_buffer in 3.11.
+#if defined(Py_LIMITED_API) && Py_LIMITED_API + 0 < 0x030B0000
+#define PyBUF_READ 0x100
+#define PyBUF_WRITE 0x200
+#endif
+
 namespace mx = mlx::core;
 namespace nb = nanobind;
 using namespace nb::literals;