Flash Linear Attention
An open-source Python library for building sequence models on your hardware, with linear attention layers and Hugging Face text generation support.
Flash Linear Attention is a Python library for researchers and developers building language models on their own hardware. It supplies reusable components for exploring linear attention, sparse attention and state space architectures.
- Models and layers: Use implementations of RetNet, GLA, DeltaNet and RWKV6, or replace standard multihead attention in your own model with individual linear attention layers.
- Hybrid architectures: Combine a model's native linear, recurrent or state space layers with standard attention, using local or full attention at different depths.
- Text generation: Use pretrained models from fla-hub and generate text through Hugging Face's text generation APIs after pretraining.
- Training components: Fused positional embedding and normalization modules target faster training, while the related flame framework builds on torchtitan to train fla models.
- Evaluation: Compare models through lm-evaluation-harness benchmarks and perplexity evaluation, with support for single and multiple GPUs.
The library is open source under the MIT license. It supports CUDA, ROCm, XPU, NPU and CPU backends, and its implementations have been verified on NVIDIA, AMD and Intel hardware. It's intended for model development in Python, so using it involves working with model code.