refactor(kda): stop 0.7 advertising unreleased KDA surface (mark prefill wrapper experimental, trim kda_kernels.__all__)

docs/api/kda_prefill.rst

@@ -156,11 +156,11 @@ records the same two kernels and preserves the workspace contract below.
 For eager packed CuTe DSL engine calls, omitting ``seq_order`` builds and
 caches a stable decreasing-length order on the host. CuTe DSL decomp retains
 the original sequence order because its CTA grid fits in one wave.
-``flashinfer.RecurrentKDAPrefillWrapper`` provides the explicit planned path
-needed for packed engine CUDA Graph capture: ``plan`` builds the order and the
-decomp ``cu_chunks`` prefix, then ``run`` consumes fixed-address buffers. The
-decomp prep kernel binary-searches this compact prefix instead of carrying a
-dense chunk-to-sequence tensor. The number of sequences, total tokens, and
+``flashinfer.RecurrentKDAPrefillWrapper`` (experimental) provides the explicit
+planned path needed for packed engine CUDA Graph capture: ``plan`` builds the
+order and the decomp ``cu_chunks`` prefix, then ``run`` consumes fixed-address
+buffers. The decomp prep kernel binary-searches this compact prefix instead of
+carrying a dense chunk-to-sequence tensor. The number of sequences, total tokens, and
 total BT=16 chunks are fixed by the first plan so the metadata and launch
 geometry remain valid across CUDA Graph replays.
 

flashinfer/kda.py

@@ -33,7 +33,7 @@
 from . import kda_prefill as _kda_prefill
 from . import kda_prefill_cute as _kda_prefill_cute
 from .jit import flash_kda_indexed as _flash_kda_indexed
-from .api_logging import flashinfer_api
+from .api_logging import flashinfer_api, flashinfer_experimental_api
 from .kda_backward import (
     RecurrentKDABackwardWorkspace as RecurrentKDABackwardWorkspace,
 )
