[Bug] recurrent_kda backend="auto" decode commits to Cake and raises instead of falling back to CuTe DSL #4935

Summary

On main, flashinfer.recurrent_kda defaults to backend="auto" (unreleased — see Environment). For T=1 decode, "auto" is documented as selecting a frozen Cake kernel for the equal-head/D128 unbounded-softplus contract and otherwise preserving CuTe DSL. It does not. Once "auto"'s coarse shape gate matches, run_recurrent_kda commits to Cake, and a failed variant selection is fatal:

ValueError: the requested Cake recurrent_kda decode contract is unsupported

The gate is coarser than the contract it commits to: it omits exclusions that the explicit backend="cake" path applies, and it omits the architecture check that variant selection performs. So the default top-level entry point raises on calls that succeed through flashinfer.kda_decode.recurrent_kda (default "cute-dsl"), and that succeeded through the top-level facade before #4535.

"auto" is not a superset of "cute-dsl". (It is not strictly narrower either — see Non-comparability — which is why the fix needs a two-way regression test.)

Environment

• main @ 1dff49bc
• Verified on NVIDIA B200, CC 10.0 (sm100a) — a fully supported architecture
• Introduced by feat(cake_kda): add native unbounded-softplus Kimi-Linear kernels #4535 (merged 2026-08-31)
• Not present in any release. v0.6.18, the current stable, was branch-cut from main on 2026-08-19 UTC and is not an ancestor of it. In the v0.6.18 tree, kda_decode.recurrent_kda and run_recurrent_kda both declare backend: Literal["cute-dsl", "cake"] = "cute-dsl", the top-level flashinfer.recurrent_kda has no backend= parameter at all, and auto_unbounded_softplus_candidate does not exist. So "auto" is unreachable in released code and no released user is affected. Nor is it reachable from any other channel: PyPI serves flashinfer-python 0.6.18 as latest, and nightly tagging stops at nightly-v0.6.18-20260819, one day before feat(kda): add CuTe DSL recurrent prefill backend #4605 added the parameter.

main is still stamped 0.6.18 with no 0.7 branch or rc tag yet, so this can be fixed before it ever ships. That is the reason the structural fix below is preferred: with no released behavior to preserve, there is no argument for keeping the fragile two-predicate arrangement that caused the bug. It is also the reason this should not be deferred — once 0.7 ships backend="auto" as the top-level default, both the default and its semantics become a compatibility constraint.

Reproduction

Dispatch-level probe. _run_flash_kda_decode is stubbed to raise a sentinel so kernel selection is observable without launching; nothing else is altered. Note flashinfer.kda_kernels.recurrent_kda is shadowed by a function of the same name, hence importlib.

import torch, importlib
R = importlib.import_module("flashinfer.kda_kernels.recurrent_kda")

class CakeSelected(Exception): pass
R._run_flash_kda_decode = lambda *a, **k: (_ for _ in ()).throw(CakeSelected())

H, D, B = 32, 128, 2
f = lambda *s: torch.zeros(*s, dtype=torch.bfloat16, device="cuda")

def args(**over):
    kw = dict(
        q=f(B,1,H,D), k=f(B,1,H,D), v=f(B,1,H,D), g=f(B,1,H,D), beta=f(B,1,H),
        A_log=torch.zeros(H, dtype=torch.float32, device="cuda"),
        dt_bias=torch.zeros(H*D, dtype=torch.float32, device="cuda"),
        use_gate_in_kernel=True, lower_bound=None,
        initial_state=f(B,H,D,D),
        ssm_state_indices=torch.arange(B, dtype=torch.int32, device="cuda"),
    )
    kw.update(over)
    return kw

