[Kimi-K3] Migrate FlashKDA to PyTorch stable ABI (#50818)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Thien Tran committed
Aug 4, 2026 at 08:18 UTC
11f88260a240e6e8ee33410f9c65f00043c3e6f7
3 files changed
+19
-7
.buildkite/check-torch-abi.py
-1
@@ -13,7 +13,6 @@ from torch_abi_audit.report import ExtensionReport, PackageReport
13
# Temporary allowlist of extensions not yet on the stable ABI.
14
# Shrink and remove over time.
15
ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = (
16
- "_flashkda_C.abi3.so",
16
"vllm_flash_attn/_vllm_fa3_C.abi3.so",
17
"third_party/deep_gemm/_C*.so",
18
)
cmake/external_projects/flashkda.cmake
+7
-1
@@ -13,7 +13,7 @@ else()
13
FetchContent_Declare(
14
flashkda
15
GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git
16
- GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f
16
+ GIT_TAG b5d11010ff01c1d4a683c0dde42e76cbeaa8107f
17
GIT_PROGRESS TRUE
18
GIT_SUBMODULES cutlass
19
)
@@ -63,6 +63,12 @@ if(FLASH_KDA_ARCHS)
63
USE_SABI 3
64
WITH_SOABI)
65
66
+ target_compile_definitions(_flashkda_C PRIVATE
67
+ TORCH_TARGET_VERSION=0x020B000000000000ULL)
68
+ if(VLLM_GPU_LANG STREQUAL "CUDA")
69
+ target_compile_definitions(_flashkda_C PRIVATE USE_CUDA)
70
+ endif()
71
+
72
target_compile_options(_flashkda_C PRIVATE
73
$<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3>
74
$<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
csrc/flashkda_registration.cpp
+12
-5
@@ -1,17 +1,24 @@
1
#include "core/registration.h"
2
#include "flash_kda.h"
3
4
-TORCH_LIBRARY(_flashkda_C, m) {
5
- m.def("get_workspace_size(int T_total, int H, int N=1) -> int",
6
- &get_workspace_size);
4
+#include <torch/csrc/stable/library.h>
5
+
6
+STABLE_TORCH_LIBRARY(_flashkda_C, m) {
7
+ m.def("get_workspace_size(int T_total, int H, int N=1) -> int");
8
m.def(
9
"fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, "
9
- "Tensor(a!) out, Tensor workspace, Tensor A_log, Tensor dt_bias, "
10
+ "Tensor(a!) out, Tensor(c!) workspace, Tensor A_log, Tensor dt_bias, "
11
"float lower_bound, "
12
"Tensor? initial_state=None, Tensor(b!)? final_state=None, "
13
"Tensor? cu_seqlens=None) -> ()");
14
}
15
15
-TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) { m.impl("fwd", &fwd); }
16
+STABLE_TORCH_LIBRARY_IMPL(_flashkda_C, CompositeExplicitAutograd, m) {
17
+ m.impl("get_workspace_size", TORCH_BOX(&get_workspace_size));
18
+}
19
+
20
+STABLE_TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) {
21
+ m.impl("fwd", TORCH_BOX(&fwd));
22
+}
23
24
REGISTER_EXTENSION(_flashkda_C)