Skip to content

Commit c9de57b

Browse files
authored
[CUDA] fix setting of CUDA architectures and enable support for NVIDIA Blackwell (#6812)
1 parent f2b959c commit c9de57b

File tree

3 files changed

+17
-10
lines changed

3 files changed

+17
-10
lines changed

.ci/check-python-dists.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ if [ "$PY_MINOR_VER" -gt 7 ]; then
3232
--inspect \
3333
--ignore 'compiled-objects-have-debug-symbols'\
3434
--ignore 'distro-too-large-compressed' \
35-
--max-allowed-size-uncompressed '100M' \
35+
--max-allowed-size-uncompressed '120M' \
3636
--max-allowed-files 800 \
3737
"$(echo "${DIST_DIR}"/*)" || exit 1
3838
elif { test "$(uname -m)" = "aarch64"; }; then

.github/workflows/cuda.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ jobs:
8686
- method: wheel
8787
compiler: gcc
8888
python_version: "3.11"
89-
cuda_version: "12.6.1"
89+
cuda_version: "12.8.0"
9090
linux_version: "ubuntu22.04"
9191
task: cuda
9292
- method: source

CMakeLists.txt

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -224,21 +224,28 @@ if(USE_CUDA)
224224
# reference for mapping of CUDA toolkit component versions to supported architectures ("compute capabilities"):
225225
# https://en.wikipedia.org/wiki/CUDA#GPUs_supported
226226
set(CUDA_ARCHS "60" "61" "62" "70" "75")
227-
if(CUDA_VERSION VERSION_GREATER_EQUAL "110")
227+
if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.0")
228228
list(APPEND CUDA_ARCHS "80")
229229
endif()
230-
if(CUDA_VERSION VERSION_GREATER_EQUAL "111")
230+
if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.1")
231231
list(APPEND CUDA_ARCHS "86")
232232
endif()
233-
if(CUDA_VERSION VERSION_GREATER_EQUAL "115")
233+
if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.5")
234234
list(APPEND CUDA_ARCHS "87")
235235
endif()
236-
if(CUDA_VERSION VERSION_GREATER_EQUAL "118")
236+
if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.8")
237237
list(APPEND CUDA_ARCHS "89")
238238
list(APPEND CUDA_ARCHS "90")
239239
endif()
240+
if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.8")
241+
list(APPEND CUDA_ARCHS "100")
242+
list(APPEND CUDA_ARCHS "120")
243+
endif()
244+
# Generate PTX for the most recent architecture for forwards compatibility
240245
list(POP_BACK CUDA_ARCHS CUDA_LAST_SUPPORTED_ARCH)
241-
list(APPEND CUDA_ARCHS "${CUDA_LAST_SUPPORTED_ARCH}+PTX")
246+
list(TRANSFORM CUDA_ARCHS APPEND "-real")
247+
list(APPEND CUDA_ARCHS "${CUDA_LAST_SUPPORTED_ARCH}-real" "${CUDA_LAST_SUPPORTED_ARCH}-virtual")
248+
message(STATUS "CUDA_ARCHITECTURES: ${CUDA_ARCHS}")
242249
if(USE_DEBUG)
243250
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -g")
244251
else()
@@ -567,22 +574,22 @@ if(USE_CUDA)
567574
set_target_properties(
568575
lightgbm_objs
569576
PROPERTIES
570-
CUDA_ARCHITECTURES ${CUDA_ARCHS}
577+
CUDA_ARCHITECTURES "${CUDA_ARCHS}"
571578
CUDA_SEPARABLE_COMPILATION ON
572579
)
573580

574581
set_target_properties(
575582
_lightgbm
576583
PROPERTIES
577-
CUDA_ARCHITECTURES ${CUDA_ARCHS}
584+
CUDA_ARCHITECTURES "${CUDA_ARCHS}"
578585
CUDA_RESOLVE_DEVICE_SYMBOLS ON
579586
)
580587

581588
if(BUILD_CLI)
582589
set_target_properties(
583590
lightgbm
584591
PROPERTIES
585-
CUDA_ARCHITECTURES ${CUDA_ARCHS}
592+
CUDA_ARCHITECTURES "${CUDA_ARCHS}"
586593
CUDA_RESOLVE_DEVICE_SYMBOLS ON
587594
)
588595
endif()

0 commit comments

Comments
 (0)