for backend in ("auto", "cute-dsl"):
    for label, over in (("l2norm=True ", {}), ("l2norm=False", {"use_qk_l2norm_in_kernel": False})):
        try:
            R.run_recurrent_kda(**args(**over), backend=backend)
            print(f"{backend:9s} {label} -> no Cake selection")
        except CakeSelected:
            print(f"{backend:9s} {label} -> CAKE selected")
        except ValueError as e:
            tag = "RAISED Cake-unsupported" if "contract is unsupported" in str(e) else f"other: {str(e)[:40]}"
            print(f"{backend:9s} {label} -> {tag}")
        except Exception as e:
            print(f"{backend:9s} {label} -> proceeded to CuTe ({type(e).__name__})")

Actual output on the B200:

auto      l2norm=True  -> CAKE selected
auto      l2norm=False -> RAISED Cake-unsupported
cute-dsl  l2norm=True  -> no Cake selection
cute-dsl  l2norm=False -> no Cake selection

Trigger classes

1. Unsupported device — any CC outside 10.0 / 10.3

The gate has no architecture check; variant selection maps only (10, 0) and (10, 3) (recurrent_kda.py:73-74) and returns None otherwise (:1487-1488). This class is an allowlist, so it covers every capability the map omits — including CC 10.7 (Rubin), which is not reachable end-to-end today: #4710 widened the four csrc/kda/*.cuh runtime guards to admit 10.7 (the sm_100f family build is valid across the SM100 line, so the check was stricter than the binary it guarded) and its scope note is explicit that the Python entry points still restrict KDA to 10.0/10.3, leaving end-to-end Rubin enablement as separate work. That is the reason the fallback matters more than the enumeration: while "auto" raises on an unmapped capability, every future architecture has to reach this map in lockstep with the kernels or "auto" regresses on it. Simulated by patching get_compute_capability:

simulated CC(9, 0)    backend='auto'      -> RAISED: Cake unsupported (no CuTe fallback)
simulated CC(9, 0)    backend='cute-dsl'  -> proceeds to CuTe
simulated CC(12, 0)   backend='auto'      -> RAISED: Cake unsupported (no CuTe fallback)
simulated CC(12, 0)   backend='cute-dsl'  -> proceeds to CuTe

The architecture map is the only device gate, so this class is exactly characterised — but it is not the only way the selector rejects on a supported device; see class 4. Confirm on real H100/SM120 before signing off the fix. This makes the production Kimi-Linear T=1 shape fail outright off Blackwell.

2. Supported device, contract field the gate does not check (verified on real B200)

Variant selection requires use_qk_l2norm_in_kernel (:1495); the gate never inspects it (:2063). See the reproduction output above. use_qk_l2norm_in_kernel=False is a documented public option, and on the flagship architecture the default entry point rejects it while "cute-dsl" accepts it.

3. Packed T=1 with explicit cu_seqlens — the paged-serving layout

This is the most serving-relevant class. The explicit Cake path has a purpose-written carve-out for it (:2015-2019):

if backend == "cake" and cu_seqlens is not None and num_spec_tokens is None:
    raise ValueError(
        "backend='cake' does not support explicit T=1 cu_seqlens; "
        "use standard decode without explicit cu_seqlens"
    )

The condition is backend == "cake" only. "auto" skips this exclusion, matches the gate anyway, and reaches the generic fatal raise instead — so a standard packed decode call fails under the default backend, with a less informative error than an explicit "cake" request would produce. Whatever the reason for the carve-out, it is conditioned on backend == "cake" alone, so the "auto" path does not inherit it and reaches the generic raise instead.

4. Supported device, ordinary contract, rejected on layout (verified on real B200)

The three classes above are all predicates the gate could in principle have carried. This one is not. With CC 10.0, use_qk_l2norm_in_kernel=True and no cu_seqlens — every predicate listed above satisfied — a q/k tensor with a padded head dimension (q.stride(-2) = 136 for head_dim = 128, a routine fused-QKV layout) makes the default entry point raise while "cute-dsl" runs. The rejection comes from :1559 (q.stride(-2) != head_dim), one of a hundred-plus predicates evaluated against tensors that do not exist at the gate site. scale=inf behaves the same way.

