Struct FlashAttention 

pub struct FlashAttention { /* private fields */ }

Flash Attention: Memory-efficient scaled dot-product attention.

This struct provides the flash attention computation without maintaining learned parameters. It is designed to be used within attention modules that handle Q, K, V projections separately.

Implementations

impl FlashAttention

pub fn new(dim_head: usize, config: FlashAttentionConfig) -> Self

Create a new Flash Attention module.

Arguments

• dim_head - Dimension of each attention head (for scaling)
• config - Flash Attention configuration

pub fn with_dim_head(dim_head: usize) -> Self

Create with default configuration.

pub fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor>

Compute flash attention.

Arguments

• q - Query tensor of shape (batch, heads, seq_q, dim_head)
• k - Key tensor of shape (batch, heads, seq_k, dim_head)
• v - Value tensor of shape (batch, heads, seq_k, dim_head)

Returns

Output tensor of shape (batch, heads, seq_q, dim_head)

Trait Implementations

impl Clone for FlashAttention

fn clone(&self) -> FlashAttention

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for FlashAttention

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.