flashinfer.kda_training.recurrent_kda_training_backward

flashinfer.kda_training.recurrent_kda_training_backward(context: RecurrentKDATrainingContext, do: Tensor, dfinal_state: Tensor, out: Sequence[Tensor] | None = None) → tuple[Tensor, ...]

Differentiate a saved route context without rerunning forward recurrence.

Parameters:

  • context (RecurrentKDATrainingContext) – Context returned by recurrent_kda_training_forward(). It must be used on the original forward stream and must not have been modified.

  • do (torch.Tensor) – Contiguous BF16 gradient of the token output, with the same shape as the forward v tensor.

  • dfinal_state (torch.Tensor) – Contiguous FP32 gradient of the final recurrent state, with the same shape as the forward initial_state tensor.

  • out (Sequence[torch.Tensor], optional) – Eight caller-owned gradient buffers for q, k, v, g, beta, A_log, dt_bias, and initial_state, in that order.

Returns:

The eight gradients in the order described by out. Q, K, V, G, and beta gradients are BF16; parameter and initial-state gradients are FP32.

Return type:

tuple[torch.Tensor, …]