Scope honestly: this was verified below the one-warp threshold (B*HV = 64), where CuTe serves the layout. At B*HV = 256 the one-warp CuTe kernel also rejects it, so class 4 is a genuine non-superset case only below the threshold. It matters less for blast radius than for fix design — it is the reason a gate-side predicate cannot be made total, and therefore the reason the fix must key on the resolved selector result.

Expected vs actual

Root cause

The speculative gate and the actual contract are two different predicates, and the gap between them is fatal.

1. :2063 — auto_unbounded_softplus_candidate checks only backend == "auto", num_spec_tokens is None, H > 0, HV == H, K == 128, V == 128, use_gate_in_kernel, lower_bound is None, and that A_log/dt_bias are present. No architecture check, no use_qk_l2norm_in_kernel check, no cu_seqlens exclusion.
2. :2262 — narrowed only by NUM_TOKENS == 1.
3. :1487-1488 and :1495 — variant selection returns None for an unmapped architecture or when use_qk_l2norm_in_kernel is false.
4. :2367-2371 — when backend == "cake" or auto_unbounded_softplus and the variant is None, raise. A speculative "auto" match and an explicit "cake" request share one fatal path.

Step 4 is the defect.

Why the obvious one-line fix is wrong

Gating the :2367 raise on backend == "cake" alone looks sufficient. It is not — it converts a loud failure into silent state-cache corruption.

auto_unbounded_softplus_candidate is also consumed ~180 lines earlier to choose the state convention (:2184-2191):

elif ssm_state_indices is not None and (
    backend == "cake" or auto_unbounded_softplus_candidate
):
    state = initial_state                      # whole pool, no gather
    ssi = ssm_state_indices.to(torch.int32).contiguous().view(-1)
elif ssm_state_indices is not None:
    state = initial_state[ssm_state_indices].contiguous()   # gather
    copy_back_indices = ssm_state_indices                   # …and copy back

Under the gate, the caller's whole pool is passed through with copy_back_indices = None — Cake's convention. But the CuTe one-warp route ignores ssm_state_indices when cu_seqlens is absent (:464-466):

else:
    seq_idx = batch_idx

So falling through to CuTe with Cake's state convention makes the kernel read and write pool rows [0 … B-1] instead of the caller's requested slots, with no copy-back to correct it. I applied this exact patch on a B200 and measured the corruption: with ssm_state_indices=[7,3,11,0,19,5,23,1] the patched "auto" wrote rows [0..7]. The state max-diff against "cute-dsl" is input-dependent and should not be quoted as a fixed number — repeated runs with the same shape and indices but different RNG gave anywhere from 12 to 27. The row sets are the reliable signature. The patched build also passes all 382 decode tests, so the existing suite provides no protection against it. It reproduces whenever B * HV >= ONE_WARP_MIN_SEQUENCE_HEADS (128, :88) — i.e. any ordinary serving batch. With a zero-initialised pool the output difference is exactly zero while the state difference is not, which is why a naive test would pass. Below that threshold the one-liner is correct: the multi-warp route does consult ssm_state_indices, so Cake's whole-pool convention happens to be the right one there and "auto" matches "cute-dsl" bit-for-bit (measured at B*HV = 64 and 120, with 128 the first corrupting size). The single failure mode is thus confined to the one-warp route — which is the ordinary serving regime, so this makes it no more acceptable.

Critically, the naive regression test ("auto succeeds wherever cute-dsl does") passes while this corruption occurs, because a zero-initialised pool makes the output tensors compare equal.

Suggested fix

The goal is that the state convention always matches the route actually taken, so that when "auto" falls through to CuTe it does so with CuTe's gather/copy-back convention intact.

