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, *, cu_seqlens_cpu: Tensor | None = None) → tuple[Tensor, Tensor, RecurrentKDATrainingContext]

Run fixed or packed KDA forward and save the selected route’s tapes.

Packed calls require cu_seqlens_cpu, a trusted CPU mirror of the CUDA cu_seqlens tensor used for host-side route and metadata planning.

Parameters:

  • q (torch.Tensor) – Contiguous BF16 queries with shape [B, T, Hqk, 128]. Packed calls use a physical batch size of one.

  • k (torch.Tensor) – Contiguous BF16 keys with the same shape as q.

  • v (torch.Tensor) – Contiguous BF16 values with shape [B, T, Hv, 128], where Hv is an integer multiple of Hqk.

  • g (torch.Tensor) – Contiguous BF16 raw gate values with the same shape as v.

  • beta (torch.Tensor) – Contiguous BF16 raw beta values with shape [B, T, Hv].

  • A_log (torch.Tensor) – Contiguous FP32 state-transition parameters with shape [Hv].

  • dt_bias (torch.Tensor) – Contiguous FP32 step-size biases with shape [Hv, 128].

  • initial_state (torch.Tensor) – Contiguous FP32 initial states with shape [num_sequences, Hv, 128, 128].

  • cu_seqlens (torch.Tensor, optional) – Contiguous CUDA int64 cumulative sequence lengths for packed input. Omit for fixed-length input.

  • scale (float, optional) – Attention scale. The training kernel currently requires 1 / sqrt(128); None selects that value.

  • lower_bound (float) – Gate lower bound. The training kernel currently requires -5.0.

  • out (torch.Tensor, optional) – Caller-owned contiguous BF16 output buffer with the same shape as v.

  • final_state_out (torch.Tensor, optional) – Caller-owned contiguous FP32 final-state buffer with the same shape as initial_state.

  • context_out (RecurrentKDATrainingContext, optional) – Same-shape context to reuse on the CUDA stream that created it.

  • cu_seqlens_cpu (torch.Tensor, optional) – Trusted contiguous CPU int64 mirror of cu_seqlens. Required for packed input and must be omitted for fixed-length input.

Returns:

Token output, final recurrent state, and the saved route context required by recurrent_kda_training_backward().

Return type:

tuple[torch.Tensor, torch.Tensor, RecurrentKDATrainingContext]