@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cmake 236 lines 12.7 KB
Raw
1 include(FetchContent)
2
3 # If FLASH_MLA_SRC_DIR is set, flash-mla is installed from that directory
4 # instead of downloading.
5 # It can be set as an environment variable or passed as a cmake argument.
6 # The environment variable takes precedence.
7 if (DEFINED ENV{FLASH_MLA_SRC_DIR})
8 set(FLASH_MLA_SRC_DIR $ENV{FLASH_MLA_SRC_DIR})
9 endif()
10
11 if(FLASH_MLA_SRC_DIR)
12 FetchContent_Declare(
13 flashmla
14 SOURCE_DIR ${FLASH_MLA_SRC_DIR}
15 CONFIGURE_COMMAND ""
16 BUILD_COMMAND ""
17 )
18 else()
19 FetchContent_Declare(
20 flashmla
21 GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
22 GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf
23 GIT_PROGRESS TRUE
24 CONFIGURE_COMMAND ""
25 BUILD_COMMAND ""
26 )
27 endif()
28
29
30 FetchContent_MakeAvailable(flashmla)
31 message(STATUS "FlashMLA is available at ${flashmla_SOURCE_DIR}")
32
33 # Vendor FlashMLA interface into vLLM with torch-ops shim.
34 set(FLASHMLA_VENDOR_DIR "${CMAKE_SOURCE_DIR}/vllm/third_party/flashmla")
35 file(MAKE_DIRECTORY "${FLASHMLA_VENDOR_DIR}")
36 file(READ "${flashmla_SOURCE_DIR}/flash_mla/flash_mla_interface.py"
37 FLASHMLA_INTERFACE_CONTENT)
38 string(REPLACE "flash_mla_cuda = torch.ops._flashmla_C"
39 "import vllm._flashmla_C\nflash_mla_cuda = torch.ops._flashmla_C"
40 FLASHMLA_INTERFACE_CONTENT
41 "${FLASHMLA_INTERFACE_CONTENT}")
42 file(WRITE "${FLASHMLA_VENDOR_DIR}/flash_mla_interface.py"
43 "${FLASHMLA_INTERFACE_CONTENT}")
44
45 # Install the generated flash_mla_interface.py to the wheel
46 # Use COMPONENT _flashmla_C to ensure it's installed with the C extension
47 install(FILES "${FLASHMLA_VENDOR_DIR}/flash_mla_interface.py"
48 DESTINATION vllm/third_party/flashmla/
49 COMPONENT _flashmla_C)
50
51 # The FlashMLA kernels only work on hopper and require CUDA 12.3 or later.
52 # Only build FlashMLA kernels if we are building for something compatible with
53 # sm90a
54
55 set(SUPPORT_ARCHS)
56 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3)
57 list(APPEND SUPPORT_ARCHS "9.0a")
58 endif()
59 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
60 # CUDA 12.9 has introduced "Family-Specific Architecture Features"
61 # this supports all compute_10x family
62 list(APPEND SUPPORT_ARCHS "10.0f")
63 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
64 list(APPEND SUPPORT_ARCHS "10.7f")
65 endif()
66 elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
67 list(APPEND SUPPORT_ARCHS "10.0a")
68 endif()
69
70
71 cuda_archs_loose_intersection(FLASH_MLA_ARCHS "${SUPPORT_ARCHS}" "${CUDA_ARCHS}")
72 if(FLASH_MLA_ARCHS)
73 message(STATUS "FlashMLA CUDA architectures: ${FLASH_MLA_ARCHS}")
74 set(VLLM_FLASHMLA_GPU_FLAGS ${VLLM_GPU_FLAGS})
75 list(APPEND VLLM_FLASHMLA_GPU_FLAGS "--expt-relaxed-constexpr" "--expt-extended-lambda" "--use_fast_math")
76
77 set(FlashMLA_SOURCES
78 ${flashmla_SOURCE_DIR}/csrc/api/api.cpp
79 ${flashmla_SOURCE_DIR}/csrc/api/sparse_prefill.cpp
80 ${flashmla_SOURCE_DIR}/csrc/api/sparse_decode.cpp
81 ${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp
82 ${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp
83
84 # Misc kernels for decoding
85 ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
86 ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu
87
88 # sm90 dense decode
89 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu
90 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu
91
92 # sm90 sparse decode
93 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu
94 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu
95 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu
96 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu
97
98 # sm90 sparse prefill
99 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu
100 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
101 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu
102 ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
103
104 # sm100 dense prefill (inference-only: dense backward stays out)
105 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
106
107 # sm100 sparse prefill
108 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu
109 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu
110 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
111 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
112 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu
113
114 # sm100 fused norm + rope + attn + rope + cast
115 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu
116 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu
117 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu
118 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu
119 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu
120 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu
121 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu
122 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu
123 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu
124 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu
125 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu
126 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu
127 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu
128 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu
129 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu
130 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu
131 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu
132 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu
133
134 # sm100 sparse decode
135 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu
136 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu
137 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu
138 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu
139 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu
140 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu
141 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu
142 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu
143 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu
144 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
145 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu
146 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu
147 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu
148 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu
149 ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu
150 )
151
152 set(FlashMLA_Extension_SOURCES
153 ${flashmla_SOURCE_DIR}/csrc/extension/torch_api.cpp
154 ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/pybind.cpp
155 ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_fp8_sm90.cu
156 ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_metadata.cu
157 )
158
159 set(FlashMLA_INCLUDES
160 ${flashmla_SOURCE_DIR}/csrc
161 ${flashmla_SOURCE_DIR}/csrc/kerutils/include
162 ${flashmla_SOURCE_DIR}/csrc/cutlass/include
163 ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
164 )
165
166 set(FlashMLA_Extension_INCLUDES
167 ${flashmla_SOURCE_DIR}/csrc
168 ${flashmla_SOURCE_DIR}/csrc/kerutils/include
169 ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
170 ${flashmla_SOURCE_DIR}/csrc/cutlass/include
171 ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
172 )
173
174 set_gencode_flags_for_srcs(
175 SRCS "${FlashMLA_SOURCES}"
176 CUDA_ARCHS "${FLASH_MLA_ARCHS}")
177
178 set_gencode_flags_for_srcs(
179 SRCS "${FlashMLA_Extension_SOURCES}"
180 CUDA_ARCHS "${FLASH_MLA_ARCHS}")
181
182 define_extension_target(
183 _flashmla_C
184 DESTINATION vllm
185 LANGUAGE ${VLLM_GPU_LANG}
186 SOURCES ${FlashMLA_SOURCES}
187 COMPILE_FLAGS ${VLLM_GPU_FLAGS}
188 ARCHITECTURES ${VLLM_GPU_ARCHES}
189 INCLUDE_DIRECTORIES ${FlashMLA_INCLUDES}
190 USE_SABI 3
191 WITH_SOABI)
192
193 # Enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
194 target_compile_options(_flashmla_C PRIVATE
195 $<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
196 $<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
197
198 # _flashmla_C is now ABI-stable torch 2.11+
199 target_compile_definitions(_flashmla_C PRIVATE
200 TORCH_TARGET_VERSION=0x020B000000000000ULL)
201 if(VLLM_GPU_LANG STREQUAL "CUDA")
202 target_compile_definitions(_flashmla_C PRIVATE USE_CUDA)
203 endif()
204
205 define_extension_target(
206 _flashmla_extension_C
207 DESTINATION vllm
208 LANGUAGE ${VLLM_GPU_LANG}
209 SOURCES ${FlashMLA_Extension_SOURCES}
210 COMPILE_FLAGS ${VLLM_FLASHMLA_GPU_FLAGS}
211 ARCHITECTURES ${VLLM_GPU_ARCHES}
212 INCLUDE_DIRECTORIES ${FlashMLA_Extension_INCLUDES}
213 USE_SABI 3
214 WITH_SOABI)
215
216 # _flashmla_extension_C is now ABI-stable w/ torch 2.11+
217 target_compile_definitions(_flashmla_extension_C PRIVATE
218 TORCH_TARGET_VERSION=0x020B000000000000ULL)
219 if(VLLM_GPU_LANG STREQUAL "CUDA")
220 target_compile_definitions(_flashmla_extension_C PRIVATE USE_CUDA)
221 endif()
222
223 # FlashMLA's sources use M_LOG2E from <cmath>. MSVC's UCRT only defines the
224 # M_* constants when _USE_MATH_DEFINES is set before <cmath> is first
225 # included, so define it for both targets on MSVC. Scoped to MSVC so the
226 # compile flags (and build caches) on other platforms stay unchanged.
227 if(MSVC)
228 target_compile_definitions(_flashmla_C PRIVATE _USE_MATH_DEFINES)
229 target_compile_definitions(_flashmla_extension_C PRIVATE _USE_MATH_DEFINES)
230 endif()
231 else()
232 message(STATUS "FlashMLA will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
233 # Create empty targets for setup.py on unsupported systems
234 add_custom_target(_flashmla_C)
235 add_custom_target(_flashmla_extension_C)
236 endif()