diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake index b70042f2f4..6ffb46e6a7 100644 --- a/cmake/external_projects/flashkda.cmake +++ b/cmake/external_projects/flashkda.cmake @@ -13,7 +13,7 @@ else() FetchContent_Declare( flashkda GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git - GIT_TAG ee0be888cd0e972f9409bf53756f8c38c6652173 + GIT_TAG 3b225bf26bb8e218928a1fe14751cb48cf31d11b GIT_PROGRESS TRUE GIT_SUBMODULES cutlass ) diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index e822ca919e..e40d2e716d 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -1083,6 +1083,56 @@ def test_fused_kda_decode_rejects_speculative_conv_state(): ) +@torch.inference_mode() +def test_flashkda_near_collinear_keys_remain_finite(): + """Guard against unstable inversion of near-collinear key blocks.""" + lower_bound = -5.0 + if not is_flashkda_supported(128, torch.bfloat16, lower_bound): + pytest.skip("FlashKDA is not supported on this platform") + + import vllm._flashkda_C # noqa: F401 + + T, H, D = 16384, 1, 128 + torch.manual_seed(0) + key = torch.randn(1, 1, H, D, dtype=torch.bfloat16, device=DEVICE) + qk = key.expand(1, T, H, D).contiguous() + value_block = torch.randn(1, 16, H, D, dtype=torch.bfloat16, device=DEVICE) + value = value_block.repeat(1, T // 16, 1, 1) + raw_gate = torch.full_like(qk, -12.0) + raw_beta = torch.full((1, T, H), 8.0, dtype=qk.dtype, device=DEVICE) + A_log = torch.zeros(H, dtype=torch.float32, device=DEVICE) + dt_bias = torch.zeros(H, D, dtype=torch.float32, device=DEVICE) + initial_state = torch.zeros(1, H, D, D, dtype=torch.float32, device=DEVICE) + final_state = torch.empty_like(initial_state) + output = torch.empty_like(value) + cu_seqlens = torch.tensor([0, T], dtype=torch.int32, device=DEVICE) + workspace = torch.empty( + torch.ops._flashkda_C.get_workspace_size(T, H, 1), + dtype=torch.uint8, + device=DEVICE, + ) + + torch.ops._flashkda_C.fwd( + qk, + qk, + value, + raw_gate, + raw_beta, + D**-0.5, + output, + workspace, + A_log, + dt_bias, + lower_bound, + initial_state, + final_state, + cu_seqlens, + ) + + assert torch.isfinite(output).all() + assert torch.isfinite(final_state).all() + + @torch.inference_mode() def test_flashkda_correctness(): if not is_flashkda_supported(128, torch.bfloat16, -3.0):