@@ -665,6 +665,13 @@ def recurrent_kda(
 class RecurrentKDAPrefillWrapper:
     """Plan-and-run wrapper for packed recurrent-KDA prefill.
 
+    .. warning::
+        ``RecurrentKDAPrefillWrapper`` is experimental: it provides no
+        compatibility guarantees and may change or be removed without
+        deprecation.  It has not appeared in a release; the plan-and-run shape
+        is expected to change as the ``recurrent_kda`` surface is unified
+        (see `#4936 <https://github.com/flashinfer-ai/flashinfer/issues/4936>`_).
+
     Compute capability 10.0 and 10.3 only.  ``run`` forces ``backend="cute-dsl"``
     and always passes the ``seq_order`` it planned, and the CC 12.0 backend
     supports neither, so a CC 12.0 caller should use
@@ -705,6 +712,7 @@ def __init__(
         self._total_chunks: Optional[int] = None
         self._lock = threading.Lock()
 
+    @flashinfer_experimental_api
     def plan(
         self,
         cu_seqlens: torch.Tensor,
@@ -807,6 +815,7 @@ def plan(
             self._workspace.__dict__["_cute_dsl_total_chunks"] = total_chunks
             self._total_tokens = offsets[-1]
 
+    @flashinfer_experimental_api
     def run(
         self,
         q: torch.Tensor,

flashinfer/kda_kernels/__init__.py

@@ -26,7 +26,11 @@
 - run_recurrent_kda: Recurrent KDA standard decode and speculative decode backend
 - run_fused_kda_decode: Fused Kimi K3 conv, recurrent KDA, and RMSNorm backend
 - run_packed_kda_decode: Packed Kimi K3 T=1 recurrent decode backend
-- run_kda_prefill_sm120: SM120a ordinary multi-token prefill backend
+
+The ``*_kda_prefill_sm120`` names below are the optional facade that
+``flashinfer.kda_prefill`` dispatches through; they are not part of this
+package's public surface. Reach that backend through ``flashinfer.kda_prefill``,
+or ``flashinfer.kda_kernels.sm120_prefill`` for its cache controls.
 """
 
 from typing import Optional
@@ -101,13 +105,10 @@
     run_kda_prefill_sm120 = None  # type: ignore
 
 __all__ = [
-    "can_implement_kda_prefill_sm120",
-    "clear_kda_prefill_sm120_caches",
     "fused_kda_decode",
     "packed_kda_decode",
     "recurrent_kda",
     "run_fused_kda_decode",
-    "run_kda_prefill_sm120",
     "run_packed_kda_decode",
     "run_recurrent_kda",
 ]

tests/experimental/test_kda_prefill_wrapper.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2026 by FlashInfer team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the experimental ``RecurrentKDAPrefillWrapper`` (see #4936).
+
+Numerics for the kernels the wrapper hands off to are covered by the stable
+lane in ``tests/kda/test_recurrent_kda_prefill.py``; these cover the wrapper's
+own contract -- the plan metadata it builds and the buffers it forwards.
+"""
+
+import importlib
+
+import pytest
+import torch
+
+from flashinfer.kda import RecurrentKDAPrefillWrapper
+
+from tests.test_helpers.kda_prefill import cpu_route_tensors
+
+kda_api = importlib.import_module("flashinfer.kda")
+
+
+@pytest.fixture
+def cuda_device():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    return torch.device("cuda")
+
+
+def test_prefill_wrapper_plan_builds_stable_device_metadata(cuda_device):
+    wrapper = RecurrentKDAPrefillWrapper(cuda_device)
+    wrapper.plan(torch.tensor([0, 0, 7, 7, 12], device=cuda_device))
+
+    cu_seqlens_ptr = wrapper._cu_seqlens_buf.data_ptr()
+    seq_order_ptr = wrapper._seq_order_buf.data_ptr()
+    cu_chunks_ptr = wrapper._cu_chunks_buf.data_ptr()
+    assert wrapper._cu_seqlens_buf.dtype == torch.int64
+    assert wrapper._cu_seqlens_buf.tolist() == [0, 0, 7, 7, 12]
+    assert wrapper._seq_order_buf.tolist() == [1, 3, 0, 2]
+    assert wrapper._cu_chunks_buf.tolist() == [0, 0, 1, 1, 2]
+    assert wrapper._workspace._cute_dsl_total_chunks == 2
+
+    wrapper.plan(torch.tensor([0, 0, 2, 2, 12], device=cuda_device))
+    assert wrapper._cu_seqlens_buf.data_ptr() == cu_seqlens_ptr
+    assert wrapper._seq_order_buf.data_ptr() == seq_order_ptr
+    assert wrapper._cu_chunks_buf.data_ptr() == cu_chunks_ptr
+    assert wrapper._seq_order_buf.tolist() == [3, 1, 0, 2]
+
+    with pytest.raises(ValueError, match="total token count is fixed"):
+        wrapper.plan(torch.tensor([0, 0, 2, 2, 13], device=cuda_device))
+
+    with pytest.raises(ValueError, match="number of sequences is fixed"):
+        wrapper.plan(torch.tensor([0, 2, 12], device=cuda_device))
+
+    chunk_wrapper = RecurrentKDAPrefillWrapper(cuda_device)
+    chunk_wrapper.plan(torch.tensor([0, 16, 16, 32], device=cuda_device))
+    with pytest.raises(ValueError, match="chunk count is fixed"):
+        chunk_wrapper.plan(torch.tensor([0, 1, 17, 32], device=cuda_device))
+
+    with pytest.raises(ValueError, match="non-decreasing"):
+        RecurrentKDAPrefillWrapper(cuda_device).plan(
+            torch.tensor([0, 2, 1, 12], device=cuda_device)
+        )
+
+
+def test_prefill_wrapper_run_forwards_planned_buffers(cuda_device, monkeypatch):
+    wrapper = RecurrentKDAPrefillWrapper(cuda_device)
+    wrapper.plan(torch.tensor([0, 1, 3], device=cuda_device))
+    calls = []
+    sentinel = (object(), object())
+    monkeypatch.setattr(
+        kda_api,
+        "recurrent_kda",
+        lambda **kwargs: calls.append(kwargs) or sentinel,
+    )
+    tensors = cpu_route_tensors(token_count=3)
+    tensors = {
+        key: value.to(cuda_device) if isinstance(value, torch.Tensor) else value
+        for key, value in tensors.items()
+    }
+
+    assert wrapper.run(**tensors) is sentinel
+    assert calls[0]["cu_seqlens"] is wrapper._cu_seqlens_buf
+    assert calls[0]["seq_order"] is wrapper._seq_order_buf
+    assert calls[0]["prefill_workspace"] is wrapper._workspace
+    assert calls[0]["backend"] == "cute-dsl"
+    assert wrapper._workspace._cute_dsl_cu_chunks is wrapper._cu_chunks_buf
+    assert wrapper._workspace._cute_dsl_total_chunks == 2

tests/kda/test_recurrent_kda_prefill.py

@@ -28,6 +28,8 @@
 from flashinfer.kda_prefill import RecurrentKDAPrefillWorkspace
 from flashinfer.utils import get_compute_capability
 
+from tests.test_helpers.kda_prefill import cpu_route_tensors
+
 kda_decode_api = importlib.import_module("flashinfer.kda_decode")
 kda_api = importlib.import_module("flashinfer.kda")
 kda_prefill_api = importlib.import_module("flashinfer.kda_prefill")
@@ -229,83 +231,6 @@ def test_cake_kda_affine_workspace_buffer_is_grow_only(monkeypatch):
         )
 
 
-def test_prefill_wrapper_plan_builds_stable_device_metadata(cuda_device):
-    wrapper = RecurrentKDAPrefillWrapper(cuda_device)
-    wrapper.plan(torch.tensor([0, 0, 7, 7, 12], device=cuda_device))
-
-    cu_seqlens_ptr = wrapper._cu_seqlens_buf.data_ptr()
-    seq_order_ptr = wrapper._seq_order_buf.data_ptr()
-    cu_chunks_ptr = wrapper._cu_chunks_buf.data_ptr()
-    assert wrapper._cu_seqlens_buf.dtype == torch.int64
-    assert wrapper._cu_seqlens_buf.tolist() == [0, 0, 7, 7, 12]
-    assert wrapper._seq_order_buf.tolist() == [1, 3, 0, 2]
-    assert wrapper._cu_chunks_buf.tolist() == [0, 0, 1, 1, 2]
-    assert wrapper._workspace._cute_dsl_total_chunks == 2
-
-    wrapper.plan(torch.tensor([0, 0, 2, 2, 12], device=cuda_device))
-    assert wrapper._cu_seqlens_buf.data_ptr() == cu_seqlens_ptr
-    assert wrapper._seq_order_buf.data_ptr() == seq_order_ptr
-    assert wrapper._cu_chunks_buf.data_ptr() == cu_chunks_ptr
-    assert wrapper._seq_order_buf.tolist() == [3, 1, 0, 2]
-
-    with pytest.raises(ValueError, match="total token count is fixed"):
-        wrapper.plan(torch.tensor([0, 0, 2, 2, 13], device=cuda_device))
-
-    with pytest.raises(ValueError, match="number of sequences is fixed"):
-        wrapper.plan(torch.tensor([0, 2, 12], device=cuda_device))
-
-    chunk_wrapper = RecurrentKDAPrefillWrapper(cuda_device)
-    chunk_wrapper.plan(torch.tensor([0, 16, 16, 32], device=cuda_device))
-    with pytest.raises(ValueError, match="chunk count is fixed"):
-        chunk_wrapper.plan(torch.tensor([0, 1, 17, 32], device=cuda_device))
-
-    with pytest.raises(ValueError, match="non-decreasing"):
-        RecurrentKDAPrefillWrapper(cuda_device).plan(
-            torch.tensor([0, 2, 1, 12], device=cuda_device)
-        )
-
-
-def test_prefill_wrapper_run_forwards_planned_buffers(cuda_device, monkeypatch):
-    wrapper = RecurrentKDAPrefillWrapper(cuda_device)
-    wrapper.plan(torch.tensor([0, 1, 3], device=cuda_device))
-    calls = []
-    sentinel = (object(), object())
-    monkeypatch.setattr(
-        kda_api,
-        "recurrent_kda",
-        lambda **kwargs: calls.append(kwargs) or sentinel,
-    )
-    tensors = _cpu_route_tensors(token_count=3)
-    tensors = {
-        key: value.to(cuda_device) if isinstance(value, torch.Tensor) else value
-        for key, value in tensors.items()
-    }
-
-    assert wrapper.run(**tensors) is sentinel
-    assert calls[0]["cu_seqlens"] is wrapper._cu_seqlens_buf
-    assert calls[0]["seq_order"] is wrapper._seq_order_buf
-    assert calls[0]["prefill_workspace"] is wrapper._workspace
-    assert calls[0]["backend"] == "cute-dsl"
-    assert wrapper._workspace._cute_dsl_cu_chunks is wrapper._cu_chunks_buf
-    assert wrapper._workspace._cute_dsl_total_chunks == 2
-
-
-def _cpu_route_tensors(token_count=2):
-    shape = (1, token_count, 1, 128)
-    return {
-        "q": torch.empty(shape, dtype=torch.bfloat16),
-        "k": torch.empty(shape, dtype=torch.bfloat16),
-        "v": torch.empty(shape, dtype=torch.bfloat16),
-        "g": torch.empty(shape, dtype=torch.bfloat16),
-        "beta": torch.empty((1, token_count, 1), dtype=torch.bfloat16),
-        "A_log": torch.empty(1, dtype=torch.float32),
-        "dt_bias": torch.empty((1, 128), dtype=torch.float32),
-        "use_gate_in_kernel": True,
-        "lower_bound": -5.0,
-        "beta_is_logit": True,
-    }
-
-
 def test_public_prefill_backend_option_routes_to_cute_dsl(monkeypatch):
     sentinel = (object(), object())
     monkeypatch.setattr(
@@ -319,7 +244,7 @@ def test_public_prefill_backend_option_routes_to_cute_dsl(monkeypatch):
         lambda **kwargs: sentinel,
     )
 
-    assert recurrent_kda(**_cpu_route_tensors(), backend="cute-dsl") is sentinel
+    assert recurrent_kda(**cpu_route_tensors(), backend="cute-dsl") is sentinel
 
 
 def test_public_prefill_auto_prefers_cute_dsl(monkeypatch):
@@ -340,7 +265,7 @@ def test_public_prefill_auto_prefers_cute_dsl(monkeypatch):
         lambda **kwargs: pytest.fail("auto should not probe Cake after a CuTe match"),
     )
 
-    assert recurrent_kda(**_cpu_route_tensors()) is sentinel
+    assert recurrent_kda(**cpu_route_tensors()) is sentinel
 
 
 def test_public_prefill_forwards_sequence_order_to_cute_dsl(monkeypatch):
@@ -360,7 +285,7 @@ def test_public_prefill_forwards_sequence_order_to_cute_dsl(monkeypatch):
     seq_order = torch.tensor([1, 0], dtype=torch.int32)
     assert (
         recurrent_kda(
-            **_cpu_route_tensors(token_count=3),
+            **cpu_route_tensors(token_count=3),
             cu_seqlens=torch.tensor([0, 1, 3], dtype=torch.int64),
             seq_order=seq_order,
         )
@@ -387,7 +312,7 @@ def test_public_prefill_auto_falls_back_to_cake(monkeypatch):
         lambda **kwargs: sentinel,
     )
 
-    assert recurrent_kda(**_cpu_route_tensors()) is sentinel
+    assert recurrent_kda(**cpu_route_tensors()) is sentinel
 
 
 def test_public_prefill_explicit_cake_skips_cute_dsl_probe_with_checkpoints(
@@ -414,7 +339,7 @@ def test_public_prefill_explicit_cake_skips_cute_dsl_probe_with_checkpoints(
     checkpoint_starts = torch.tensor([0, 1], dtype=torch.int64)
     assert (
         recurrent_kda(
-            **_cpu_route_tensors(),
+            **cpu_route_tensors(),
             state_checkpoints=checkpoint_state,
             checkpoint_cu_starts=checkpoint_starts,
             checkpoint_every_n_tokens=32,
@@ -442,7 +367,7 @@ def test_public_prefill_auto_routes_supported_checkpoints_to_cute_dsl(monkeypatc
     starts = torch.tensor([0, 1], dtype=torch.int64)
     assert (
         recurrent_kda(
-            **_cpu_route_tensors(),
+            **cpu_route_tensors(),
             state_checkpoints=checkpoints,
             checkpoint_cu_starts=starts,
             checkpoint_every_n_tokens=32,
@@ -462,7 +387,7 @@ def test_public_prefill_cake_backend_is_strict(monkeypatch):
     )
 
     with pytest.raises(ValueError, match="backend='cake' does not support"):
-        recurrent_kda(**_cpu_route_tensors(), backend="cake")
+        recurrent_kda(**cpu_route_tensors(), backend="cake")
 
 
 def test_public_decode_backend_option_forwards_to_decode_layer(monkeypatch):
@@ -474,15 +399,13 @@ def run(**kwargs):
         return sentinel
 
     monkeypatch.setattr(kda_decode_api, "_run_recurrent_kda", run)
-    assert (
-        recurrent_kda(**_cpu_route_tensors(token_count=1), backend="cake") is sentinel
-    )
+    assert recurrent_kda(**cpu_route_tensors(token_count=1), backend="cake") is sentinel
     assert calls[0]["backend"] == "cake"
 
 
 def test_public_backend_option_rejects_unknown_value():
     with pytest.raises(ValueError, match="backend must be"):
-        recurrent_kda(**_cpu_route_tensors(), backend="unknown")
+        recurrent_kda(**cpu_route_tensors(), backend="unknown")
 
 
 def test_cute_dsl_prefill_adapter_preserves_indexed_in_place_state_semantics(
@@ -519,7 +442,7 @@ def get_compiled(**kwargs):
         torch.cuda, "current_stream", lambda device=None: SimpleNamespace(cuda_stream=7)
     )
 
-    inputs = _cpu_route_tensors()
+    inputs = cpu_route_tensors()
     state = torch.empty((3, 1, 128, 128), dtype=torch.bfloat16)
     state_indices = torch.tensor([2], dtype=torch.int32)
     output = torch.empty_like(inputs["q"])
@@ -595,7 +518,7 @@ def __call__(self, *args, **kwargs):
         torch.cuda, "current_stream", lambda device=None: SimpleNamespace(cuda_stream=7)
     )
 
-    inputs = _cpu_route_tensors()
+    inputs = cpu_route_tensors()
     output = torch.empty_like(inputs["q"])
     cu_seqlens = torch.tensor([0, 1, 2], dtype=torch.int64)
     seq_order = torch.tensor([1, 0], dtype=torch.int32) if explicit_order else None
@@ -675,7 +598,7 @@ def test_cute_dsl_unplanned_packed_engine_rejects_graph_capture(monkeypatch):
             "has_state_indices": False,
         },
     )
-    inputs = _cpu_route_tensors()
+    inputs = cpu_route_tensors()
     cu_seqlens = torch.tensor([0, 1, 2], dtype=torch.int64)
 
     with pytest.raises(RuntimeError, match=r"Wrapper\.plan\(\)"):