fix: derive CUDA_CORE_BUILD_MAJOR from headers instead of bindings version

Fixes build failures when cuda-bindings reports major version 13 but
CUDA headers are version 12, causing missing enum errors for
CU_MEM_LOCATION_TYPE_NONE and CU_MEM_ALLOCATION_TYPE_MANAGED.

The new _get_cuda_core_build_major_version() function prioritizes:
1. Explicit CUDA_CORE_BUILD_MAJOR env var (CI override)
2. CUDA_VERSION from cuda.h headers (matches compile target)
3. nvidia-smi driver-reported version (fallback)
4. cuda-bindings major version (last resort)

Adds tests for the version detection logic in test_build_hooks.py.

Author: Andy-Jost

cuda_core/build_hooks.py

@@ -54,6 +54,42 @@ def _get_proper_cuda_bindings_major_version() -> str:
     return "13"
 
 
+@functools.cache
+def _get_cuda_core_build_major_version() -> str:
+    # Explicit overwrite, e.g. in CI.
+    cuda_major = os.environ.get("CUDA_CORE_BUILD_MAJOR")
+    if cuda_major is not None:
+        return cuda_major
+
+    # Try to derive from the CUDA headers (preferred; matches what we compile against).
+    CUDA_PATH = os.environ.get("CUDA_PATH", os.environ.get("CUDA_HOME", None))
+    if CUDA_PATH:
+        for root in CUDA_PATH.split(os.pathsep):
+            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.
+                            return str(v // 1000)
+            except OSError:
+                continue
+
+    # Fall back to driver-reported CUDA version if available.
+    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):
+        pass
+
+    # Last resort: align to cuda-bindings major.
+    return _get_proper_cuda_bindings_major_version()
+
+
 # used later by setup()
 _extensions = None
 
@@ -104,7 +140,7 @@ def get_cuda_paths():
     )
 
     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(_get_cuda_core_build_major_version())}
     compiler_directives = {"embedsignature": True, "warn.deprecated.IF": False, "freethreading_compatible": True}
     if COMPILE_FOR_COVERAGE:
         compiler_directives["linetrace"] = True

cuda_core/tests/test_build_hooks.py

@@ -0,0 +1,115 @@
+# 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 _get_cuda_core_build_major_version() function which
+derives the CUDA major version from headers (preferred) rather than
+cuda-bindings version to avoid mismatches.
+
+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
+"""
+
+import importlib.util
+import os
+import tempfile
+from pathlib import Path
+from unittest import mock
+
+import pytest
+
+
+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_core_build_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._get_cuda_core_build_major_version()
+            assert result == expected_major
+
+
+class TestGetCudaCoreBuildMajorVersion:
+    """Tests for _get_cuda_core_build_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_core_build_major_version.cache_clear()
+        with mock.patch.dict(os.environ, {"CUDA_CORE_BUILD_MAJOR": version}, clear=False):
+            result = build_hooks._get_cuda_core_build_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")