flashinfer.kda_decode.recurrent_kda

flashinfer.kda_decode.recurrent_kda(q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor, A_log: Tensor | None = None, dt_bias: Tensor | None = None, scale: float | None = None, initial_state: Tensor | None = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = True, use_gate_in_kernel: bool = False, lower_bound: float | None = None, cu_seqlens: Tensor | None = None, ssm_state_indices: Tensor | None = None, num_spec_tokens: int | None = None, num_accepted_tokens: Tensor | None = None, output: Tensor | None = None, initial_state_source: Tensor | None = None, initial_state_indices: Tensor | None = None, beta_is_logit: bool = False, *, disable_state_update: bool = False, correction_cache: Tensor | None = None, kg_cache: Tensor | None = None, backend: Literal['cute-dsl', 'cake', 'auto'] = 'cute-dsl') → tuple[Tensor, Tensor | None]

Recurrent KDA (Kimi Delta Attention) decode kernel.

This public API supports the existing CuTe DSL implementation and an explicit exported Cake backend in flashinfer.kda_kernels.recurrent_kda. It supports single-token decode, fused speculative decode, GQA, optional cu_seqlens packing, and the same gate modes as the selected backend implementation.

Parameters:

  • q (torch.Tensor) – Current query of shape [B, 1, H, K], or [1, total_tokens, H, K] when using cu_seqlens. Must be bfloat16.

  • k (torch.Tensor) – Current key of shape [B, 1, H, K]. Must be bfloat16.

  • v (torch.Tensor) – Current value of shape [B, 1, HV, V]. Must be bfloat16. GQA is applied when HV != H.

  • g (torch.Tensor) – Per-K-dimension gate of shape [B, 1, HV, K]. Must be bfloat16. Log-space if pre-computed, raw input if use_gate_in_kernel=True.

  • beta (torch.Tensor) – Delta-rule learning rate of shape [B, 1, HV]. Must be bfloat16. Pre-sigmoided unless beta_is_logit=True.

  • A_log (Optional[torch.Tensor]) – Log decay parameter of shape [H]. Must be float32. Required when use_gate_in_kernel=True.

  • dt_bias (Optional[torch.Tensor]) – Per-head-K decay bias of shape [H*K]. Must be float32.

  • scale (Optional[float]) – Scale factor for queries. If None, defaults to 1 / sqrt(K).

  • initial_state (Optional[torch.Tensor]) – Initial state of shape [N, HV, V, K]. Must be bfloat16. If None, zero-initialized. Updated in-place. For batched spec decode without cu_seqlens, N is the packed checkpoint-slot count B * (1 + num_spec_tokens) when ssm_state_indices is omitted.

  • output_final_state (bool) – Whether to return the final state. Default: False.

  • use_qk_l2norm_in_kernel (bool) – Whether to apply L2 normalization to Q and K. Default: True.

  • use_gate_in_kernel (bool) – Whether to compute the gate inside the kernel from A_log and g. Default: False.

  • lower_bound (Optional[float]) – If set, uses lower_bound * sigmoid(exp(A_log) * (g + dt_bias)) gate formula. If None, uses -exp(A_log) * softplus(g + dt_bias). A supplied bound must be negative.

  • cu_seqlens (Optional[torch.Tensor]) – Cumulative sequence lengths of shape [N+1]. Must be int32.

  • ssm_state_indices (Optional[torch.Tensor]) – State cache indices. Shape [N] int32 for standard decode, or [N, 1+S] int32 for spec decode (num_spec_tokens must also be set).

  • num_spec_tokens (Optional[int]) – Number of speculative tokens (S). When set, processes 1+S tokens in a single fused kernel launch. Must be >= 1.

  • num_accepted_tokens (Optional[torch.Tensor]) – Per-sequence accepted token count from the previous spec decode round. Shape [N] int32. If None, initial state is loaded from ssm_state_indices[n, 0]. Values above 1+S are clamped to the final checkpoint slot.

  • output (Optional[torch.Tensor]) – Pre-allocated output tensor. Shape [B, 1, HV, V] for standard decode, [1, N*(1+S), HV, V] for spec decode with cu_seqlens. If None, a new tensor is allocated.

  • initial_state_source (Optional[torch.Tensor]) – Optional read-only committed state pool [N0, HV, V, K]. When provided, token 0 is loaded from this pool instead of initial_state.

  • initial_state_indices (Optional[torch.Tensor]) – Source slot per sequence, shape [N] int32. Required together with initial_state_source.

  • beta_is_logit (bool) – If True, apply sigmoid to beta inside the recurrent kernel.

  • disable_state_update (bool) – Frozen / speculative-verify mode (mirrors GDN’s gated_delta_rule_mtp flag): compute the outputs for up to 16 tokens per sequence from the committed state and never write any state back; final_state is always None. Dispatches internally to the WY-parallel tensor-core kernel or a grouped register recurrence by problem size. Supports the batched [B, T, ...] form directly and the packed cu_seqlens form (ragged per-sequence lengths). backend="cake" raises in this mode (no frozen-state Cake kernels; no silent fallback), and output_final_state=True is rejected. Requires K == V == 128 and a bf16 state pool.

  • correction_cache (Optional[torch.Tensor]) – Only with disable_state_update=True. Slot-indexed float32 buffer [num_slots, HV, T_max, V] receiving the per-token delta-rule corrections sigmoid-or-raw(beta) * (v - u) for a downstream commit/recovery kernel (the analog of GDN’s slot-indexed intermediate_states_buffer). Rows past each sequence’s length and null slots are left untouched.

  • kg_cache (Optional[torch.Tensor]) – Only with disable_state_update=True. Slot-indexed bf16 buffer [num_slots, HV, T_max, 2*K] receiving the raw (unnormalized) key in [..., :K] and the raw gate in [..., K:] per token, matching the vLLM RecoverSSM cache convention.

  • backend (Literal["cute-dsl", "cake", "auto"]) – Implementation backend. "cute-dsl" preserves the existing FlashInfer implementation. "cake" strictly selects an exported Cake kernel and raises when the call does not match one of its supported contracts. "auto" selects Cake only for its equal-head/D128/T1 unbounded-softplus contract, preserving CuTe DSL for every other decode surface. Default: "cute-dsl".

Returns:

Tuple of (output, final_state) where final_state is None when output_final_state=False. See flashinfer.kda_kernels.recurrent_kda.run_recurrent_kda() for the backend implementation.