fix: derive CUDA major version from headers for build (#1395)

* fix: derive CUDA major version from headers for build

Fixes build failures when cuda-bindings reports a different major version
than the CUDA headers being compiled against.

The new _get_cuda_major_version() function is used for both:
1. Determining which cuda-bindings version to install as a build dependency
2. Setting CUDA_CORE_BUILD_MAJOR for Cython compile-time conditionals

Version is derived from (in order of priority):
1. CUDA_CORE_BUILD_MAJOR env var (explicit override, e.g. in CI)
2. CUDA_VERSION macro in cuda.h from CUDA_PATH or CUDA_HOME

Since CUDA_PATH or CUDA_HOME is required for the build anyway, the
cuda.h header should always be available, ensuring consistency between
the installed cuda-bindings and the compile-time conditionals.

* Minor refactoring.

Author: Andy-Jost

cuda_core/build_hooks.py

@@ -11,7 +11,6 @@
 import glob
 import os
 import re
-import subprocess
 
 from Cython.Build import cythonize
 from setuptools import Extension
@@ -26,32 +25,60 @@
 
 
 @functools.cache
-def _get_proper_cuda_bindings_major_version() -> str:
-    # for local development (with/without build isolation)
-    try:
-        import cuda.bindings
+def _get_cuda_paths() -> list[str]:
+    CUDA_PATH = os.environ.get("CUDA_PATH", os.environ.get("CUDA_HOME", None))
+    if not CUDA_PATH:
+        raise RuntimeError("Environment variable CUDA_PATH or CUDA_HOME is not set")
+    CUDA_PATH = CUDA_PATH.split(os.pathsep)
+    print("CUDA paths:", CUDA_PATH)
+    return CUDA_PATH
 
-        return cuda.bindings.__version__.split(".")[0]
-    except ImportError:
-        pass
 
-    # for custom overwrite, e.g. in CI
+@functools.cache
+def _determine_cuda_major_version() -> str:
+    """Determine the CUDA major version for building cuda.core.
+
+    This version is used for two purposes:
+    1. Determining which cuda-bindings version to install as a build dependency
+    2. Setting CUDA_CORE_BUILD_MAJOR for Cython compile-time conditionals
+
+    The version is derived from (in order of priority):
+    1. CUDA_CORE_BUILD_MAJOR environment variable (explicit override, e.g. in CI)
+    2. CUDA_VERSION macro in cuda.h from CUDA_PATH or CUDA_HOME
+
+    Since CUDA_PATH or CUDA_HOME is required for the build (to provide include
+    directories), the cuda.h header should always be available.
+    """
+    # Explicit override, e.g. in CI.
     cuda_major = os.environ.get("CUDA_CORE_BUILD_MAJOR")
     if cuda_major is not None:
+        print("CUDA MAJOR VERSION:", cuda_major)
         return cuda_major
 
-    # also for local development
-    try:
-        out = subprocess.run("nvidia-smi", env=os.environ, capture_output=True, check=True)  # noqa: S603, S607
-        m = re.search(r"CUDA Version:\s*([\d\.]+)", out.stdout.decode())
-        if m:
-            return m.group(1).split(".")[0]
-    except (FileNotFoundError, subprocess.CalledProcessError):
-        # the build machine has no driver installed
-        pass
-
-    # default fallback
-    return "13"
+    # Derive from the CUDA headers (the authoritative source for what we compile against).
+    cuda_path = _get_cuda_paths()
+    for root in cuda_path:
+        cuda_h = os.path.join(root, "include", "cuda.h")
+        try:
+            with open(cuda_h, encoding="utf-8") as f:
+                for line in f:
+                    m = re.match(r"^#\s*define\s+CUDA_VERSION\s+(\d+)\s*$", line)
+                    if m:
+                        v = int(m.group(1))
+                        # CUDA_VERSION is e.g. 12020 for 12.2.
+                        cuda_major = str(v // 1000)
+                        print("CUDA MAJOR VERSION:", cuda_major)
+                        return cuda_major
+        except OSError:
+            continue
+
+    # CUDA_PATH or CUDA_HOME is required for the build, so we should not reach here
+    # in normal circumstances. Raise an error to make the issue clear.
+    raise RuntimeError(
+        "Cannot determine CUDA major version. "
+        "Set CUDA_CORE_BUILD_MAJOR environment variable, or ensure CUDA_PATH or CUDA_HOME "
+        "points to a valid CUDA installation with include/cuda.h."
+    )
 
 
 # used later by setup()
@@ -68,25 +95,12 @@ def _build_cuda_core():
 
     # It seems setuptools' wildcard support has problems for namespace packages,
     # so we explicitly spell out all Extension instances.
-    root_module = "cuda.core"
-    root_path = f"{os.path.sep}".join(root_module.split(".")) + os.path.sep
-    ext_files = glob.glob(f"{root_path}/**/*.pyx", recursive=True)
-
-    def strip_prefix_suffix(filename):
-        return filename[len(root_path) : -4]
-
-    module_names = (strip_prefix_suffix(f) for f in ext_files)
-
-    @functools.cache
-    def get_cuda_paths():
-        CUDA_PATH = os.environ.get("CUDA_PATH", os.environ.get("CUDA_HOME", None))
-        if not CUDA_PATH:
-            raise RuntimeError("Environment variable CUDA_PATH or CUDA_HOME is not set")
-        CUDA_PATH = CUDA_PATH.split(os.pathsep)
-        print("CUDA paths:", CUDA_PATH)
-        return CUDA_PATH
+    def module_names():
+        root_path = os.path.sep.join(["cuda", "core", ""])
+        for filename in glob.glob(f"{root_path}/**/*.pyx", recursive=True):
+            yield filename[len(root_path) : -4]
 
-    all_include_dirs = list(os.path.join(root, "include") for root in get_cuda_paths())
+    all_include_dirs = list(os.path.join(root, "include") for root in _get_cuda_paths())
     extra_compile_args = []
     if COMPILE_FOR_COVERAGE:
         # CYTHON_TRACE_NOGIL indicates to trace nogil functions.  It is not
@@ -101,11 +115,11 @@ def get_cuda_paths():
             language="c++",
             extra_compile_args=extra_compile_args,
         )
-        for mod in module_names
+        for mod in module_names()
     )
 
     nthreads = int(os.environ.get("CUDA_PYTHON_PARALLEL_LEVEL", os.cpu_count() // 2))
-    compile_time_env = {"CUDA_CORE_BUILD_MAJOR": int(_get_proper_cuda_bindings_major_version())}
+    compile_time_env = {"CUDA_CORE_BUILD_MAJOR": int(_determine_cuda_major_version())}
     compiler_directives = {"embedsignature": True, "warn.deprecated.IF": False, "freethreading_compatible": True}
     if COMPILE_FOR_COVERAGE:
         compiler_directives["linetrace"] = True
@@ -132,7 +146,7 @@ def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
 
 
 def _get_cuda_bindings_require():
-    cuda_major = _get_proper_cuda_bindings_major_version()
+    cuda_major = _determine_cuda_major_version()
     return [f"cuda-bindings=={cuda_major}.*"]
 
 

cuda_core/tests/test_build_hooks.py

@@ -0,0 +1,131 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Tests for build_hooks.py build infrastructure.
+
+These tests verify the CUDA version detection logic used during builds,
+particularly the _determine_cuda_major_version() function which derives the
+CUDA major version from headers.
+
+Note: These tests do NOT require cuda.core to be built/installed since they
+test build-time infrastructure. Run with --noconftest to avoid loading
+conftest.py which imports cuda.core modules:
+
+    pytest tests/test_build_hooks.py -v --noconftest
+
+These tests require Cython to be installed (build_hooks.py imports it).
+"""
+
+import importlib.util
+import os
+import tempfile
+from pathlib import Path
+from unittest import mock
+
+import pytest
+
+# build_hooks.py imports Cython at the top level, so skip if not available
+pytest.importorskip("Cython")
+
+
+def _load_build_hooks():
+    """Load build_hooks module from source without permanently modifying sys.path.
+
+    build_hooks.py is a PEP 517 build backend, not an installed module.
+    We use importlib to load it directly from source to avoid polluting
+    sys.path with the cuda_core/ directory (which contains cuda/core/ source
+    that could shadow the installed package).
+    """
+    build_hooks_path = Path(__file__).parent.parent / "build_hooks.py"
+    spec = importlib.util.spec_from_file_location("build_hooks", build_hooks_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+# Load the module once at import time
+build_hooks = _load_build_hooks()
+
+
+def _check_version_detection(
+    cuda_version, expected_major, *, use_cuda_path=True, use_cuda_home=False, cuda_core_build_major=None
+):
+    """Test version detection with a mock cuda.h.
+
+    Args:
+        cuda_version: CUDA_VERSION to write in mock cuda.h (e.g., 12080)
+        expected_major: Expected return value (e.g., "12")
+        use_cuda_path: If True, set CUDA_PATH to the mock headers directory
+        use_cuda_home: If True, set CUDA_HOME to the mock headers directory
+        cuda_core_build_major: If set, override with this CUDA_CORE_BUILD_MAJOR env var
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        include_dir = Path(tmpdir) / "include"
+        include_dir.mkdir()
+        cuda_h = include_dir / "cuda.h"
+        cuda_h.write_text(f"#define CUDA_VERSION {cuda_version}\n")
+
+        build_hooks._get_cuda_paths.cache_clear()
+        build_hooks._determine_cuda_major_version.cache_clear()
+
+        mock_env = {
+            k: v
+            for k, v in {
+                "CUDA_CORE_BUILD_MAJOR": cuda_core_build_major,
+                "CUDA_PATH": tmpdir if use_cuda_path else None,
+                "CUDA_HOME": tmpdir if use_cuda_home else None,
+            }.items()
+            if v is not None
+        }
+
+        with mock.patch.dict(os.environ, mock_env, clear=True):
+            result = build_hooks._determine_cuda_major_version()
+            assert result == expected_major
+
+
+class TestGetCudaMajorVersion:
+    """Tests for _determine_cuda_major_version()."""
+
+    @pytest.mark.parametrize("version", ["11", "12", "13", "14"])
+    def test_env_var_override(self, version):
+        """CUDA_CORE_BUILD_MAJOR env var override works with various versions."""
+        build_hooks._get_cuda_paths.cache_clear()
+        build_hooks._determine_cuda_major_version.cache_clear()
+        with mock.patch.dict(os.environ, {"CUDA_CORE_BUILD_MAJOR": version}, clear=False):
+            result = build_hooks._determine_cuda_major_version()
+            assert result == version
+
+    @pytest.mark.parametrize(
+        ("cuda_version", "expected_major"),
+        [
+            (11000, "11"),  # CUDA 11.0
+            (11080, "11"),  # CUDA 11.8
+            (12000, "12"),  # CUDA 12.0
+            (12020, "12"),  # CUDA 12.2
+            (12080, "12"),  # CUDA 12.8
+            (13000, "13"),  # CUDA 13.0
+            (13010, "13"),  # CUDA 13.1
+        ],
+        ids=["11.0", "11.8", "12.0", "12.2", "12.8", "13.0", "13.1"],
+    )
+    def test_cuda_headers_parsing(self, cuda_version, expected_major):
+        """CUDA_VERSION is correctly parsed from cuda.h headers."""
+        _check_version_detection(cuda_version, expected_major)
+
+    def test_cuda_home_fallback(self):
+        """CUDA_HOME is used if CUDA_PATH is not set."""
+        _check_version_detection(12050, "12", use_cuda_path=False, use_cuda_home=True)
+
+    def test_env_var_takes_priority_over_headers(self):
+        """Env var override takes priority even when headers exist."""
+        _check_version_detection(12080, "11", cuda_core_build_major="11")
+
+    def test_missing_cuda_path_raises_error(self):
+        """RuntimeError is raised when CUDA_PATH/CUDA_HOME not set and no env var override."""
+        build_hooks._get_cuda_paths.cache_clear()
+        build_hooks._determine_cuda_major_version.cache_clear()
+        with (
+            mock.patch.dict(os.environ, {}, clear=True),
+            pytest.raises(RuntimeError, match="CUDA_PATH or CUDA_HOME"),
+        ):
+            build_hooks._determine_cuda_major_version()