| 1 | # vLLM flash attention requires VLLM_GPU_ARCHES to contain the set of target |
| 2 | # arches in the CMake syntax (75-real, 89-virtual, etc), since we clear the |
| 3 | # arches in the CUDA case (and instead set the gencodes on a per file basis) |
| 4 | # we need to manually set VLLM_GPU_ARCHES here. |
| 5 | if(VLLM_GPU_LANG STREQUAL "CUDA") |
| 6 | foreach(_ARCH ${CUDA_ARCHS}) |
| 7 | string(REPLACE "." "" _ARCH "${_ARCH}") |
| 8 | list(APPEND VLLM_GPU_ARCHES "${_ARCH}-real") |
| 9 | endforeach() |
| 10 | endif() |
| 11 | |
| 12 | # |
| 13 | # Build vLLM flash attention from source |
| 14 | # |
| 15 | # IMPORTANT: This has to be the last thing we do, because vllm-flash-attn uses the same macros/functions as vLLM. |
| 16 | # Because functions all belong to the global scope, vllm-flash-attn's functions overwrite vLLMs. |
| 17 | # They should be identical but if they aren't, this is a massive footgun. |
| 18 | # |
| 19 | # The vllm-flash-attn install rules are nested under vllm to make sure the library gets installed in the correct place. |
| 20 | # To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2), --component _vllm_fa3_C (for FA3), |
| 21 | # or --component _vllm_fa4_cutedsl_C (for FA4 CuteDSL Python files). |
| 22 | # If no component is specified, vllm-flash-attn is still installed. |
| 23 | |
| 24 | # If VLLM_FLASH_ATTN_SRC_DIR is set, vllm-flash-attn is installed from that directory instead of downloading. |
| 25 | # This is to enable local development of vllm-flash-attn within vLLM. |
| 26 | # It can be set as an environment variable or passed as a cmake argument. |
| 27 | # The environment variable takes precedence. |
| 28 | if (DEFINED ENV{VLLM_FLASH_ATTN_SRC_DIR}) |
| 29 | set(VLLM_FLASH_ATTN_SRC_DIR $ENV{VLLM_FLASH_ATTN_SRC_DIR}) |
| 30 | endif() |
| 31 | |
| 32 | if(VLLM_FLASH_ATTN_SRC_DIR) |
| 33 | FetchContent_Declare( |
| 34 | vllm-flash-attn SOURCE_DIR |
| 35 | ${VLLM_FLASH_ATTN_SRC_DIR} |
| 36 | BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn |
| 37 | ) |
| 38 | else() |
| 39 | FetchContent_Declare( |
| 40 | vllm-flash-attn |
| 41 | GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git |
| 42 | GIT_TAG 9cd61de38763d712bb6ce56e2a02cc2bf718c89f |
| 43 | GIT_PROGRESS TRUE |
| 44 | # Don't share the vllm-flash-attn build between build types |
| 45 | BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn |
| 46 | ) |
| 47 | endif() |
| 48 | |
| 49 | # Make sure vllm-flash-attn install rules are nested under vllm/ |
| 50 | # ALL_COMPONENTS ensures the save/modify/restore runs exactly once regardless |
| 51 | # of how many components are being installed, avoiding double-append of /vllm/. |
| 52 | install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" ALL_COMPONENTS) |
| 53 | install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS) |
| 54 | install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" ALL_COMPONENTS) |
| 55 | |
| 56 | # Fetch the vllm-flash-attn library |
| 57 | FetchContent_MakeAvailable(vllm-flash-attn) |
| 58 | message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}") |
| 59 | |
| 60 | # Restore the install prefix after FA's install rules |
| 61 | install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS) |
| 62 | install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) |
| 63 | |
| 64 | # Install shared Python files for both FA2 and FA3 components |
| 65 | foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C) |
| 66 | # Ensure the vllm/vllm_flash_attn directory exists before installation |
| 67 | install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")" |
| 68 | COMPONENT ${_FA_COMPONENT}) |
| 69 | |
| 70 | # Copy vllm_flash_attn python files (except __init__.py and flash_attn_interface.py |
| 71 | # which are source-controlled in vllm) |
| 72 | install( |
| 73 | DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/ |
| 74 | DESTINATION vllm/vllm_flash_attn |
| 75 | COMPONENT ${_FA_COMPONENT} |
| 76 | FILES_MATCHING PATTERN "*.py" |
| 77 | PATTERN "__init__.py" EXCLUDE |
| 78 | PATTERN "flash_attn_interface.py" EXCLUDE |
| 79 | ) |
| 80 | |
| 81 | endforeach() |
| 82 | |
| 83 | # |
| 84 | # FA4 CuteDSL component |
| 85 | # This is a Python-only component that copies the flash_attn/cute directory |
| 86 | # and transforms imports to match our package structure. |
| 87 | # |
| 88 | add_custom_target(_vllm_fa4_cutedsl_C) |
| 89 | |
| 90 | # Install flash_attn/cute directory (needed for FA4). |
| 91 | # When using a local source dir (VLLM_FLASH_ATTN_SRC_DIR), create a symlink |
| 92 | # so edits to cute-dsl Python files take effect immediately without rebuilding. |
| 93 | # Otherwise, copy files and transform flash_attn.cute imports to |
| 94 | # vllm.vllm_flash_attn.cute to match our package structure. |
| 95 | if(VLLM_FLASH_ATTN_SRC_DIR) |
| 96 | install(CODE " |
| 97 | set(LINK_TARGET \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\") |
| 98 | set(LINK_NAME \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute\") |
| 99 | file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\") |
| 100 | file(REMOVE_RECURSE \"\${LINK_NAME}\") |
| 101 | file(CREATE_LINK \"\${LINK_TARGET}\" \"\${LINK_NAME}\" SYMBOLIC) |
| 102 | " COMPONENT _vllm_fa4_cutedsl_C) |
| 103 | else() |
| 104 | install(CODE " |
| 105 | file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\") |
| 106 | foreach(SRC_FILE \${CUTE_PY_FILES}) |
| 107 | file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE}) |
| 108 | set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\") |
| 109 | get_filename_component(DST_DIR \${DST_FILE} DIRECTORY) |
| 110 | file(MAKE_DIRECTORY \${DST_DIR}) |
| 111 | file(READ \${SRC_FILE} FILE_CONTENTS) |
| 112 | string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\") |
| 113 | file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\") |
| 114 | endforeach() |
| 115 | " COMPONENT _vllm_fa4_cutedsl_C) |
| 116 | endif() |