Function flash_attention

pub fn flash_attention(
    q: &Tensor,
    k: &Tensor,
    v: &Tensor,
    opts: impl FlashAttentionOpts,
) -> Tensor