Add CUDA 13 runtime headers

.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

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<std::string>& include_path_args() {
   static std::vector<std::string> cached_args = []() {
     std::vector<std::string> 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<std::string>& 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<std::filesystem::path> 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.");

mlx/backend/cuda/runtime_headers.h

@@ -0,0 +1,49 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include <filesystem>
+#include <span>
+#include <string>
+
+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<const std::filesystem::path> 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

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}")

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

tests/cuda_runtime_headers_tests.cpp

@@ -0,0 +1,99 @@
+// Copyright © 2026 Apple Inc.
+
+#include <filesystem>
+#include <fstream>
+#include <random>
+#include <stdexcept>
+#include <vector>
+
+#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<std::filesystem::path> 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());
+  }
+}