| 1 | include(FetchContent) |
| 2 | |
| 3 | if(DEFINED ENV{FLASH_KDA_SRC_DIR}) |
| 4 | set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR}) |
| 5 | endif() |
| 6 | |
| 7 | if(FLASH_KDA_SRC_DIR) |
| 8 | FetchContent_Declare( |
| 9 | flashkda |
| 10 | SOURCE_DIR ${FLASH_KDA_SRC_DIR} |
| 11 | ) |
| 12 | else() |
| 13 | FetchContent_Declare( |
| 14 | flashkda |
| 15 | GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git |
| 16 | GIT_TAG 17a037d98da546deb4591e967cf961a43c034d8b |
| 17 | GIT_PROGRESS TRUE |
| 18 | GIT_SUBMODULES cutlass |
| 19 | ) |
| 20 | endif() |
| 21 | |
| 22 | FetchContent_MakeAvailable(flashkda) |
| 23 | message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}") |
| 24 | |
| 25 | set(FLASH_KDA_SUPPORT_ARCHS) |
| 26 | if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0) |
| 27 | list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a") |
| 28 | endif() |
| 29 | if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) |
| 30 | list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f") |
| 31 | elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) |
| 32 | list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a") |
| 33 | endif() |
| 34 | |
| 35 | cuda_archs_loose_intersection( |
| 36 | FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}") |
| 37 | |
| 38 | if(FLASH_KDA_ARCHS) |
| 39 | message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}") |
| 40 | |
| 41 | set(FLASH_KDA_SOURCES |
| 42 | csrc/flashkda_registration.cpp |
| 43 | ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp |
| 44 | ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu) |
| 45 | set(FLASH_KDA_INCLUDES |
| 46 | ${flashkda_SOURCE_DIR}/csrc |
| 47 | ${flashkda_SOURCE_DIR}/cutlass/include |
| 48 | ${flashkda_SOURCE_DIR}/cutlass/examples/common |
| 49 | ${flashkda_SOURCE_DIR}/cutlass/tools/util/include) |
| 50 | |
| 51 | set_gencode_flags_for_srcs( |
| 52 | SRCS "${FLASH_KDA_SOURCES}" |
| 53 | CUDA_ARCHS "${FLASH_KDA_ARCHS}") |
| 54 | |
| 55 | define_extension_target( |
| 56 | _flashkda_C |
| 57 | DESTINATION vllm |
| 58 | LANGUAGE ${VLLM_GPU_LANG} |
| 59 | SOURCES ${FLASH_KDA_SOURCES} |
| 60 | COMPILE_FLAGS ${VLLM_GPU_FLAGS} |
| 61 | ARCHITECTURES ${VLLM_GPU_ARCHES} |
| 62 | INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES} |
| 63 | USE_SABI 3 |
| 64 | WITH_SOABI) |
| 65 | |
| 66 | target_compile_definitions(_flashkda_C PRIVATE |
| 67 | TORCH_TARGET_VERSION=0x020B000000000000ULL) |
| 68 | if(VLLM_GPU_LANG STREQUAL "CUDA") |
| 69 | target_compile_definitions(_flashkda_C PRIVATE USE_CUDA) |
| 70 | endif() |
| 71 | |
| 72 | target_compile_options(_flashkda_C PRIVATE |
| 73 | $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3> |
| 74 | $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>) |
| 75 | else() |
| 76 | message(STATUS |
| 77 | "FlashKDA will not compile: CUDA >=12.0 and a supported architecture " |
| 78 | "(SM90, SM10x, or SM12x) are required") |
| 79 | add_custom_target(_flashkda_C) |
| 80 | endif() |