Module flash_attention 

Flash Attention - Memory-Efficient Attention Mechanism

Implements FlashAttention-2 algorithm for O(N) memory instead of O(N²). Uses tiled matrix multiplication with online softmax normalization.

Structs

FlashAttention

Flash Attention layer with tiled computation and online softmax.

FlashAttentionGrads

Gradients from Flash Attention backward pass.

GroupedQueryAttention

Grouped Query Attention (GQA) - Efficient attention with fewer KV heads

KVCache

KV-Cache for efficient autoregressive generation

MultiQueryAttention

Multi-Query Attention (MQA) - Extreme efficiency with single KV head

SlidingWindowAttention

Sliding Window Attention for efficient long-context processing

Enums

AttentionMask

Attention mask for custom attention patterns