Move to nanobind

.clang-format

@@ -0,0 +1,22 @@
+BasedOnStyle: Google
+Language: Cpp
+
+# Make braces stay on the same line (like your diffs)
+BreakBeforeBraces: Attach
+AllowShortFunctionsOnASingleLine: None
+
+# Compact/“binpack” parameter lists (what produced your earlier diffs)
+BinPackParameters: true
+BinPackArguments: true
+
+# Typical CUDA/C++ ergonomics
+IndentWidth: 2
+ColumnLimit: 100
+PointerAlignment: Left
+DerivePointerAlignment: false
+
+# Don’t reorder #includes if you don’t want surprise churn
+SortIncludes: false
+
+# Optional: make templates break more aggressively
+AlwaysBreakTemplateDeclarations: Yes

.github/workflows/publish.yml

@@ -1,36 +1,152 @@
-# This workflow will upload a Python Package using Twine when a release is created
-# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries
+# https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries
+# https://github.com/pypa/cibuildwheel/blob/main/examples/github-deploy.yml
 
-# This workflow uses actions that are not certified by GitHub.
-# They are provided by a third-party and are governed by
-# separate terms of service, privacy policy, and support
-# documentation.
-
-name: Upload Python Package
+name: Build and upload to PyPI
 
 on:
+  workflow_dispatch:
+  pull_request:
+  push:
+    branches: [main]
   release:
     types: [published]
 
 jobs:
-  deploy:
+  build_wheels:
+    name: Build wheels for ${{ matrix.os }}-cu${{ matrix.cuda_major }}
+    runs-on: ${{ matrix.runs-on }}
+    strategy:
+      matrix:
+        include:
+          # CUDA 12 wheels (compiled with 12.2 for runtime 12.2+ compatibility)
+          - os: linux-intel
+            cuda_major: "12"
+            runs-on: ubuntu-latest
+            cibw_image: "ghcr.io/scverse/rapids_singlecell:manylinux_2_28_x86_64_cuda12.2"
+            dockerfile: "docker/manylinux_2_28_x86_64_cuda12.2.Dockerfile"
+            cuda_archs: "75-real;80-real;86-real;89-real;90"
+          - os: linux-arm
+            cuda_major: "12"
+            runs-on: ubuntu-24.04-arm
+            cibw_image: "ghcr.io/scverse/rapids_singlecell:manylinux_2_28_aarch64_cuda12.2"
+            dockerfile: "docker/manylinux_2_28_aarch64_cuda12.2.Dockerfile"
+            cuda_archs: "75-real;80-real;86-real;89-real;90"
+          # CUDA 13 wheels (native Blackwell support)
+          - os: linux-intel
+            cuda_major: "13"
+            runs-on: ubuntu-latest
+            cibw_image: "ghcr.io/scverse/rapids_singlecell:manylinux_2_28_x86_64_cuda13.0"
+            dockerfile: "docker/manylinux_2_28_x86_64_cuda13.0.Dockerfile"
+            cuda_archs: "75-real;80-real;86-real;89-real;90-real;100-real;120"
+          - os: linux-arm
+            cuda_major: "13"
+            runs-on: ubuntu-24.04-arm
+            cibw_image: "ghcr.io/scverse/rapids_singlecell:manylinux_2_28_aarch64_cuda13.0"
+            dockerfile: "docker/manylinux_2_28_aarch64_cuda13.0.Dockerfile"
+            cuda_archs: "75-real;80-real;86-real;89-real;90-real;100-real;120"
 
-    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
 
