Skip to content

Commit 10d67e1

Browse files
committed
Redesigned CMake build system for flashinfer.
1 parent 7be6f1d commit 10d67e1

14 files changed

+1175
-548
lines changed

CMakeLists.txt

+94-409
Large diffs are not rendered by default.

cmake/Components.cmake

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# Define the component structure
2+
set(FLASHINFER_COMPONENTS "Headers")
3+
4+
if(FLASHINFER_BUILD_KERNELS)
5+
list(APPEND FLASHINFER_COMPONENTS "Kernels")
6+
7+
if(FLASHINFER_TVM_BINDING)
8+
list(APPEND FLASHINFER_COMPONENTS "TVMBinding")
9+
endif()
10+
endif()
11+
12+
if(FLASHINFER_DISTRIBUTED)
13+
list(APPEND FLASHINFER_COMPONENTS "Distributed")
14+
endif()
15+
16+
# Setup component-specific build flags
17+
macro(add_component_flags component)
18+
add_definitions(-DFLASHINFER_COMPONENT_${component})
19+
endmacro()
20+
21+
# For each enabled component, add compile-time flags
22+
foreach(comp ${FLASHINFER_COMPONENTS})
23+
add_component_flags(${comp})
24+
endforeach()

cmake/Dependencies.cmake

