Struct FlashAttention 

pub struct FlashAttention { /* private fields */ }

Flash attention with block-wise computation

Computes attention in tiles to minimize memory usage while maintaining numerical stability.

Implementations

impl FlashAttention

pub fn new(dim: usize, block_size: usize) -> Self

Create new flash attention

pub fn causal(dim: usize, block_size: usize) -> Self

Create with causal masking

Trait Implementations

impl Attention for FlashAttention

fn compute( &self, query: &[f32], keys: &[&[f32]], values: &[&[f32]], ) -> AttentionResult<Vec<f32>>

Computes attention over the given query, keys, and values.

fn compute_with_mask( &self, query: &[f32], keys: &[&[f32]], values: &[&[f32]], mask: Option<&[bool]>, ) -> AttentionResult<Vec<f32>>

Computes attention with optional mask.

fn dim(&self) -> usize

Returns the model dimension.

fn num_heads(&self) -> usize

Returns the number of attention heads (1 for single-head attention).