A literal hoist of variant selection above :2184 does not achieve this, and would make things worse. _select_flash_kda_decode_variant is written against post-normalization values: frozen_q is a torch.as_strided packed re-view built at :2271 (itself gated on copy_back_indices is None, an output of the normalization), frozen_out is the out_buf allocated during normalization, and frozen_cu_seqlens / frozen_ssi are synthesized at :2328-2336 when the caller passed none. The selector then rejects on exactly those: cu_seqlens is None → return None (:1491), ssm_state_indices is None → return None (:1492), and q.shape[0] != 1 (:1523), which the caller's [B, 1, H, D] fails. Hoisted verbatim, the selector returns None where the real call site returns d128_t1_unbounded_softplus_direct_split16 — silently disabling Cake "auto" altogether instead of giving it a fallback.

Reordering cannot break the dependency either, because it is genuinely circular: the state convention at :2184 writes copy_back_indices, the frozen-normalization block at :2268 is gated on copy_back_indices is None, and the selector consumes the frozen tensors that block produces. Nor is splitting the selector into a contract half and a layout half sufficient on its own: the layout guards can still reject on a fully supported device (see trigger class 4 above), so a late fallback would still be required. That makes the late fallback the core of any correct fix, and the reordering optional polish on top of it.

Recommended fix: drive the fallback off the resolved selector result, in place. Keep the gate and the ordering; when the variant comes back None under "auto", re-normalize state to CuTe's convention before falling through, and restrict the raise to backend == "cake":

if flash_kda_decode_variant is None and auto_unbounded_softplus and backend != "cake":
    if cu_seqlens_i32 is None and ssm_state_indices is not None:
        state = initial_state[ssm_state_indices].contiguous()
        copy_back_indices = ssm_state_indices
        ssi = None  # CuTe's convention is copy_back_indices with ssi unset; the Cake branch set it
if backend == "cake" and flash_kda_decode_variant is None:
    raise ValueError("the requested Cake recurrent_kda decode contract is unsupported")

This is total on every trigger class and every witness above, because the fallback is driven by the selector's actual return value rather than by a second predicate that must be kept in step — the property the alternative below fails to deliver. One residual: the fallback body is guarded on ssm_state_indices is not None, so it is a no-op on the no-indices arm, where the gate assigned state = initial_state (:2192) rather than initial_state.contiguous() (:2195). Restoring .contiguous() there closes it; see the fourth variable below. It preserves "cake" strictness and the Cake happy path, and it costs one gather on a cold path. Implemented on a B200, this gives exact output and state-pool equality against "cute-dsl" on every trigger class including class 4, with backend="cake" still raising and Cake still launching on the happy path, and the KDA suite unchanged against baseline.

A second shape satisfies the same invariant and is cleaner if the extra lines are acceptable: sink state normalization below variant resolution rather than patching up after it. That is roughly 20 lines, confined to the dense branch — the cu_seqlens branch never gated its state convention on the flag in the first place — and it also measures bit-identical output and state with 382 passed / 5 skipped, identical to baseline. Either shape is fine; the rule to enforce in review is the thing to state: no state convention may reach the kernel launch before the variant is resolved. The recommended shape satisfies this by repairing an early commitment; the second satisfies it by never committing early. What is forbidden is a second predicate deciding the convention independently of the selector's return value — the arrangement that produced this bug.

Worth noting for the follow-up rather than the fix: the repo already has a decorator that enforces this shape by construction. @backend_requirement (utils.py:1185) takes per-backend eligibility checkers plus an "auto" heuristic, computes the set of suitable backends first — filtering on compute capability and problem size — and raises only when that set is empty. It is used at 44 sites across 15 modules and at zero sites in KDA, which is why KDA hand-rolls the commit-then-check ordering that produced this bug. It is not a drop-in replacement here — its checkers run before normalization, so they can carry the contract predicates but not the layout ones — but the ordering is the part worth copying.

