From 9d88ceb02694c6e3df182ec3c87201284f50e2c6 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Fri, 11 Sep 2026 19:53:24 -0400 Subject: [PATCH] [KDA] Update flashKDA to support bf16 checkpoint state (#56485) Signed-off-by: wzhao18 --- cmake/external_projects/flashkda.cmake | 2 +- tests/models/kimi_k3/test_kda.py | 21 ++++++++++++++------- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake index 6ffb46e6a7..5ac4f4b415 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 3b225bf26bb8e218928a1fe14751cb48cf31d11b + GIT_TAG b59532f1f464fbd536272780e30df5bf6a2ccc02 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 1fd4982e76..84a5e7eed4 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -1421,14 +1421,21 @@ def test_flashinfer_kda_prefill_breakable_graph_cross_stream(): torch.testing.assert_close(graph_value, torch.full_like(graph_value, 2)) +@pytest.mark.parametrize( + ("state_dtype", "tolerance"), + [ + pytest.param(torch.bfloat16, 0.03, id="bf16"), + pytest.param(torch.float32, 0.01, id="fp32"), + ], +) @torch.inference_mode() -def test_flashkda_checkpoint_correctness(): +def test_flashkda_checkpoint_correctness(state_dtype: torch.dtype, tolerance: float): lower_bound = -3.0 - _require_kda_prefill_backend("flashkda", torch.float32, lower_bound) + _require_kda_prefill_backend("flashkda", state_dtype, lower_bound) import vllm._flashkda_C # noqa: F401 - inputs = _make_kda_prefill_inputs(torch.float32, lower_bound=lower_bound) + inputs = _make_kda_prefill_inputs(state_dtype, lower_bound=lower_bound) q, k, v = inputs.q, inputs.k, inputs.v raw_g, raw_beta = inputs.raw_g, inputs.raw_beta A_log, dt_bias = inputs.A_log, inputs.dt_bias @@ -1480,12 +1487,12 @@ def test_flashkda_checkpoint_correctness(): checkpoint_offsets=checkpoint_offsets, ) - assert_close("checkpoint_o", expected_out, checkpoint_out, 0.01) - assert_close("checkpoint_ht", expected_state, checkpoint_final_state, 0.01) - assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], 0.01) + assert_close("checkpoint_o", expected_out, checkpoint_out, tolerance) + assert_close("checkpoint_ht", expected_state, checkpoint_final_state, tolerance) + assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], tolerance) conv_state = torch.zeros(2, H * D, 3, dtype=q.dtype, device=DEVICE) - recurrent_storage = torch.zeros(2, H * D * D + 8, device=DEVICE) + recurrent_storage = torch.zeros(2, H * D * D + 8, dtype=state_dtype, device=DEVICE) recurrent_state = recurrent_storage[:, : H * D * D].view(2, H, D, D) conv_input = q[0].flatten(1) checkpoint_state_indices = torch.tensor(