[KDA] Update flashKDA to support bf16 checkpoint state (#56485)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Sep 11, 2026 at 19:53 UTC
9d88ceb02694c6e3df182ec3c87201284f50e2c6
2 files changed
+15
-8
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 3b225bf26bb8e218928a1fe14751cb48cf31d11b
16
+ GIT_TAG b59532f1f464fbd536272780e30df5bf6a2ccc02
17
GIT_PROGRESS TRUE
18
GIT_SUBMODULES cutlass
19
)
tests/models/kimi_k3/test_kda.py
+14
-7
@@ -1421,14 +1421,21 @@ def test_flashinfer_kda_prefill_breakable_graph_cross_stream():
1421
torch.testing.assert_close(graph_value, torch.full_like(graph_value, 2))
1422
1423
1424
+@pytest.mark.parametrize(
1425
+ ("state_dtype", "tolerance"),
1426
+ [
1427
+ pytest.param(torch.bfloat16, 0.03, id="bf16"),
1428
+ pytest.param(torch.float32, 0.01, id="fp32"),
1429
+ ],
1430
+)
1431
@torch.inference_mode()
1425
-def test_flashkda_checkpoint_correctness():
1432
+def test_flashkda_checkpoint_correctness(state_dtype: torch.dtype, tolerance: float):
1433
lower_bound = -3.0
1427
- _require_kda_prefill_backend("flashkda", torch.float32, lower_bound)
1434
+ _require_kda_prefill_backend("flashkda", state_dtype, lower_bound)
1435
1436
import vllm._flashkda_C # noqa: F401
1437
1431
- inputs = _make_kda_prefill_inputs(torch.float32, lower_bound=lower_bound)
1438
+ inputs = _make_kda_prefill_inputs(state_dtype, lower_bound=lower_bound)
1439
q, k, v = inputs.q, inputs.k, inputs.v
1440
raw_g, raw_beta = inputs.raw_g, inputs.raw_beta
1441
A_log, dt_bias = inputs.A_log, inputs.dt_bias
@@ -1480,12 +1487,12 @@ def test_flashkda_checkpoint_correctness():
1487
checkpoint_offsets=checkpoint_offsets,
1488
)
1489
1483
- assert_close("checkpoint_o", expected_out, checkpoint_out, 0.01)
1484
- assert_close("checkpoint_ht", expected_state, checkpoint_final_state, 0.01)
1485
- assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], 0.01)
1490
+ assert_close("checkpoint_o", expected_out, checkpoint_out, tolerance)
1491
+ assert_close("checkpoint_ht", expected_state, checkpoint_final_state, tolerance)
1492
+ assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], tolerance)
1493
1494
conv_state = torch.zeros(2, H * D, 3, dtype=q.dtype, device=DEVICE)
1488
- recurrent_storage = torch.zeros(2, H * D * D + 8, device=DEVICE)
1495
+ recurrent_storage = torch.zeros(2, H * D * D + 8, dtype=state_dtype, device=DEVICE)
1496
recurrent_state = recurrent_storage[:, : H * D * D].view(2, H, D, D)
1497
conv_input = q[0].flatten(1)
1498
checkpoint_state_indices = torch.tensor(