A third and a fourth coupled assignment. The gate at :2184-2195 writes three variables, not two. Its Cake branch sets state = initial_state and ssi = ssm_state_indices.to(torch.int32).contiguous().view(-1); its CuTe branch sets state = initial_state[ssm_state_indices].contiguous() and copy_back_indices, leaving ssi unset. So there are three state conventions, not two: Cake, CuTe one-warp, and CuTe multi-warp. The one-warp route ignores ssm_state_indices (:464-466), which is the mechanism behind the corruption above; the multi-warp route does consult it. A patch that gathers state and sets copy_back_indices but leaves ssi set writes zero rows on the multi-warp route — a new failure mode introduced by the fix, measured. Non-finite output accompanies it on some runs, from an out-of-bounds read, so the reliable signature is the empty mutated-row set rather than the NaN. Hence the ssi = None line, and hence the requirement below to test both sides of the B*HV = 128 threshold.

The no-indices arm carries a fourth divergence, in contiguity rather than in indexing: :2192 passes initial_state through where :2195 would have called .contiguous(). It is reachable only through the guard carve-out discussed under Non-comparability, and the observable cost is a worse error rather than a wrong answer — on a transposed pool the fall-through raises an opaque Mismatched mH.strides[2] from the CuTe launch instead of the guard's own message. These coupled variables are easy to miss — the third and the fourth each surfaced only after a patch that looked correct had been measured — which is the best available evidence about how risky a larger refactor here would be under deadline.

Because "auto" is unreleased (see Environment) a larger redesign would also be permissible, but it is not warranted: splitting the selector's contract guards from its layout guards is a reasonable follow-up on main after the cut, unblocked by the release, and it is not release-critical once the fallback is total.

Fixing this before the 0.7 cut is preferable to fixing it after: "auto" is the top-level default, so once it ships, changing its behaviour presumably owes users a deprecation cycle rather than being a straight correction.

The commonly suggested cheap alternative does not work — this is measured, not argued. Extending the gate at :2063 to mirror the selector's preconditions — architecture, use_qk_l2norm_in_kernel, the cu_seqlens exclusion from :2015 — was implemented and leaves at least five surviving witnesses (seven counting shape variants) where "auto" still raises and "cute-dsl" succeeds:

The reason is structural: _select_flash_kda_decode_variant is 259 lines whose eligibility test is a single 112-line boolean carrying over a hundred predicates spanning dtypes, every stride, 8- and 16-byte pointer alignment, nine int32-overflow products plus six raw stride bounds, and tensor-aliasing analysis. Mirroring it is reimplementing it, and most of its subjects are tensors that do not exist at the gate site until :2263-2337. So this option should be ruled out as non-viable rather than merely disfavoured on maintainability grounds.

Regression test requirements. Assert both directions and check state, not just the absence of an exception:

• Extend the existing test_t1_unbounded_softplus_auto_route_* (test_recurrent_kda_decode_export.py:1685) with ineligible-input cases rather than writing a new test. It already does most of what is prescribed here — a non-zero random pool, non-identity indices, a two-way "auto"-versus-"cute-dsl" comparison, state-pool equality, and bit-exact untouched-slot checks — and misses the bug only because every parametrization it runs is Cake-servable. The gap is negative contracts, not harness.
• For every input where "cute-dsl" succeeds, "auto" must succeed and produce numerically equal output and an equal mutated state pool. Non-identity ssm_state_indices are the load-bearing requirement: a state-pool assertion catches the corruption even with a zero-initialised pool, whereas identity indices hide it regardless of pool contents. A non-zero pool is belt-and-braces — it is what makes the output differ too.
• Exercise both sides of the B*HV = 128 one-warp threshold. The two CuTe routes have different state conventions, and a fix that is correct on one can produce NaN on the other.
• Cover at least one supported and one unsupported architecture, plus each trigger class above.
• Because the accepted sets are not nested (below), do not assume "auto" ⊇ "cute-dsl" as an invariant to test in one direction only.

Non-comparability

"auto" is not merely narrower than "cute-dsl"; the accepted input sets are incomparable. The non-contiguous-initial_state guard is skipped for Cake and for the "auto" gate (:2160-2165, note and not auto_unbounded_softplus_candidate), so a non-contiguous pool without cu_seqlens skips a check under "auto" that "cute-dsl" enforces. The witness has to be chosen with care, because the guard skip only helps if the selector then accepts the layout: a leading-dim-strided pool (inner dims still contiguous, e.g. a slice of a larger pool) is accepted under "auto" and raises ValueError: non-contiguous initial_state requires cu_seqlens under "cute-dsl". A transposed pool is not a witness: it skips the guard under "auto" but is then rejected on strides by the selector, so both backends raise.

