Avoids float-casting of CUDA version strings to preemptively prevent CUDA extension install bugs

setup.py

@@ -17,20 +17,34 @@
 import glob
 from setuptools import setup
 
+import re
 import torch
 from torch.utils.cpp_extension import CUDAExtension, BuildExtension
 
 
-def cuda_extension():
-    cuda_version = float(torch.version.cuda)
+def cuda_extension() -> CUDAExtension:
+    # Parse CUDA version string assuming SemVer (handles formats like "12.8-rc.2", "12.10", etc.)
+    if torch.version.cuda:
+        version_pattern = re.match(
+            r"^(\d+)(?:\.(\d+))?(?:\.(\d+))?", torch.version.cuda
+        )
+    else:  # Occurs on CPU-only PyTorch installations
+        raise RuntimeError(f"CUDA version not found: {torch.version.cuda=}")
+    if not version_pattern:
+        raise RuntimeError(f"Unable to parse CUDA version: {torch.version.cuda=}")
+
+    # Extract components, defaulting to 0 if not present
+    cuda_major = int(version_pattern.group(1))
+    cuda_minor = int(version_pattern.group(2) or 0)
+
     nvcc_args = [
         "-gencode=arch=compute_75,code=sm_75",
     ]
-    if cuda_version >= 11:
+    if cuda_major >= 11:
         nvcc_args.append("-gencode=arch=compute_80,code=sm_80")
-    if cuda_version >= 11.1:
+    if (cuda_major == 11 and cuda_minor >= 1) or (cuda_major > 11):
         nvcc_args.append("-gencode=arch=compute_86,code=sm_86")
-    if cuda_version >= 12:
+    if cuda_major >= 12:
         nvcc_args.append("-gencode=arch=compute_90,code=sm_90")
     if cuda_version >= 12.8:
         nvcc_args.append("-gencode=arch=compute_100,code=sm_100")