flashinfer.kda.recurrent_kda

flashinfer.kda.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, seq_order: Tensor | None = None, prefill_workspace: RecurrentKDAPrefillWorkspace | None = None) → tuple[Tensor, Tensor | None]

Recurrent KDA (Kimi Delta Attention) decode and prefill kernel.

This is the public API layer for the CuTe DSL implementation 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 backend implementation. On SM100a (B200/GB200) and SM103a (B300/GB300), the FlashKDA-compatible subset of ordinary multi-token prefill is dispatched to the corresponding frozen SM100-family kernel. All existing decode and speculative-decode calls retain the CuTe DSL backend.

Parameters:

  • q (torch.Tensor) – Query of shape [B, T, H, K], or [1, total_tokens, H, K] when using cu_seqlens. Must be bfloat16. T=1 selects decode; eligible T>1 calls may select the frozen prefill backend.

  • k (torch.Tensor) – Key with the same shape as q. Must be bfloat16.

  • v (torch.Tensor) – Value of shape [B, T, HV, V], or [1, total_tokens, HV, V] when packed. Must be bfloat16. GQA is applied when HV != H.

  • g (torch.Tensor) – Per-K-dimension gate of shape [B, T, HV, K], or [1, total_tokens, HV, K] when packed. 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, T, HV], or [1, total_tokens, HV] when packed. 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] or [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 instead of softplus. Must be negative.

  • cu_seqlens (Optional[torch.Tensor]) – Contiguous CUDA cumulative sequence lengths of shape [N+1]. May be int32 or int64. Frozen prefill converts int32 offsets to int64 outside graph capture; graph capture requires caller-provided int64 offsets. For frozen prefill, values must start at zero, be strictly increasing, and end at the total token count. This value contract is not host-validated to avoid a device synchronization.

  • 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, T, HV, V] for fixed layout, or the corresponding packed/speculative shape when using cu_seqlens. If None, a new tensor is allocated. Frozen prefill requires storage disjoint from Q, K, V, G, beta, and initial_state.

  • 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.

  • seq_order (Optional[torch.Tensor]) – Optional packed-prefill sequence order, as a contiguous CUDA int32 permutation of shape [N]. Sorting by descending sequence length improves tail utilization. It is only consumed by the frozen FlashKDA prefill backend; prepare it before CUDA graph capture or timed launches. Fixed-layout prefill and decode calls must leave it as None.

  • prefill_workspace (Optional[RecurrentKDAPrefillWorkspace]) – Caller-owned workspace for the frozen SM100-family prefill backend. It is optional for eager execution and required for CUDA graph capture. Warm it eagerly with the exact tensors on the capture stream before capture. Use one workspace per captured recurrent_kda invocation.

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.