[Perf] Integrate FlashInfer KDA kernels (#55364)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

tests/models/kimi_k3/test_kda.py

@@ -13,6 +13,7 @@ import torch
 import torch.nn.functional as F
 
 from vllm import _custom_ops as ops
+from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
 from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update
 from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
     gather_initial_states,
@@ -22,8 +23,11 @@ from vllm.models.kimi_k3.amd.ops.third_party.kda import (
 )
 from vllm.models.kimi_k3.nvidia import kda as nvidia_kda
 from vllm.models.kimi_k3.nvidia.kda import (
+    _flashinfer_kda_prefill,
     _flashkda_prefill,
     _store_cache_checkpoints_kernel,
+    is_flashinfer_fused_kda_decode_supported,
+    is_flashinfer_recurrent_kda_prefill_supported,
     is_flashkda_supported,
     is_fused_kda_decode_supported,
 )
@@ -43,6 +47,7 @@ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
 )
 from vllm.platforms import current_platform
 from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
+from vllm.utils.flashinfer import flashinfer_fused_kda_decode
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 
 DEVICE = current_platform.device_type
@@ -87,6 +92,7 @@ def test_kda_recoverssm_config_state_layout():
         ),
         cache_config=SimpleNamespace(
             mamba_cache_dtype="auto",
+            mamba_ssm_cache_dtype="auto",
             use_kda_recoverssm=True,
         ),
         parallel_config=SimpleNamespace(tensor_parallel_size=1),
@@ -909,6 +915,7 @@ def test_kda_recoverssm_verify_and_group_commit(
         (96, 1, -5.0, True, "DS"),
     ],
 )
+@pytest.mark.parametrize("decode_backend", ["native", "flashinfer"])
 @torch.inference_mode()
 def test_fused_kda_decode_correctness(
     num_heads: int,
@@ -916,17 +923,36 @@ def test_fused_kda_decode_correctness(
     lower_bound: float | None,
     fuse_output_norm: bool,
     conv_layout: str,
+    decode_backend: str,
 ):
     D, W = 128, 4
-    if not is_fused_kda_decode_supported(
-        num_heads,
-        D,
-        W,
-        num_spec=0,
-        input_dtype=torch.bfloat16,
-        conv_state_dtype=torch.bfloat16,
-    ):
-        pytest.skip("Fused KDA decode is not supported on this platform")
+    state_dtype = torch.bfloat16 if decode_backend == "flashinfer" else torch.float32
+    if decode_backend == "flashinfer":
+        if conv_layout == "DS":
+            pytest.skip("FlashInfer fused decode requires SD conv-state layout")
+        if not fuse_output_norm:
+            pytest.skip("FlashInfer's fused decode always applies output norm")
+        supported = is_flashinfer_fused_kda_decode_supported(
+            num_heads,
+            D,
+            W,
+            num_spec=0,
+            input_dtype=torch.bfloat16,
+            conv_state_dtype=torch.bfloat16,
+            recurrent_state_dtype=state_dtype,
+        )
+    else:
+        supported = is_fused_kda_decode_supported(
+            num_heads,
+            D,
+            W,
+            num_spec=0,
+            input_dtype=torch.bfloat16,
+            conv_state_dtype=torch.bfloat16,
+            recurrent_state_dtype=state_dtype,
+        )
+    if not supported:
+        pytest.skip(f"{decode_backend} fused KDA decode is not supported")
     torch.manual_seed(967 + num_heads + num_seqs + (conv_layout == "DS"))
     dim = num_heads * D
     slots = num_seqs + 2
@@ -995,7 +1021,7 @@ def test_fused_kda_decode_correctness(
         D,
         dtype=torch.float32,
         device=DEVICE,
-    )
+    ).to(state_dtype)
 
     conv_ref = conv_seed.clone()
     state_ref = state_seed.clone()
