From 399982e34f6be114b720afd01a00e5c2c24992f7 Mon Sep 17 00:00:00 2001 From: jasp Date: Mon, 27 Jul 2026 10:40:28 +0200 Subject: [PATCH] Add CUDA 13 runtime headers --- .github/actions/build-wheel/action.yml | 65 +++++++++++++++++ mlx/backend/cuda/jit_module.cpp | 23 +++--- mlx/backend/cuda/runtime_headers.h | 49 +++++++++++++ setup.py | 1 + tests/CMakeLists.txt | 1 + tests/cuda_runtime_headers_tests.cpp | 99 ++++++++++++++++++++++++++ 6 files changed, 228 insertions(+), 10 deletions(-) create mode 100644 mlx/backend/cuda/runtime_headers.h create mode 100644 tests/cuda_runtime_headers_tests.cpp diff --git a/.github/actions/build-wheel/action.yml b/.github/actions/build-wheel/action.yml index 96cca9eadc..fcd8ea7db9 100644 --- a/.github/actions/build-wheel/action.yml +++ b/.github/actions/build-wheel/action.yml @@ -79,6 +79,71 @@ runs: echo "::group::Build backend package" python setup.py clean --all MLX_BUILD_STAGE=2 python -m build -w + python - <<'PY' + import os + import re + import subprocess + from email.parser import Parser + from pathlib import Path + from zipfile import ZipFile + + from packaging.requirements import Requirement + from packaging.specifiers import SpecifierSet + from packaging.utils import canonicalize_name + + wheels = list(Path("dist").glob("mlx_cuda_13-*.whl")) + build_cuda = "MLX_BUILD_CUDA=ON" in os.environ.get("CMAKE_ARGS", "") + if build_cuda: + nvcc_version = subprocess.check_output(["nvcc", "--version"], text=True) + match = re.search(r"release (\d+)", nvcc_version) + if not match: + raise RuntimeError( + f"Could not determine the CUDA version from: {nvcc_version}" + ) + cuda_major = int(match.group(1)) + else: + cuda_major = None + + if cuda_major == 13 and len(wheels) != 1: + raise RuntimeError(f"Expected one CUDA 13 wheel, found: {wheels}") + + for wheel in wheels: + with ZipFile(wheel) as archive: + metadata_files = [ + name for name in archive.namelist() + if name.endswith(".dist-info/METADATA") + ] + if len(metadata_files) != 1: + raise RuntimeError( + f"Expected one METADATA file in {wheel}, " + f"found: {metadata_files}" + ) + metadata = Parser().parsestr( + archive.read(metadata_files[0]).decode() + ) + requirements = [ + Requirement(value) + for value in metadata.get_all("Requires-Dist", []) + ] + runtime_requirements = [ + requirement + for requirement in requirements + if canonicalize_name(requirement.name) + == "nvidia-cuda-runtime" + ] + if ( + len(runtime_requirements) != 1 + or runtime_requirements[0].specifier + != SpecifierSet("==13.*") + or runtime_requirements[0].marker is not None + or runtime_requirements[0].extras + or runtime_requirements[0].url is not None + ): + raise RuntimeError( + f"{wheel} is missing the CUDA 13 runtime dependency: " + f"{requirements}" + ) + PY echo "::endgroup::" - name: Post-process backend package diff --git a/mlx/backend/cuda/jit_module.cpp b/mlx/backend/cuda/jit_module.cpp index 3de1ddb018..d712e4a227 100644 --- a/mlx/backend/cuda/jit_module.cpp +++ b/mlx/backend/cuda/jit_module.cpp @@ -2,6 +2,7 @@ #include "mlx/backend/cuda/jit_module.h" #include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/runtime_headers.h" #include "mlx/version.h" #include "cuda_jit_sources.h" @@ -54,10 +55,7 @@ const std::vector& include_path_args() { static std::vector cached_args = []() { std::vector args; // Add path to bundled headers. - auto root_dir = current_binary_dir(); -#if !defined(_WIN32) - root_dir = root_dir.parent_path(); -#endif + auto root_dir = detail::mlx_package_root(current_binary_dir()); auto path = root_dir / "include"; if (std::filesystem::exists(path)) { args.push_back(fmt::format("--include-path={}", path.string())); @@ -72,17 +70,22 @@ const std::vector& include_path_args() { } // Add path to CUDA runtime headers, try local-installed python package // first and then system-installed headers. - path = root_dir.parent_path() / "nvidia" / "cuda_runtime" / "include"; - if (!std::filesystem::exists(path)) { + std::vector toolkit_roots; + path = detail::find_cuda_runtime_include_dir( + root_dir, CUDA_VERSION / 1000, toolkit_roots); + if (path.empty()) { const char* home = std::getenv("CUDA_HOME"); if (!home) { home = std::getenv("CUDA_PATH"); } - path = home ? std::filesystem::path(home) : default_cuda_toolkit_path(); - if (!path.empty()) { - path = path / "include"; + if (home) { + toolkit_roots.emplace_back(home); + } else if (!default_cuda_toolkit_path().empty()) { + toolkit_roots.push_back(default_cuda_toolkit_path()); } - if (path.empty() || !std::filesystem::exists(path)) { + path = detail::find_cuda_runtime_include_dir( + root_dir, CUDA_VERSION / 1000, toolkit_roots); + if (path.empty()) { throw std::runtime_error( "Can not find locations of CUDA headers, please set environment " "variable CUDA_HOME or CUDA_PATH."); diff --git a/mlx/backend/cuda/runtime_headers.h b/mlx/backend/cuda/runtime_headers.h new file mode 100644 index 0000000000..be2b8a0fff --- /dev/null +++ b/mlx/backend/cuda/runtime_headers.h @@ -0,0 +1,49 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include +#include +#include + +namespace mlx::core::cu::detail { + +inline std::filesystem::path mlx_package_root( + const std::filesystem::path& binary_dir) { +#if defined(_WIN32) + return binary_dir; +#else + return binary_dir.parent_path(); +#endif +} + +inline bool has_cuda_runtime_headers(const std::filesystem::path& path) { + return std::filesystem::exists(path / "cuda.h") && + std::filesystem::exists(path / "cuda_runtime.h"); +} + +inline std::filesystem::path find_cuda_runtime_include_dir( + const std::filesystem::path& root_dir, + int cuda_major_version, + std::span toolkit_roots) { + auto path = root_dir.parent_path() / "nvidia"; + if (cuda_major_version >= 13) { + path /= "cu" + std::to_string(cuda_major_version); + } else { + path /= "cuda_runtime"; + } + path /= "include"; + if (has_cuda_runtime_headers(path)) { + return path; + } + + for (const auto& root : toolkit_roots) { + path = root / "include"; + if (has_cuda_runtime_headers(path)) { + return path; + } + } + return {}; +} + +} // namespace mlx::core::cu::detail diff --git a/setup.py b/setup.py index 3c2f138048..553a0b2f3c 100644 --- a/setup.py +++ b/setup.py @@ -323,6 +323,7 @@ def get_tag(self) -> tuple[str, str, str]: "nvidia-cublas", "nvidia-cufft", "nvidia-cuda-nvrtc", + "nvidia-cuda-runtime==13.*", ] else: raise ValueError(f"Unknown toolkit {toolkit}") diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2a4a41c6b6..ae0d5e4d57 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -21,6 +21,7 @@ target_sources( blas_tests.cpp compile_tests.cpp custom_vjp_tests.cpp + cuda_runtime_headers_tests.cpp creations_tests.cpp device_tests.cpp einsum_tests.cpp diff --git a/tests/cuda_runtime_headers_tests.cpp b/tests/cuda_runtime_headers_tests.cpp new file mode 100644 index 0000000000..1d12ae5984 --- /dev/null +++ b/tests/cuda_runtime_headers_tests.cpp @@ -0,0 +1,99 @@ +// Copyright © 2026 Apple Inc. + +#include +#include +#include +#include +#include + +#include "doctest/doctest.h" + +#include "mlx/backend/cuda/runtime_headers.h" + +namespace { + +class TempDirectory { + public: + TempDirectory() { + std::mt19937_64 random(std::random_device{}()); + for (int attempt = 0; attempt < 100; ++attempt) { + path_ = std::filesystem::temp_directory_path() / + ("mlx-cuda-runtime-headers-" + std::to_string(random())); + if (std::filesystem::create_directory(path_)) { + return; + } + } + throw std::runtime_error("Failed to create a temporary directory"); + } + + ~TempDirectory() { + std::error_code error; + std::filesystem::remove_all(path_, error); + } + + const std::filesystem::path& path() const { + return path_; + } + + private: + std::filesystem::path path_; +}; + +void create_cuda_headers(const std::filesystem::path& include_dir) { + std::filesystem::create_directories(include_dir); + std::ofstream(include_dir / "cuda.h"); + std::ofstream(include_dir / "cuda_runtime.h"); +} + +} // namespace + +using namespace mlx::core::cu::detail; + +TEST_CASE("test CUDA runtime header package layouts") { + TempDirectory temp; + auto mlx_root = temp.path() / "site-packages" / "mlx"; + + SUBCASE("binary directory resolves to the MLX package root") { +#if defined(_WIN32) + auto binary_dir = mlx_root; +#else + auto binary_dir = mlx_root / "lib"; +#endif + CHECK_EQ(mlx_package_root(binary_dir), mlx_root); + } + + SUBCASE("CUDA 13 uses the cu13 package") { + auto expected = + temp.path() / "site-packages" / "nvidia" / "cu13" / "include"; + create_cuda_headers(expected); + + CHECK_EQ(find_cuda_runtime_include_dir(mlx_root, 13, {}), expected); + } + + SUBCASE("CUDA 12 uses the legacy cuda_runtime package") { + auto expected = + temp.path() / "site-packages" / "nvidia" / "cuda_runtime" / "include"; + create_cuda_headers(expected); + + CHECK_EQ(find_cuda_runtime_include_dir(mlx_root, 12, {}), expected); + } + + SUBCASE("a toolkit root is used when the package is incomplete") { + auto package_include = + temp.path() / "site-packages" / "nvidia" / "cu13" / "include"; + std::filesystem::create_directories(package_include); + std::ofstream(package_include / "cuda_runtime.h"); + + auto toolkit_root = temp.path() / "toolkit"; + create_cuda_headers(toolkit_root / "include"); + std::vector toolkit_roots{toolkit_root}; + + CHECK_EQ( + find_cuda_runtime_include_dir(mlx_root, 13, toolkit_roots), + toolkit_root / "include"); + } + + SUBCASE("no candidates returns an empty path") { + CHECK(find_cuda_runtime_include_dir(mlx_root, 13, {}).empty()); + } +}