flash-attention

name	flash-attention
description	Enables fused transformer attention (PyTorch scaled_dot_product_attention, the flash-attn library, optional H100 FP8 and sliding windows) so long sequences use less GPU memory and less wall time. Use for training or inference past about 512 tokens, attention OOMs, or Ampere-or-newer SDPA backends. Not for CNN or UNet blocks, CPU-only training, or float32 attention. Never assume Volta V100 kernels exist.
version	1.0.1
author	Orchestra Research
license	MIT
dependencies	["flash-attn","torch","transformers"]
platforms	["linux","macos","windows"]
metadata	{"hermes":{"tags":["Optimization","Flash Attention","Attention Optimization","Memory Efficiency","Speed Optimization","Long Context","PyTorch","SDPA","H100","FP8","Transformers"]}}

Overview

Flash Attention provides 2-4x speedup and 10-20x memory reduction for transformer attention through IO-aware tiling and recomputation. It supports PyTorch native SDPA, the flash-attn library, H100 FP8, and sliding window attention.

When to Use

• Training transformers with sequences >512 tokens.
• Running inference with long context (>2K tokens).
• GPU memory constrained (OOM with standard attention).
• Need 2-4x speedup without accuracy loss.
• Using PyTorch 2.2+ or can install flash-attn.

Prerequisites

• GPU: NVIDIA Ampere+ (A100, A10, A30) or AMD MI200+. Turing (T4) is supported. Volta (V100) is NOT supported.
• VRAM: Same as standard attention (Flash Attention doesn't increase memory).
• CUDA: 12.0+ (11.8 minimum).
• PyTorch: 2.2+ for native SDPA support.

Procedure

Workflow 1: Enable in existing PyTorch model (Native SDPA)

1. Check PyTorch version (≥2.2):

python -c "import torch; print(torch.__version__)"

If <2.2, upgrade:

pip install --upgrade torch

2. Enable Flash Attention backend: Replace standard attention:

# Before (standard attention)
attn_weights = torch.softmax(q @ k.transpose(-2, -1) / math.sqrt(d_k), dim=-1)
out = attn_weights @ v

# After (Flash Attention)
import torch.nn.functional as F
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)

Force Flash Attention backend:

with torch.backends.cuda.sdp_kernel(
    enable_flash=True,
    enable_math=False,
    enable_mem_efficient=False
):
    out = F.scaled_dot_product_attention(q, k, v)

Workflow 2: Use flash-attn library for advanced features

1. Install flash-attn library:

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