@@ -1030,7 +1056,7 @@ def test_fused_kda_decode_correctness(
     conv_slot_elements = 3 * dim * (W - 1)
     state_slot_elements = num_heads * D * D
     conv_slot_bytes = conv_slot_elements * torch.bfloat16.itemsize
-    page_bytes = conv_slot_bytes + state_slot_elements * torch.float32.itemsize
+    page_bytes = conv_slot_bytes + state_slot_elements * state_dtype.itemsize
     cache_storage = torch.empty(slots * page_bytes, dtype=torch.uint8, device=DEVICE)
     conv_actual = torch.as_strided(
         cache_storage.view(torch.bfloat16),
@@ -1042,30 +1068,57 @@ def test_fused_kda_decode_correctness(
         ),
     )
     state_actual = torch.as_strided(
-        cache_storage.view(torch.float32),
+        cache_storage.view(state_dtype),
         size=(slots, num_heads, D, D),
-        stride=(page_bytes // torch.float32.itemsize, D * D, D, 1),
-        storage_offset=conv_slot_bytes // torch.float32.itemsize,
+        stride=(page_bytes // state_dtype.itemsize, D * D, D, 1),
+        storage_offset=conv_slot_bytes // state_dtype.itemsize,
     )
     conv_actual.copy_(conv_seed)
     state_actual.copy_(state_seed)
     fused_weight = weight.reshape(3, dim, W).transpose(1, 2).contiguous()
-    actual = ops.fused_kda_decode(
-        x=packed_x,
-        weight=fused_weight,
-        bias=None,
-        conv_state=conv_actual,
-        raw_g=raw_g,
-        raw_beta=raw_beta,
-        A_log=A_log,
-        dt_bias=dt_bias,
-        state_indices=state_indices,
-        state=state_actual,
-        lower_bound=lower_bound,
-        output_gate=output_gate if fuse_output_norm else None,
-        norm_weight=norm_weight if fuse_output_norm else None,
-        norm_eps=norm_eps,
-    )
+    if decode_backend == "flashinfer":
+        output = torch.empty(
+            1,
+            num_seqs,
+            num_heads,
+            D,
+            dtype=torch.bfloat16,
+            device=DEVICE,
+        )
+
+        actual = flashinfer_fused_kda_decode(
+            x=packed_x,
+            weight=fused_weight,
+            conv_state=conv_actual,
+            raw_gate=raw_g,
+            raw_beta=raw_beta,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            state_indices=state_indices,
+            state=state_actual,
+            output_gate=output_gate,
+            norm_weight=norm_weight,
+            lower_bound=lower_bound,
+            norm_eps=norm_eps,
+            output=output,
+        )
+    else:
+        actual = ops.fused_kda_decode(
+            x=packed_x,
+            weight=fused_weight,
+            bias=None,
+            conv_state=conv_actual,
+            raw_g=raw_g,
+            raw_beta=raw_beta,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            state_indices=state_indices,
+            state=state_actual,
+            lower_bound=lower_bound,
+            output_gate=output_gate if fuse_output_norm else None,
+            norm_weight=norm_weight if fuse_output_norm else None,
+            norm_eps=norm_eps,
+        )
 
     torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2)
     torch.testing.assert_close(conv_actual, conv_ref, atol=0, rtol=0)
@@ -1080,17 +1133,171 @@ def test_fused_kda_decode_rejects_speculative_conv_state():
         num_spec=2,
         input_dtype=torch.bfloat16,
         conv_state_dtype=torch.bfloat16,
+        recurrent_state_dtype=torch.float32,
     )
 
 
+def _make_kda_prefill_inputs(
+    state_dtype: torch.dtype,
+    *,
+    lower_bound: float,
+) -> SimpleNamespace:
+    B, T, H, D = 1, 48, 2, 128
+    torch.manual_seed(11)
+    q, k, v, raw_g = [
+        torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4)
+    ]
+    raw_beta = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE)
+    A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5
+    dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1
+    initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE).to(
+        state_dtype
+    )
+    cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE)
+    return SimpleNamespace(
+        q=q,
+        k=k,
+        v=v,
+        raw_g=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        initial_state=initial_state,
+        cu_seqlens=cu_seqlens,
+        lower_bound=lower_bound,
+    )
+
+
+def _require_kda_prefill_backend(
+    backend: str,
+    state_dtype: torch.dtype,
+    lower_bound: float,
+) -> None:
+    if backend == "flashinfer":
+        supported = is_flashinfer_recurrent_kda_prefill_supported(
+            128, torch.bfloat16, state_dtype, lower_bound
+        )
+    else:
+        assert backend == "flashkda"
+        supported = is_flashkda_supported(128, torch.bfloat16, state_dtype, lower_bound)
+    if not supported:
+        pytest.skip(f"{backend} KDA prefill is not supported on this platform")
+
+
+def _run_kda_prefill_backend(
+    backend: str,
+    *,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    initial_state: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    lower_bound: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    initial_state = initial_state.clone()
+    output = torch.empty_like(v)
+    if backend == "flashinfer":
+        flashinfer_query_start_loc = cu_seqlens.to(torch.int64)
+        seq_order = None
+        if q.shape[1] > initial_state.shape[0]:
+            seq_order = torch.argsort(
+                flashinfer_query_start_loc.diff(), descending=True
+            ).to(torch.int32)
+        return _flashinfer_kda_prefill(
+            q=q,
+            k=k,
+            v=v,
+            raw_g=raw_g,
+            raw_beta=raw_beta,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            lower_bound=lower_bound,
+            initial_state=initial_state,
+            cu_seqlens=flashinfer_query_start_loc,
+            out=output,
+            seq_order=seq_order,
+        )
+
+    assert backend == "flashkda"
+    import vllm._flashkda_C  # noqa: F401
+
+    final_state = torch.empty_like(initial_state)
+    workspace = torch.empty(
+        torch.ops._flashkda_C.get_workspace_size(
+            q.shape[1], q.shape[2], cu_seqlens.numel() - 1
+        ),
+        dtype=torch.uint8,
+        device=q.device,
+    )
+    return _flashkda_prefill(
+        q=q,
+        k=k,
+        v=v,
+        g=raw_g,
+        beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        lower_bound=lower_bound,
+        initial_state=initial_state,
+        cu_seqlens=cu_seqlens,
+        out=output,
+        final_state=final_state,
+        workspace=workspace,
+    )
+
+
+def _kda_prefill_reference(
+    inputs: SimpleNamespace,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    gate = inputs.lower_bound * torch.sigmoid(
+        inputs.A_log.exp()[None, None, :, None]
+        * (inputs.raw_g.float() + inputs.dt_bias[None, None])
+    )
+    beta = inputs.raw_beta.float().sigmoid()
+    q_norm = l2norm_fwd(inputs.q.contiguous())
+    k_norm = l2norm_fwd(inputs.k.contiguous())
+    expected_outputs = []
+    expected_states = []
+    for i, (start, end) in enumerate(
+        zip(inputs.cu_seqlens[:-1].tolist(), inputs.cu_seqlens[1:].tolist())
+    ):
+        output, final_state = naive_recurrent_kda(
+            q_norm[:, start:end],
+            k_norm[:, start:end],
+            inputs.v[:, start:end],
+            gate[:, start:end],
+            beta[:, start:end],
+            initial_state=inputs.initial_state[i : i + 1].transpose(-1, -2),
+            output_final_state=True,
+        )
+        assert final_state is not None
+        expected_outputs.append(output)
+        expected_states.append(final_state)
+    return (
+        torch.cat(expected_outputs, dim=1),
+        torch.cat(expected_states).transpose(-1, -2).contiguous(),
+    )
+
+
+@pytest.mark.parametrize(
+    ("backend", "state_dtype"),
+    [
+        pytest.param("flashkda", torch.float32, id="flashkda"),
+        pytest.param("flashinfer", torch.bfloat16, id="flashinfer"),
+    ],
+)
 @torch.inference_mode()
-def test_flashkda_near_collinear_keys_remain_finite():
+def test_kda_prefill_near_collinear_keys_remain_finite(
+    backend: str,
+    state_dtype: torch.dtype,
+):
     """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
+    _require_kda_prefill_backend(backend, state_dtype, lower_bound)
 
     T, H, D = 16384, 1, 128
     torch.manual_seed(0)
@@ -1102,81 +1309,131 @@ def test_flashkda_near_collinear_keys_remain_finite():
     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)
+    initial_state = torch.zeros(1, H, D, D, dtype=state_dtype, device=DEVICE)
     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,
+    output, final_state = _run_kda_prefill_backend(
+        backend,
+        q=qk,
+        k=qk,
+        v=value,
+        raw_g=raw_gate,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        initial_state=initial_state,
+        cu_seqlens=cu_seqlens,
+        lower_bound=lower_bound,
     )
 
     assert torch.isfinite(output).all()
     assert torch.isfinite(final_state).all()
 
 
+@pytest.mark.parametrize(
+    ("backend", "state_dtype", "tolerance"),
+    [
+        pytest.param("flashinfer", torch.bfloat16, 0.03, id="flashinfer-bf16"),
+        pytest.param("flashkda", torch.bfloat16, 0.03, id="flashkda-bf16"),
+        pytest.param("flashkda", torch.float32, 0.01, id="flashkda-fp32"),
+    ],
+)
 @torch.inference_mode()
-def test_flashkda_correctness():
-    if not is_flashkda_supported(128, torch.bfloat16, -3.0):
-        pytest.skip("FlashKDA is not supported on this platform")
+def test_kda_prefill_correctness(
+    backend: str,
+    state_dtype: torch.dtype,
+    tolerance: float,
+):
+    lower_bound = -5.0
+    _require_kda_prefill_backend(backend, state_dtype, lower_bound)
+    inputs = _make_kda_prefill_inputs(state_dtype, lower_bound=lower_bound)
+    expected_out, expected_state = _kda_prefill_reference(inputs)
+    actual_out, actual_state = _run_kda_prefill_backend(
+        backend,
+        **vars(inputs),
+    )
 
-    import vllm._flashkda_C  # noqa: F401
+    assert_close("o", expected_out, actual_out, tolerance)
+    assert_close("ht", expected_state, actual_state, tolerance)
 
-    B, T, H, D = 1, 48, 2, 128
-    torch.manual_seed(11)
+
+@torch.inference_mode()
+def test_flashinfer_kda_prefill_breakable_graph_cross_stream():
+    if not is_flashinfer_recurrent_kda_prefill_supported(
+        128,
+        torch.bfloat16,
+        torch.bfloat16,
+        -5.0,
+    ):
+        pytest.skip("FlashInfer KDA prefill is not supported on this platform")
+
+    B, T, H, D = 1, 8, 12, 128
     q, k, v, raw_g = [
         torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4)
     ]
-    beta_logits = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE)
-    A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5
-    dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1
-    initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE)
-    cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE)
-    lower_bound = -3.0
+    raw_beta = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE)
+    initial_state = torch.randn(
+        B,
+        H,
+        D,
+        D,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    kwargs = {
+        "q": q,
+        "k": k,
+        "v": v,
+        "raw_g": raw_g,
+        "raw_beta": raw_beta,
+        "A_log": torch.randn(H, dtype=torch.float32, device=DEVICE),
+        "dt_bias": torch.randn(H, D, dtype=torch.float32, device=DEVICE),
+        "lower_bound": -5.0,
+        "initial_state": initial_state,
+        "cu_seqlens": torch.tensor([0, T], dtype=torch.int64, device=DEVICE),
+        "out": torch.empty_like(q),
+        "seq_order": torch.zeros(B, dtype=torch.int32, device=DEVICE),
+    }
 
+    capture_stream = torch.Stream(device=DEVICE)
+    original_stream = torch.accelerator.current_stream()
+    torch.accelerator.set_stream(capture_stream)
+    try:
+        graph_value = torch.zeros(1, device=DEVICE)
+        capture = BreakableCUDAGraphCapture()
+        with capture:
+            graph_value.add_(1)
+            capture.add_eager(lambda: _flashinfer_kda_prefill(**kwargs))
+            graph_value.add_(1)
+        capture_stream.synchronize()
+    finally:
+        torch.accelerator.set_stream(original_stream)
+
+    torch.testing.assert_close(graph_value, torch.zeros_like(graph_value))
+    capture.replay()
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(graph_value, torch.full_like(graph_value, 2))
+
+
+@torch.inference_mode()
+def test_flashkda_checkpoint_correctness():
+    lower_bound = -3.0
+    _require_kda_prefill_backend("flashkda", torch.float32, lower_bound)
+
+    import vllm._flashkda_C  # noqa: F401
+
+    inputs = _make_kda_prefill_inputs(torch.float32, 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
+    initial_state, cu_seqlens = inputs.initial_state, inputs.cu_seqlens
+    _, T, H, D = q.shape
+    expected_out, expected_state = _kda_prefill_reference(inputs)
     gate = lower_bound * torch.sigmoid(
         A_log.exp()[None, None, :, None] * (raw_g.float() + dt_bias[None, None, :, :])
     )
-    beta = beta_logits.float().sigmoid()
+    beta = raw_beta.float().sigmoid()
     q_norm = l2norm_fwd(q.contiguous())
     k_norm = l2norm_fwd(k.contiguous())
-
-    expected_outputs = []
-    expected_states = []
-    for i, (start, end) in enumerate(
-        zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist())
-    ):
-        output, final_state = naive_recurrent_kda(
-            q_norm[:, start:end],
-            k_norm[:, start:end],
-            v[:, start:end],
-            gate[:, start:end],
-            beta[:, start:end],
-            initial_state=initial_state[i].transpose(-1, -2),
-            output_final_state=True,
-        )
-        expected_outputs.append(output)
-        expected_states.append(final_state)
-    expected_out = torch.cat(expected_outputs, dim=1)
-    expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous()
     _, expected_checkpoint = naive_recurrent_kda(
         q_norm[:, :16],
         k_norm[:, :16],
@@ -1189,33 +1446,11 @@ def test_flashkda_correctness():
     assert expected_checkpoint is not None
     expected_checkpoint = expected_checkpoint.transpose(-1, -2).contiguous()
 
-    actual_out = torch.empty_like(v)
-    actual_state = torch.empty_like(initial_state)
     workspace = torch.empty(
         torch.ops._flashkda_C.get_workspace_size(T, H, cu_seqlens.numel() - 1),
         dtype=torch.uint8,
         device=DEVICE,
     )
-    torch.ops._flashkda_C.fwd(
-        q,
-        k,
-        v,
-        raw_g,
-        beta_logits,
-        D**-0.5,
-        actual_out,
-        workspace,
-        A_log,
-        dt_bias,
-        lower_bound,
-        initial_state,
-        actual_state,
-        cu_seqlens,
-    )
-
-    assert_close("o", expected_out, actual_out, 0.01)
-    assert_close("ht", expected_state, actual_state, 0.01)
-
     checkpoint_out = torch.empty_like(v)
     checkpoint_final_state = torch.empty_like(initial_state)
     checkpoint_state = torch.empty_like(initial_state)
@@ -1225,7 +1460,7 @@ def test_flashkda_correctness():
         k=k,
         v=v,
         g=raw_g,
-        beta=beta_logits,
+        beta=raw_beta,
         A_log=A_log,
         dt_bias=dt_bias,
         lower_bound=lower_bound,

tests/models/kimi_k3/test_kda_metadata.py

@@ -49,6 +49,8 @@ PRUNED_METADATA_FIELDS = {
     "prefill_state_indices",
     "prefill_has_initial_state",
     "spec_sequence_masks",
+    "flashinfer_prefill_query_start_loc",
+    "flashinfer_prefill_seq_order",
 }
 
 
@@ -59,10 +61,10 @@ def _assert_matches_shared_gdn(
     assert actual.recoverssm_context is None
     for field in fields(GDNAttentionMetadata):
         actual_value = getattr(actual, field.name)
-        expected_value = getattr(reference, field.name)
         if field.name in PRUNED_METADATA_FIELDS:
             assert actual_value is None
             continue
+        expected_value = getattr(reference, field.name)
         if (
             field.name in {"spec_token_indx", "non_spec_token_indx"}
             and actual.num_spec_decodes > 0
@@ -157,6 +159,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
         ),
         cache_config=SimpleNamespace(
             mamba_cache_dtype="auto",
+            mamba_ssm_cache_dtype="auto",
             use_kda_recoverssm=True,
         ),
         parallel_config=SimpleNamespace(tensor_parallel_size=1),

vllm/engine/arg_utils.py

@@ -772,7 +772,10 @@ class EngineArgs:
 
     fail_on_environ_validation: bool = False
     gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None
-    kda_prefill_backend: Literal["auto", "triton", "flashkda"] | None = None
+    kda_prefill_backend: Literal["auto", "triton", "flashkda", "flashinfer"] | None = (
+        None
+    )
+    kda_decode_backend: Literal["auto", "native", "flashinfer", "triton"] | None = None
 
     def __post_init__(self):
         # support `EngineArgs(compilation_config={...})`
@@ -1747,10 +1750,17 @@ class EngineArgs:
         parser.add_argument(
             "--kda-prefill-backend",
             dest="kda_prefill_backend",
-            choices=["auto", "triton", "flashkda"],
+            choices=["auto", "triton", "flashkda", "flashinfer"],
             default=None,
             help="Select KDA prefill backend.",
         )
+        parser.add_argument(
+            "--kda-decode-backend",
+            dest="kda_decode_backend",
+            choices=["auto", "native", "flashinfer", "triton"],
+            default=None,
+            help="Select KDA decode backend.",
+        )
         return parser
 
     @classmethod
@@ -2567,6 +2577,8 @@ class EngineArgs:
             self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend
         if self.kda_prefill_backend is not None:
             self.additional_config["kda_prefill_backend"] = self.kda_prefill_backend
+        if self.kda_decode_backend is not None:
+            self.additional_config["kda_decode_backend"] = self.kda_decode_backend
 
         config = VllmConfig(
             model_config=model_config,

vllm/model_executor/layers/mamba/mamba_utils.py

@@ -134,9 +134,19 @@ class MambaStateDtypeCalculator:
         cls,
         model_dtype: ModelDType | torch.dtype,
         mamba_cache_dtype: MambaDType,
+        mamba_ssm_cache_dtype: MambaDType = "auto",
     ) -> tuple[torch.dtype, torch.dtype]:
-        state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
-        return (state_dtype, torch.float32)
+        conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
+        if mamba_ssm_cache_dtype == "auto":
+            recurrent_state_dtype = torch.float32
+        elif mamba_ssm_cache_dtype in ("float32", "bfloat16"):
+            recurrent_state_dtype = STR_DTYPE_TO_TORCH_DTYPE[mamba_ssm_cache_dtype]
+        else:
+            raise ValueError(
+                "KDA recurrent state supports only auto, float32, and bfloat16; "
+                f"got {mamba_ssm_cache_dtype}"
+            )
+        return (conv_state_dtype, recurrent_state_dtype)
 
     @classmethod
     def append_kda_recoverssm_record(

vllm/models/kimi_k3/nvidia/kda.py

@@ -50,6 +50,12 @@ from vllm.platforms import current_platform
 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
 from vllm.triton_utils import tl, triton
+from vllm.utils.flashinfer import (
+    flashinfer_fused_kda_decode,
+    flashinfer_recurrent_kda,
+    has_flashinfer_fused_kda_decode,
+    has_flashinfer_recurrent_kda,
+)
 from vllm.v1.attention.backend import AttentionBackend
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.kv_cache_interface import MambaSpec
@@ -155,6 +161,7 @@ def is_fused_kda_decode_supported(
     num_spec: int,
     input_dtype: torch.dtype,
     conv_state_dtype: torch.dtype,
+    recurrent_state_dtype: torch.dtype,
 ) -> bool:
     # The fused kernel handles both conv-state cache layouts (SD and DS); the
     # inner strides are selected from the tensor at launch time.
@@ -165,6 +172,7 @@ def is_fused_kda_decode_supported(
         or num_spec != 0
         or input_dtype != torch.bfloat16
         or conv_state_dtype != torch.bfloat16
+        or recurrent_state_dtype != torch.float32
         or not hasattr(torch.ops._C, "fused_kda_decode")
     ):
         return False
@@ -176,9 +184,93 @@ def is_fused_kda_decode_supported(
     )
 
 
+def is_flashinfer_fused_kda_decode_supported(
+    num_heads: int,
+    head_dim: int,
+    conv_width: int,
+    num_spec: int,
+    input_dtype: torch.dtype,
+    conv_state_dtype: torch.dtype,
+    recurrent_state_dtype: torch.dtype,
+) -> bool:
+    if not has_flashinfer_fused_kda_decode():
+        return False
+    capability = current_platform.get_device_capability()
+    if capability is None:
+        return False
+    compute_capability = (capability.major, capability.minor)
+    return (
+        compute_capability in ((10, 0), (10, 3))
+        and num_heads in (12, 24, 32, 48, 96)
+        and head_dim == 128
+        and conv_width == 4
+        and num_spec == 0
+        and input_dtype == torch.bfloat16
+        and conv_state_dtype == torch.bfloat16
+        and recurrent_state_dtype in (torch.float32, torch.bfloat16)
+        and not is_conv_state_dim_first()
+    )
+
+
+def resolve_kda_decode_backend(
+    backend: str,
+    num_heads: int,
+    head_dim: int,
+    conv_width: int,
+    num_spec: int,
+    input_dtype: torch.dtype,
+    conv_state_dtype: torch.dtype,
+    recurrent_state_dtype: torch.dtype,
+) -> str:
+    if backend not in ("auto", "native", "flashinfer", "triton"):
+        raise ValueError(f"Unsupported KDA decode backend: {backend}")
+    if backend == "triton":
+        return "triton"
+
+    native_supported = is_fused_kda_decode_supported(
+        num_heads,
+        head_dim,
+        conv_width,
+        num_spec,
+        input_dtype,
+        conv_state_dtype,
+        recurrent_state_dtype,
+    )
+    if native_supported and backend in ("auto", "native"):
+        logger.info_once("Using native fused KDA decode backend.")
+        return "native"
+    if backend == "native":
+        raise RuntimeError(
+            "Native fused KDA decode requires a supported CUDA architecture, "
+            "bfloat16 activations and convolution state, float32 recurrent "
+            "state, head_dim=128, convolution width 4, and no speculation."
+        )
+
+    flashinfer_supported = is_flashinfer_fused_kda_decode_supported(
+        num_heads,
+        head_dim,
+        conv_width,
+        num_spec,
+        input_dtype,
+        conv_state_dtype,
+        recurrent_state_dtype,
+    )
+    if flashinfer_supported and backend == "flashinfer":
+        logger.info_once("Using FlashInfer fused KDA decode backend.")
+        return "flashinfer"
+    if backend == "flashinfer":
+        raise RuntimeError(
+            "FlashInfer fused KDA decode requires CUDA SM100 or SM103, "
+            "bfloat16 activations and convolution state, bfloat16 or float32 "
+            "recurrent state, head_dim=128, convolution width 4, no speculation."
+        )
+    return "triton"
+
+
 def is_flashkda_supported(
     head_dim: int,
-    dtype: torch.dtype,
+    input_dtype: torch.dtype,
+    recurrent_state_dtype: torch.dtype,
     lower_bound: float | None,
 ) -> bool:
     if not current_platform.is_cuda():
@@ -188,7 +280,29 @@ def is_flashkda_supported(
         capability is not None
         and capability.major in (9, 10, 12)
         and head_dim == 128
-        and dtype == torch.bfloat16
+        and input_dtype == torch.bfloat16
+        and recurrent_state_dtype in (torch.bfloat16, torch.float32)
+        and lower_bound is not None
+    )
+
+
+def is_flashinfer_recurrent_kda_prefill_supported(
+    head_dim: int,
+    input_dtype: torch.dtype,
+    recurrent_state_dtype: torch.dtype,
+    lower_bound: float | None,
+) -> bool:
+    if not current_platform.is_cuda() or not has_flashinfer_recurrent_kda():
+        return False
+    capability = current_platform.get_device_capability()
+    if capability is None:
+        return False
+    compute_capability = (capability.major, capability.minor)
+    return (
+        compute_capability in ((10, 0), (10, 3))
+        and head_dim == 128
+        and input_dtype == torch.bfloat16
+        and recurrent_state_dtype == torch.bfloat16
         and lower_bound is not None
     )
 
@@ -237,6 +351,44 @@ def _flashkda_prefill(
     return out, final_state
 
 
+def _flashinfer_kda_prefill(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    lower_bound: float,
+    initial_state: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    out: torch.Tensor,
+    seq_order: torch.Tensor | None = None,
+    prefill_workspace: object | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    output, _ = flashinfer_recurrent_kda(
+        q=q.contiguous(),
+        k=k.contiguous(),
+        v=v.contiguous(),
+        g=raw_g.contiguous(),
+        beta=raw_beta.contiguous(),
+        A_log=A_log.contiguous(),
+        dt_bias=dt_bias.contiguous(),
+        scale=q.shape[-1] ** -0.5,
+        initial_state=initial_state.contiguous(),
+        output_final_state=False,
+        use_qk_l2norm_in_kernel=True,
+        use_gate_in_kernel=True,
+        lower_bound=lower_bound,
+        cu_seqlens=cu_seqlens.contiguous(),
+        output=out,
+        beta_is_logit=True,
+        seq_order=seq_order,
+        prefill_workspace=prefill_workspace,
+    )
+    return output, initial_state
+
+
 @triton.jit
 def _store_cache_checkpoints_kernel(
     x_ptr,
@@ -303,18 +455,40 @@ def _store_cache_checkpoints_kernel(
 def resolve_kda_prefill_backend(
     backend: str,
     head_dim: int,
-    dtype: torch.dtype,
+    input_dtype: torch.dtype,
+    recurrent_state_dtype: torch.dtype,
     lower_bound: float | None,
 ) -> str:
-    if backend not in ("auto", "triton", "flashkda"):
+    if backend not in ("auto", "triton", "flashkda", "flashinfer"):
         raise ValueError(f"Unsupported KDA prefill backend: {backend}")
-    supported = is_flashkda_supported(head_dim, dtype, lower_bound)
-    if backend == "flashkda" and not supported:
+    flashinfer_supported = is_flashinfer_recurrent_kda_prefill_supported(
+        head_dim,
+        input_dtype,
+        recurrent_state_dtype,
+        lower_bound,
+    )
+    flashkda_supported = is_flashkda_supported(
+        head_dim,
+        input_dtype,
+        recurrent_state_dtype,
+        lower_bound,
+    )
+    if backend == "flashinfer" and not flashinfer_supported:
         raise RuntimeError(
-            "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, "
-            "head_dim=128, and a bounded KDA gate."
+            "FlashInfer KDA prefill requires CUDA SM100 or SM103, bfloat16 "
+            "activations and recurrent state, head_dim=128, a bounded KDA "
+            "gate, and flashinfer-python 0.6.18 or newer."
         )
-    if supported and backend != "triton":
+    if backend == "flashkda" and not flashkda_supported:
+        raise RuntimeError(
+            "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16 input, "
+            "bfloat16 or float32 recurrent state, head_dim=128, and a bounded "
+            "KDA gate."
+        )
+    if flashinfer_supported and backend == "flashinfer":
+        logger.info_once("Using FlashInfer KDA prefill backend.")
+        return "flashinfer"
+    if flashkda_supported and backend in ("auto", "flashkda"):
         logger.info_once("Using FlashKDA KDA prefill backend.")
         return "flashkda"
     return "triton"
@@ -369,7 +543,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
         if self.model_config is None or self.cache_config is None:
             raise ValueError("model_config and cache_config must be set")
         base_dtypes = MambaStateDtypeCalculator.kda_state_dtype(
-            self.model_config.dtype, self.cache_config.mamba_cache_dtype
+            self.model_config.dtype,
+            self.cache_config.mamba_cache_dtype,
+            self.cache_config.mamba_ssm_cache_dtype,
         )
         if self.cache_config.use_kda_recoverssm:
             return MambaStateDtypeCalculator.append_kda_recoverssm_record(
@@ -486,17 +662,25 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
         self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
         # Keep a width-major copy for fused decode without changing the layout
         # consumed by the prefill and fallback decode kernels.
-        conv_state_dtype = self.get_state_dtype()[0]
-        decode_conv1d_weight = None
-        if is_fused_kda_decode_supported(
+        conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()[:2]
+        additional_config = vllm_config.additional_config
+        decode_backend = (
+            additional_config.get("kda_decode_backend", "auto")
+            if isinstance(additional_config, dict)
+            else "auto"
+        )
+        self.kda_decode_backend = resolve_kda_decode_backend(
+            decode_backend,
             self.local_num_heads,
             self.head_dim,
             self.conv_size,
             self.num_spec,
             vllm_config.model_config.dtype,
             conv_state_dtype,
-        ):
-            logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.")
+            recurrent_state_dtype,
+        )
+        decode_conv1d_weight = None
+        if self.kda_decode_backend != "triton":
             decode_conv1d_weight = torch.empty(
                 3,
                 self.conv_size,
@@ -533,7 +717,6 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                 f"Got {self.gate_lower_bound}."
             )
 
-        additional_config = vllm_config.additional_config
         backend = (
             additional_config.get("kda_prefill_backend", "auto")
             if isinstance(additional_config, dict)
@@ -543,11 +726,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
             backend,
             self.head_dim,
             vllm_config.model_config.dtype,
+            recurrent_state_dtype,
             self.gate_lower_bound,
         )
         self._flashkda_buffer_specs: (
             tuple[tuple[tuple[int, ...], torch.dtype], ...] | None
         ) = None
+        self._flashinfer_kda_output_spec: tuple[tuple[int, ...], torch.dtype] | None = (
+            None
+        )
         if self.kda_prefill_backend == "flashkda":
             T = vllm_config.scheduler_config.max_num_batched_tokens
             N = vllm_config.scheduler_config.max_num_seqs
@@ -561,6 +748,13 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                 ((N, H, D, D), self.get_state_dtype()[1]),
                 ((workspace_size,), torch.uint8),
             )
+        elif self.kda_prefill_backend == "flashinfer":
+            T = vllm_config.scheduler_config.max_num_batched_tokens
+            H, D = self.local_num_heads, self.head_dim
+            self._flashinfer_kda_output_spec = (
+                (1, T, H, D),
+                self.model_config.dtype,
+            )
 
         self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
         decode_norm_weight = None
@@ -725,30 +919,50 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
             conv_state = conv_state.transpose(-1, -2)
 
         if (
-            self.decode_conv1d_weight is not None
+            self.kda_decode_backend != "triton"
+            and self.decode_conv1d_weight is not None
             and self.decode_norm_weight is not None
             and not has_spec_decode
             and m.num_prefills == 0
             and m.num_decodes > 0
         ):
             assert non_spec_state_indices_tensor is not None
-            ops.fused_kda_decode(
-                x=mixed_qkv,
-                weight=self.decode_conv1d_weight,
-                bias=self.conv1d.bias,
-                conv_state=conv_state,
-                raw_g=g1,
-                raw_beta=beta,
-                A_log=self.A_log,
-                dt_bias=self.dt_bias,
-                state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
-                state=recurrent_state,
-                out=core_attn_out[:, :num_actual_tokens],
-                lower_bound=self.gate_lower_bound,
-                output_gate=g2[:num_actual_tokens],
-                norm_weight=self.decode_norm_weight,
-                norm_eps=self.o_norm.eps,
-            )
+            state_indices = non_spec_state_indices_tensor[:num_actual_tokens]
+            if self.kda_decode_backend == "flashinfer":
+                flashinfer_fused_kda_decode(
+                    x=mixed_qkv,
+                    weight=self.decode_conv1d_weight,
+                    conv_state=conv_state,
+                    raw_gate=g1,
+                    raw_beta=beta,
+                    A_log=self.A_log,
+                    dt_bias=self.dt_bias,
+                    state_indices=state_indices,
+                    state=recurrent_state,
+                    output_gate=g2[:num_actual_tokens],
+                    norm_weight=self.decode_norm_weight,
+                    lower_bound=self.gate_lower_bound,
+                    norm_eps=self.o_norm.eps,
+                    output=core_attn_out[:, :num_actual_tokens],
+                )
+            else:
+                ops.fused_kda_decode(
+                    x=mixed_qkv,
+                    weight=self.decode_conv1d_weight,
+                    bias=self.conv1d.bias,
+                    conv_state=conv_state,
+                    raw_g=g1,
+                    raw_beta=beta,
+                    A_log=self.A_log,
+                    dt_bias=self.dt_bias,
+                    state_indices=state_indices,
+                    state=recurrent_state,
+                    out=core_attn_out[:, :num_actual_tokens],
+                    lower_bound=self.gate_lower_bound,
+                    output_gate=g2[:num_actual_tokens],
+                    norm_weight=self.decode_norm_weight,
+                    norm_eps=self.o_norm.eps,
+                )
             return
 
         conv_weights = self.conv1d.weight.view(
@@ -991,6 +1205,35 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                             final_state=final_state[: initial_state.shape[0]],
                             workspace=workspace,
                         )
+                elif self.kda_prefill_backend == "flashinfer":
+                    assert self.gate_lower_bound is not None
+                    assert m.flashinfer_prefill_query_start_loc is not None
+                    if q_ns.shape[1] > initial_state.shape[0]:
+                        assert m.flashinfer_prefill_seq_order is not None
+                    flashinfer_out = core_attn_out[:, : q_ns.shape[1]]
+                    if has_spec_decode:
+                        assert self._flashinfer_kda_output_spec is not None
+                        (workspace_out,) = current_workspace_manager().get_simultaneous(
+                            self._flashinfer_kda_output_spec
+                        )
+                        flashinfer_out = workspace_out[:, : q_ns.shape[1]]
+                    (
+                        core_attn_out_non_spec,
+                        last_recurrent_state,
+                    ) = _flashinfer_kda_prefill(
+                        q=q_ns,
+                        k=k_ns,
+                        v=v_ns,
+                        raw_g=g1_ns,
+                        raw_beta=beta_ns,
+                        A_log=self.A_log,
+                        dt_bias=self.dt_bias,
+                        lower_bound=self.gate_lower_bound,
+                        initial_state=initial_state,
+                        cu_seqlens=m.flashinfer_prefill_query_start_loc,
+                        out=flashinfer_out,
+                        seq_order=m.flashinfer_prefill_seq_order,
+                    )
                 else:
                     (
                         core_attn_out_non_spec,
@@ -1009,7 +1252,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                         use_qk_l2norm_in_kernel=True,
                         cu_seqlens=non_spec_query_start_loc,
                     )
-                recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
+                recurrent_state[non_spec_state_indices_tensor] = (
+                    last_recurrent_state.to(recurrent_state.dtype)
+                )
             else:
                 # Pure non-speculative decode.
                 assert non_spec_state_indices_tensor is not None
@@ -1046,7 +1291,10 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
             core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
             core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
         elif core_attn_out_non_spec is not None:
-            if self.kda_prefill_backend != "flashkda" or m.num_prefills == 0:
+            if (
+                self.kda_prefill_backend not in ("flashkda", "flashinfer")
+                or m.num_prefills == 0
+            ):
                 # TODO: decode kernels write directly to core_attn_out
                 core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
                     0, :num_actual_tokens

vllm/models/kimi_k3/nvidia/kda_metadata.py

@@ -270,6 +270,8 @@ class KDACheckpointMetadata:
 
 @dataclass
 class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
+    flashinfer_prefill_query_start_loc: torch.Tensor | None = None
+    flashinfer_prefill_seq_order: torch.Tensor | None = None
     recoverssm_commit: KDARecoverSSMCommitMetadata | None = None
     recoverssm_context: "KDARecoverSSMCommitContext | None" = field(
         default=None, repr=False, compare=False
@@ -318,6 +320,11 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
         device: torch.device,
     ) -> None:
         super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        additional_config = vllm_config.additional_config
+        self.use_flashinfer_prefill = (
+            isinstance(additional_config, dict)
+            and additional_config.get("kda_prefill_backend") == "flashinfer"
+        )
         self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm
         self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1
         self.recoverssm_num_accepted_tokens: torch.Tensor | None = None
@@ -721,6 +728,20 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
                 align=align,
             )
 
+        flashinfer_prefill_query_start_loc = None
+        flashinfer_prefill_seq_order = None
+        if self.use_flashinfer_prefill and num_prefills > 0:
+            assert non_spec_query_start_loc is not None
+            flashinfer_prefill_query_start_loc = non_spec_query_start_loc.to(
+                torch.int64
+            )
+            num_non_spec_requests = non_spec_query_start_loc.shape[0] - 1
+            num_non_spec_tokens = num_prefill_tokens + num_decode_tokens
+            if num_non_spec_tokens > num_non_spec_requests:
+                flashinfer_prefill_seq_order = torch.argsort(
+                    flashinfer_prefill_query_start_loc.diff(), descending=True
+                ).to(torch.int32)
+
         return KimiK3KDAMetadata(
             num_prefills=num_prefills,
             num_prefill_tokens=num_prefill_tokens,
@@ -738,6 +759,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
             spec_token_indx=spec_token_indx,
             non_spec_token_indx=non_spec_token_indx,
             num_accepted_tokens=num_accepted_tokens,
+            flashinfer_prefill_query_start_loc=flashinfer_prefill_query_start_loc,
+            flashinfer_prefill_seq_order=flashinfer_prefill_seq_order,
             recoverssm_commit=recoverssm_commit,
             recoverssm_context=(
                 self._get_recoverssm_context()

vllm/models/kimi_k3/nvidia/model.py

@@ -1649,7 +1649,9 @@ class KimiLinearForCausalLM(
         vllm_config: "VllmConfig",
     ) -> tuple[torch.dtype, ...]:
         dtypes = MambaStateDtypeCalculator.kda_state_dtype(
-            vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
+            vllm_config.model_config.dtype,
+            vllm_config.cache_config.mamba_cache_dtype,
+            vllm_config.cache_config.mamba_ssm_cache_dtype,
         )
         if vllm_config.cache_config.use_kda_recoverssm:
             dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record(

vllm/utils/flashinfer.py

@@ -264,6 +264,14 @@ flashinfer_xqa_batch_decode_with_kv_cache = _lazy_import_wrapper(
     "flashinfer.decode",
     "xqa_batch_decode_with_kv_cache",
 )
+flashinfer_recurrent_kda = _lazy_import_wrapper(
+    "flashinfer.kda",
+    "recurrent_kda",
+)
+flashinfer_fused_kda_decode = _lazy_import_wrapper(
+    "flashinfer.kda_decode",
+    "fused_kda_decode",
+)
 
 
 # Special case for autotune since it returns a context manager
@@ -399,6 +407,28 @@ def has_flashinfer_bf16_fp4() -> bool:
     )
 
 
+@functools.cache
+def has_flashinfer_recurrent_kda() -> bool:
+    """Return whether FlashInfer recurrent KDA prefill is available."""
+    if not has_flashinfer():
+        return False
+    mod = _get_submodule("flashinfer.kda")
+    return mod is not None and callable(getattr(mod, "recurrent_kda", None))
+
+
+@functools.cache
+def has_flashinfer_fused_kda_decode() -> bool:
+    """Return whether FlashInfer fused KDA decode is available."""
+    if not has_flashinfer():
+        return False
+    mod = _get_submodule("flashinfer.kda_decode")
+    return (
+        mod is not None
+        and bool(getattr(mod, "_FUSED_KDA_DECODE_AVAILABLE", False))
+        and callable(getattr(mod, "fused_kda_decode", None))
+    )
+
+
 @functools.cache
 def has_flashinfer_trtllm_fused_moe() -> bool:
     """Return `True` if FlashInfer TRTLLM fused MoE is available."""
@@ -1261,6 +1291,8 @@ __all__ = [
     "flashinfer_trtllm_batch_decode_with_kv_cache_mla",
     "flashinfer_trtllm_batch_decode_sparse_mla_dsv4",
     "flashinfer_xqa_batch_decode_with_kv_cache",
+    "flashinfer_recurrent_kda",
+    "flashinfer_fused_kda_decode",
     "autotune",
     "has_flashinfer_moe",
     "has_flashinfer_comm",
@@ -1268,6 +1300,8 @@ __all__ = [
     "has_flashinfer_nvlink_one_sided",
     "has_flashinfer_cutlass_fused_moe",
     "has_flashinfer_cutedsl_grouped_gemm_nt_masked",
+    "has_flashinfer_recurrent_kda",
+    "has_flashinfer_fused_kda_decode",
     "has_flashinfer_cutedsl_moe_nvfp4",
     "has_flashinfer_bf16_fp4",
     "has_flashinfer_b12x_moe",