-    environment: publish
+      - name: Set package name and extras for CUDA ${{ matrix.cuda_major }}
+        run: |
+          python3 - ${{ matrix.cuda_major }} <<'SCRIPT'
+          import sys, pathlib
+
+          cuda = sys.argv[1]
+          other = "13" if cuda == "12" else "12"
+          path = pathlib.Path("pyproject.toml")
+          text = path.read_text()
+
+          # Rename package
+          text = text.replace(
+              'name = "rapids_singlecell"',
+              f'name = "rapids_singlecell-cu{cuda}"',
+          )
+          # Rename matching extra to "rapids", remove the other
+          text = text.replace(f'rapids-cu{cuda} =', 'rapids =')
+          # Remove the other CUDA extra line entirely
+          lines = text.splitlines(keepends=True)
+          text = "".join(l for l in lines if f'rapids-cu{other}' not in l)
+
+          path.write_text(text)
+
+          # Verify
+          for line in path.read_text().splitlines():
+              if "name" in line or "rapids" in line.lower():
+                  print(line)
+          SCRIPT
+
+      - name: Build CUDA manylinux image
+        run: |
+          docker build -t "${{ matrix.cibw_image }}" -f "${{ matrix.dockerfile }}" docker
+
+      - name: Build wheels
+        uses: pypa/cibuildwheel@v3.1.4
+        env:
+          CIBW_SKIP: '*-musllinux*'
+          CIBW_MANYLINUX_X86_64_IMAGE: ${{ matrix.os == 'linux-intel' && matrix.cibw_image || '' }}
+          CIBW_MANYLINUX_AARCH64_IMAGE: ${{ matrix.os == 'linux-arm' && matrix.cibw_image || '' }}
+          CIBW_ENVIRONMENT: >
+            CUDA_PATH=/usr/local/cuda
+            LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
+            PATH=/usr/local/cuda/bin:$PATH
+            SKBUILD_CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=${{ matrix.cuda_archs }}"
+          CIBW_BEFORE_BUILD: >
+            python -m pip install -U pip
+            scikit-build-core cmake ninja nanobind
+          CIBW_TEST_SKIP: "*"
+          CIBW_TEST_COMMAND: ""
+          CIBW_REPAIR_WHEEL_COMMAND: "auditwheel repair -w {dest_dir} {wheel}"
+          CIBW_BUILD_VERBOSITY: "1"
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cibw-wheels-${{ matrix.os }}-cu${{ matrix.cuda_major }}-${{ strategy.job-index }}
+          path: ./wheelhouse/*.whl
 
+  build_sdist:
+    name: Build source distribution
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
+      - name: Build sdist
+        run: pipx run build --sdist
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cibw-sdist
+          path: dist/*.tar.gz
+
+  upload_wheels:
+    needs: [build_wheels]
+    runs-on: ubuntu-latest
+    environment: publish-cu${{ matrix.cuda_major }}
     permissions:
       id-token: write
+    if: github.event_name == 'release' && github.event.action == 'published'
+    strategy:
+      matrix:
+        cuda_major: ["12", "13"]
+    steps:
+      - uses: actions/download-artifact@v5
+        with:
+          pattern: cibw-wheels-*-cu${{ matrix.cuda_major }}-*
+          path: dist
+          merge-multiple: true
+
+      - uses: pypa/gh-action-pypi-publish@release/v1
 
+  upload_sdist:
+    needs: [build_sdist]
+    runs-on: ubuntu-latest
+    environment: publish
+    permissions:
+      id-token: write
+    if: github.event_name == 'release' && github.event.action == 'published'
     steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python
-      uses: actions/setup-python@v5
-      with:
-        python-version: '3.x'
-    - name: Install CLI tool
-      run: pip install build
-    - name: Build package
-      run: python -m build
-    - name: Publish package
-      uses: pypa/gh-action-pypi-publish@release/v1
+      - uses: actions/download-artifact@v5
+        with:
+          name: cibw-sdist
+          path: dist
+
+      - uses: pypa/gh-action-pypi-publish@release/v1

.github/workflows/test-gpu.yml

@@ -83,6 +83,11 @@ jobs:
         with:
           python-version: ${{ matrix.env.python }}
 
+      - name: Add CUDA to PATH
+        run: |
+          echo "/usr/local/cuda/bin" >> $GITHUB_PATH
+          echo "LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" >> $GITHUB_ENV
+
       - name: Install dependencies
         run: uvx hatch -v env create ${{ matrix.env.name }}
 
@@ -140,6 +145,17 @@ jobs:
         with:
           python-version: ${{ matrix.env.python }}
 
+      - name: Add CUDA to PATH
+        run: |
+          echo "/usr/local/cuda/bin" >> $GITHUB_PATH
+          echo "LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" >> $GITHUB_ENV
+
+      - name: Check CUDA version
+        run: nvcc --version
+
+      - name: Check NVIDIA SMI
+        run: nvidia-smi
+
       - name: Install dependencies
         run: uvx hatch -v env create ${{ matrix.env.name }}
 

.gitignore

@@ -31,8 +31,10 @@ __pycache__/
 /_skbuild/
 /*.egg-info/
 
-# Compiled CUDA extension copied for editable installs
+# Compiled CUDA extensions and stubs copied for editable installs
 *.so
+*.pyi
+py.typed
 
 # Test/coverage caches
 /.pytest_cache/

.pre-commit-config.yaml

@@ -36,3 +36,9 @@ repos:
           (?x)(
               docs/references.bib|
           )
+-   repo: https://github.com/pre-commit/mirrors-clang-format
+    rev: v18.1.8
+    hooks:
+    -   id: clang-format
+        args: [--style=file, -i]
+        types_or: [c, c++, cuda]

.readthedocs.yml

@@ -16,5 +16,5 @@ build:
       - asdf global uv latest
     build:
       html:
-        - uvx hatch run docs:build
+        - CMAKE_ARGS="-DRSC_BUILD_EXTENSIONS=OFF" uvx hatch run docs:build
         - mv docs/_build $READTHEDOCS_OUTPUT

CMakeLists.txt

@@ -0,0 +1,96 @@
+cmake_minimum_required(VERSION 3.24)
+
+project(rapids_singlecell_cuda LANGUAGES CXX)
+
+# Option to disable building compiled extensions (for docs/RTD)
+option(RSC_BUILD_EXTENSIONS "Build CUDA/C++ extensions" ON)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+if (RSC_BUILD_EXTENSIONS)
+  enable_language(CUDA)
+  find_package(Python REQUIRED COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT})
+  find_package(nanobind CONFIG REQUIRED)
+  find_package(CUDAToolkit REQUIRED)
+  message(STATUS "Building for CUDA architectures: ${CMAKE_CUDA_ARCHITECTURES}")
+else()
+  message(STATUS "RSC_BUILD_EXTENSIONS=OFF -> skipping compiled extensions for docs")
+endif()
+
+# Helper to declare a nanobind CUDA module uniformly
+function(add_nb_cuda_module target src)
+  if (RSC_BUILD_EXTENSIONS)
+    nanobind_add_module(${target} STABLE_ABI LTO
+        ${src}
+    )
+    target_link_libraries(${target} PRIVATE CUDA::cudart)
+    set_target_properties(${target} PROPERTIES
+        CUDA_SEPARABLE_COMPILATION ON
+    )
+    install(TARGETS ${target} LIBRARY DESTINATION rapids_singlecell/_cuda)
+    # Generate type stubs at install time (for wheel installs)
+    nanobind_add_stub(${target}_stub
+        MODULE ${target}
+        OUTPUT rapids_singlecell/_cuda/${target}.pyi
+        PYTHON_PATH $<TARGET_FILE_DIR:${target}>
+        DEPENDS ${target}
+        INSTALL_TIME
+        MARKER_FILE rapids_singlecell/_cuda/py.typed
+    )
+    # Generate type stubs at build time (for editable installs)
+    nanobind_add_stub(${target}_stub_dev
+        MODULE ${target}
+        OUTPUT ${target}.pyi
+        PYTHON_PATH $<TARGET_FILE_DIR:${target}>
+        DEPENDS ${target}
+    )
+    # Copy built module + stub into source tree for editable installs
+    add_custom_command(TARGET ${target}_stub_dev POST_BUILD
+        COMMAND ${CMAKE_COMMAND} -E copy
+            ${CMAKE_CURRENT_BINARY_DIR}/${target}.pyi
+            ${PROJECT_SOURCE_DIR}/src/rapids_singlecell/_cuda/${target}.pyi
+        COMMAND ${CMAKE_COMMAND} -E touch
+            ${PROJECT_SOURCE_DIR}/src/rapids_singlecell/_cuda/py.typed
+    )
+    add_custom_command(TARGET ${target} POST_BUILD
+        COMMAND ${CMAKE_COMMAND} -E copy
+            $<TARGET_FILE:${target}>
+            ${PROJECT_SOURCE_DIR}/src/rapids_singlecell/_cuda/$<TARGET_FILE_NAME:${target}>
+    )
+  endif()
+endfunction()
+
+if (RSC_BUILD_EXTENSIONS)
+  # CUDA modules
+  add_nb_cuda_module(_mean_var_cuda     src/rapids_singlecell/_cuda/mean_var/mean_var.cu)
+  add_nb_cuda_module(_sparse2dense_cuda src/rapids_singlecell/_cuda/sparse2dense/sparse2dense.cu)
+  add_nb_cuda_module(_scale_cuda        src/rapids_singlecell/_cuda/scale/scale.cu)
+  add_nb_cuda_module(_qc_cuda           src/rapids_singlecell/_cuda/qc/qc.cu)
+  add_nb_cuda_module(_qc_dask_cuda      src/rapids_singlecell/_cuda/qc_dask/qc_kernels_dask.cu)
+  add_nb_cuda_module(_bbknn_cuda        src/rapids_singlecell/_cuda/bbknn/bbknn.cu)
+  add_nb_cuda_module(_norm_cuda         src/rapids_singlecell/_cuda/norm/norm.cu)
+  add_nb_cuda_module(_pr_cuda           src/rapids_singlecell/_cuda/pr/pr.cu)
+  add_nb_cuda_module(_nn_descent_cuda   src/rapids_singlecell/_cuda/nn_descent/nn_descent.cu)
+  add_nb_cuda_module(_aucell_cuda       src/rapids_singlecell/_cuda/aucell/aucell.cu)
+  add_nb_cuda_module(_nanmean_cuda      src/rapids_singlecell/_cuda/nanmean/nanmean.cu)
+  add_nb_cuda_module(_autocorr_cuda     src/rapids_singlecell/_cuda/autocorr/autocorr.cu)
+  add_nb_cuda_module(_cooc_cuda         src/rapids_singlecell/_cuda/cooc/cooc.cu)
+  add_nb_cuda_module(_aggr_cuda         src/rapids_singlecell/_cuda/aggr/aggr.cu)
+  add_nb_cuda_module(_spca_cuda         src/rapids_singlecell/_cuda/spca/spca.cu)
+  add_nb_cuda_module(_ligrec_cuda       src/rapids_singlecell/_cuda/ligrec/ligrec.cu)
+  add_nb_cuda_module(_pv_cuda           src/rapids_singlecell/_cuda/pv/pv.cu)
+  add_nb_cuda_module(_edistance_cuda    src/rapids_singlecell/_cuda/edistance/edistance.cu)
+  add_nb_cuda_module(_hvg_cuda          src/rapids_singlecell/_cuda/hvg/hvg.cu)
+  add_nb_cuda_module(_wilcoxon_cuda     src/rapids_singlecell/_cuda/wilcoxon/wilcoxon.cu)
+  # Harmony CUDA modules
+  add_nb_cuda_module(_harmony_scatter_cuda   src/rapids_singlecell/_cuda/harmony/scatter/scatter.cu)
+  add_nb_cuda_module(_harmony_outer_cuda     src/rapids_singlecell/_cuda/harmony/outer/outer.cu)
+  add_nb_cuda_module(_harmony_colsum_cuda    src/rapids_singlecell/_cuda/harmony/colsum/colsum.cu)
+  add_nb_cuda_module(_harmony_kmeans_cuda    src/rapids_singlecell/_cuda/harmony/kmeans/kmeans.cu)
+  add_nb_cuda_module(_harmony_normalize_cuda src/rapids_singlecell/_cuda/harmony/normalize/normalize.cu)
+  add_nb_cuda_module(_harmony_pen_cuda       src/rapids_singlecell/_cuda/harmony/pen/pen.cu)
+  # Wilcoxon binned histogram CUDA module
+  add_nb_cuda_module(_wilcoxon_binned_cuda   src/rapids_singlecell/_cuda/wilcoxon_binned/wilcoxon_binned.cu)
+endif()

conda/rsc_rapids_25.12_cuda12.yml

@@ -18,4 +18,4 @@ dependencies:
     - omnipath
     - gdown
     - wget
-    - rapids-singlecell
+    - rapids-singlecell-cu12

conda/rsc_rapids_25.12_cuda13.yml

@@ -18,4 +18,4 @@ dependencies:
     - omnipath
     - gdown
     - wget
-    - rapids-singlecell
+    - rapids-singlecell-cu13

conda/rsc_rapids_26.02_cuda12.yml

@@ -17,4 +17,4 @@ dependencies:
     - gdown
     - wget
     - scikit-misc
-    - rapids-singlecell
+    - rapids-singlecell-cu12

conda/rsc_rapids_26.02_cuda13.yml

@@ -17,4 +17,4 @@ dependencies:
     - gdown
     - wget
     - scikit-misc
-    - rapids-singlecell
+    - rapids-singlecell-cu13

docker/Dockerfile

@@ -5,9 +5,10 @@ ARG GIT_ID=main
 SHELL ["/bin/bash", "-euo", "pipefail", "-c"]
 
 ENV PATH=/opt/conda/bin:$PATH
+ENV SKBUILD_CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=75-real;80-real;86-real;89-real;90-real;100-real;120"
 
 RUN <<EOF
-# install rapids_singlecell from source
+# install rapids_singlecell from source (compiled for all supported GPU architectures)
 set -x
 mkdir /src
 cd /src