Note this single carve-out is what makes the sets incomparable. It is not, however, unsound in the way the guard's message suggests: the guard protects the path that calls .contiguous() on the pool (:2195), and the Cake and "auto" arms (:2187, :2193) pass initial_state straight through, so a fall-through writes into the caller's tensor and the update is not lost — measured against a contiguous-pool "cute-dsl" reference to within bf16 rounding, with the interleaved rows outside the view untouched. What the carve-out does cost is a layout the selector may still reject, which then reaches CuTe un-normalized. The right disposition is to decide whether the guard should have been skipped at all — not to treat the acceptance as a compatibility constraint on unreleased surface.

Incorrect documentation

Both describe a fallback that does not exist, and should be corrected with the fix:

• flashinfer/kda_decode.py docstring: "auto" "selects Cake only for its equal-head/D128/T1 unbounded-softplus contract, preserving CuTe DSL for every other decode surface."
• docs/api/kda_decode.rst: "backend=\"auto\" selects Cake only for the equal-head D128 T1 unbounded-softplus contract and preserves CuTe-DSL for other decode modes."

Blast radius

• Scope is source builds of main only. No stable release contains "auto", and no tagged nightly does either: nightly tagging stops at nightly-v0.6.18-20260819, whose tree predates feat(kda): add CuTe DSL recurrent prefill backend #4605 and has no "auto" in kda.py. Severity for users on the current stable is therefore nil. Severity for anyone tracking main, and for the next release, is high: this is the default path, and the same source that worked against v0.6.18 — a defaulted top-level call, which then reached run_recurrent_kda's "cute-dsl" default — now raises on these inputs.
• flashinfer.recurrent_kda is the documented top-level entry point and defaults to "auto", so anyone building from main hits it without opting in.
• flashinfer.kda_decode.recurrent_kda defaults to "cute-dsl" (kda_decode.py:79) and is unaffected, so the two public spellings of one operation disagree — see the API audit in #4936.
• Trigger class 1 fails the production Kimi-Linear T=1 shape on every device outside CC 10.0/10.3 — including CC 12.0 and 12.1, which are Blackwell but not datacentre Blackwell; class 3 fails the paged-serving layout everywhere.
• Why CI misses it: not for lack of coverage — for lack of a negative test. Instrumenting the "auto" gate across tests/kda/ shows it matches on 6 of 190 run_recurrent_kda calls, all six from test_recurrent_kda_decode_export.py:1685's test_t1_unbounded_softplus_auto_route_* and all six with an explicit backend="auto". That test asserts frozen_calls == [expected_variant] — it pins Cake selection, so by construction it cannot detect a missing fallback. Meanwhile the three defaulted top-level call sites in test_packed_kda_decode_cute.py (eleven invocations) all miss the gate: two on lower_bound=-5.0 (:728, :837) and the third (:902) on use_gate_in_kernel, which kda.py:66 defaults to False — it passes no lower_bound at all. Five of the eleven never reach the gate site at all, exiting at the T=1 fast path (:1933-1964) first. So the gate is never entered through a public default, and the one test that does enter it only checks the happy path.

Related

• #4936 — KDA public API unification audit (documents the resulting facade divergence); supersedes #4483
• feat(cake_kda): add native unbounded-softplus Kimi-Linear kernels #4535 — introduced "auto" decode selection
• [Bug][v0.6.17] tests.trace.test_template_registry::test_registration_module_inventory_is_complete AssertionError: assert {'flashinfer....cat_ops', ...} == {'flashinfer....cat_ops', ...} #4367 — trace-registry inventory failure from the same duplicate-facade split; already fixed on main and should be closed