bugfix: Fix cuda-python v13.0 import compatibility (#1455)

<!-- .github/pull_request_template.md -->

## 📌 Description

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

cc:  @yzh119  @paul841029

Author: yongwww

flashinfer/comm/mnnvl.py

@@ -517,7 +517,15 @@ def __init__(
         checkCudaErrors(cuda.cuCtxSetCurrent(primary_ctx))
 
         # Set CUDA device
-        import cuda.cudart as cudart
+        # Check if cuda.cudart is available and import accordingly
+        from flashinfer.utils import has_cuda_cudart
+
+        if has_cuda_cudart():
+            # cuda-python <= 12.9
+            import cuda.cudart as cudart
+        else:
+            # cuda-python >= 13.0
+            import cuda.bindings.runtime as cudart
 
         checkCudaErrors(cudart.cudaSetDevice(device_idx))
 

flashinfer/cuda_utils.py

@@ -14,10 +14,20 @@
 limitations under the License.
 """
 
-import cuda.bindings.driver as driver
-import cuda.bindings.runtime as runtime
-import cuda.cudart as cudart
-import cuda.nvrtc as nvrtc
+from flashinfer.utils import has_cuda_cudart
+
+# Check if cuda.cudart module is available and import accordingly
+if 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
+    import cuda.nvrtc as nvrtc
+else:
+    # cuda-python >= 13.0 (no cuda.cudart, use runtime as cudart)
+    from cuda.bindings import driver, nvrtc, runtime
+
+    cudart = runtime  # Alias runtime as cudart for compatibility
 
 
 def _cudaGetErrorEnum(error):

flashinfer/utils.py

@@ -417,6 +417,18 @@ def version_at_least(version: str, base_version: str) -> bool:
     return pkg_version.parse(version) >= pkg_version.parse(base_version)
 
 
+def has_cuda_cudart() -> bool:
+    """
+    Check if cuda.cudart module is available (cuda-python <= 12.9).
+
+    Returns:
+        True if cuda.cudart exists, False otherwise
+    """
+    import importlib.util
+
+    return importlib.util.find_spec("cuda.cudart") is not None
+
+
 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")