[CUDA] fix setting of CUDA architectures and enable support for NVIDI…

…A Blackwell (#6812)
microsoft · Feb 2, 2025 · c9de57b · c9de57b
1 parent f2b959c
commit c9de57b
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 10 deletions.
diff --git a/.ci/check-python-dists.sh b/.ci/check-python-dists.sh
@@ -32,7 +32,7 @@ if [ "$PY_MINOR_VER" -gt 7 ]; then
             --inspect \
             --ignore 'compiled-objects-have-debug-symbols'\
             --ignore 'distro-too-large-compressed' \
-            --max-allowed-size-uncompressed '100M' \
+            --max-allowed-size-uncompressed '120M' \
             --max-allowed-files 800 \
             "$(echo "${DIST_DIR}"/*)" || exit 1
     elif { test "$(uname -m)" = "aarch64"; }; then

diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml
@@ -86,7 +86,7 @@ jobs:
           - method: wheel
             compiler: gcc
             python_version: "3.11"
-            cuda_version: "12.6.1"
+            cuda_version: "12.8.0"
             linux_version: "ubuntu22.04"
             task: cuda
           - method: source

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -224,21 +224,28 @@ if(USE_CUDA)
     # reference for mapping of CUDA toolkit component versions to supported architectures ("compute capabilities"):
     # https://en.wikipedia.org/wiki/CUDA#GPUs_supported
     set(CUDA_ARCHS "60" "61" "62" "70" "75")
-    if(CUDA_VERSION VERSION_GREATER_EQUAL "110")
+    if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.0")
         list(APPEND CUDA_ARCHS "80")
     endif()
-    if(CUDA_VERSION VERSION_GREATER_EQUAL "111")
+    if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.1")
         list(APPEND CUDA_ARCHS "86")
     endif()
-    if(CUDA_VERSION VERSION_GREATER_EQUAL "115")
+    if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.5")
         list(APPEND CUDA_ARCHS "87")
     endif()
-    if(CUDA_VERSION VERSION_GREATER_EQUAL "118")
+    if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.8")
         list(APPEND CUDA_ARCHS "89")
         list(APPEND CUDA_ARCHS "90")
     endif()
+    if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.8")
+        list(APPEND CUDA_ARCHS "100")
+        list(APPEND CUDA_ARCHS "120")
+    endif()
+    # Generate PTX for the most recent architecture for forwards compatibility
     list(POP_BACK CUDA_ARCHS CUDA_LAST_SUPPORTED_ARCH)
-    list(APPEND CUDA_ARCHS "${CUDA_LAST_SUPPORTED_ARCH}+PTX")
+    list(TRANSFORM CUDA_ARCHS APPEND "-real")
+    list(APPEND CUDA_ARCHS "${CUDA_LAST_SUPPORTED_ARCH}-real" "${CUDA_LAST_SUPPORTED_ARCH}-virtual")
+    message(STATUS "CUDA_ARCHITECTURES: ${CUDA_ARCHS}")
     if(USE_DEBUG)
       set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -g")
     else()
@@ -567,22 +574,22 @@ if(USE_CUDA)
   set_target_properties(
     lightgbm_objs
     PROPERTIES
-      CUDA_ARCHITECTURES ${CUDA_ARCHS}
+      CUDA_ARCHITECTURES "${CUDA_ARCHS}"
       CUDA_SEPARABLE_COMPILATION ON
   )
 
   set_target_properties(
     _lightgbm
     PROPERTIES
-      CUDA_ARCHITECTURES ${CUDA_ARCHS}
+      CUDA_ARCHITECTURES "${CUDA_ARCHS}"
       CUDA_RESOLVE_DEVICE_SYMBOLS ON
   )
 
   if(BUILD_CLI)
     set_target_properties(
       lightgbm
       PROPERTIES
-        CUDA_ARCHITECTURES ${CUDA_ARCHS}
+        CUDA_ARCHITECTURES "${CUDA_ARCHS}"
         CUDA_RESOLVE_DEVICE_SYMBOLS ON
     )
   endif()