fix: fix cannot import name 'cuda' from 'cuda' in CUDA13 (#1764)

Author: LuYanFCP

flashinfer/comm/mnnvl.py

@@ -25,12 +25,18 @@
 import torch
 
 try:
-    from cuda import cuda
-except ImportError as e:
-    raise ImportError(
-        "Could not import the 'cuda' module. "
-        "Please install cuda-python that matches your CUDA version."
-    ) from e
+    # cuda-python >= 12.9 (has cuda.bindings.driver)
+    from cuda.bindings import driver as cuda
+except ImportError:
+    try:
+        # cuda-python < 12.9 (no cuda.bindings.driver, use cuda as driver)
+        # from cuda import cuda is not available in cuda-python >= 13.0
+        from cuda import cuda
+    except ImportError as e:
+        raise ImportError(
+            "Could not import the 'cuda' module. "
+            "Please install cuda-python that matches your CUDA version."
+        ) from e
 
 from ..cuda_utils import checkCudaErrors
 from .dlpack_utils import create_dlpack_capsule, pack_strided_memory

flashinfer/cuda_utils.py

@@ -19,7 +19,7 @@
 try:
     # Check if cuda.cudart module is available and import accordingly
     if has_cuda_cudart():
-        # cuda-python <= 12.9 (has cuda.cudart)
+        # cuda-python < 12.9 (has cuda.cudart)
         import cuda.bindings.driver as driver
         import cuda.bindings.runtime as runtime
         import cuda.cudart as cudart

flashinfer/cute_dsl/gemm_allreduce_two_shot.py

@@ -2,7 +2,20 @@
 
 import torch
 import torch.distributed as dist
-from cuda import cuda
+
+try:
+    # cuda-python >= 12.9 (has cuda.bindings.driver)
+    from cuda.bindings import driver as cuda
+except ImportError:
+    try:
+        # cuda-python < 12.9 (no cuda.bindings.driver, use cuda as driver)
+        # from cuda import cuda is not available in cuda-python >= 13.0
+        from cuda import cuda
+    except ImportError as e:
+        raise ImportError(
+            "Could not import the 'cuda' module. "
+            "Please install cuda-python that matches your CUDA version."
+        ) from e
 
 import cutlass
 import cutlass.cute as cute

flashinfer/utils.py

@@ -453,6 +453,12 @@ def has_cuda_cudart() -> bool:
     return importlib.util.find_spec("cuda.cudart") is not None
 
 
+def get_cuda_python_version() -> str:
+    import cuda
+
+    return cuda.__version__
+
+
 def is_sm90a_supported(device: torch.device) -> bool:
     major, _ = get_compute_capability(device)
     return major == 9 and version_at_least(torch.version.cuda, "12.3")

tests/test_cute_dsl_gemm_allreduce_two_shot.py

@@ -4,7 +4,20 @@
 import socket
 from typing import Any, Tuple, Type
 
-from cuda import cuda
+try:
+    # cuda-python >= 12.9 (has cuda.bindings.driver)
+    from cuda.bindings import driver as cuda
+except ImportError:
+    try:
+        # cuda-python <= 12.9 (no cuda.bindings.driver, use cuda as driver)
+        # from cuda import cuda is not available in cuda-python >= 13.0
+        from cuda import cuda
+    except ImportError as e:
+        raise ImportError(
+            "Could not import the 'cuda' module. "
+            "Please install cuda-python that matches your CUDA version."
+        ) from e
+
 import cutlass
 import cutlass.cute as cute
 import cutlass.cute.testing as testing