Skip to content

[KDA] Unify the public recurrent_kda API, and settle the unreleased KDA surface before FI v0.7 #4936

Description

@kahyunnam

TL;DR. Two public functions named recurrent_kda disagree on their default backend, and the top-level default is not a safe superset of the other. The unification itself is not proposed for 0.7 — it touches released API and belongs on a normal deprecation schedule. Four things do want attention before the cut:

  1. A bug on the default path: #4935backend="auto" decode raises instead of falling back, contradicting its own docstring, which says "auto" selects Cake for its native contract and "otherwise preserves CuTe DSL" (recurrent_kda.py:1873-1878). Fixing it is a restoration of the documented contract, not a behaviour change. The only release-critical engineering item here; fix proposed in #5037.
  2. A holding action: six of the ten top-level KDA symbols have never been in a release. 0.7 should not advertise them as stable API. The plan is to mark the unreleased surface experimental, which says plainly that these signatures may change as the unification below lands. See Holding action.
  3. One comment on an open PR: #4779 adds a fourth backend literal to one facade only. The cheapest intervention on this list, and it is open now.
  4. One housekeeping close: #4367 is already fixed on main.

The rest is background and a proposed sequence.

Context

Status as of 2026-09-03 on FlashInfer main (7a3c04f0).

This supersedes #4483 ([KDA] Unify flashinfer.recurrent_kda API, 2026-08-12). That audit is still worth reading as the historical record, but enough has changed in three weeks — a backend= parameter landed, the top-level surface grew from 3 symbols to 10, and a release was cut in the middle of it — that patching it in place would have been more confusing than restating it. Relationship to #4483 summarises what moved.

FlashInfer has two decorated functions named recurrent_kda registering the same trace template. Since #4483, backend= has been added to the top-level facade (#4605), which resolves the original discoverability gap. But the two facades now ship different defaults, and the top-level default is not a safe superset of the other: backend="auto" decode commits to a frozen Cake kernel and raises instead of falling back to CuTe DSL. That defect is filed separately as #4935; this issue covers the public Python API contract only.

The KDA public surface also grew from the 3 top-level symbols #4483 audited to 10 — of which v0.6.18 shipped 4 — adding a training/backward family and a plan-and-run wrapper class that #4483 did not cover.

Six of those ten symbols, and the entire backend= parameter on the top-level facade, have never been in a release — see Release status. That is the most important input to the plan below: the parts of this design that are still wrong can mostly be fixed outright, but only until 0.7 is cut. Since that cut is close, Holding action proposes narrowing the advertised surface back to what v0.6.18 shipped rather than rushing the consolidation.

This is a design / tracking issue. Related CAKE progress is tracked in #4254. The same problem class on the sibling GDN op is already tracked in #2687, with an in-flight implementation in #2706 — that pair is the closest precedent for the plan below and should be aligned with rather than duplicated.

Suggested implementation sequence

Steps are ordered by dependency, but the release window is the scheduling constraint: anything touching unreleased surface is cheap now and locked after 0.7.

