flashinfer.kda_training.recurrent_kda_training_forward¶
- flashinfer.kda_training.recurrent_kda_training_forward(q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor, A_log: Tensor, dt_bias: Tensor, initial_state: Tensor, cu_seqlens: Tensor | None = None, scale: float | None = None, lower_bound: float = -5.0, out: Tensor | None = None, final_state_out: Tensor | None = None, context_out: RecurrentKDATrainingContext | None = None) tuple[Tensor, Tensor, RecurrentKDATrainingContext]¶
Run fixed or packed KDA forward and save the selected route’s tapes.