# Flash Attention — Fast Memory-Efficient Exact Attention

> A CUDA kernel library that computes exact attention 2-4x faster than standard implementations while using 5-20x less GPU memory, enabling training of longer sequences on existing hardware.

## Install

Save in your project root:

# Flash Attention — Fast Memory-Efficient Exact Attention

## Quick Use
```bash
pip install flash-attn --no-build-isolation
```
```python
from flash_attn import flash_attn_func

# q, k, v: (batch, seqlen, nheads, headdim)
output = flash_attn_func(q, k, v, causal=True)
```

## Introduction
Flash Attention is an IO-aware exact attention algorithm that restructures the computation to minimize memory reads and writes between GPU HBM and SRAM. By tiling the attention computation and fusing operations into a single kernel, it achieves significant speedups and memory savings without any approximation of the attention output.

## What Flash Attention Does
- Computes exact multi-head attention 2-4x faster than standard PyTorch implementations
- Reduces memory usage from O(N^2) to O(N) in sequence length, enabling much longer contexts
- Provides fused kernels for forward and backward passes with automatic gradient computation
- Supports causal masking, variable-length sequences, and grouped-query attention (GQA/MQA)
- Integrates with PyTorch's autograd for seamless use in any transformer model

## Architecture Overview
Flash Attention uses a tiling strategy that loads blocks of Q, K, V from slow GPU HBM into fast on-chip SRAM, computes partial attention scores locally, and accumulates results using the online softmax trick. This avoids materializing the full N x N attention matrix in HBM. The backward pass recomputes attention from saved O and logsumexp statistics rather than storing the attention matrix, trading minimal extra compute for massive memory savings.

## Self-Hosting & Configuration
- Requires NVIDIA GPU with compute capability 8.0+ (A100, H100, RTX 3090/4090)
- Install via pip with CUDA toolkit matching your PyTorch installation
- Drop-in replacement for standard attention by swapping the attention function call
- Configure block sizes automatically tuned for your GPU architecture
- Supports both FP16 and BF16 data types for mixed-precision training

## Key Features
- Exact attention output (not an approximation) with IO-optimized memory access patterns
- Kernel fusion eliminates intermediate memory allocation for the attention matrix
- Native support for causal, sliding window, and custom attention masks
- Flash Attention 2 further improves parallelism over sequence length
- Used by default in Hugging Face Transformers, PyTorch 2.0 SDPA, and most LLM training stacks

## Comparison with Similar Tools
- **PyTorch SDPA** — wraps Flash Attention as one backend; direct usage gives more control over options
- **xFormers** — provides memory-efficient attention with approximations; Flash Attention is exact
- **FlexAttention (PyTorch)** — compiler-based approach for custom masks; Flash Attention uses hand-tuned CUDA kernels
- **Triton-based attention** — portable across GPUs; Flash Attention's CUDA kernels are faster on NVIDIA hardware
- **Ring Attention** — distributes attention across devices; Flash Attention optimizes single-device computation

## FAQ
**Q: Does Flash Attention produce the same results as standard attention?**
A: Yes. It computes mathematically exact attention with bit-level equivalence (modulo floating-point reordering).

**Q: What GPUs are supported?**
A: NVIDIA Ampere (A100, RTX 3090) and Hopper (H100) architectures. Older GPUs like V100 are not supported.

**Q: How much memory does it save?**
A: Memory scales linearly with sequence length instead of quadratically. For a 4K sequence, typical savings are 10-20x on the attention computation.

**Q: Can I use it with Hugging Face models?**
A: Yes. Most Hugging Face models automatically use Flash Attention 2 when installed and enabled via model config.

## Sources
- https://github.com/Dao-AILab/flash-attention
- https://arxiv.org/abs/2205.14135

---
Source: https://tokrepo.com/en/workflows/asset-7d1f0b81
Author: AI Open Source