Step Scope Compatibility Before 0.7? Status
0 (done) Add backend= to flashinfer.kda.recurrent_kda; wire Cake decode and prefill eligibility through it Additive Landed in #4605 / #4535, not yet released
0b Comment on #4779 and #4968, asking that each new backend literal cover both facades or go in as an experimental backend under #4880 Two comments; done Yes Highest leverage per unit effort: both are open now and each would widen the divergence this issue exists to stop. #4968 adds "cudnn" to kda.py's enum and leaves kda_decode.py untouched — the same shape as #4779's "vibecuda". Raised on #4968 and #4779; #4779 also AOT-registers its backend, which #4880's policy forbids for an experimental one
0c Housekeeping: close #4367 as already fixed (verified passing on main) None no
1 Prerequisite (#5037, open): fix #4935 so the fallback is driven by the resolved selector result rather than any duplicated predicate, with a negative test asserting output and state-pool equality on both sides of the one-warp threshold Bug fix on unreleased surface Yes Blocks step 5. Neither reordering nor a gate-side predicate works — see #4935
2 Document "auto" as the top-level default, keeping the decode facade's released "cute-dsl" default untouched Free choice while unreleased Yes Depends on step 1. No value is safe for a phase-neutral facade until "auto" decode is total — see remedy 2 above
2b Confirm the "auto" decode kernel swap against v0.6.18's effective "cute-dsl" routing on the intended serving path Unreleased Yes Blocks ratifying "auto" in step 2. Measured: Cake wins at every batch size under CUDA-graph capture; the eager small-batch loss is host dispatch
3 Give the training/backward family real trace templates (or state the omission), and reconsider the _kda_training_impl split Free while unreleased no Problem 8. Disappears entirely if #4965 lands; otherwise is_experimental keeps them out of the stable inventory, so 3b-i removes the urgency
3b-i Holding action (#5040, open): mark the unreleased top-level symbols experimental under #4880RecurrentKDAPrefillWrapper per-method on plan / run, plus the five training/backward symbols only if #4965 has not landed Non-breaking for callers (warns once) Yes No kernel or wheel change. #5040 also relocates the wrapper's own two tests to tests/experimental/, leaving the three CuTe-DSL-numerics tests that merely drive the wrapper in the stable lane
3b-ii Containment, only if #4965 does not land: relocate kda_backward.py + _kda_training_impl.py behind thin gated entry points, move their tests to tests/experimental/, de-register the ten aot.py sites, file the owner / tracking issue / graduation plan Wheel contents change; ~91 s first-call JIT; 1,522 test lines leave the default lanes no — first release after 0.7 Superseded by #4965's removal, which deletes the same code instead of moving it
3c Settle the SM120a tranche. The __all__ half is done in #5040 (open), which drops the three unreleased kda_kernels.__all__ exports and restores __all__ to byte-identical with v0.6.18; a zero-consumer change, since dispatch reads them as attributes and every doc, test and benchmark uses the deeper kda_kernels.sm120_prefill path. Still open: the explicit backend= literal and gating the automatic branch with experimental_auto_backends_allowed() Unreleased throughout Yes for the __all__ half; the gate can follow Independent of 3b-i. The gate is deliberately not in #5040: it changes "auto" prefill routing on CC 12.0 and mostly by raising (see remedy 2 above), so it wants its own PR validated on a CC 12.0 device
4 Factor a shared _dispatch_recurrent_kda_decode used by both facades; decide which facade owns the recurrent_kda trace Refactor, no behavior change no
5 Point kda_decode.recurrent_kda at that dispatcher with omitted→"cute-dsl" + deprecate Soft deprecation of released API no Needs steps 1, 4
6 Give the two recurrent_kda traces distinguishable labels; extend the trace template and retag stage Additive no
7 Retarget decode benchmarks, test_recurrent_kda.py, Cake export tests, and fi_trace to the canonical facade; add top-level fi_api trace coverage Test-only (touches many call sites) no
8 Deprecate then remove kda_kernels noun aliases (un-shadows submodules); consolidate docs pages; document or internalize FLASHINFER_KDA_T1_FAST_PATH Soft deprecation of released API no

The time-critical items are 0b, 1, 2, 2b, 3b-i and 3c; all are confined to main-only surface. Everything else concerns released API and can proceed on a normal deprecation schedule.

Stated plainly: if step 3b-i lands, the release-critical engineering reduces to two items — step 1, and step 3c's disposal of the three unreleased kda_kernels SM120a exports, because the top-level backend= parameter and its default would then be the sole unreleased inference surface that a 0.7 cut freezes — and step 2 cannot be settled without it, since no backend value is safe across both phases until "auto" decode is total. Step 0b is one comment; step 3's subject disappears with 3b-i.

The benchmark in step 2b is not a formality. On a B200 (H=32, K=V=128, T=1, eager, 200 iterations) "auto" decode selects Cake and crosses over against "cute-dsl" at about B = 35: below it "auto" is up to 3.1x slower, above it up to 7.5x faster (0.32x at B=4, 1.90x at B=64, 7.53x at B=256). Cake's flat ~108 µs is host dispatch rather than device time, and under CUDA-graph capture the picture inverts completely: Cake wins at every batch size, by 3.7x to 7.9x. So the eager small-batch loss is ~75 µs of Python-side enqueue that capture erases, and no size condition on the gate is warranted — the default is right for the graph-captured serving path it exists for. What remains worth confirming is that the intended consumers do capture.

If step 1 slips, the contingency is not to ship "auto" as-is. Disabling the speculative "auto" → Cake decode gate is a one-line change to auto_unbounded_softplus_candidate (recurrent_kda.py:2063-2064) that makes "auto" decode identical to "cute-dsl" — bit-exact on every trigger class with a non-zero pool and non-identity indices, backend="cake" strictness intact, and exactly the routing a defaulted v0.6.18 call took. It costs the Cake decode win at serving batch sizes for one release and requires retargeting six parametrizations of test_t1_unbounded_softplus_auto_route_* to backend="cake". Land step 1; keep this in reserve.

Current public surface (main)

Full inventory: both facades, the ten top-level symbols, and the docs split
Import path Defined in @flashinfer_api Notable kwargs Routing
flashinfer.recurrent_kda flashinfer/kda.py yes (recurrent_kda_trace) backend (default "auto"), seq_order, prefill_workspace, state_checkpoints, checkpoint_cu_starts, checkpoint_every_n_tokens Classifies phase by shape first (_is_plain_multi_token_prefill). Plain multi-token prefill: under "auto" tries SM120a CuTe → SM100-family CuTe DSL BT=16 → frozen Cake fallback. Otherwise decode/spec → run_recurrent_kda via raw _run_recurrent_kda (still bypasses decorated kda_decode.recurrent_kda), forwarding backend=
flashinfer.kda_decode.recurrent_kda flashinfer/kda_decode.py yes (same trace template) backend (default "cute-dsl") No phase classification; always run_recurrent_kda(backend=…). Default "cute-dsl" → CuTe DSL, never Cake
flashinfer.packed_kda_decode flashinfer/kda_decode.py yes (packed_kda_decode_trace) packed QKV + raw gate/beta Serving-native Kimi K3 T=1, H=12, K=V=128; strict Cake (#4445)
flashinfer.fused_kda_decode flashinfer/kda_decode.py yes (fused_kda_decode_trace) packed QKV + conv + RMSNorm Separate Kimi K3 fused op
flashinfer.RecurrentKDAPrefillWrapper flashinfer/kda.py n/a (class) plan() / run() New. CC 10.0/10.3 packed prefill; run forces backend="cute-dsl"
flashinfer.RecurrentKDAPrefillWorkspace flashinfer/kda_prefill.py n/a (type) graph-capture workspace Top-level export
flashinfer.recurrent_kda_backward, RecurrentKDABackwardWorkspace flashinfer/kda_backward.py bare, no trace / n/a training backward New. Frozen SM100a/SM103a
flashinfer.recurrent_kda_training_forward / _backward, RecurrentKDATrainingContext flashinfer/kda_training.py (re-export of private _kda_training_impl.py) bare, no trace / n/a paired training API New. Own paired-function convention

Top-level __init__.py re-exports 10 KDA symbols (__init__.py:165-184): recurrent_kda and RecurrentKDAPrefillWrapper from kda; fused_kda_decode and packed_kda_decode from kda_decode; RecurrentKDAPrefillWorkspace from kda_prefill; recurrent_kda_backward and RecurrentKDABackwardWorkspace from kda_backward; and recurrent_kda_training_forward, recurrent_kda_training_backward, RecurrentKDATrainingContext from kda_training.

flashinfer.kda is reachable as an attribute (the package-root from .kda import … binds the submodule) but is not an intentional re-export.

flashinfer.kda_kernels still re-exports three noun aliases (recurrent_kda, fused_kda_decode, and now packed_kda_decode) alongside the run_* backends. Its __all__ has grown from 6 entries to 9 since v0.6.18; all three additions are the SM120a prefill helpers, which are unreleased. Two of the three are not merely redundant — they shadow a submodule of the same name, so attribute access cannot reach the module:

K.recurrent_kda      -> function   (shadows kda_kernels/recurrent_kda.py)
K.fused_kda_decode   -> function   (shadows kda_kernels/fused_kda_decode.py)
K.packed_kda_decode  -> function   (no same-named module; impl is cake_packed_kda_decode.py)
importlib.import_module("flashinfer.kda_kernels.recurrent_kda") -> module

In-repo consumers that need the shadowed modules work around this with importlib.import_module. That makes the alias cleanup a correctness fix for the first two, not cosmetics. Separately, kda_kernels/__init__.py does now carry an explicit note that packed_kda_decode_cute is an internal implementation module — the desired precedent, applied inconsistently with the shadowing noun aliases directly above it.

Docs are split across five pages: docs/api/kda.rst, kda_decode.rst, kda_prefill.rst, and the new kda_backward.rst and kda_training.rst (both deleted by #4965). kda_decode.rst still autosummarizes recurrent_kda, packed_kda_decode, and fused_kda_decode. kda_prefill.rst documents RecurrentKDAPrefillWrapper, which is defined in kda.py.

In-repo usage today: Prefill tests/benchmarks use the top-level / kda facade (bench_recurrent_kda_prefill.py, tests/kda/test_recurrent_kda_prefill.py). Decode paths — including Cake export tests and the main CuTe DSL correctness suite tests/kda/test_recurrent_kda.py, plus bench_recurrent_kda.py, bench_recurrent_kda_decode_export.py, and bench_recurrent_kda_decode_split_sweep.py — import flashinfer.kda_decode.recurrent_kda. The newer tests/kda/test_packed_kda_decode_cute.py uses top-level flashinfer.recurrent_kda, so the codebase now exercises both facades for decode-shaped calls. fi_trace fixtures still pin only the decode facade.

The split is also why CI does not catch #4935, though the mechanism is subtler than "the top-level default is untested". Across tests/kda/, the "auto" gate matches on 6 of 190 run_recurrent_kda calls, all six with an explicit backend="auto" and all from one test that asserts Cake is selected — so it cannot detect a missing fallback. The three defaulted top-level calls in test_packed_kda_decode_cute.py all miss the gate, two on lower_bound=-5.0 and the third on use_gate_in_kernel. The gap is a missing negative test, not missing coverage, which makes it cheap to close.

Release status: what is actually shipped

Most of the surface this issue describes has never been in a release. That widens the design freedom considerably, and it also puts a clock on it.

v0.6.18 (published 2026-08-29, tag commit 69ff11fc) is the current stable release. It is not an ancestor of main: the release branch was cut from main at 61a6c651 on 2026-08-19 (author date 01:03Z; 08-18 in US-Pacific, which is how it appears in git log locally), then took cherry-picks through 08-28. So "did it ship?" cannot be inferred from merge dates — the classification below comes from reading the v0.6.18 tree directly.

Top-level symbol v0.6.18
recurrent_kda (phase-neutral facade, no backend=) shipped
fused_kda_decode shipped
packed_kda_decode shipped
RecurrentKDAPrefillWorkspace shipped
RecurrentKDAPrefillWrapper unreleased
recurrent_kda_backward, RecurrentKDABackwardWorkspace unreleased
recurrent_kda_training_forward / _backward, RecurrentKDATrainingContext unreleased

So v0.6.18 exports 4 KDA symbols where main exports 10.

No other channel carries the newer surface either. PyPI serves flashinfer-python 0.6.18 as latest (plus a 0.6.18rc10 prerelease, checked and clean); the release-v0.6.18 branch head is the tag itself; there are no 0.6.18.post* releases, which is informative rather than accidental since v0.6.16 had four. Nightly tagging stops at nightly-v0.6.18-20260819 — one day before #4605 — and nightlies track main rather than the release branch, which is what makes that check meaningful. No nightly contains a top-level backend= or the auto_unbounded_softplus_candidate gate.

Per-item breakdown: what is changeable outright vs. what carries a compatibility constraint

Unreleased — changeable outright, no deprecation cycle

  • The backend= parameter on the top-level facade, in its entirety. Released flashinfer.recurrent_kda has no such parameter; the only backend= in the released kda.py is an internal hardcoded backend="cake" prefill call. "auto" as the top-level default is therefore unreleased.
  • "auto" as a literal anywhere. Released kda_decode.recurrent_kda and run_recurrent_kda both declare Literal["cute-dsl", "cake"] = "cute-dsl", and the auto_unbounded_softplus_candidate gate is absent from the released kernel module. [Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935 has never shipped, so it can be fixed by redesign rather than by a compat-preserving patch.
  • RecurrentKDAPrefillWrapper and the entire training/backward family (five symbols), with their modules kda_backward.py, kda_training.py, _kda_training_impl.py, _kda_training_dispatch.py, kda_prefill_cute.py and docs pages kda_backward.rst, kda_training.rst. The bare-decorator/no-trace choice and the private _*_impl split (problem 8) can simply be corrected. The unreleased footprint is wider than the public symbols suggest: 15 KDA-related Python modules exist on main and not in v0.6.18, the rest being kda_kernels/kda_chunked_bt16.py, the four kda_kernels/sm120_prefill/ files, and five jit/ modules (_kda_jit_common.py, cake_kda.py, cake_kda_decode.py, flash_kda_backward.py, flash_kda_training.py).
  • The CuTe DSL prefill backends and the version-dependent routing they introduced (feat(kda): add CuTe DSL recurrent prefill backend #4605 / feat(kda): add SM120a CuTe DSL prefill backend #4633 / fix(kda): fall back to Cake when CuTe DSL predates cutlass.experimental #4667). One precision: fix(kda): fall back to Cake when CuTe DSL predates cutlass.experimental #4667 was partially cherry-picked onto the release branch as 21a95d73, but only its shared is_cute_dsl_experimental_available probe (14 lines in cute_dsl/utils.py) — the KDA routing did not ship, and the release engineer said so in the commit message: "0.6.18 has no kda_prefill_cute / feat(kda): add CuTe DSL recurrent prefill backend #4605 auto CuTe prefill path (public recurrent_kda already uses Cake), so only the shared is_cute_dsl_experimental_available probe is cherry-picked." That is upstream confirming this section's thesis in the release branch's own history.
  • The whole SM120a prefill tranche. kda_kernels.__all__ grew from 6 entries to 9, and all three additions are SM120a: can_implement_kda_prefill_sm120, clear_kda_prefill_sm120_caches, run_kda_prefill_sm120. Behind them sits a new flashinfer/kda_kernels/sm120_prefill/ subpackage (decomp.py, fused.py, runtime.py), the _sm120_kda_prefill_* eligibility and rejection machinery in kda_prefill.py, RecurrentKDAPrefillWorkspace._sm120_state, and the "auto" prefill route that tries SM120a first (kda.py:302-306). None of it is in v0.6.18.

Released — real compatibility constraints

Note that the core complaint of this issue is in the released set, so consolidation still needs a deprecation story:

  • Both recurrent_kda facades shipped under the same name, and both were already decorated with recurrent_kda_trace — the duplicate name and the trace-registry collision (problems 1 and 2) are released.
  • Problem 3(b) is released: _is_plain_multi_token_prefill and the top-level facade's use of it both shipped, while the decode facade never classified phase.
  • kda_decode.recurrent_kda's keyword-only backend: Literal["cute-dsl", "cake"] = "cute-dsl". main has since widened that enum to include "auto" (feat(cake_kda): add native unbounded-softplus Kimi-Linear kernels #4535), which is purely additive; changing the default would be a breaking behavior change.
  • seq_order, prefill_workspace, and the checkpoint kwargs on the top-level facade.
  • The kda_kernels noun aliases and their submodule shadowing (shim item 4).
  • FLASHINFER_KDA_T1_FAST_PATH (problem 6).
  • The three-page docs split.

Problem 3(a) — the default divergence — is the one part of problem 3 that is unreleased, precisely because the released top-level facade had no backend at all.

Timing

main is still stamped 0.6.18, and no 0.7 branch or rc tag exists yet. Recent stable cadence is roughly two to three weeks (v0.6.17 2026-08-11, v0.6.18 2026-08-29), so the next cut is plausibly mid-September. Once 0.7 ships backend="auto" as the top-level default, changing it presumably owes users a deprecation cycle rather than being a straight correction, and the cheap fixes here turn into expensive ones. The unreleased-surface decisions should land before the next release cut.

Holding action: what 0.7 should advertise

The plan is not to advertise the new KDA surface as stable in 0.7. Six of the ten top-level symbols have never shipped, and several are ones this issue proposes to change. Marking them experimental states that directly: they remain usable, but their signatures may change as the unification below lands, and nobody acquires a compatibility claim on them in the meantime.

Update (2026-09-04): five of those six are now proposed for removal rather than marking. Per @YangXu1990uiuc's comment, FlashInfer and cuDNN Frontend are splitting KDA by phase — inference here, training there, where cudnn.linear_attention.ops.kimi_delta_attention already provides fwd+bwd with autograd and the CAKE training kernels will land as an engine under it. #4965 implements the removal of recurrent_kda_training_forward, recurrent_kda_training_backward, RecurrentKDATrainingContext, recurrent_kda_backward and RecurrentKDABackwardWorkspace. Since none has shipped and a GitHub-wide search finds no external caller, removal before the cut needs no deprecation cycle, and it disposes of problem 8 and the aot.py registrations outright instead of relocating them. That leaves RecurrentKDAPrefillWrapper as the only unreleased top-level symbol needing the experimental treatment. #4965 is still a draft pending its cuDNN Frontend counterpart (NVIDIA/cudnn-frontend#912), so the marking below remains the fallback for the training family if the cut lands first. The rest of this section is written for that fallback; ordinary prefill, including the m128_n16_checkpoint variant and the checkpoint kwargs that serve SGLang prefix caching, is untouched either way.

The mechanism is #4880, which merged on 2026-09-04. Marking a symbol is non-breaking: @flashinfer_experimental_api emits an ExperimentalWarning once rather than raising, and its docs are explicit that calling an experimental API is itself the opt-in, requiring no environment variable. What the env var FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1 gates is automatic selection — whether backend="auto" may pick a backend marked @experimental_backend. trace is optional for experimental APIs and is_experimental excludes them from the stable trace inventory, which incidentally sanctions the training family's bare-decorator state (problem 8) rather than leaving it an inconsistency.

Applying it means following #4880's own guidelines rather than bolting a decorator onto the current layout. Concretely, for the KDA surface:

Decorate methods, not classes. @flashinfer_experimental_api returns a function, so applying it to a class silently replaces the class with a function and breaks isinstance, subclassing and Sphinx autoclass — and docs/api/kda_prefill.rst autosummarizes RecurrentKDAPrefillWrapper, so that would be a visible docs regression. Per-method decoration avoids this and is already the codebase's convention for class-based public ops: 98 indented @flashinfer_api sites, including sparse.py:1224 where BlockSparseAttentionWrapper.run carries a full trace template, against zero applications to a class. That covers all six symbols:

Unreleased symbol Kind Disposition
RecurrentKDAPrefillWrapper class mark — per-method on plan / run
recurrent_kda_backward function remove (#4965); markable as fallback
recurrent_kda_training_forward function remove (#4965); markable as fallback
recurrent_kda_training_backward function remove (#4965); markable as fallback
RecurrentKDABackwardWorkspace class remove (#4965); as fallback, transitively covered — inert without recurrent_kda_backward
RecurrentKDATrainingContext class remove (#4965); as fallback, transitively covered — a returned tape dataclass, never user-constructed

Worth adding an isinstance(f, type) guard to the decorator so the degradation is an error rather than silent; raised on #4880. Note that per-method decoration warns but does not document: RecurrentKDABackwardWorkspace and RecurrentKDATrainingContext have their own autosummary entries and would render as ordinary stable pages, so each needs an explicit .. warning:: in its class docstring.

Marking and containment are separate pieces of work, and only marking belongs before the cut. Decorating the symbols is hours of work with no kernel, wheel or CI-lane change. Containment is not: #4880's CONTRIBUTING.md requires experimental tests to live in tests/experimental/, and its pytest.ini sets norecursedirs = test_helpers experimental, so complying moves tests/kda/test_recurrent_kda_backward.py and test_recurrent_kda_training.py — 1,522 lines of Blackwell coverage — out of the default lanes days before a cut. De-registering also breaks test_aot_registers_flash_kda_backward, which pins the spec list exactly (tests/jit/test_flash_kda_backward_jit.py:206-213), and the real footprint is ten aot.py sites plus three tests, not four entries: imports at :98-99, capability probes at :553-557, the four appends at :674-679, and detect_sm_capabilities at :1235-1242. The relocation is also not a file move — the containment rule wants a thin core entry point, and flashinfer/kda_training.py is currently a bare re-export, so the 15-parameter signature has to be restated in core, where tests/kda/test_recurrent_kda_training.py:86-121 pins it verbatim and a *args, **kwargs handoff fails. First-call JIT for training users is roughly 91 s per cache directory, measured. None of that is an argument against containment — it is the agreed destination — but nothing in the 0.7 contract depends on file placement, so it should land in the first release after. RecurrentKDAPrefillWrapper is exempt either way: it is host-side planning that ends in a call to the stable recurrent_kda facade, so it has no backend implementation to contain.

Relocate and de-register, per the containment rule — the fallback path only, superseded if #4965 lands. #4880's README is explicit that a core entry point may include only the signature, shared validation, the gate check, backend selection and a handoff, and that "if the API is experimental, its implementation lives here too, regardless of the implementation's maturity — this keeps removal a one-directory deletion". So kda_backward.py and _kda_training_impl.py (~3,000 lines) move under flashinfer/experimental/ behind thin gated entry points, and the ten aot.py sites enumerated above come out, leaving the training and backward kernels JIT-only as the policy requires. Two consequences to expect rather than to debate: wheel contents change, and training users take first-call JIT latency. Both are the normal price of the experimental track, and both reverse cleanly when these APIs graduate.

One genuine convergence, worth aligning deliberately. FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS exists precisely so that backend="auto" declines to select experimental backends — which is exactly the shape of the SM120a prefill question. The unreleased surface splits by architecture, and the two halves want different halves of the policy. The six top-level symbols are Blackwell-datacentre only: recurrent_kda_backward and the training trio gate on _SUPPORTED_COMPUTE_CAPABILITIES = ((10, 0), (10, 3)) (kda_backward.py:35, _kda_training_impl.py:22), and RecurrentKDAPrefillWrapper is documented "Compute capability 10.0 and 10.3 only". Those are API surface. The SM120a prefill tranche is backend surface reached through a stable API — the shape @experimental_backend is built for. It cannot use that decorator yet: @experimental_backend marks a @backend_requirement checker, and KDA has none, so it only becomes available after the migration in remedy 2. Until then the hand-rolled equivalent is warn_experimental_backend_once on an explicit branch plus experimental_auto_backends_allowed() guarding the automatic one at kda.py:302not require_experimental_auto_backends, which raises rather than falling through. Gating SM120a also changes "auto" prefill routing on CC 12.0, and mostly by raising. Of three prefill spellings that work on CC 10.0, only the packed one still succeeds once SM120a is gated off: a fixed-T call (q=[B, T>1, H, 128], no cu_seqlens) falls through to the generic path and then raises Decode only supports T=1 without cu_seqlens, and a call passing prefill_workspace= raises ValueError because SM120a is the only CC-12.0 backend that accepts a workspace (prefill_workspace at kda_prefill.py:3627, the raise at kda.py:473). Neither is a regression against v0.6.18, which had no CC-12.0 prefill backend at all, but the second contradicts a "falls through rather than raises" reading. Worth noting #4880's stated default here is to raise with a message naming the variable, so a fall-through is a deliberate divergence. There is no SM121 question: the gate is exactly _SM120_KDA_SUPPORTED_COMPUTE_CAPABILITIES = {(12, 0)} (kda_prefill.py:3594, mirrored by kda.py:325), so CC 12.1 is served by no KDA prefill backend, released or not.

The architecture predicate is enumerated in eight Python places, four of which independently spell the same pair. ((10, 0), (10, 3)) appears verbatim at kda_backward.py:35, _kda_training_impl.py:22, kda_prefill_cute.py:34 and kda_prefill.py:43 — the last two being the prefill entry-point gates. Alongside them: _FLASH_KDA_DECODE_ARCH_BY_COMPUTE_CAPABILITY (recurrent_kda.py:70-75) for decode, and (12, 0) three times over for SM120a — _SM120_KDA_SUPPORTED_COMPUTE_CAPABILITIES (kda_prefill.py:3594), SM120_CAPABILITY (kda_kernels/sm120_prefill/runtime.py:138) and DEVICE_CC (sm120_prefill/fused.py:99). #4710 is a fair illustration rather than a counterexample: it widened the four csrc/kda/*.cuh runtime guards to admit CC 10.7 and refreshed the JIT cache keys that edit invalidates, and its scope note states plainly that the Python entry points still restrict KDA to 10.0/10.3 so end-to-end Rubin support is separate work. The observation is only that the separate work is spread across five sites, and that while "auto" raises on an unmapped capability (#4935) they have to move in lockstep with the kernels.

Marking does not fix #4935. "auto" is the top-level default and the defect is in its decode routing, so that needs the change in step 1 regardless.

The marking and the bug fix are independent, and only the bug fix is release-critical. Marking is worth doing before the cut because it is what keeps the unification below cheap, but nothing breaks if it slips. #4935 is different: it is the behaviour of the default path, so shipping it means shipping a broken default.

How we got here

Timeline of the merged PRs that produced the current shape
PR / issue Effect on API
#4262 Introduced phase-neutral flashinfer.kda.recurrent_kda and switched the top-level export to it; automatic frozen prefill dispatch
#4279 Added backend="cake" to kda_decode.recurrent_kda and the run_recurrent_kda backend — updating the older decode facade, not the new top-level one
#4243 Added fused_kda_decode (reasonable separate op: conv + recurrence + norm)
#4313 / #4314 SM100-family kernel sharing around the same timeframe (facade split adjacent, not the root cause)
#4417 (merged 2026-08-14) CuTe packed decode landed internal, with an explicit "not public API" note — the guidance below was followed
#4445 (merged 2026-08-16) Added packed_kda_decode as a public noun-form facade (trace + top-level export + docs), superseding #4378. Also added checkpoint kwargs to the top-level facade
#4605 (merged 2026-08-20) Added backend="auto"|"cute-dsl"|"cake" to flashinfer.kda.recurrent_kda and wired CuTe DSL recurrent prefill. This is Step 1 of #4483's proposed sequence
#4633 (merged 2026-08-22) SM120a CuTe DSL prefill backend, tried first under "auto"/"cute-dsl"
#4667 (merged 2026-08-26) Made "auto" fall back to Cake when CuTe DSL predates cutlass.experimental (<4.7), and made explicit "cute-dsl" raise ImportError there — introduced version-dependent routing. The strict path is currently red in CI (#4911)
#4535 (merged 2026-08-31) Added "auto" to kda_decode.recurrent_kda and run_recurrent_kda, selecting Cake for the equal-head/D128/T=1 unbounded-softplus contract. Source of #4935
#4182 (open) Indexed paging / dispatch heuristics; continues editing the decode path without consolidating facades
#4378 (closed unmerged) Cake serving-native packed decode; superseded by #4445

The v0.6.18 branch cut falls between #4445 (08-16, released) and #4605 (08-20, unreleased). Everything from #4605 onward in that table — the top-level backend= parameter, both CuTe DSL prefill backends, the version gate, and "auto" decode — is main-only.

Concrete problems

  1. Duplicate canonical name. Two @flashinfer_api wrappers share the name recurrent_kda and the same trace template, and are not the same object (flashinfer.kda.recurrent_kda is flashinfer.kda_decode.recurrent_kdaFalse). The top-level facade still bypasses the decode facade by calling raw _run_recurrent_kda, so decode through the top-level path is never seen by the decode wrapper's own logging/tracing. (It does now forward backend=, so the original "never sets backend" finding no longer applies.)

  2. Trace registry collision. Both facades register in _TRACE_REGISTRY under the label 'recurrent_kda'. The count depends on import scope — plain import flashinfer shows two colliding labels, and importing every submodule shows four:

    plain import flashinfer:   {'cute_dsl_fused_moe': 2, 'recurrent_kda': 2}
    all submodules imported:   {'cute_dsl_fused_moe': 2, 'recurrent_kda': 2,
                                'block_sparse_run': 2, 'gdn_decode': 2}
    

    (Reproducing the second line needs a sweep that tolerates import failures; a naive import-everything loop raises before finishing.)

    These are not all the same defect, and the distinction matters for the fix:

    Label Shape of the collision
    recurrent_kda Two same-named functions in different modules (kda, kda_decode) — the duplicate-facade problem this issue is about
    gdn_decode Two differently-named functions in one module sharing a template — the sibling op, tracked in [RFC] Unified GDN Decode/Prefill API #2687
    cute_dsl_fused_moe Two differently-named functions (…_fused_moe, …_fused_moe_nvfp4) in one module sharing a template
    block_sparse_run Two differently-named functions in different modules sharing a label

    Only recurrent_kda is two same-named public facades; the other three are deliberate template reuse that merely produces a non-unique label. Any registry-uniqueness fix should distinguish these cases rather than treat them alike.

    No test catches this. The registry's duplicate key is (func.__module__, func.__qualname__, label) (tests/trace/template_registry.py:100-103), and the two facades differ in __module__, so the shared label never trips the guard. The collision is not a failing test — it is unpoliced, which is the weaker-looking but more durable argument for fixing it.

    recurrent_kda_trace has since gained ssm_state_indices, state_checkpoints, checkpoint_cu_starts, and checkpoint_every_n_tokens, but still does not model backend, seq_order, or prefill_workspace, and is still tagged stage:decode despite dispatching prefill. fi_trace fixtures still pin fi_api:flashinfer.kda_decode.recurrent_kda only, so top-level fi_api strings remain untested.

  3. The two facades disagree on the default, and the top-level default is not a superset. Both delegate decode to the same run_recurrent_kda, so for a fixed backend value the decode kwargs are identical. Two things still diverge:

    (a) The default value. Verified by intercepting the forwarded value:

    top-level backend='auto'      -> decode backend received: 'auto'
    top-level backend='cake'      -> decode backend received: 'cake'
    top-level backend omitted     -> decode backend received: 'auto'
    
    Entry point backend default Argument-free T=1 decode
    flashinfer.recurrent_kda "auto" Cake when the equal-head/D128 unbounded-softplus gate matches — or ValueError if Cake cannot serve it ([Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935)
    flashinfer.kda_decode.recurrent_kda "cute-dsl" CuTe DSL always

    Because "auto" commits to Cake without falling back (#4935), the top-level default is not a superset of the decode facade's — and the two accepted sets are in fact incomparable, since one guard is skipped under "auto" that "cute-dsl" enforces. The same call can succeed through one import path and raise through the other, in both directions. This is the strongest argument for consolidation, but it is a defect to fix in [Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935, not something to design around here.

    (b) Phase classification runs before delegation. flashinfer.kda.recurrent_kda calls _is_plain_multi_token_prefill first, which classifies on shape only (kda_prefill.py:315-327): no num_spec_tokens, and either fixed T > 1 or packed total_tokens > num_sequences. There is no head-dim, dtype, or architecture guard. So a packed multi-token decode call is claimed by the prefill branch:

    packed q=[1,4,H,128], cu_seqlens=[0,2,4], num_spec_tokens=None, precomputed gate:
      kda_decode.recurrent_kda (today)      -> OK (reaches decode)
      kda.recurrent_kda backend='cute-dsl'  -> ValueError: does not support this prefill contract
      kda.recurrent_kda backend='cake'      -> ValueError: does not support this prefill contract
      kda.recurrent_kda backend='auto'      -> OK (reaches decode, prefill backends ineligible)
    

    The decode facade never runs this classifier. The "auto" outcome here is gate-mode dependent — with the raw Kimi-Linear softplus gate (use_gate_in_kernel=True, lower_bound=None) the same call raises instead, because it then matches the "auto" Cake gate discussed in [Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935. Either way the conclusion holds and in fact strengthens: no single backend value reproduces the decode facade's current behavior. This is the constraint that shapes the shim design below.

  4. Discoverability. Resolved by feat(kda): add CuTe DSL recurrent prefill backend #4605. flashinfer.recurrent_kda(..., backend="cake") now exists and reaches the strict Cake decode path.

  5. Docs / traces lag the facade. Five Sphinx pages now describe pieces of one recurrent API; kda_decode.rst still autosummarizes recurrent_kda; kda_prefill.rst documents a class defined in kda.py. Both the kda_decode.py docstring and docs/api/kda_decode.rst describe an "auto" → CuTe-DSL decode fallback that does not exist (see [Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935). kda_decode.rst documents the Cake export set thoroughly — the exported bodies, gate modes, per-T value-row splits and per-architecture split policies — but not the layout boundary: whether a given call is accepted depends on a 259-line selector over strides, pointer alignment and int32-overflow products that no user can predict from the prose. Once "auto" falls back correctly this is a performance-predictability problem rather than a correctness one, but it should be stated somewhere. The trace template gaps and stage:decode tag are covered in problem 2.

  6. An environment variable participates in public dispatch, documented only for contributors. FLASHINFER_KDA_T1_FAST_PATH (default "1", recurrent_kda.py:139) selects a different T=1 decode route. It appears in CLAUDE.md:607 but nowhere in docs/api/, so users of the public API cannot discover it. Which kernel a given public backend= value resolves to therefore depends on process environment. Any unified contract needs to either document this as public API or make it internal-only.

  7. Fragmentation is now happening, not hypothetical. Open PRs are each adding surface in a different style, concurrently editing the same two facades:

    PR New convention introduced
    #4779 Adds flashinfer/kda_vibecuda.py and widens the top-level enum to Literal["auto","cute-dsl","cake","vibecuda"] on kda.py only — a fourth backend literal that the decode facade will not accept, pushing the two signatures further apart
    #4765 Exports generated Blackwell recurrent-KDA specializations, fp32 state and long context; edits kda.py and kda_prefill.py and adds a jit/flash_kda_nvrtc.py compilation path
    #4809 Adds top-level fused_kda_decode_packed — a fourth *_kda_decode noun
    #4886, #4845, #4846 Further kda.py / kda_prefill.py prefill backends
    perf(kda): optimize indexed recurrent decode dispatch and paging #4182 Decode dispatch/paging edits in kda_decode.py

    feat(vibecuda_kda): add SM100/SM103 recurrent KDA prefill backend #4779 is the immediate instance: it adds "vibecuda" to kda.py's enum only, which is the divergence this issue is about — no comment on the kernel work, just the one line that would need to cover both facades.

    Worth citing as the counterexample: #4709 adds disable_state_update with backend-specific raises but updates both kda.py and kda_decode.py, keeping the two signatures in step. That is the behavior this issue wants to make the default expectation.

  8. New surface outside the original scope. kda_backward.py and kda_training.py introduced a paired-function training family, and RecurrentKDAPrefillWrapper introduced a plan-and-run class convention. Neither is covered by the "noun-form facade vs internal run_*" rule below, so that rule needs cases for stateful plan/run wrappers and paired forward/backward ops. Two further inconsistencies came along with them:

    • All three new ops use bare @flashinfer_api with no trace template, unlike every recurrent_kda/*_kda_decode facade. So "public op" currently means "decorated" in two incompatible ways, and the training family is invisible to fi_trace.
    • flashinfer/kda_training.py is a 23-line re-export of private flashinfer/_kda_training_impl.py — a module-layout convention used nowhere else in KDA (kda_backward.py keeps its implementation inline). There is a stated reason: the module also defines _get_training_module, an "internal indirection seam used by the training API tests". That justifies the seam; whether it also justifies the split is worth a sentence from the author. Either way it is cosmetic next to the missing trace templates alongside it.

    Adjacent: #4896 (production KDA prefill for SGLang prefix caching).

Non-goals

  • Do not merge fused_kda_decode or packed_kda_decode into recurrent_kda — different serving contracts (packed QKV, conv, RMSNorm; different return types).
  • Do not fold the training/backward family into the inference facade.
  • Do not require open kernel PRs to refactor all of KDA; they should either follow the target facade pattern or stay explicitly internal.
  • Do not change numerical contracts of existing CuTe / Cake kernels in the unification PR.
  • Do not fix [Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935 here. It is a behavior defect with its own regression test; this issue depends on it but should not absorb it.

Proposed target (backwards-compatible)

Canonical API

One public recurrent inference entrypoint — this now exists and already accepts backend=:

flashinfer.recurrent_kda(..., backend="auto", ...)
# alias: flashinfer.kda.recurrent_kda

Current backend semantics, as actually implemented on main:

backend Behavior
"auto" (default) Prefill: SM120a CuTe → SM100-family CuTe DSL BT=16 → frozen Cake fallback (Cake also used when nvidia-cutlass-dsl<4.7). Decode: commits to Cake when the equal-head/D128/T=1 unbounded-softplus gate matches, raising rather than falling back if Cake cannot serve it (#4935)
"cute-dsl" CuTe DSL strictly; no frozen-prefill dispatch and no Cake decode. Raises ValueError on an unsupported plain-prefill contract, ImportError when the DSL predates 4.7
"cake" Frozen Cake strictly; raises on unsupported prefill contracts and on unexported decode shapes. No silent CuTe fallback

Target semantics: "auto" becomes a true superset — never raising where "cute-dsl" or "cake" would have succeeded. The plan below assumes that fix has landed.

Because none of this parameter has shipped, the fix should be the structural one described in #4935 — resolve backend eligibility before state normalization, so one resolved decision drives both the state convention and the launch — rather than the cheaper "keep the speculative gate and hand-synchronise a second predicate" variant. There is no released behavior to preserve, so there is no reason to accept two predicates that must be kept in step by hand; that fragility is what produced the defect.

The remaining work is not adding this parameter — it is reconciling the second facade against it. Keep seq_order / prefill_workspace / checkpoint kwargs on this facade only (all three are released there).

Compatibility shims

  1. flashinfer.kda_decode.recurrent_kda cannot simply forward to flashinfer.kda.recurrent_kda. No single backend value preserves today's behavior, because the two facades differ in both the default and the phase classifier (problem 3):

    Existing decode call Preserved by "cute-dsl" Preserved by "auto"
    T=1 Kimi-Linear equal-head/D128 yes no — commits to Cake, may raise ([Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935)
    packed multi-token, num_spec_tokens=None no — claimed by prefill classifier, raises yes

    So the shim must bypass phase classification and target a shared decode dispatcher directly, rather than re-entering the phase-neutral facade:

    • Factor the decode tail of kda.recurrent_kda into one internal _dispatch_recurrent_kda_decode(...), used by both facades, giving decode a single contract regardless of entry point. Note this does not by itself fix tracing: an internal dispatcher is undecorated by construction, so problem 1 additionally requires deciding which facade owns the recurrent_kda trace (see the docs/traces section).
    • kda_decode.recurrent_kda forwards to that dispatcher with omitted backend"cute-dsl", preserving today's kernel selection exactly.
    • Accept and forward "cake" and "auto" unchanged.
    • Emit DeprecationWarning pointing at flashinfer.recurrent_kda (soft; keep for ≥1 minor release).
    • Do not accept seq_order / prefill_workspace on the shim — those are prefill-only and the shim deliberately does not reach prefill.
  2. Decide the top-level default explicitly, before 0.7. This is a policy choice separate from [Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935: even once "auto" stops raising, it selects a different decode kernel than "cute-dsl" would. Today, no single value is a safe default for a phase-neutral facade — and that is the sharpest evidence for this whole issue. The two phases have opposite totality properties:

    backend="auto" backend="cute-dsl"
    decode raises on four input classes (#4935) accepts classes 1-3, and accepts class 4's padded layout only below the one-warp threshold; separately raises on a non-contiguous initial_state pool that "auto" accepts
    prefill falls back to Cake when CuTe DSL is ineligible raises ImportError below nvidia-cutlass-dsl 4.7 (kda.py:366-372) — the second of the two failures in [Bug][v0.6.18] test_recurrent_kda_prefill.py ModuleNotFoundError: No module named 'cutlass.experimental', ImportError: backend='cute-dsl' requires nvidia-cutlass-dsl>=4.7.0 (cutlass.experimental); #4911

    Neither value is total on either phase, and on decode the two accepted sets are incomparable rather than nested (see #4935, Non-comparability); "cake" is strictly narrower than both. The prefill raise says so in its own message: "backend='cute-dsl' requires nvidia-cutlass-dsl>=4.7.0 (cutlass.experimental); backend='auto' falls back to Cake." So defaulting the top-level facade to "cute-dsl" would break defaulted prefill calls on older CuTe DSL, and defaulting it to "auto" ships the known-broken decode default. A single backend knob on a phase-neutral facade is being asked to encode two independent phase-specific decisions, which is problem 3(b) restated as a concrete release blocker.

    That resolves the choice rather than blocking it:

    • "auto" is the right top-level default, but only once step 1 lands. It is already the only value with fallback semantics on both phases; step 1 makes those semantics true on decode as well as prefill. Without step 1 there is no safe default at all, which is the strongest reason to treat it as release-critical rather than merely desirable.

    • The top-level default is unreleased, so it can be set either way with no deprecation — but note that in v0.6.18 the top-level facade had no backend at all, so a defaulted top-level decode call reached run_recurrent_kda's own "cute-dsl" default. That call has already changed kernel between the release and main — and on the inputs in #4935 it now raises where the v0.6.18 spelling of the same call worked, which makes step 1 a regression fix against the last release rather than merely a correction to unreleased surface. Ratifying "auto" is a silent kernel swap for existing top-level callers, not a neutral blessing of the status quo, and it should be benchmarked rather than assumed.

    • Adopt @backend_requirement rather than hand-rolling the dispatch. utils.py:1185 already provides the repo's standard mechanism: per-backend eligibility checkers plus an "auto" heuristic, resolved by computing the set of suitable backends first — filtered on compute capability and problem size — and raising only when that set is empty. It is applied at 44 sites across 15 modules (gemm/, grouped_mm/, fused_moe/, quantization/, deep_gemm.py) and at zero sites in KDA, and it composes with @flashinfer_api and a trace template at 43 of those 44 (e.g. attn_scores.py:1020-1021). Its resolve-then-dispatch ordering is what makes commit-then-check hard to write by accident, which is why it is the right target shape rather than a bespoke selector, and it gives the phase asymmetry above a natural home as two decorated entry points with their own checker sets.

      Two limits to plan around, though, which is why this is post-0.7 and not a wholesale replacement. Checkers run before normalization (utils.py:1412-1459 validates and then calls; there is no post-normalization hook), so they can host KDA's contract predicates — trigger classes 1-3 — but not its layout predicates, which are evaluated against re-viewed and synthesized tensors that do not exist until :2263-2337 and which produce trigger class 4. The decorator also injects a public skip_check kwarg whose True path bypasses all checking. So the realistic shape is contract predicates in checkers and a late fallback for layout, not one mechanism for both.

    • The decode facade's "cute-dsl" default is released and is decode-only, so no prefill concern touches it and it should not be changed at all. Its enum already accepts "auto" (kda_decode.py:79, added by feat(cake_kda): add native unbounded-softplus Kimi-Linear kernels #4535), so nothing needs adding there. The two facades converge on "auto" only after step 1 — and if step 1 slips, they should not be converged at this cut.

  3. flashinfer.fused_kda_decode / packed_kda_decode: unchanged.

  4. flashinfer.kda_kernels.run_*: document as backends only; remove the three noun aliases (recurrent_kda, fused_kda_decode, packed_kda_decode) from kda_kernels.__all__. This un-shadows the same-named submodules and lets in-repo consumers drop their importlib.import_module workarounds. Extend the existing packed_kda_decode_cute "not public API" note into a module-wide statement. These aliases shipped in v0.6.18, so the removal needs the same soft-deprecation window as item 1 even though it is a correctness fix.

Docs / traces

  • Single "start here" page for recurrent KDA (docs/api/kda.rst); kda_decode.rst focuses on fused_kda_decode + packed_kda_decode plus a compatibility note for the old recurrent_kda import; kda_prefill.rst remains workspace/contract detail; consider moving RecurrentKDAPrefillWrapper onto the page for its defining module (Sphinx resolves it correctly today via .. currentmodule:: flashinfer.kda, so this is page organisation rather than breakage); kda_backward.rst / kda_training.rst stay as the training pair.
  • Correct the "auto" fallback description in kda_decode.py and docs/api/kda_decode.rst (tracked in [Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935).
  • Document or internalize FLASHINFER_KDA_T1_FAST_PATH.
  • Extend recurrent_kda_trace for backend, seq_order, and prefill_workspace, retag the stage now that the template covers prefill, and give the two facades distinguishable registry labels (or collapse them once the shim lands). Treat the gdn_decode / cute_dsl_fused_moe / block_sparse_run collisions as a separate, lower-priority cleanup — they are template reuse, not duplicate facades.
  • Add fi_trace coverage for the top-level facade's fi_api strings, then point decode benches/tests (CuTe + Cake) at the canonical facade.

Rule for new kernels (immediate guidance)

New KDA work should do one of:

  1. Public op: noun-form facade (@flashinfer_api with a trace template + top-level/docs export); run_* stays in kda_kernels as backend. State explicitly if a trace template is deliberately omitted, rather than defaulting to a bare decorator as the training ops do.
  2. Public stateful op: plan-and-run class next to its facade, documented in its defining module's page (as RecurrentKDAPrefillWrapper does), or a paired forward/backward module (as kda_training). Keep the implementation in the public module unless there is a stated reason for a private _*_impl split.
  3. Internal/experimental: do not add to public __all__ / top-level; state that clearly in the PR — the precedent set by feat(kda): add packed-input CuTe decode kernel #4417. With #4880 landing in 0.7 this becomes a formal track: @flashinfer_experimental_api for an unstable public API, and flashinfer/experimental/ plus @experimental_backend for an opt-in backend reached through an explicit backend= literal. Prefer that over widening a stable enum.

Additionally:

  • Do not add a new backend literal to one facade only (cf. feat(vibecuda_kda): add SM100/SM103 recurrent KDA prefill backend #4779). A new backend value is part of the shared contract.
  • Do not add a new top-level *_kda_decode noun without stating why it cannot be a backend= value or a kwarg on an existing facade.
  • Do not introduce a further convention such as "only kda_kernels.run_foo / kda_kernels.<submod>.run_foo is the product API."

Why this is backwards compatible

Compatibility argument, item by item

The compatibility baseline is v0.6.18, not current main.

Related

Relationship to #4483

What #4483 got right, what it got wrong, and what is new since

#4483 raised this problem on 2026-08-12 and was closed on 2026-09-03 in favour of this issue. What changed in the three weeks between:

Resolved. Discoverability — #4483's fourth problem — is fixed: backend= landed on the top-level facade in #4605 and is forwarded to decode, so flashinfer.recurrent_kda(..., backend="cake") now reaches the strict Cake path. Two pieces of #4483's guidance were also followed: #4417 merged internal-only, and #4378 was closed in favour of #4445's public packed_kda_decode noun facade.

Superseded by events. #4483 described a single "opposite backend policy" between the two facades. That framing no longer fits. The facades now share one run_recurrent_kda and differ in their defaults, and the top-level default is not a safe superset — "auto" decode commits to Cake and raises rather than falling back, which is #4935. Its prefill description has also inverted, but by events rather than by error: #4483 described "auto" prefill as preferring frozen FlashKDA, which was correct on 2026-08-12 and is still what released v0.6.18 does (kda.py:260 hardcodes backend="cake"). CuTe-DSL prefill only arrived with #4605 on 08-20. Today "auto" prefers CuTe DSL with Cake as the fallback, and that fallback is nvidia-cutlass-dsl version-dependent (#4667).

New since #4483. The backend= parameter itself; the "auto" literal and the defect behind it; a training/backward family and a plan-and-run wrapper class (six top-level symbols, all decorated-but-untraced or classes, plus a private _kda_training_impl split); an SM120a prefill tranche that added three kda_kernels exports and a new subpackage; two more docs pages; and concrete fragmentation across seven open PRs, with #4709 as a counterexample that keeps both facades in step. Critically, none of that has shipped — see Release status — which is what makes the plan above cheap now and expensive after 0.7.

Unchanged and still reproducing. The duplicate recurrent_kda name and shared trace label; the top-level facade bypassing the decode wrapper; the stage:decode tag on a template that dispatches prefill; missing backend / seq_order / prefill_workspace trace fields; kda_decode.rst autosummarising recurrent_kda; the kda_kernels noun aliases; decode-only fi_trace coverage; and the prefill/decode test import split.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Type

No type

Projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions