flashinfer.kda_backward.recurrent_kda_backward

flashinfer.kda_backward.recurrent_kda_backward(q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor, A_log: Tensor, dt_bias: Tensor, initial_state: Tensor, do: Tensor, dfinal_state: Tensor, cu_seqlens: Tensor | None = None, scale: float | None = None, lower_bound: float = -5.0, workspace: RecurrentKDABackwardWorkspace | None = None, out: Sequence[Tensor] | None = None) → tuple[Tensor, ...]

Compute all gradients of the recurrent KDA training recurrence.

The kernel differentiates BF16 Q/K/V/raw-gate/raw-beta token inputs and FP32 parameters and recurrent state. Q and K are L2-normalized with epsilon 1e-6; decay is exp(lower_bound * sigmoid(exp(A_log) * (g + dt_bias))) and beta is passed through a sigmoid. The loss adjoints are do for token output and dfinal_state for final recurrent state.

This frozen implementation is specialized for SM100a/SM103a, head/key/value dimension 128, FP32 state, and the eight shapes listed in flashinfer.kda_backward. It returns (dq, dk, dv, dg, dbeta, dA_log, ddt_bias, dinitial_state) in that order.

Parameters:

  • q – Contiguous BF16 [1,T,H,128] query tensor.

  • k – Contiguous BF16 [1,T,H,128] key tensor.

  • v – Contiguous BF16 [1,T,H,128] value tensor.

  • g – Contiguous BF16 [1,T,H,128] raw gate tensor.

  • beta – Contiguous BF16 [1,T,H] raw beta logits.

  • A_log – Contiguous FP32 [H] decay parameter.

  • dt_bias – Contiguous FP32 [H,128] decay bias.

  • initial_state – Contiguous FP32 value-first [N,H,V,K] state, with V=K=128.

  • do – Contiguous BF16 token-output adjoint matching q.

  • dfinal_state – Contiguous FP32 value-first final-state adjoint matching initial_state.

  • cu_seqlens – Exact contiguous CUDA int64 packed offsets for one of the documented packed shapes, otherwise None.

  • scale – Fixed output scale 1 / sqrt(128).

  • lower_bound – Fixed safe-gate lower bound -5.0.

  • workspace – Reusable scratch. Required, eagerly warmed, for CUDA Graph capture.

  • out – Eight preallocated gradient tensors in return order. Required for CUDA Graph capture.