vision_rs::models::yolo::kernels::attention

Module flash_attn2

Expand description

Flash Attention 2 — forward and backward kernels.

Layout: all tensors are stored as [BATCH * N_HEADS, N_CTX, HEAD_DIM] row-major (contiguous). The caller is responsible for reshaping [B, H, N, D] PyTorch tensors to this flat 3-D layout before calling.

Algorithm: each CTA processes one (batch, head, q_row) triple. The kernel iterates over all N_CTX_K key/value rows with the online softmax recurrence (Flash Attention paper, Dao et al. 2022/2023), so the full N_CTX_Q × N_CTX_K attention matrix is never materialised. Memory is O(N_CTX × HEAD_DIM) per CTA rather than O(N_CTX²).

HEAD_DIM must be a power of two and is a compile-time const so the inner vector loads are always fully unmasked.

Structs§

FlashAttention2BackwardDkv: Flash Attention 2 backward: computes dK and dV for one key row.
FlashAttention2BackwardDq: Flash Attention 2 backward: computes dQ.
FlashAttention2Forward: Flash Attention 2 forward pass.
FlashAttention2Op: Bundles the forward and backward Flash Attention 2 kernels for a given head dimension.

Functions§

flash_attention2_backward_dkv: Flash Attention 2 backward: computes dK and dV for one key row.
flash_attention2_backward_dq: Flash Attention 2 backward: computes dQ.
flash_attention2_forward: Flash Attention 2 forward pass.