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) → tuple[Tensor, Tensor | None]

Recurrent KDA (Key-Driven Attention) decode 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.

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.

  • 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 instead of softplus. 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].

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

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.