flash-attn

What it is and what it does

FlashAttention is a GPU-accelerated library that implements fast, memory-efficient exact attention mechanisms for transformer models. It provides optimized kernels for scaled dot-product attention that reduce memory I/O and improve computational efficiency on modern GPUs. The package includes FlashAttention-2 (optimized for Ampere, Ada, and Hopper architectures) and a beta FlashAttention-3 for Hopper GPUs, with support for both NVIDIA CUDA and AMD ROCm backends.

The library integrates with PyTorch and einops, exposing functions like flash_attn_func and flash_attn_qkvpacked_func that drop in as replacements for standard attention. It supports various attention patterns including causal masking, sliding window attention, multi-query and grouped query attention, dropout, rotary embeddings, and ALiBi. Installation requires compilation against your GPU toolkit, which adds setup complexity but is necessary to generate hardware-specific optimized code.

Use it for:

• Accelerate training and inference of large language models and transformers on NVIDIA A100, RTX 4090, H100, or AMD MI200/MI300 GPUs
• Reduce memory consumption during attention computation, enabling longer sequences or larger batch sizes on the same hardware
• Implement causal attention for autoregressive language model generation with improved efficiency
• Deploy transformer models with lower latency by replacing standard PyTorch attention with optimized FlashAttention kernels

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Provides optimized GPU implementations of scaled dot-product attention (FlashAttention and FlashAttention-2) that reduce memory usage and improve speed compared to standard attention, with support for NVIDIA CUDA and AMD ROCm.

Yes, if you have a supported GPU (Ampere, Ada, Hopper for CUDA; MI200/MI300 for ROCm) and are training or deploying transformer models where attention is a bottleneck. The high install friction (compilation, GPU toolkit requirements) is justified by significant speed and memory gains. Not worth installing for CPU-only workflows or unsupported GPU architectures (e.g., older Turing GPUs without FlashAttention 1.x).

Install

flash-attn on PyPI

pip install flash-attn

uv add flash-attn

poetry add flash-attn

Installing flash-attn

Before you install

High install friction: requires CUDA 12.0+ or ROCm 6.0+, PyTorch 2.2+, ninja build tool, and compilation from source. Compilation takes 3–5 minutes on well-equipped machines but can exhaust RAM on systems with many CPU cores and <96GB memory; MAX_JOBS environment variable can mitigate this. Package is actively maintained with recent releases.

License in practice

BSD License (permissive): you may use, modify, and distribute the package freely, including in commercial projects, provided you retain the license notice and credit the authors as requested in the repository.

Quickstart

pip install flash-attn --no-build-isolation

from flash_attn import flash_attn_func
# flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)

Requires CUDA 12.0+ or ROCm 6.0+, PyTorch 2.2+, ninja build tool, and Linux (Windows support experimental as of v2.3.2). Compilation requires a GPU and appropriate toolkit installed.

Verify before relying

• Exact performance gains over standard attention on specific hardware and sequence lengths
• Compatibility with all transformer architectures and attention variants beyond those documented
• Windows support stability and completeness as of current version

Package facts

Evidence: flash_attn-2.8.3.post1.tar.gz

Tags

More Artificial Intelligence packages

Further reading