Build with ABI stable FlashMLA (#48174)
Signed-off-by: Jane Xu <janeyx@meta.com> Signed-off-by: Shengqi Chen <i@harrychen.xyz> Co-authored-by: Shengqi Chen <i@harrychen.xyz>
Jane (Yuan) Xu committed
Jul 14, 2026 at 23:29 UTC
37aa52821df831ea118ebfe3ec5b967227935733
1 file changed
+18
-13
cmake/external_projects/flashmla.cmake
+18
-13
@@ -19,7 +19,7 @@ else()
19
FetchContent_Declare(
20
flashmla
21
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
22
- GIT_TAG b70aff3d110a2b1a037e62eac295166b5143643a
22
+ GIT_TAG a8f794d1251cbfd88a5011445dd5582289c727e4
23
GIT_PROGRESS TRUE
24
CONFIGURE_COMMAND ""
25
BUILD_COMMAND ""
@@ -35,7 +35,7 @@ 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 "import flash_mla.cuda as flash_mla_cuda"
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}")
@@ -72,7 +72,7 @@ if(FLASH_MLA_ARCHS)
72
list(APPEND VLLM_FLASHMLA_GPU_FLAGS "--expt-relaxed-constexpr" "--expt-extended-lambda" "--use_fast_math")
73
74
set(FlashMLA_SOURCES
75
- ${flashmla_SOURCE_DIR}/csrc/torch_api.cpp
75
+ ${flashmla_SOURCE_DIR}/csrc/api/api.cpp
76
77
# Misc kernels for decoding
78
${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
@@ -128,6 +128,7 @@ if(FLASH_MLA_ARCHS)
128
129
set(FlashMLA_Extension_INCLUDES
130
${flashmla_SOURCE_DIR}/csrc
131
+ ${flashmla_SOURCE_DIR}/csrc/kerutils/include
132
${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
133
${flashmla_SOURCE_DIR}/csrc/cutlass/include
134
${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
@@ -152,15 +153,18 @@ if(FLASH_MLA_ARCHS)
153
USE_SABI 3
154
WITH_SOABI)
155
155
- # Keep Stable ABI for the module, but *not* for CUDA/C++ files.
156
- # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
157
- # Also enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
156
+ # Enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
157
target_compile_options(_flashmla_C PRIVATE
159
- $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
160
- $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
158
$<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
159
$<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
160
161
+ # _flashmla_C is now ABI-stable torch 2.11+
162
+ target_compile_definitions(_flashmla_C PRIVATE
163
+ TORCH_TARGET_VERSION=0x020B000000000000ULL)
164
+ if(VLLM_GPU_LANG STREQUAL "CUDA")
165
+ target_compile_definitions(_flashmla_C PRIVATE USE_CUDA)
166
+ endif()
167
+
168
define_extension_target(
169
_flashmla_extension_C
170
DESTINATION vllm
@@ -172,11 +176,12 @@ if(FLASH_MLA_ARCHS)
176
USE_SABI 3
177
WITH_SOABI)
178
175
- # Keep Stable ABI for the module, but *not* for CUDA/C++ files.
176
- # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
177
- target_compile_options(_flashmla_extension_C PRIVATE
178
- $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
179
- $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
179
+ # _flashmla_extension_C is now ABI-stable w/ torch 2.11+
180
+ target_compile_definitions(_flashmla_extension_C PRIVATE
181
+ TORCH_TARGET_VERSION=0x020B000000000000ULL)
182
+ if(VLLM_GPU_LANG STREQUAL "CUDA")
183
+ target_compile_definitions(_flashmla_extension_C PRIVATE USE_CUDA)
184
+ endif()
185
else()
186
message(STATUS "FlashMLA will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
187
# Create empty targets for setup.py on unsupported systems