fix binding version check

Author: leofang

cuda_core/cuda/core/experimental/_launcher.py

@@ -11,7 +11,7 @@
 from cuda.core.experimental._kernel_arg_handler import ParamHolder
 from cuda.core.experimental._module import Kernel
 from cuda.core.experimental._stream import Stream
-from cuda.core.experimental._utils import CUDAError, check_or_create_options, handle_return
+from cuda.core.experimental._utils import CUDAError, check_or_create_options, get_binding_version, handle_return
 
 # TODO: revisit this treatment for py313t builds
 _inited = False
@@ -25,7 +25,7 @@ def _lazy_init():
 
     global _use_ex
     # binding availability depends on cuda-python version
-    _py_major_minor = tuple(int(v) for v in (importlib.metadata.version("cuda-python").split(".")[:2]))
+    _py_major_minor = get_binding_version()
     _driver_ver = handle_return(cuda.cuDriverGetVersion())
     _use_ex = (_driver_ver >= 11080) and (_py_major_minor >= (11, 8))
     _inited = True

cuda_core/cuda/core/experimental/_module.py

@@ -5,7 +5,7 @@
 import importlib.metadata
 
 from cuda import cuda
-from cuda.core.experimental._utils import handle_return, precondition
+from cuda.core.experimental._utils import get_binding_version, handle_return, precondition
 
 _backend = {
     "old": {
@@ -30,7 +30,7 @@ def _lazy_init():
 
     global _py_major_ver, _driver_ver, _kernel_ctypes
     # binding availability depends on cuda-python version
-    _py_major_ver = int(importlib.metadata.version("cuda-python").split(".")[0])
+    _py_major_ver, _ = get_binding_version()
     if _py_major_ver >= 12:
         _backend["new"] = {
             "file": cuda.cuLibraryLoadFromFile,

cuda_core/cuda/core/experimental/_utils.py

@@ -134,3 +134,11 @@ def get_device_from_ctx(ctx_handle) -> int:
         assert ctx_handle == handle_return(cuda.cuCtxPopCurrent())
         handle_return(cuda.cuCtxPushCurrent(prev_ctx))
     return device_id
+
+
+def get_binding_version():
+    try:
+        major_minor = importlib.metadata.version("cuda-bindings").split(".")[:2]
+    except importlib.metadata.PackageNotFoundError:
+        major_minor = importlib.metadata.version("cuda-python").split(".")[:2]
+    return tuple(int(v) for v in major_minor)