feat(kda): add CuTe DSL recurrent prefill backend

benchmarks/bench_recurrent_kda_prefill.py

@@ -22,8 +22,10 @@
 ``recurrent_kda`` API. ``--candidate-route dispatcher`` measures the natural
 device/shape policy, while ``nonpersistent`` supplies the same explicit
 workspace and packed sequence order used by the historical benchmark to keep
-B200 on the direct schedule family. The resolved module and exact target are
-recorded by observing the real dispatcher during untimed warmup. With
+B200 on the direct schedule family. ``--backend`` selects one public API
+backend per invocation; compare auto, CuTe DSL, and Cake with separate commands
+over the same case set. The resolved backend, schedule variant, and target are
+recorded during untimed warmup. With
 ``--flash-kda-peer``, two commit-verified MoonshotAI/FlashKDA measurements are
 reported:
 
@@ -266,6 +268,7 @@ def _make_case(
     *,
     state_rotations: int,
     candidate_route: str,
+    candidate_backend: str,
     flash_kda=None,
 ) -> PreparedCase:
     total_tokens = sum(case.seq_lens)
@@ -356,6 +359,7 @@ def candidate_run():
             beta_is_logit=True,
             seq_order=seq_order,
             prefill_workspace=candidate_workspace,
+            backend=candidate_backend,
         )
 
     peer_raw_run = None
@@ -449,25 +453,48 @@ def reset_state_pools() -> None:
     # untimed warmup. This avoids duplicating dispatcher policy in the evidence
     # harness while keeping route logging out of every timed call.
     kda_prefill_module = import_module("flashinfer.kda_prefill")
+    kda_prefill_cute_module = import_module("flashinfer.kda_prefill_cute")
     original_get_module = kda_prefill_module._get_flash_kda_prefill_module
-    resolved_routes = []
+    original_cute_run = kda_prefill_cute_module._run_cute_dsl_kda_prefill
+    resolved_cake_routes = []
+    resolved_backends = []
 
     def recording_get_module(variant, target):
-        resolved_routes.append((variant, target))
+        resolved_cake_routes.append((variant, target))
         return original_get_module(variant, target)
 
+    def recording_cute_run(**kwargs):
+        resolved_backends.append("cute-dsl")
+        return original_cute_run(**kwargs)
+
     kda_prefill_module._get_flash_kda_prefill_module = recording_get_module
+    kda_prefill_cute_module._run_cute_dsl_kda_prefill = recording_cute_run
     try:
         candidate_run()
         torch.cuda.synchronize()
     finally:
         kda_prefill_module._get_flash_kda_prefill_module = original_get_module
+        kda_prefill_cute_module._run_cute_dsl_kda_prefill = original_cute_run
         reset_state_pools()
-    if len(resolved_routes) != 1:
+    if resolved_backends:
+        if resolved_backends != ["cute-dsl"] or resolved_cake_routes:
+            raise RuntimeError(
+                "expected exactly one CuTe DSL route during warmup, got "
+                f"backends={resolved_backends}, cake={resolved_cake_routes}"
+            )
+        resolved_backend = "cute-dsl"
+        decomp_ctas = len(case.seq_lens) * case.num_heads * 2
+        sm_count = torch.cuda.get_device_properties(q.device).multi_processor_count
+        resolved_variant = "decomp" if decomp_ctas <= sm_count else "engine"
+        resolved_target = "bt16"
+    elif len(resolved_cake_routes) == 1:
+        resolved_backend = "cake"
+        resolved_variant, resolved_target = resolved_cake_routes[0]
+    else:
         raise RuntimeError(
-            f"expected one FlashKDA prefill route during warmup, got {resolved_routes}"
+            "expected one recurrent-KDA prefill route during warmup, got "
+            f"backends={resolved_backends}, cake={resolved_cake_routes}"
         )
-    resolved_variant, resolved_target = resolved_routes[0]
 
     metadata = {
         "name": case.name,
@@ -478,6 +505,8 @@ def recording_get_module(variant, target):
         "variant": resolved_variant,
         "target": resolved_target,
         "candidate_route": candidate_route,
+        "requested_backend": candidate_backend,
+        "resolved_backend": resolved_backend,
         "seed": case.seed,
         "state_rotation_capacity": state_rotations,
     }
@@ -592,6 +621,15 @@ def main() -> None:
             "non-persistent direct/M64 family with an explicit workspace."
         ),
     )
+    parser.add_argument(
+        "--backend",
+        choices=("auto", "cute-dsl", "cake"),
+        default="auto",
+        help=(
+            "Select one backend for this invocation of the public recurrent_kda "
+            "API; run separate commands to compare backends."
+        ),
+    )
     parser.add_argument(
         "--flash-kda-peer",
         action="store_true",
@@ -680,6 +718,7 @@ def main() -> None:
             case,
             state_rotations=state_rotations,
             candidate_route=args.candidate_route,
+            candidate_backend=args.backend,
             flash_kda=flash_kda,
         )
         result = {**prepared.metadata, "hardware": hardware}
@@ -775,7 +814,8 @@ def main() -> None:
         results.append(result)
         if prepared.peer_raw_run is None:
             print(
-                f"{result['name']:<18} {result['variant']:<4} "
+                f"{result['name']:<18} {result['resolved_backend']:<8} "
+                f"{result['variant']:<10} "
                 f"{result['median_us']:10.3f} us"
             )
         else:

docs/api/kda_prefill.rst

