[Kimi-K3] Bump FlashKDA to fix unstable inverse (#54859)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
Sep 2, 2026 at 15:16 UTC
f4e61361462882f82d91f60dfc4133807ef00822
2 files changed
+51
-1
cmake/external_projects/flashkda.cmake
+1
-1
@@ -13,7 +13,7 @@ else()
13
FetchContent_Declare(
14
flashkda
15
GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git
16
- GIT_TAG ee0be888cd0e972f9409bf53756f8c38c6652173
16
+ GIT_TAG 3b225bf26bb8e218928a1fe14751cb48cf31d11b
17
GIT_PROGRESS TRUE
18
GIT_SUBMODULES cutlass
19
)
tests/models/kimi_k3/test_kda.py
+50
@@ -1083,6 +1083,56 @@ def test_fused_kda_decode_rejects_speculative_conv_state():
1083
)
1084
1085
1086
+@torch.inference_mode()
1087
+def test_flashkda_near_collinear_keys_remain_finite():
1088
+ """Guard against unstable inversion of near-collinear key blocks."""
1089
+ lower_bound = -5.0
1090
+ if not is_flashkda_supported(128, torch.bfloat16, lower_bound):
1091
+ pytest.skip("FlashKDA is not supported on this platform")
1092
+
1093
+ import vllm._flashkda_C # noqa: F401
1094
+
1095
+ T, H, D = 16384, 1, 128
1096
+ torch.manual_seed(0)
1097
+ key = torch.randn(1, 1, H, D, dtype=torch.bfloat16, device=DEVICE)
1098
+ qk = key.expand(1, T, H, D).contiguous()
1099
+ value_block = torch.randn(1, 16, H, D, dtype=torch.bfloat16, device=DEVICE)
1100
+ value = value_block.repeat(1, T // 16, 1, 1)
1101
+ raw_gate = torch.full_like(qk, -12.0)
1102
+ raw_beta = torch.full((1, T, H), 8.0, dtype=qk.dtype, device=DEVICE)
1103
+ A_log = torch.zeros(H, dtype=torch.float32, device=DEVICE)
1104
+ dt_bias = torch.zeros(H, D, dtype=torch.float32, device=DEVICE)
1105
+ initial_state = torch.zeros(1, H, D, D, dtype=torch.float32, device=DEVICE)
1106
+ final_state = torch.empty_like(initial_state)
1107
+ output = torch.empty_like(value)
1108
+ cu_seqlens = torch.tensor([0, T], dtype=torch.int32, device=DEVICE)
1109
+ workspace = torch.empty(
1110
+ torch.ops._flashkda_C.get_workspace_size(T, H, 1),
1111
+ dtype=torch.uint8,
1112
+ device=DEVICE,
1113
+ )
1114
+
1115
+ torch.ops._flashkda_C.fwd(
1116
+ qk,
1117
+ qk,
1118
+ value,
1119
+ raw_gate,
1120
+ raw_beta,
1121
+ D**-0.5,
1122
+ output,
1123
+ workspace,
1124
+ A_log,
1125
+ dt_bias,
1126
+ lower_bound,
1127
+ initial_state,
1128
+ final_state,
1129
+ cu_seqlens,
1130
+ )
1131
+
1132
+ assert torch.isfinite(output).all()
1133
+ assert torch.isfinite(final_state).all()
1134
+
1135
+
1136
@torch.inference_mode()
1137
def test_flashkda_correctness():
1138
if not is_flashkda_supported(128, torch.bfloat16, -3.0):