Use cpm and move flash_attention to external

CMakeLists.txt

@@ -116,6 +116,7 @@ list(APPEND CMAKE_MODULE_PATH ${oneflow_cmake_dir})
 
 include(threading)
 include(util)
+include(CPM)
 include(proto2cpp)
 
 if(NOT DEFINED USE_CXX11_ABI)

cmake/CPM.cmake

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: MIT
+#
+# SPDX-FileCopyrightText: Copyright (c) 2019-2023 Lars Melchior and contributors
+
+set(CPM_DOWNLOAD_VERSION 0.39.0)
+set(CPM_HASH_SUM "66639bcac9dd2907b2918de466783554c1334446b9874e90d38e3778d404c2ef")
+
+if(CPM_SOURCE_CACHE)
+  set(CPM_DOWNLOAD_LOCATION "${CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+elseif(DEFINED ENV{CPM_SOURCE_CACHE})
+  set(CPM_DOWNLOAD_LOCATION "$ENV{CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+else()
+  set(CPM_DOWNLOAD_LOCATION "${CMAKE_BINARY_DIR}/cmake/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+endif()
+
+# Expand relative path. This is important if the provided path contains a tilde (~)
+get_filename_component(CPM_DOWNLOAD_LOCATION ${CPM_DOWNLOAD_LOCATION} ABSOLUTE)
+
+set(CPM_CMAKE_URL
+  https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake)
+use_mirror(VARIABLE CPM_CMAKE_URL URL ${CPM_CMAKE_URL})
+
+file(DOWNLOAD ${CPM_CMAKE_URL} ${CPM_DOWNLOAD_LOCATION} EXPECTED_HASH SHA256=${CPM_HASH_SUM})
+
+include(${CPM_DOWNLOAD_LOCATION})

cmake/oneflow.cmake

@@ -632,11 +632,11 @@ if(BUILD_CPP_API)
     checkdirandappendslash(DIR ${TRT_FLASH_ATTENTION_LIBRARY_DIR} OUTPUT
                            TRT_FLASH_ATTENTION_LIBRARY_DIR_APPENDED)
     list(APPEND LIBONEFLOW_THIRD_PARTY_DIRS ${TRT_FLASH_ATTENTION_LIBRARY_DIR_APPENDED})
-    if(CUDA_VERSION VERSION_GREATER_EQUAL "11.7")
-      checkdirandappendslash(DIR ${FLASH_ATTENTION_LIBRARY_DIR} OUTPUT
-                             FLASH_ATTENTION_LIBRARY_DIR_APPENDED)
-      list(APPEND LIBONEFLOW_THIRD_PARTY_DIRS ${FLASH_ATTENTION_LIBRARY_DIR_APPENDED})
-    endif()
+    # if(CUDA_VERSION VERSION_GREATER_EQUAL "11.7")
+    #   checkdirandappendslash(DIR ${FLASH_ATTENTION_LIBRARY_DIR} OUTPUT
+    #                          FLASH_ATTENTION_LIBRARY_DIR_APPENDED)
+    #   list(APPEND LIBONEFLOW_THIRD_PARTY_DIRS ${FLASH_ATTENTION_LIBRARY_DIR_APPENDED})
+    # endif()
     if(WITH_CUTLASS)
       checkdirandappendslash(DIR ${CUTLASS_LIBRARY_DIR} OUTPUT CUTLASS_LIBRARY_DIR_APPENDED)
       list(APPEND LIBONEFLOW_THIRD_PARTY_DIRS ${CUTLASS_LIBRARY_DIR_APPENDED})

cmake/third_party.cmake

@@ -145,9 +145,9 @@ if(BUILD_CUDA)
   include(nccl)
   include(cutlass)
   include(trt_flash_attention)
-  if(CUDA_VERSION VERSION_GREATER_EQUAL "11.7")
-    include(flash_attention)
-  endif()
+  # if(CUDA_VERSION VERSION_GREATER_EQUAL "11.7")
+  #   include(flash_attention)
+  # endif()
 
   list(APPEND oneflow_third_party_libs ${NCCL_LIBRARIES})
   list(APPEND oneflow_third_party_libs ${CUDNN_LIBRARIES})
@@ -167,11 +167,11 @@ if(BUILD_CUDA)
   list(APPEND oneflow_third_party_dependencies trt_flash_attention)
   list(APPEND oneflow_third_party_libs ${TRT_FLASH_ATTENTION_LIBRARIES})
   list(APPEND ONEFLOW_THIRD_PARTY_INCLUDE_DIRS ${TRT_FLASH_ATTENTION_INCLUDE_DIR})
-  if(CUDA_VERSION VERSION_GREATER_EQUAL "11.7")
-    list(APPEND oneflow_third_party_dependencies flash_attention)
-    list(APPEND oneflow_third_party_libs ${FLASH_ATTENTION_LIBRARIES})
-    list(APPEND ONEFLOW_THIRD_PARTY_INCLUDE_DIRS ${FLASH_ATTENTION_INCLUDE_DIR})
-  endif()
+  # if(CUDA_VERSION VERSION_GREATER_EQUAL "11.7")
+  #   list(APPEND oneflow_third_party_dependencies flash_attention)
+  #   list(APPEND oneflow_third_party_libs ${FLASH_ATTENTION_LIBRARIES})
+  #   list(APPEND ONEFLOW_THIRD_PARTY_INCLUDE_DIRS ${FLASH_ATTENTION_INCLUDE_DIR})
+  # endif()
 endif()
 
 if(BUILD_RDMA)

external/CMakeLists.txt

@@ -20,8 +20,8 @@ set(KINETO_MD5 f9b550591b3899fb267270c19484933f)
 set(EXTERNAL_TARGETS)
 
 if(WITH_TBB) # set(WITH_${threading_runtime_item} ON) in threading.cmake
-  add_subdirectory(onetbb)
-  list(APPEND EXTERNAL_TARGETS tbb)
+    add_subdirectory(onetbb)
+    list(APPEND EXTERNAL_TARGETS tbb)
 endif()
 
 add_subdirectory(robin-hood-hashing)
@@ -33,6 +33,11 @@ list(APPEND EXTERNAL_TARGETS fmt)
 add_subdirectory(kineto)
 list(APPEND EXTERNAL_TARGETS kineto)
 
+if(CUDA_VERSION VERSION_GREATER_EQUAL "11.7")
+    add_subdirectory(flash_attention)
+    list(APPEND EXTERNAL_TARGETS flash_attention)
+endif()
+
 mark_targets_as_system(${EXTERNAL_TARGETS})
 
 set_property(GLOBAL PROPERTY EXTERNAL_TARGETS ${EXTERNAL_TARGETS})

external/flash_attention/CMakeLists.txt

@@ -0,0 +1,9 @@
+set(FLASH_ATTENTION_URL https://github.com/Oneflow-Inc/flash-attention-v2/archive/eed2e82b880e06237af3e50ceac4cf6728b15645.zip)
+use_mirror(VARIABLE FLASH_ATTENTION_URL URL ${FLASH_ATTENTION_URL})
+set(FLASH_ATTENTION_MD5 cff17aa2a95ad754c9889bb3a03d5c51)
+
+CPMAddPackage(
+    NAME flash_attention
+    URL ${FLASH_ATTENTION_URL}
+    URL_HASH MD5=${FLASH_ATTENTION_MD5}
+)