gpuastar-cu

gpuastar-cu is the CUDA solver backend for GPUAStar. The search kernels in csrc/bwas_kernel.cu are compiled into an extension and exposed as PyTorch custom operators. The search loop in src/gpuastar_cu/search.py launches them.

Install through the root project extra:

python -m pip install "gpuastar[cu]"

The package owns the gpuastar_cu import package and registers the cu backend through the gpuastar.backends entry-point group. The base gpuastar wheel contains no CUDA extension files. gpuastar_cu.load_library() probes the compiled extension independently, while gpuastar_cu.is_available() requires both a loadable extension and CUDA hardware detected by Torch.

Release wheels target glibc-based Linux on x86_64 and aarch64. Official PyPI Torch wheels for Windows AMD64 are CPU-only, and macOS has no NVIDIA CUDA runtime. The root gpuastar[cu] marker therefore selects this package only on the two supported Linux architectures. The CUDA package CMake project can still install its Python sources without the extension for source-tree development and import checks. That source-only mode is not a published target of the root extra. On any platform, gpuastar.backends.get_backend("cu") raises BackendUnavailableError when Torch cannot detect CUDA or the compiled extension cannot load.

Linux x86_64 wheels for CPython 3.10 through 3.13 use Torch 2.6 and CUDA 12.4. Linux x86_64 wheels for CPython 3.14 and 3.14t use Torch 2.11 and CUDA 13.0. All Linux aarch64 wheels use Torch 2.11 and CUDA 13.0 because the earlier official PyPI aarch64 Torch wheels are CPU-only.

The default binary contains SASS for compute capabilities 7.5, 8.0, 8.6, 8.7, 8.9, and 9.0. Compute capability 7.5 is the supported floor. This covers Turing GPUs and Jetson Orin while following the CUDA 13 library floor. The compute 9.0 entry also contains PTX for driver JIT compilation on later architectures.

Run the backend through the GPUAStar CLI:

gpuastar command=astar_search backend.language=cu backend.device=cuda:0 input.env=cube3 input.states=states.pkl input.model=model_state_dict.pt input.results_dir=results/cu

The input.* options can be omitted when the gpuastar[experiments] extra is installed, which provides Rubik's Cube example defaults.

In a source checkout, the extension is built with pixi run -e dev cu-build and the Rubik's Cube smoke test goes through the same CLI entry point:

CUDA_VISIBLE_DEVICES=1 pixi run -e cu-experiments cu-cube3-smoke