chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,218 @@
|
||||
include(FetchContent)
|
||||
|
||||
# If DEEPGEMM_SRC_DIR is set, DeepGEMM is built from that directory
|
||||
# instead of downloading.
|
||||
# It can be set as an environment variable or passed as a cmake argument.
|
||||
# The environment variable takes precedence.
|
||||
if (DEFINED ENV{DEEPGEMM_SRC_DIR})
|
||||
set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR})
|
||||
endif()
|
||||
|
||||
# Local tree: set deepgemm_SOURCE_DIR directly (no FetchContent download).
|
||||
# Upstream git: use FetchContent_Populate with explicit options (CMP0169 NEW
|
||||
# disallows one-argument Populate(dep) after Declare; MakeAvailable would run
|
||||
# DeepGEMM's top-level CMakeLists.txt, which vLLM must not load).
|
||||
if(DEEPGEMM_SRC_DIR)
|
||||
# cmake_path(ABSOLUTE_PATH <var> ...) reads the path from <var>; NORMALIZE is a
|
||||
# flag (no trailing path argument). Resolve relative paths against vLLM root.
|
||||
set(_deepgemm_user_src "${DEEPGEMM_SRC_DIR}")
|
||||
cmake_path(ABSOLUTE_PATH _deepgemm_user_src
|
||||
BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
|
||||
NORMALIZE)
|
||||
set(DEEPGEMM_SRC_DIR "${_deepgemm_user_src}")
|
||||
if(NOT IS_DIRECTORY "${DEEPGEMM_SRC_DIR}")
|
||||
message(FATAL_ERROR
|
||||
"DEEPGEMM_SRC_DIR is not an existing directory: '${DEEPGEMM_SRC_DIR}'")
|
||||
endif()
|
||||
set(deepgemm_SOURCE_DIR "${DEEPGEMM_SRC_DIR}")
|
||||
message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
|
||||
else()
|
||||
# Keep in sync with tools/install_deepgemm.sh
|
||||
set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
|
||||
# NOTE: This is currently targeting nv-dev branch due to sm120 support
|
||||
set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
|
||||
|
||||
set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
|
||||
if(NOT _deepgemm_fc_root)
|
||||
set(_deepgemm_fc_root "${CMAKE_BINARY_DIR}/_deps")
|
||||
endif()
|
||||
set(_deepgemm_src "${_deepgemm_fc_root}/deepgemm-src")
|
||||
set(_deepgemm_bin "${_deepgemm_fc_root}/deepgemm-build")
|
||||
set(_deepgemm_sub "${_deepgemm_fc_root}/deepgemm-subbuild")
|
||||
|
||||
if(EXISTS "${_deepgemm_src}/csrc/python_api.cpp")
|
||||
set(deepgemm_SOURCE_DIR "${_deepgemm_src}")
|
||||
set(deepgemm_BINARY_DIR "${_deepgemm_bin}")
|
||||
else()
|
||||
FetchContent_Populate(
|
||||
deepgemm
|
||||
SUBBUILD_DIR "${_deepgemm_sub}"
|
||||
SOURCE_DIR "${_deepgemm_src}"
|
||||
BINARY_DIR "${_deepgemm_bin}"
|
||||
GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}"
|
||||
GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}"
|
||||
GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
|
||||
GIT_PROGRESS TRUE
|
||||
)
|
||||
endif()
|
||||
message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
|
||||
endif()
|
||||
|
||||
# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 (official upstream),
|
||||
# and 12.8+ for SM120 / SM12x. CUDA 13+ can use the family-specific SM12x
|
||||
# arch; CUDA 12.x builds the arch-specific SM120/SM121 variants.
|
||||
set(DEEPGEMM_SUPPORT_ARCHS)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3)
|
||||
list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
|
||||
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
|
||||
else()
|
||||
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0f")
|
||||
else()
|
||||
list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0a" "12.1a")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
cuda_archs_loose_intersection(DEEPGEMM_ARCHS
|
||||
"${DEEPGEMM_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
|
||||
|
||||
if(DEEPGEMM_ARCHS)
|
||||
message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")
|
||||
|
||||
#
|
||||
# DeepGEMM integration notes
|
||||
# --------------------------
|
||||
# We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle a
|
||||
# `_C.cpython-X.Y-*.so` for every CPython in `requires-python`. The
|
||||
# per-Python build is delegated to tools/build_deepgemm_C.py.
|
||||
#
|
||||
# Why per-Python: DeepGEMM's binding uses PYBIND11_MODULE, which links
|
||||
# private CPython symbols — a single `_C.abi3.so` is not viable today
|
||||
# (see #41476 / #41512 for the failed attempt).
|
||||
#
|
||||
# TODOs (tracked in vllm-project/vllm#42431):
|
||||
# - Replace DeepGEMM's pybind11 binding with a TORCH_LIBRARY + shim
|
||||
# binding (cf. vllm-flash-attention/csrc/common/pytorch_shim.h) to
|
||||
# collapse to one `_C.abi3.so`. Needs either an upstream change or
|
||||
# a maintained binding fork in vLLM.
|
||||
# - AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop
|
||||
# the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime
|
||||
# requirement.
|
||||
#
|
||||
|
||||
# DEEPGEMM_PYTHON_INTERPRETERS: ":"-separated target Python paths.
|
||||
# Empty/unset → fall back to the build interpreter (editable installs).
|
||||
# (Empty-but-set env vars test as DEFINED in cmake — treat as unset.)
|
||||
if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "")
|
||||
string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}")
|
||||
else()
|
||||
set(_dg_pythons "${Python_EXECUTABLE}")
|
||||
endif()
|
||||
message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")
|
||||
|
||||
# add_custom_command does no implicit header scanning; glob explicitly so
|
||||
# header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild.
|
||||
file(GLOB_RECURSE _dg_headers
|
||||
"${deepgemm_SOURCE_DIR}/csrc/*.h"
|
||||
"${deepgemm_SOURCE_DIR}/csrc/*.hpp"
|
||||
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h"
|
||||
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp"
|
||||
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh")
|
||||
|
||||
set(_dg_markers)
|
||||
set(_dg_seen_soabis)
|
||||
foreach(_pybin IN LISTS _dg_pythons)
|
||||
execute_process(
|
||||
COMMAND "${_pybin}" -c
|
||||
"import sysconfig; print(sysconfig.get_config_var('SOABI'))"
|
||||
OUTPUT_VARIABLE _dg_soabi
|
||||
OUTPUT_STRIP_TRAILING_WHITESPACE
|
||||
COMMAND_ERROR_IS_FATAL ANY)
|
||||
# Dedup interpreters that resolve to the same CPython.
|
||||
if(_dg_soabi IN_LIST _dg_seen_soabis)
|
||||
continue()
|
||||
endif()
|
||||
list(APPEND _dg_seen_soabis "${_dg_soabi}")
|
||||
set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C_${_dg_soabi}")
|
||||
set(_dg_marker "${_dg_dir}/.built")
|
||||
add_custom_command(
|
||||
OUTPUT "${_dg_marker}"
|
||||
COMMAND "${Python_EXECUTABLE}"
|
||||
"${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
|
||||
"${deepgemm_SOURCE_DIR}" "${_dg_dir}" "${_pybin}"
|
||||
COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}"
|
||||
DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
|
||||
"${deepgemm_SOURCE_DIR}/csrc/python_api.cpp"
|
||||
${_dg_headers}
|
||||
COMMENT "Building DeepGEMM _C for ${_pybin}"
|
||||
VERBATIM)
|
||||
list(APPEND _dg_markers "${_dg_marker}")
|
||||
install(DIRECTORY "${_dg_dir}/"
|
||||
DESTINATION vllm/third_party/deep_gemm
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "_C.cpython-*.so")
|
||||
endforeach()
|
||||
add_custom_target(_deep_gemm_C ALL DEPENDS ${_dg_markers})
|
||||
|
||||
#
|
||||
# Vendor DeepGEMM Python package files
|
||||
#
|
||||
install(FILES
|
||||
"${deepgemm_SOURCE_DIR}/deep_gemm/__init__.py"
|
||||
DESTINATION vllm/third_party/deep_gemm
|
||||
COMPONENT _deep_gemm_C)
|
||||
|
||||
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/utils/"
|
||||
DESTINATION vllm/third_party/deep_gemm/utils
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "*.py")
|
||||
|
||||
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/testing/"
|
||||
DESTINATION vllm/third_party/deep_gemm/testing
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "*.py")
|
||||
|
||||
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/legacy/"
|
||||
DESTINATION vllm/third_party/deep_gemm/legacy
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "*.py")
|
||||
|
||||
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/"
|
||||
DESTINATION vllm/third_party/deep_gemm/mega
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "*.py")
|
||||
|
||||
# Generate envs.py (normally generated by DeepGEMM's setup.py build step)
|
||||
file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
|
||||
"# Pre-installed environment variables\npersistent_envs = dict()\n")
|
||||
install(FILES "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
|
||||
DESTINATION vllm/third_party/deep_gemm
|
||||
RENAME envs.py
|
||||
COMPONENT _deep_gemm_C)
|
||||
|
||||
#
|
||||
# Install include files needed for JIT compilation at runtime.
|
||||
# The JIT compiler finds these relative to the package directory.
|
||||
#
|
||||
|
||||
# DeepGEMM's own CUDA headers
|
||||
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/include/"
|
||||
DESTINATION vllm/third_party/deep_gemm/include
|
||||
COMPONENT _deep_gemm_C)
|
||||
|
||||
# CUTLASS and CuTe headers (vendored for JIT, separate from vLLM's CUTLASS)
|
||||
install(DIRECTORY "${deepgemm_SOURCE_DIR}/third-party/cutlass/include/"
|
||||
DESTINATION vllm/third_party/deep_gemm/include
|
||||
COMPONENT _deep_gemm_C)
|
||||
|
||||
else()
|
||||
message(STATUS "DeepGEMM will not compile: "
|
||||
"unsupported CUDA architecture ${CUDA_ARCHS}")
|
||||
# Create empty target so setup.py doesn't fail on unsupported systems
|
||||
add_custom_target(_deep_gemm_C)
|
||||
endif()
|
||||
@@ -0,0 +1,186 @@
|
||||
include(FetchContent)
|
||||
|
||||
# If FLASH_MLA_SRC_DIR is set, flash-mla is installed from that directory
|
||||
# instead of downloading.
|
||||
# It can be set as an environment variable or passed as a cmake argument.
|
||||
# The environment variable takes precedence.
|
||||
if (DEFINED ENV{FLASH_MLA_SRC_DIR})
|
||||
set(FLASH_MLA_SRC_DIR $ENV{FLASH_MLA_SRC_DIR})
|
||||
endif()
|
||||
|
||||
if(FLASH_MLA_SRC_DIR)
|
||||
FetchContent_Declare(
|
||||
flashmla
|
||||
SOURCE_DIR ${FLASH_MLA_SRC_DIR}
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
)
|
||||
else()
|
||||
FetchContent_Declare(
|
||||
flashmla
|
||||
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
|
||||
GIT_TAG b70aff3d110a2b1a037e62eac295166b5143643a
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
)
|
||||
endif()
|
||||
|
||||
|
||||
FetchContent_MakeAvailable(flashmla)
|
||||
message(STATUS "FlashMLA is available at ${flashmla_SOURCE_DIR}")
|
||||
|
||||
# Vendor FlashMLA interface into vLLM with torch-ops shim.
|
||||
set(FLASHMLA_VENDOR_DIR "${CMAKE_SOURCE_DIR}/vllm/third_party/flashmla")
|
||||
file(MAKE_DIRECTORY "${FLASHMLA_VENDOR_DIR}")
|
||||
file(READ "${flashmla_SOURCE_DIR}/flash_mla/flash_mla_interface.py"
|
||||
FLASHMLA_INTERFACE_CONTENT)
|
||||
string(REPLACE "import flash_mla.cuda as flash_mla_cuda"
|
||||
"import vllm._flashmla_C\nflash_mla_cuda = torch.ops._flashmla_C"
|
||||
FLASHMLA_INTERFACE_CONTENT
|
||||
"${FLASHMLA_INTERFACE_CONTENT}")
|
||||
file(WRITE "${FLASHMLA_VENDOR_DIR}/flash_mla_interface.py"
|
||||
"${FLASHMLA_INTERFACE_CONTENT}")
|
||||
|
||||
# Install the generated flash_mla_interface.py to the wheel
|
||||
# Use COMPONENT _flashmla_C to ensure it's installed with the C extension
|
||||
install(FILES "${FLASHMLA_VENDOR_DIR}/flash_mla_interface.py"
|
||||
DESTINATION vllm/third_party/flashmla/
|
||||
COMPONENT _flashmla_C)
|
||||
|
||||
# The FlashMLA kernels only work on hopper and require CUDA 12.3 or later.
|
||||
# Only build FlashMLA kernels if we are building for something compatible with
|
||||
# sm90a
|
||||
|
||||
set(SUPPORT_ARCHS)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3)
|
||||
list(APPEND SUPPORT_ARCHS "9.0a")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
|
||||
# CUDA 12.9 has introduced "Family-Specific Architecture Features"
|
||||
# this supports all compute_10x family
|
||||
list(APPEND SUPPORT_ARCHS "10.0f")
|
||||
elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
|
||||
list(APPEND SUPPORT_ARCHS "10.0a")
|
||||
endif()
|
||||
|
||||
|
||||
cuda_archs_loose_intersection(FLASH_MLA_ARCHS "${SUPPORT_ARCHS}" "${CUDA_ARCHS}")
|
||||
if(FLASH_MLA_ARCHS)
|
||||
message(STATUS "FlashMLA CUDA architectures: ${FLASH_MLA_ARCHS}")
|
||||
set(VLLM_FLASHMLA_GPU_FLAGS ${VLLM_GPU_FLAGS})
|
||||
list(APPEND VLLM_FLASHMLA_GPU_FLAGS "--expt-relaxed-constexpr" "--expt-extended-lambda" "--use_fast_math")
|
||||
|
||||
set(FlashMLA_SOURCES
|
||||
${flashmla_SOURCE_DIR}/csrc/torch_api.cpp
|
||||
|
||||
# Misc kernels for decoding
|
||||
${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu
|
||||
|
||||
# sm90 dense decode
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu
|
||||
|
||||
# sm90 sparse decode
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu
|
||||
|
||||
# sm90 sparse prefill
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
|
||||
|
||||
# sm100 dense prefill & backward
|
||||
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
|
||||
|
||||
# sm100 sparse prefill
|
||||
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
|
||||
|
||||
# sm100 sparse decode
|
||||
${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
|
||||
)
|
||||
|
||||
set(FlashMLA_Extension_SOURCES
|
||||
${flashmla_SOURCE_DIR}/csrc/extension/torch_api.cpp
|
||||
${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/pybind.cpp
|
||||
${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_fp8_sm90.cu
|
||||
${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_metadata.cu
|
||||
)
|
||||
|
||||
set(FlashMLA_INCLUDES
|
||||
${flashmla_SOURCE_DIR}/csrc
|
||||
${flashmla_SOURCE_DIR}/csrc/kerutils/include
|
||||
${flashmla_SOURCE_DIR}/csrc/sm90
|
||||
${flashmla_SOURCE_DIR}/csrc/cutlass/include
|
||||
${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
|
||||
)
|
||||
|
||||
set(FlashMLA_Extension_INCLUDES
|
||||
${flashmla_SOURCE_DIR}/csrc
|
||||
${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
|
||||
${flashmla_SOURCE_DIR}/csrc/cutlass/include
|
||||
${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
|
||||
)
|
||||
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${FlashMLA_SOURCES}"
|
||||
CUDA_ARCHS "${FLASH_MLA_ARCHS}")
|
||||
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${FlashMLA_Extension_SOURCES}"
|
||||
CUDA_ARCHS "${FLASH_MLA_ARCHS}")
|
||||
|
||||
define_extension_target(
|
||||
_flashmla_C
|
||||
DESTINATION vllm
|
||||
LANGUAGE ${VLLM_GPU_LANG}
|
||||
SOURCES ${FlashMLA_SOURCES}
|
||||
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
|
||||
ARCHITECTURES ${VLLM_GPU_ARCHES}
|
||||
INCLUDE_DIRECTORIES ${FlashMLA_INCLUDES}
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
# Keep Stable ABI for the module, but *not* for CUDA/C++ files.
|
||||
# This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
|
||||
# Also enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
|
||||
target_compile_options(_flashmla_C PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
|
||||
|
||||
define_extension_target(
|
||||
_flashmla_extension_C
|
||||
DESTINATION vllm
|
||||
LANGUAGE ${VLLM_GPU_LANG}
|
||||
SOURCES ${FlashMLA_Extension_SOURCES}
|
||||
COMPILE_FLAGS ${VLLM_FLASHMLA_GPU_FLAGS}
|
||||
ARCHITECTURES ${VLLM_GPU_ARCHES}
|
||||
INCLUDE_DIRECTORIES ${FlashMLA_Extension_INCLUDES}
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
# Keep Stable ABI for the module, but *not* for CUDA/C++ files.
|
||||
# This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
|
||||
target_compile_options(_flashmla_extension_C PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
|
||||
else()
|
||||
message(STATUS "FlashMLA will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
|
||||
# Create empty targets for setup.py on unsupported systems
|
||||
add_custom_target(_flashmla_C)
|
||||
add_custom_target(_flashmla_extension_C)
|
||||
endif()
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
include(FetchContent)
|
||||
|
||||
# If FMHA_SM100_SRC_DIR is set, fmha_sm100 is installed from that directory
|
||||
# instead of downloading. This is useful for local MSA development.
|
||||
if(DEFINED ENV{FMHA_SM100_SRC_DIR})
|
||||
set(FMHA_SM100_SRC_DIR $ENV{FMHA_SM100_SRC_DIR})
|
||||
endif()
|
||||
|
||||
if(FMHA_SM100_SRC_DIR)
|
||||
FetchContent_Declare(
|
||||
fmha_sm100
|
||||
SOURCE_DIR ${FMHA_SM100_SRC_DIR}
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
)
|
||||
else()
|
||||
FetchContent_Declare(
|
||||
fmha_sm100
|
||||
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
|
||||
GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
)
|
||||
endif()
|
||||
|
||||
FetchContent_GetProperties(fmha_sm100)
|
||||
if(NOT fmha_sm100_POPULATED)
|
||||
FetchContent_Populate(fmha_sm100)
|
||||
endif()
|
||||
message(STATUS "fmha_sm100 is available at ${fmha_sm100_SOURCE_DIR}")
|
||||
|
||||
add_custom_target(fmha_sm100)
|
||||
|
||||
set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")
|
||||
|
||||
install(FILES
|
||||
"${FMHA_SM100_PY_ROOT}/__init__.py"
|
||||
"${FMHA_SM100_PY_ROOT}/api.py"
|
||||
"${FMHA_SM100_PY_ROOT}/bench_utils.py"
|
||||
"${FMHA_SM100_PY_ROOT}/jit.py"
|
||||
"${FMHA_SM100_PY_ROOT}/sparse.py"
|
||||
"${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py"
|
||||
DESTINATION vllm/third_party/fmha_sm100
|
||||
COMPONENT fmha_sm100)
|
||||
|
||||
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/"
|
||||
DESTINATION vllm/third_party/fmha_sm100/csrc
|
||||
COMPONENT fmha_sm100
|
||||
PATTERN "__pycache__" EXCLUDE
|
||||
PATTERN "*.pyc" EXCLUDE
|
||||
PATTERN ".git*" EXCLUDE)
|
||||
|
||||
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
|
||||
DESTINATION vllm/third_party/fmha_sm100/cute
|
||||
COMPONENT fmha_sm100
|
||||
PATTERN "__pycache__" EXCLUDE
|
||||
PATTERN "*.pyc" EXCLUDE
|
||||
PATTERN ".git*" EXCLUDE)
|
||||
|
||||
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/include/"
|
||||
DESTINATION vllm/third_party/fmha_sm100/cutlass/include
|
||||
COMPONENT fmha_sm100
|
||||
PATTERN "__pycache__" EXCLUDE
|
||||
PATTERN "*.pyc" EXCLUDE
|
||||
PATTERN ".git*" EXCLUDE)
|
||||
|
||||
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/tools/util/include/"
|
||||
DESTINATION vllm/third_party/fmha_sm100/cutlass/tools/util/include
|
||||
COMPONENT fmha_sm100
|
||||
PATTERN "__pycache__" EXCLUDE
|
||||
PATTERN "*.pyc" EXCLUDE
|
||||
PATTERN ".git*" EXCLUDE)
|
||||
@@ -0,0 +1,160 @@
|
||||
include(FetchContent)
|
||||
|
||||
set(CUTLASS_INCLUDE_DIR "${CUTLASS_INCLUDE_DIR}" CACHE PATH "Path to CUTLASS include/ directory")
|
||||
|
||||
if(DEFINED ENV{QUTLASS_SRC_DIR})
|
||||
set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR})
|
||||
endif()
|
||||
|
||||
# CMP0169 NEW: one-argument FetchContent_Populate(name) after Declare is invalid.
|
||||
# Use explicit Populate(...) for git, or set SOURCE_DIR for local trees.
|
||||
if(QUTLASS_SRC_DIR)
|
||||
set(_qutlass_user_src "${QUTLASS_SRC_DIR}")
|
||||
cmake_path(ABSOLUTE_PATH _qutlass_user_src
|
||||
BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
|
||||
NORMALIZE)
|
||||
set(QUTLASS_SRC_DIR "${_qutlass_user_src}")
|
||||
if(NOT IS_DIRECTORY "${QUTLASS_SRC_DIR}")
|
||||
message(FATAL_ERROR
|
||||
"[QUTLASS] QUTLASS_SRC_DIR is not an existing directory: '${QUTLASS_SRC_DIR}'")
|
||||
endif()
|
||||
set(qutlass_SOURCE_DIR "${QUTLASS_SRC_DIR}")
|
||||
set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
|
||||
else()
|
||||
set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
|
||||
set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
|
||||
|
||||
set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
|
||||
if(NOT _qutlass_fc_root)
|
||||
set(_qutlass_fc_root "${CMAKE_BINARY_DIR}/_deps")
|
||||
endif()
|
||||
set(_qutlass_src "${_qutlass_fc_root}/qutlass-src")
|
||||
set(_qutlass_bin "${_qutlass_fc_root}/qutlass-build")
|
||||
set(_qutlass_sub "${_qutlass_fc_root}/qutlass-subbuild")
|
||||
|
||||
if(EXISTS "${_qutlass_src}/qutlass/csrc/bindings.cpp")
|
||||
set(qutlass_SOURCE_DIR "${_qutlass_src}")
|
||||
set(qutlass_BINARY_DIR "${_qutlass_bin}")
|
||||
else()
|
||||
FetchContent_Populate(
|
||||
qutlass
|
||||
SUBBUILD_DIR "${_qutlass_sub}"
|
||||
SOURCE_DIR "${_qutlass_src}"
|
||||
BINARY_DIR "${_qutlass_bin}"
|
||||
GIT_REPOSITORY "${_QUTLASS_UPSTREAM_REPO}"
|
||||
GIT_TAG "${_QUTLASS_UPSTREAM_TAG}"
|
||||
GIT_PROGRESS TRUE
|
||||
)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if(NOT qutlass_SOURCE_DIR)
|
||||
message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.")
|
||||
endif()
|
||||
message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
|
||||
# QUTLASS uses TARGET_CUDA_ARCH as a single preprocessor selector for all its
|
||||
# sources. Do not compile a mixed SM100/SM120 arch list with one selector; prefer
|
||||
# SM100 when both families are requested because that is the primary deployed
|
||||
# target for this extension today.
|
||||
if(QUTLASS_SM100_ARCHS)
|
||||
set(QUTLASS_ARCHS "${QUTLASS_SM100_ARCHS}")
|
||||
set(QUTLASS_TARGET_CC 100)
|
||||
if(QUTLASS_SM120_ARCHS)
|
||||
message(WARNING
|
||||
"[QUTLASS] Both SM100 and SM120 archs were requested; selecting SM100 "
|
||||
"because TARGET_CUDA_ARCH is a single compile-time selector.")
|
||||
endif()
|
||||
elseif(QUTLASS_SM120_ARCHS)
|
||||
set(QUTLASS_ARCHS "${QUTLASS_SM120_ARCHS}")
|
||||
set(QUTLASS_TARGET_CC 120)
|
||||
else()
|
||||
set(QUTLASS_ARCHS)
|
||||
endif()
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
set(QUTLASS_SOURCES
|
||||
csrc/qutlass_registration.cpp
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/fused_quantize_mx.cu
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/fused_quantize_nv.cu
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/fused_quantize_mx_sm100.cu
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/fused_quantize_nv_sm100.cu
|
||||
)
|
||||
|
||||
set(QUTLASS_INCLUDES
|
||||
${qutlass_SOURCE_DIR}
|
||||
${qutlass_SOURCE_DIR}/qutlass
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/include
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/include/cutlass_extensions
|
||||
)
|
||||
|
||||
if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
|
||||
list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
|
||||
if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
|
||||
EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
|
||||
list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
|
||||
else()
|
||||
get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
|
||||
if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
|
||||
list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
|
||||
endif()
|
||||
endif()
|
||||
elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
|
||||
list(APPEND QUTLASS_INCLUDES
|
||||
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
|
||||
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
|
||||
message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
|
||||
else()
|
||||
message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
|
||||
"Set -DCUTLASS_INCLUDE_DIR=/path/to/cutlass/include")
|
||||
endif()
|
||||
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${QUTLASS_SOURCES}"
|
||||
CUDA_ARCHS "${QUTLASS_ARCHS}"
|
||||
)
|
||||
|
||||
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
|
||||
# Keep it as its own extension (registers torch.ops._qutlass_C).
|
||||
define_extension_target(
|
||||
_qutlass_C
|
||||
DESTINATION vllm
|
||||
LANGUAGE ${VLLM_GPU_LANG}
|
||||
SOURCES ${QUTLASS_SOURCES}
|
||||
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
|
||||
ARCHITECTURES ${VLLM_GPU_ARCHES}
|
||||
INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
target_compile_definitions(_qutlass_C PRIVATE
|
||||
QUTLASS_DISABLE_PYBIND=1
|
||||
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
|
||||
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
|
||||
)
|
||||
|
||||
else()
|
||||
if("${CMAKE_CUDA_COMPILER_VERSION}" VERSION_LESS "12.8")
|
||||
message(STATUS
|
||||
"[QUTLASS] Skipping build: CUDA 12.8 or newer is required (found ${CMAKE_CUDA_COMPILER_VERSION}).")
|
||||
else()
|
||||
message(STATUS
|
||||
"[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
|
||||
"CUDA_ARCHS='${CUDA_ARCHS}'.")
|
||||
endif()
|
||||
add_custom_target(_qutlass_C)
|
||||
endif()
|
||||
@@ -0,0 +1,53 @@
|
||||
# Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels
|
||||
|
||||
set(DEFAULT_TRITON_KERNELS_TAG "v3.5.1")
|
||||
|
||||
# Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to
|
||||
# be directly set to the triton_kernels python directory.
|
||||
if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
|
||||
message(STATUS "[triton_kernels] Fetch from $ENV{TRITON_KERNELS_SRC_DIR}")
|
||||
FetchContent_Declare(
|
||||
triton_kernels
|
||||
SOURCE_DIR $ENV{TRITON_KERNELS_SRC_DIR}
|
||||
)
|
||||
|
||||
else()
|
||||
set(TRITON_GIT "https://github.com/triton-lang/triton.git")
|
||||
message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${DEFAULT_TRITON_KERNELS_TAG}")
|
||||
FetchContent_Declare(
|
||||
triton_kernels
|
||||
# TODO (varun) : Fetch just the triton_kernels directory from Triton
|
||||
GIT_REPOSITORY https://github.com/triton-lang/triton.git
|
||||
GIT_TAG ${DEFAULT_TRITON_KERNELS_TAG}
|
||||
GIT_PROGRESS TRUE
|
||||
SOURCE_SUBDIR python/triton_kernels/triton_kernels
|
||||
)
|
||||
endif()
|
||||
|
||||
# Fetch content
|
||||
FetchContent_MakeAvailable(triton_kernels)
|
||||
|
||||
if (NOT triton_kernels_SOURCE_DIR)
|
||||
message (FATAL_ERROR "[triton_kernels] Cannot resolve triton_kernels_SOURCE_DIR")
|
||||
endif()
|
||||
|
||||
if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
|
||||
set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/")
|
||||
else()
|
||||
set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/python/triton_kernels/triton_kernels/")
|
||||
endif()
|
||||
|
||||
message (STATUS "[triton_kernels] triton_kernels is available at ${TRITON_KERNELS_PYTHON_DIR}")
|
||||
|
||||
add_custom_target(triton_kernels)
|
||||
|
||||
# Ensure the vllm/third_party directory exists before installation
|
||||
install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/triton_kernels\")")
|
||||
|
||||
## Copy .py files to install directory.
|
||||
install(DIRECTORY
|
||||
${TRITON_KERNELS_PYTHON_DIR}
|
||||
DESTINATION
|
||||
vllm/third_party/triton_kernels/
|
||||
COMPONENT triton_kernels
|
||||
FILES_MATCHING PATTERN "*.py")
|
||||
@@ -0,0 +1,116 @@
|
||||
# vLLM flash attention requires VLLM_GPU_ARCHES to contain the set of target
|
||||
# arches in the CMake syntax (75-real, 89-virtual, etc), since we clear the
|
||||
# arches in the CUDA case (and instead set the gencodes on a per file basis)
|
||||
# we need to manually set VLLM_GPU_ARCHES here.
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
foreach(_ARCH ${CUDA_ARCHS})
|
||||
string(REPLACE "." "" _ARCH "${_ARCH}")
|
||||
list(APPEND VLLM_GPU_ARCHES "${_ARCH}-real")
|
||||
endforeach()
|
||||
endif()
|
||||
|
||||
#
|
||||
# Build vLLM flash attention from source
|
||||
#
|
||||
# IMPORTANT: This has to be the last thing we do, because vllm-flash-attn uses the same macros/functions as vLLM.
|
||||
# Because functions all belong to the global scope, vllm-flash-attn's functions overwrite vLLMs.
|
||||
# They should be identical but if they aren't, this is a massive footgun.
|
||||
#
|
||||
# The vllm-flash-attn install rules are nested under vllm to make sure the library gets installed in the correct place.
|
||||
# To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2), --component _vllm_fa3_C (for FA3),
|
||||
# or --component _vllm_fa4_cutedsl_C (for FA4 CuteDSL Python files).
|
||||
# If no component is specified, vllm-flash-attn is still installed.
|
||||
|
||||
# If VLLM_FLASH_ATTN_SRC_DIR is set, vllm-flash-attn is installed from that directory instead of downloading.
|
||||
# This is to enable local development of vllm-flash-attn within vLLM.
|
||||
# It can be set as an environment variable or passed as a cmake argument.
|
||||
# The environment variable takes precedence.
|
||||
if (DEFINED ENV{VLLM_FLASH_ATTN_SRC_DIR})
|
||||
set(VLLM_FLASH_ATTN_SRC_DIR $ENV{VLLM_FLASH_ATTN_SRC_DIR})
|
||||
endif()
|
||||
|
||||
if(VLLM_FLASH_ATTN_SRC_DIR)
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn SOURCE_DIR
|
||||
${VLLM_FLASH_ATTN_SRC_DIR}
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
)
|
||||
else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG bb9a72e7dde0dc614ffc663e052cd6a19ce73a42
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
)
|
||||
endif()
|
||||
|
||||
# Make sure vllm-flash-attn install rules are nested under vllm/
|
||||
# ALL_COMPONENTS ensures the save/modify/restore runs exactly once regardless
|
||||
# of how many components are being installed, avoiding double-append of /vllm/.
|
||||
install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" ALL_COMPONENTS)
|
||||
install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
|
||||
install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" ALL_COMPONENTS)
|
||||
|
||||
# Fetch the vllm-flash-attn library
|
||||
FetchContent_MakeAvailable(vllm-flash-attn)
|
||||
message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}")
|
||||
|
||||
# Restore the install prefix after FA's install rules
|
||||
install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
|
||||
install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
|
||||
|
||||
# Install shared Python files for both FA2 and FA3 components
|
||||
foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
|
||||
# Ensure the vllm/vllm_flash_attn directory exists before installation
|
||||
install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")"
|
||||
COMPONENT ${_FA_COMPONENT})
|
||||
|
||||
# Copy vllm_flash_attn python files (except __init__.py and flash_attn_interface.py
|
||||
# which are source-controlled in vllm)
|
||||
install(
|
||||
DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
|
||||
DESTINATION vllm/vllm_flash_attn
|
||||
COMPONENT ${_FA_COMPONENT}
|
||||
FILES_MATCHING PATTERN "*.py"
|
||||
PATTERN "__init__.py" EXCLUDE
|
||||
PATTERN "flash_attn_interface.py" EXCLUDE
|
||||
)
|
||||
|
||||
endforeach()
|
||||
|
||||
#
|
||||
# FA4 CuteDSL component
|
||||
# This is a Python-only component that copies the flash_attn/cute directory
|
||||
# and transforms imports to match our package structure.
|
||||
#
|
||||
add_custom_target(_vllm_fa4_cutedsl_C)
|
||||
|
||||
# Install flash_attn/cute directory (needed for FA4).
|
||||
# When using a local source dir (VLLM_FLASH_ATTN_SRC_DIR), create a symlink
|
||||
# so edits to cute-dsl Python files take effect immediately without rebuilding.
|
||||
# Otherwise, copy files and transform flash_attn.cute imports to
|
||||
# vllm.vllm_flash_attn.cute to match our package structure.
|
||||
if(VLLM_FLASH_ATTN_SRC_DIR)
|
||||
install(CODE "
|
||||
set(LINK_TARGET \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\")
|
||||
set(LINK_NAME \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute\")
|
||||
file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")
|
||||
file(REMOVE_RECURSE \"\${LINK_NAME}\")
|
||||
file(CREATE_LINK \"\${LINK_TARGET}\" \"\${LINK_NAME}\" SYMBOLIC)
|
||||
" COMPONENT _vllm_fa4_cutedsl_C)
|
||||
else()
|
||||
install(CODE "
|
||||
file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\")
|
||||
foreach(SRC_FILE \${CUTE_PY_FILES})
|
||||
file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
|
||||
set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\")
|
||||
get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
|
||||
file(MAKE_DIRECTORY \${DST_DIR})
|
||||
file(READ \${SRC_FILE} FILE_CONTENTS)
|
||||
string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\")
|
||||
file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
|
||||
endforeach()
|
||||
" COMPONENT _vllm_fa4_cutedsl_C)
|
||||
endif()
|
||||
Reference in New Issue
Block a user