build: added build for blackwell (#459)

Author: guocuimi

CMakeLists.txt

@@ -75,7 +75,7 @@ endif()
 # Build TORCH_CUDA_ARCH_LIST
 set(TORCH_CUDA_ARCH_LIST "")
 foreach(CUDA_ARCH IN LISTS CMAKE_CUDA_ARCHITECTURES)
-  if(CUDA_ARCH MATCHES "^([0-9])([0-9])*$")
+  if(CUDA_ARCH MATCHES "^([1-9][0-9]*)([0-9]a?)$")
     set(TORCH_ARCH "${CMAKE_MATCH_1}.${CMAKE_MATCH_2}")
   elseif(CUDA_ARCH STREQUAL "native")
     set(TORCH_ARCH "Auto")

scripts/build_scalellm.sh

@@ -11,11 +11,11 @@ command -v ccache >/dev/null && ccache -M 25Gi
 command -v ccache >/dev/null && ccache -z
 
 # build
-cmake -G Ninja -S . -B build -DCMAKE_CUDA_ARCHITECTURES="80;89;90"
+cmake -G Ninja -S . -B build -DCMAKE_CUDA_ARCHITECTURES="80;89;90a;100a;120a"
 cmake --build build --target scalellm --config Release -j$(nproc)
 
 # show ccache statistics if ccache is installed
 command -v ccache >/dev/null && ccache -vs
 
 # install
-cmake --install build --prefix ./app
+cmake --install build --prefix ./app

setup.py

@@ -9,8 +9,9 @@
 import sysconfig
 from pathlib import Path
 from typing import List
-from jinja2 import Template
 
+from jinja2 import Template
+from packaging.version import Version
 from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
 
@@ -33,6 +34,17 @@ def get_nccl_root():
         return None
 
 
+def get_cuda_version():
+    import torch.utils.cpp_extension as torch_cpp_ext
+
+    if torch_cpp_ext.CUDA_HOME is None:
+        nvcc = "nvcc"
+    else:
+        nvcc = os.path.join(torch_cpp_ext.CUDA_HOME, "bin/nvcc")
+    txt = subprocess.check_output([nvcc, "--version"], text=True)
+    return Version(re.findall(r"release (\d+\.\d+),", txt)[0])
+
+
 def get_base_dir():
     return os.path.abspath(os.path.dirname(__file__))
 
@@ -55,7 +67,7 @@ def get_scalellm_version():
 
     if not version:
         raise RuntimeError("Unable to find version string.")
-    
+
     version_suffix = os.getenv("SCALELLM_VERSION_SUFFIX")
     if version_suffix:
         version += version_suffix
@@ -158,8 +170,11 @@ def build_extension(self, ext: CMakeExtension):
         debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug
         build_type = "Debug" if debug else "Release"
 
-        # python directories
-        cuda_architectures = "80;89;90"
+        cuda_version = get_cuda_version()
+        cuda_architectures = "80;89;90a"
+        if cuda_version >= Version("12.8"):
+            # blackwell needs cuda 12.8
+            cuda_architectures += ";100a;120a"
         cmake_args = [
             "-G",
             "Ninja",  # Ninja is much faster than make
@@ -247,7 +262,6 @@ def build_extension(self, ext: CMakeExtension):
             "License :: OSI Approved :: Apache Software License",
             "Topic :: Scientific/Engineering",
             "Topic :: Scientific/Engineering :: Artificial Intelligence",
-            
         ],
         packages=["scalellm", "scalellm/serve", "scalellm/_C", "examples"],
         ext_modules=[CMakeExtension("_C", "scalellm/")],