Add sm89 by default for rtx 4090

li-plus · Mar 16, 2024 · 1ad2c59 · 1ad2c59
1 parent f7a2457
commit 1ad2c59
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -26,11 +26,18 @@ add_subdirectory(third_party/sentencepiece)
 if (GGML_CUBLAS)
     add_compile_definitions(GGML_USE_CUBLAS)
     enable_language(CUDA)
+    # ref: https://stackoverflow.com/questions/28932864/which-compute-capability-is-supported-by-which-cuda-versions
+    set(CUDA_ARCH_LIST "52;61;70;75")
+    if (CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "11.0")
+        set(CUDA_ARCH_LIST "${CUDA_ARCH_LIST};80")
+    endif ()
+    if (CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "11.1")
+        set(CUDA_ARCH_LIST "${CUDA_ARCH_LIST};86")
+    endif ()
     if (CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "11.8")
-        set(CUDA_ARCHITECTURES "52;61;70;75;80;86;90" CACHE STRING "chatglm: cuda architectures to compile")
-    else ()
-        set(CUDA_ARCHITECTURES "52;61;70;75;80;86" CACHE STRING "chatglm: cuda architectures to compile")
+        set(CUDA_ARCH_LIST "${CUDA_ARCH_LIST};89;90")
     endif ()
+    set(CUDA_ARCHITECTURES ${CUDA_ARCH_LIST} CACHE STRING "chatglm: cuda architectures to compile")
     set_property(TARGET ggml PROPERTY CUDA_ARCHITECTURES ${CUDA_ARCHITECTURES})
 endif ()