Reproduce “Faster Causal Attention Over Large Sequences Through Sparse Flash Attention” on OpenResearch

About this paper

Matteo Pagliardini, Daniele Paliotta, Martin Jaggi, François Fleuret

Researchers from EPFL and the University of Geneva developed Sparse Causal Flash Attention (SCFA), a new GPU kernel that enables efficient and exact implementation of dynamic sparse attention patterns for Transformers. The method achieves multi-fold training speedups for models processing long sequences (e.g., 1.8x-3.5x for 8k-16k tokens) without compromising model performance, by extending the I/O-aware principles of FlashAttention.

Key insights

• FlashAttention's memory-efficient tiling and block-wise computation can be extended to handle dynamic and irregular causal sparsity by incorporating original sequence indices to determine local causal masks.
• Dynamically reordering and compacting queries and keys based on sparsity patterns (e.g., hash buckets or selected elements) enables SCFA to perform efficient attention calculations only on relevant sub-blocks.
• The initial linear overhead for preprocessing (sorting/compacting) is quickly outweighed by the quadratic savings in attention computation for longer sequences, making dynamic sparsity practically viable.

Summary by alphaXiv.