@@ -4,8 +4,8 @@ flashinfer.kda_prefill
 ======================
 
 Optimized recurrent Kimi Delta Attention (KDA) prefill support. The
-:func:`flashinfer.kda.recurrent_kda` facade dispatches a strict ordinary
-multi-token prefill subset to frozen FlashKDA-compatible SM100-family kernels.
+:func:`flashinfer.kda.recurrent_kda` facade exposes frozen Cake and source-level
+CuTe DSL implementations for a strict ordinary multi-token prefill subset.
 
 .. currentmodule:: flashinfer.kda_prefill
 
@@ -14,10 +14,34 @@ multi-token prefill subset to frozen FlashKDA-compatible SM100-family kernels.
 
     RecurrentKDAPrefillWorkspace
 
+.. currentmodule:: flashinfer.kda
+
+.. autosummary::
+    :toctree: ../generated
+
+    RecurrentKDAPrefillWrapper
+
+Backend selection
+-----------------
+
+``backend="auto"`` selects the source-level CuTe DSL backend for eligible
+ordinary multi-token prefill and falls back to the frozen Cake backend for
+unsupported contracts. Decode retains the existing KDA decode routing.
+``backend="cake"`` and ``backend="cute-dsl"`` select a backend strictly and
+raise when its contract is unsupported.
+
+For multi-token prefill, ``backend="cute-dsl"`` selects a BT=16 CuTe DSL kernel.
+It supports contiguous BF16 Q, K, V, G, and beta with one shared head count and
+head dimension 128, the in-kernel lower-bound gate, fixed or packed-varlen
+layout, BF16 recurrent state, explicit ``seq_order``, and the same checkpoint
+contract as Cake. ``checkpoint_cu_starts`` must always be int64. Packed
+``cu_seqlens`` must be int64 during CUDA graph capture. The CuTe DSL schedule
+is non-persistent.
+
 Optimized Blackwell prefill subset
 -----------------------------------
 
-``flashinfer.kda.recurrent_kda`` uses the frozen prefill backend only when
+The strict Cake backend is available only when
 every condition below holds:
 
 * the device has compute capability 10.0 (SM100a; B200/GB200) or 10.3
@@ -36,8 +60,7 @@ every condition below holds:
   features are not enabled. Plain int32 ``ssm_state_indices`` and native
   prefill checkpoints are supported by direct M128.
 
-Calls outside that subset retain the existing CuTe-DSL path. In particular,
-T=1 decode and speculative decode are not rerouted.
+T=1 decode and speculative decode are not handled by either prefill backend.
 
 CUDA 12.8 predates the family target, so CC 10.0 uses legacy exact
 ``sm_100a`` modules. With CUDA 12.9 or newer, JIT and AOT compile one
@@ -66,10 +89,11 @@ Fixed input omits ``cu_seqlens``. Packed input has ``B=1`` and accepts a
 contiguous CUDA int32 or int64 ``cu_seqlens``. The frozen binding consumes
 int64 offsets; pass int64 directly for CUDA graph capture to avoid an
 in-capture conversion allocation. Offset values are a caller contract:
-``cu_seqlens[0] == 0``, entries are strictly increasing (every sequence is
-non-empty), and ``cu_seqlens[-1] == total_tokens``. FlashInfer does not
-synchronize the device to inspect these values; invalid offsets may cause
-out-of-bounds device access.
+``cu_seqlens[0] == 0``, entries are non-decreasing, and
+``cu_seqlens[-1] == total_tokens``. CuTe DSL accepts equal adjacent offsets for
+zero-length sequences; Cake requires every sequence to be non-empty.
+FlashInfer does not synchronize the device to inspect these values; invalid
+offsets may cause out-of-bounds device access.
 
 Packed scheduling
 -----------------
@@ -80,12 +104,23 @@ is a permutation of ``[0, N)``. Ordering sequences by decreasing length
 reduces the final partial wave. FlashInfer validates dtype, device, rank, and
 size without synchronizing the device to inspect permutation values.
 
-When ``seq_order=None``, a cached identity order is used. H12 selects the
-dedicated M128 schedule with a 16-token recurrence chunk for both fixed and
-packed layouts. Fixed ``B=1,H=64`` selects the two-CTA M64 value-split kernel;
-the fixed small-BH region described above selects its eight-CTA owner/helper
-schedule; all remaining eligible inputs select the general 32-token M128
-schedule.
+For Cake, omitting ``seq_order`` uses its cached eager scheduling metadata. H12
+selects the dedicated M128 schedule with a 16-token recurrence chunk for both
+fixed and packed layouts. Fixed ``B=1,H=64`` selects the two-CTA M64
+value-split kernel; the fixed small-BH region described above selects its
+eight-CTA owner/helper schedule; all remaining eligible inputs select the
+general 32-token M128 schedule.
+
+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
+total BT=16 chunks are fixed by the first plan so the metadata and launch
+geometry remain valid across CUDA Graph replays.
 
 State and graph semantics
 -------------------------

flashinfer/__init__.py

@@ -130,6 +130,7 @@
 from .kda_prefill import (
     RecurrentKDAPrefillWorkspace as RecurrentKDAPrefillWorkspace,
 )
+from .kda import RecurrentKDAPrefillWrapper as RecurrentKDAPrefillWrapper
 from .kda import recurrent_kda as recurrent_kda
 from .kda_decode import fused_kda_decode as fused_kda_decode
 from .kda_decode import packed_kda_decode as packed_kda_decode