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, state_checkpoints: Tensor | None = None, checkpoint_cu_starts: Tensor | None = None, checkpoint_every_n_tokens: int = 0, *, backend: Literal['auto', 'cute-dsl', 'cake'] = 'auto') → tuple[Tensor, Tensor | None] | tuple[Tensor, Tensor | None, Tensor]

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 SM120a, eligible ordinary multi-token prefill uses the architecture-specific CuTe DSL backend. On SM100a (B200/GB200) and SM103a (B300/GB300), the FlashKDA-compatible subset can use either the frozen Cake schedules or the source-level CuTe DSL BT=16 kernel. backend="auto" prefers CuTe DSL for supported plain prefill contracts and keeps Cake as the feature-complete fallback.

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. Eligible frozen prefill accepts non-overlapping token-row-strided storage with a unit head stride, including a view into a fused projection.

  • 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. For eligible frozen prefill with ssm_state_indices, this is a state pool [N_pool, H, 128, 128] whose inner slots are contiguous; padding between pool slots is allowed.

  • 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 non-decreasing, and end at the total token count. This value contract is not normally host-validated. Eager calls without an explicit workspace or seq_order read these values once per unchanged offsets tensor to schedule longer sequences first on Cake. Eager packed CuTe DSL engine calls also cache a longest-sequence-first order; CuTe DSL decomp calls retain the original order because their CTA grid fits in one wave. Eligible 148-SM B200 and 152-SM GB200 Cake calls additionally cache persistent worker task bins.

  • 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). Eligible frozen packed prefill accepts contiguous CUDA int32 [N_seq] indices and updates the selected initial_state pool slots directly.

  • 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]. For eager CuTe DSL packed engine calls, omitting it builds and caches a longest-sequence-first order; CuTe DSL decomp keeps the original order because its CTA grid fits in one wave. CUDA Graph capture of a packed CuTe DSL engine call requires an explicit plan prepared with RecurrentKDAPrefillWrapper. Cake constructs and caches its own eager host metadata. On Cake, supplying an order keeps the direct schedule so caller-owned ordering is not replaced by persistent task bins. Fixed-layout prefill and decode calls must leave it as None.

  • prefill_workspace (Optional[RecurrentKDAPrefillWorkspace]) – Caller-owned workspace for SM100-family and SM120 prefill backends. 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. Explicit workspaces and CUDA Graph capture use direct/M64 schedules; persistent task planning is an eager-only B200/GB200 route because its bins depend on host-visible sequence lengths.

  • state_checkpoints (Optional[torch.Tensor]) – Caller-owned BF16 checkpoint output [C, H, 128, 128] for frozen prefill. Row zero for each sequence is its initial state; later rows are the states before token blocks beginning at N, 2N, .... Required when checkpoint_every_n_tokens > 0.

  • checkpoint_cu_starts (Optional[torch.Tensor]) – Contiguous CUDA int64 cumulative checkpoint counts [N_seq+1]. Each count must equal ceil(seq_len / checkpoint_every_n_tokens).

  • checkpoint_every_n_tokens (int) – Checkpoint interval. Zero disables checkpoints; a positive value must be divisible by 32, except that the SM100-family exact-N16 frozen route also accepts multiples of 16. SGLang normally uses 64 or a larger cache-page-aligned multiple.

  • backend (Literal["auto", "cute-dsl", "cake"]) – Implementation backend. "auto" selects the architecture- appropriate CuTe DSL kernel for supported ordinary multi-token prefill, including the SM120 backend and SM100-family state checkpoints, and otherwise falls back to an exported frozen Cake specialization. "cake" and "cute-dsl" select those backends strictly.

Returns:

Tuple of (output, final_state) where final_state is None when output_final_state=False. When checkpointing is enabled, a triple (output, final_state, state_checkpoints) is returned. See flashinfer.kda_kernels.recurrent_kda.run_recurrent_kda() for the backend implementation.