+146
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
# === Required Dependencies for Core Functionality ===
2+
find_package(CUDAToolkit REQUIRED)
3+
find_package(Python3 REQUIRED)
4+
if(NOT Python3_FOUND)
5+
message(
6+
FATAL_ERROR
7+
"Python3 not found it is required to generate the kernel sources.")
8+
endif()
9+
10+
find_package(Thrust REQUIRED)
11+
12+
# === Test Dependencies ===
13+
if(FLASHINFER_UNITTESTS)
14+
include(FetchContent)
15+
16+
# Google Test for unit testing
17+
FetchContent_Declare(
18+
googletest
19+
GIT_REPOSITORY https://github.com/google/googletest.git
20+
GIT_TAG 6910c9d9165801d8827d628cb72eb7ea9dd538c5 # release-1.16.0
21+
FIND_PACKAGE_ARGS NAMES GTest)
22+
FetchContent_MakeAvailable(googletest)
23+
endif()
24+
25+
# === Benchmark Dependencies ===
26+
if(FLASHINFER_CXX_BENCHMARKS)
27+
include(FetchContent)
28+
29+
# NVBench for GPU benchmarking
30+
FetchContent_Declare(
31+
nvbench
32+
GIT_REPOSITORY https://github.com/NVIDIA/nvbench.git
33+
GIT_TAG c03033b50e46748207b27685b1cdfcbe4a2fec59)
34+
FetchContent_MakeAvailable(nvbench)
35+
endif()
36+
37+
# === Boost Dependency for FP16 QK Reductions ===
38+
if(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS)
39+
include(FetchContent)
40+
set(BOOST_ENABLE_CMAKE ON)
41+
FetchContent_Declare(boost_math
42+
GIT_REPOSITORY https://github.com/boostorg/math.git)
43+
FetchContent_MakeAvailable(boost_math)
44+
45+
set(USE_FP16_QK_REDUCTIONS "true")
46+
message(STATUS "USE_FP16_QK_REDUCTIONS=${USE_FP16_QK_REDUCTIONS}")
47+
else()
48+
set(USE_FP16_QK_REDUCTIONS "false")
49+
message(STATUS "USE_FP16_QK_REDUCTIONS=${USE_FP16_QK_REDUCTIONS}")
50+
endif()
51+
52+
# === Distributed component dependencies ===
53+
if(FLASHINFER_DISTRIBUTED OR FLASHINFER_DIST_UNITTESTS)
54+
include(FetchContent)
55+
FetchContent_Declare(
56+
mscclpp
57+
GIT_REPOSITORY https://github.com/microsoft/mscclpp.git
58+
GIT_TAG 11e62024d3eb190e005b4689f8c8443d91a6c82e)
59+
FetchContent_MakeAvailable(mscclpp)
60+
61+
# Create alias for distributed component
62+
if(NOT TARGET flashinfer::mscclpp)
63+
add_library(flashinfer::mscclpp ALIAS mscclpp)
64+
endif()
65+
66+
# Fetch spdlog for distributed tests (header-only usage)
67+
FetchContent_Declare(
68+
spdlog
69+
GIT_REPOSITORY https://github.com/gabime/spdlog.git
70+
GIT_TAG f355b3d58f7067eee1706ff3c801c2361011f3d5 # release-1.15.1
71+
FIND_PACKAGE_ARGS NAMES spdlog)
72+
73+
# Use Populate instead of MakeAvailable since we only need the headers
74+
FetchContent_Populate(spdlog)
75+
76+
# Set the include directory for later use
77+
set(SPDLOG_INCLUDE_DIR "${spdlog_SOURCE_DIR}/include")
78+
message(STATUS "Using spdlog from ${SPDLOG_INCLUDE_DIR}")
79+
80+
find_package(MPI REQUIRED)
81+
endif()
82+
83+
# === FP8 Dependencies ===
84+
if(FLASHINFER_FP8_TESTS OR FLASHINFER_FP8_BENCHMARKS)
85+
# Verify CUDA architecture is SM90 or higher
86+
if(NOT CMAKE_CUDA_ARCHITECTURES STREQUAL "90"
87+
AND NOT CMAKE_CUDA_ARCHITECTURES STREQUAL "90a")
88+
message(
89+
FATAL_ERROR "FP8 tests/benchmarks require SM90 or higher architecture")
90+
endif()
91+
92+
# Find PyTorch which is required for FP8 features
93+
find_package(Torch REQUIRED)
94+
if(NOT Torch_FOUND)
95+
message(
96+
FATAL_ERROR "PyTorch is required for FP8 tests/benchmarks but not found")
97+
endif()
98+
message(STATUS "Found PyTorch: ${TORCH_INCLUDE_DIRS}")
99+
100+
# Fetch Flash Attention repository with specific commit
101+
include(FetchContent)
102+
FetchContent_Declare(
103+
flash_attention
104+
GIT_REPOSITORY https://github.com/Dao-AILab/flash-attention.git
105+
GIT_TAG 29ef580560761838c0e9e82bc0e98d04ba75f949)
106+
FetchContent_Populate(flash_attention)
107+
108+
# Set Flash Attention 3 include directory
109+
set(FA3_INCLUDE_DIR "${flash_attention_SOURCE_DIR}/csrc/flash_attn/hopper")
110+
message(STATUS "Flash Attention 3 source directory: ${FA3_INCLUDE_DIR}")
111+
112+
# Compile Flash Attention 3 kernel library
113+
file(GLOB FA3_IMPL_FILES "${FA3_INCLUDE_DIR}/flash_fwd_*.cu")
114+
endif()
115+
116+
# === TVM Binding dependencies ===
117+
if(FLASHINFER_TVM_BINDING)
118+
# Resolve TVM source directory
119+
if(NOT FLASHINFER_TVM_SOURCE_DIR STREQUAL "")
120+
set(TVM_SOURCE_DIR_SET ${FLASHINFER_TVM_SOURCE_DIR})
121+
elseif(DEFINED ENV{TVM_SOURCE_DIR})
122+
set(TVM_SOURCE_DIR_SET $ENV{TVM_SOURCE_DIR})
123+
elseif(DEFINED ENV{TVM_HOME})
124+
set(TVM_SOURCE_DIR_SET $ENV{TVM_HOME})
125+
else()
126+
message(
127+
FATAL_ERROR
128+
"TVM source directory not found. Set FLASHINFER_TVM_SOURCE_DIR.")
129+
endif()
130+
endif()
131+
132+
# === CUTLASS Configuration ===
133+
if(FLASHINFER_CUTLASS_DIR)
134+
list(APPEND CMAKE_PREFIX_PATH ${FLASHINFER_CUTLASS_DIR})
135+
endif()
136+
137+
if(FLASHINFER_CUTLASS_DIR)
138+
# Add CUTLASS include directories directly
139+
include_directories(${FLASHINFER_CUTLASS_DIR}/include)
140+
include_directories(${FLASHINFER_CUTLASS_DIR}/tools/util/include)
141+
142+
message(STATUS "Using CUTLASS from ${FLASHINFER_CUTLASS_DIR}")
143+
else()
144+
message(
145+
FATAL_ERROR "FLASHINFER_CUTLASS_DIR must be set to the path of CUTLASS")
146+
endif()

cmake/Options.cmake

+86
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
# cmake-format: off
2+
# NOTE:
3+
# a) Do not modify this file to change option values. Options should be
4+
# configured using either a config.cmake file (refer the default file
5+
# inside the cmake folder), or by setting the required -DFLASHINFER_XXX
6+
# option through command-line.
7+
#
8+
# b) This file should only contain option definitions and should not contain
9+
# any other CMake commands.
10+
#
11+
# c) All new options should be defined here with a default value and a short
12+
# description.
13+
#
14+
# d) Add new options under the appropriate section.
15+
16+
# === COMPONENT OPTIONS ===
17+
flashinfer_option(FLASHINFER_BUILD_KERNELS "Build and install kernel libraries" OFF)
18+
flashinfer_option(FLASHINFER_TVM_BINDING "Build TVM binding support" OFF)
19+
flashinfer_option(FLASHINFER_DISTRIBUTED "Build distributed support" OFF)
20+
21+
# === DATA TYPE OPTIONS ===
22+
flashinfer_option(FLASHINFER_ENABLE_FP8 "Enable FP8 data type support" ON)
23+
flashinfer_option(FLASHINFER_ENABLE_FP8_E4M3 "Enable FP8 E4M3 format specifically" ON)
24+
flashinfer_option(FLASHINFER_ENABLE_FP8_E5M2 "Enable FP8 E5M2 format specifically" ON)
25+
flashinfer_option(FLASHINFER_ENABLE_F16 "Enable F16 data type support" ON)
26+
flashinfer_option(FLASHINFER_ENABLE_BF16 "Enable BF16 data type support" ON)
27+
28+
# === CODE GENERATION OPTIONS ===
29+
flashinfer_option(FLASHINFER_GEN_HEAD_DIMS "Head dimensions to enable" 64 128 256)
30+
flashinfer_option(FLASHINFER_GEN_POS_ENCODING_MODES "Position encoding modes to enable" 0 1 2)
31+
flashinfer_option(FLASHINFER_GEN_MASK_MODES "Mask modes to enable" 0 1 2)
32+
# FIXME: CAUTION!!! Turning on this option will cause build failures. (refer #806, #936)
33+
flashinfer_option(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS "Use FP16 for QK reductions" OFF)
34+
35+
# === BUILD TYPE OPTIONS ===
36+
flashinfer_option(FLASHINFER_UNITTESTS "Build unit tests" OFF)
37+
flashinfer_option(FLASHINFER_CXX_BENCHMARKS "Build benchmarks" OFF)
38+
flashinfer_option(FLASHINFER_DIST_UNITTESTS "Build distributed unit tests" OFF)
39+
40+
# === FEATURE-SPECIFIC TESTS/BENCHMARKS ===
41+
flashinfer_option(FLASHINFER_FP8_TESTS "Build FP8 tests" OFF)
42+
flashinfer_option(FLASHINFER_FP8_BENCHMARKS "Build FP8 benchmarks" OFF)
43+
44+
# === ARCHITECTURE OPTIONS ===
45+
flashinfer_option(FLASHINFER_CUDA_ARCHITECTURES "CUDA architectures to compile for" "")
46+
47+
# === PATH OPTIONS ===
48+
flashinfer_option(FLASHINFER_CUTLASS_DIR "Path to CUTLASS installation" "")
49+
flashinfer_option(FLASHINFER_TVM_SOURCE_DIR "Path to TVM source directory" "")
50+
51+
# === AUTO-DERIVED OPTIONS ===
52+
# Handle CUDA architectures
53+
if(FLASHINFER_CUDA_ARCHITECTURES)
54+
message(STATUS "CMAKE_CUDA_ARCHITECTURES set to ${FLASHINFER_CUDA_ARCHITECTURES}.")
55+
set(CMAKE_CUDA_ARCHITECTURES ${FLASHINFER_CUDA_ARCHITECTURES})
56+
endif()
57+
58+
# Handle automatic enabling of dependent features
59+
if(FLASHINFER_FP8_TESTS)
60+
set(FLASHINFER_UNITTESTS ON CACHE BOOL "Tests enabled for FP8" FORCE)
61+
endif()
62+
63+
if(FLASHINFER_FP8_BENCHMARKS)
64+
set(FLASHINFER_CXX_BENCHMARKS ON CACHE BOOL "Benchmarks enabled for FP8" FORCE)
65+
endif()
66+
67+
if(FLASHINFER_DIST_UNITTESTS)
68+
set(FLASHINFER_UNITTESTS ON CACHE BOOL "Tests enabled for distributed" FORCE)
69+
endif()
70+
71+
if(FLASHINFER_TVM_BINDING AND NOT FLASHINFER_BUILD_KERNELS)
72+
message(FATAL_ERROR "TVM binding requires FLASHINFER_BUILD_KERNELS to be ON")
73+
endif()
74+
75+
if(FLASHINFER_ENABLE_FP8)
76+
# Enable both FP8 formats when FP8 is enabled
77+
set(FLASHINFER_ENABLE_FP8_E4M3 ON CACHE BOOL "Enable FP8 E4M3 format" FORCE)
78+
set(FLASHINFER_ENABLE_FP8_E5M2 ON CACHE BOOL "Enable FP8 E5M2 format" FORCE)
79+
endif()
80+
81+
# Ensure FP8 is enabled for FP8 tests/benchmarks
82+
if(FLASHINFER_FP8_TESTS OR FLASHINFER_FP8_BENCHMARKS)
83+
set(FLASHINFER_ENABLE_FP8 ON CACHE BOOL "FP8 enabled for tests/benchmarks" FORCE)
84+
set(FLASHINFER_ENABLE_FP8_E4M3 ON CACHE BOOL "FP8_E4M3 enabled for tests/benchmarks" FORCE)
85+
endif()
86+
# cmake-format: on

cmake/flashinferConfig.cmake.in

+72
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
@PACKAGE_INIT@
2+
3+
include(CMakeFindDependencyMacro)
4+
# Required dependencies for FlashInfer headers
5+
find_dependency(CUDAToolkit REQUIRED)
6+
7+
# Optional dependencies based on components
8+
if("Distributed" IN_LIST flashinfer_FIND_COMPONENTS)
9+
find_dependency(MPI REQUIRED)
10+
find_package(mscclpp REQUIRED)
11+
endif()
12+
13+
# Define available components passed from main CMakeLists.txt
14+
set(_flashinfer_available_components @FLASHINFER_COMPONENTS@)
15+
16+
# Initialize component found status
17+
foreach(_comp Headers Kernels TVMBinding Distributed)
18+
set(flashinfer_${_comp}_FOUND FALSE)
19+
endforeach()
20+
21+
# Mark available components as found
22+
foreach(_comp ${_flashinfer_available_components})
23+
set(flashinfer_${_comp}_FOUND TRUE)
24+
endforeach()
25+
26+
# Check for requested components
27+
foreach(_comp ${flashinfer_FIND_COMPONENTS})
28+
if(NOT _comp IN_LIST _flashinfer_available_components)
29+
set(flashinfer_FOUND False)
30+
set(flashinfer_NOT_FOUND_MESSAGE "Requested component: ${_comp} is not available")
31+
return()
32+
endif()
33+
endforeach()
34+
35+
# Headers component is always available and included
36+
set(flashinfer_Headers_FOUND TRUE)
37+
38+
# Include kernel targets if available and requested
39+
if(flashinfer_Kernels_FOUND AND
40+
("Kernels" IN_LIST flashinfer_FIND_COMPONENTS OR NOT flashinfer_FIND_COMPONENTS))
41+
include("${CMAKE_CURRENT_LIST_DIR}/libflashinferTargets.cmake" OPTIONAL)
42+
endif()
43+
44+
# Include TVM binding targets if available and requested
45+
if(flashinfer_TVMBinding_FOUND AND
46+
"TVMBinding" IN_LIST flashinfer_FIND_COMPONENTS)
47+
include("${CMAKE_CURRENT_LIST_DIR}/flashinferTVMBindingTargets.cmake" OPTIONAL)
48+
endif()
49+
50+
# Include Distributed targets if available and requested
51+
if(flashinfer_Distributed_FOUND AND
52+
"Distributed" IN_LIST flashinfer_FIND_COMPONENTS)
53+
# Create aliases to make usage consistent
54+
if(NOT TARGET flashinfer::dist)
55+
add_library(flashinfer::dist INTERFACE IMPORTED)
56+
set_target_properties(flashinfer::dist PROPERTIES
57+
INTERFACE_LINK_LIBRARIES mscclpp)
58+
endif()
59+
60+
include("${CMAKE_CURRENT_LIST_DIR}/flashinferDistTargets.cmake")
61+
endif()
62+
63+
# Set include directories
64+
set(FLASHINFER_INCLUDE_DIRS "${PACKAGE_PREFIX_DIR}/include")
65+
66+
# Print status message
67+
if(NOT flashinfer_FIND_QUIETLY)
68+
message(STATUS "Found flashinfer: ${PACKAGE_PREFIX_DIR} (version: ${flashinfer_VERSION})")
69+
message(STATUS "Available components: ${_flashinfer_available_components}")
70+
endif()
71+
72+
check_required_components(flashinfer)

0 commit comments

Comments
 (0)