Flash Attention Reference Guide

Flash Attention accelerates transformer attention and reduces memory usage through IO-aware tiling and recomputation. It achieves a 2-4x speedup and 10-20x memory reduction compared to standard attention. The technique was introduced in 2022 and has seen widespread adoption by 2024.

When to Use Flash Attention

Use Flash Attention when:

• Training transformers with sequences longer than 512 tokens
• Running inference with long context (over 2K tokens)
• GPU memory is constrained and standard attention causes out-of-memory errors
• You need a 2-4x speedup without accuracy loss
• Using PyTorch 2.2+ or can install the flash-attn library

Flash Attention is not beneficial for sequences under 256 tokens, where standard attention may be faster due to overhead. It requires a GPU with compute capability 7.5 or higher (Turing+). Volta (V100) GPUs are not supported. Ampere (A100, A10) and Turing (T4) GPUs are fully supported.

Capabilities

• 2-4x speedup over standard attention
• 10-20x memory reduction
• Automatic use via PyTorch's F.scaled_dot_product_attention (PyTorch 2.2+)
• Dropout support via the dropout_p parameter
• Causal masking via causal=True
• Sliding window attention via the window_size parameter
• Multi-query attention (MQA) – automatically handles fewer KV heads
• FP8 forward-only on H100 (via FlashAttention-3 hopper build)
• FP16/BF16 forward+backward (via FlashAttention-3 hopper build)
• Softmax scaling via the softmax_scale parameter (auto-scales if None)

Prerequisites

• NVIDIA Ampere+ GPU (A100, A10, A30) or AMD MI200+ (for FA2); Hopper (H100/H800) for FA3 FP8
• CUDA 12.0+ (minimum 11.8)
• PyTorch 2.2+ for native support
• For flash-attn library: pip install flash-attn --no-build-isolation
• For FA3: build from source (hopper/ directory) with CUDA toolchain and Hopper GPU

Enabling Flash Attention in Existing PyTorch Models

Step 1: Check PyTorch Version

Verify your PyTorch version is 2.2.0 or higher:

python -c "import torch; print(torch.__version__)"
# Should be ≥2.2.0

If the version is below 2.2.0, upgrade:

pip install --upgrade torch

Step 2: Replace Standard Attention Code

Replace standard attention code with F.scaled_dot_product_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)

The function automatically uses Flash Attention if available. Here is a complete example:

import torch
import torch.nn.functional as F

q = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)  # [batch, heads, seq, dim]
k = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)

# Automatically uses Flash Attention if available
out = F.scaled_dot_product_attention(q, k, v)

Step 3: Force Flash Attention Backend (Optional)

To force the Flash Attention backend, use SDPBackend and sdpa_kernel:

from torch.nn.attention import SDPBackend, sdpa_kernel

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v)

Note that torch.nn.attention.sdpa_kernel is the modern API, while torch.backends.cuda.sdp_kernel is an older equivalent.

Step 4: Verify Speedup

Benchmark Flash Attention against standard attention:

import torch.utils.benchmark as benchmark

def test_attention(use_flash):
    q, k, v = [torch.randn(2, 8, 2048, 64, device='cuda', dtype=torch.float16) for _ in range(3)]

    if use_flash:
        from torch.nn.attention import SDPBackend, sdpa_kernel
        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
            return F.scaled_dot_product_attention(q, k, v)
    else:
        attn = (q @ k.transpose(-2, -1) / 8.0).softmax(dim=-1)
        return attn @ v

# Benchmark
t_flash = benchmark.Timer(stmt='test_attention(True)', globals=globals())
t_standard = benchmark.Timer(stmt='test_attention(False)', globals=globals())

print(f"Flash: {t_flash.timeit(100).mean:.3f}s")
print(f"Standard: {t_standard.timeit(100).mean:.3f}s")

Step 5: Test Accuracy

Verify that Flash Attention output matches standard attention within tolerance:

# Compare outputs
q, k, v = [torch.randn(1, 8, 512, 64, device='cuda', dtype=torch.float16) for _ in range(3)]

# Flash Attention
out_flash = F.scaled_dot_product_attention(q, k, v)

# Standard attention
attn_weights = torch.softmax(q @ k.transpose(-2, -1) / 8.0, dim=-1)
out_standard = attn_weights @ v

# Check difference
diff = (out_flash - out_standard).abs().max()
print(f"Max difference: {diff:.6f}")
# Should be <1e-3 for float16

The maximum difference should be less than 1e-3 for float16. Accuracy degradation can occur if the dtype is float32, which is not supported. Cast to float16 or bfloat16:

q = q.to(torch.float16)  # Or torch.bfloat16

Using the flash-attn Library for Advanced Features

Step 1: Install the Library

Install the flash-attn library:

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

For NVIDIA GPUs with CUDA 12.0+:

# NVIDIA GPUs (CUDA 12.0+)
pip install flash-attn --no-build-isolation

# Verify installation
python -c "from flash_attn import flash_attn_func; print('Success')"

If you encounter issues, install the CUDA toolkit first:

conda install cuda -c nvidia
pip install flash-attn --no-build-isolation

Step 2: Modify Attention Code

The flash_attn_func expects inputs in [batch, seq, heads, dim] format. Transpose if needed:

from flash_attn import flash_attn_func

# Input: [batch_size, seq_len, num_heads, head_dim]
# Transpose from [batch, heads, seq, dim] if needed
q = q.transpose(1, 2)  # [batch, seq, heads, dim]
k = k.transpose(1, 2)
v = v.transpose(1, 2)

out = flash_attn_func(
    q, k, v,
    dropout_p=0.1,
    causal=True,  # For autoregressive models
    window_size=(-1, -1),  # No sliding window
    softmax_scale=None  # Auto-scale
)

out = out.transpose(1, 2)  # Back to [batch, heads, seq, dim]

Step 3: Enable Advanced Features

Multi-Query Attention (MQA): Pass fewer KV heads than query heads:

from flash_attn import flash_attn_func

# q: [batch, seq, num_q_heads, dim]
# k, v: [batch, seq, num_kv_heads, dim]  # Fewer KV heads
out = flash_attn_func(q, k, v)  # Automatically handles MQA

Sliding Window Attention: Set window_size to a tuple (left, right):

# Only attend to window of 256 tokens before/after
out = flash_attn_func(
    q, k, v,
    window_size=(256, 256),  # (left, right) window
    causal=True
)

Step 4: Benchmark Performance

Benchmark the flash-attn library:

import torch
from flash_attn import flash_attn_func
import time

q, k, v = [torch.randn(4, 4096, 32, 64, device='cuda', dtype=torch.float16) for _ in range(3)]

# Warmup
for _ in range(10):
    _ = flash_attn_func(q, k, v)

# Benchmark
torch.cuda.synchronize()
start = time.time()
for _ in range(100):
    out = flash_attn_func(q, k, v)
    torch.cuda.synchronize()
end = time.time()

print(f"Time per iteration: {(end-start)/100*1000:.2f}ms")
print(f"Memory allocated: {torch.cuda.max_memory_allocated()/1e9:.2f}GB")

H100 FP8 Optimization (FlashAttention-3)

Step 1: Verify Hopper GPU

Check that you have an H100 or H800 GPU:

nvidia-smi --query-gpu=name --format=csv
# Should show "H100" or "H800"

Step 2: Build and Install FlashAttention-3

FlashAttention-3 is a separate beta build from source, not included in the pip package:

git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install
# (compilation is heavy and requires a CUDA toolchain + Hopper GPU)

Step 3: Use the FA3 Interface

Import from flash_attn_interface instead of flash_attn:

import torch
from flash_attn_interface import flash_attn_func  # FA3 (hopper build), not `flash_attn`

# q, k, v: [batch, seqlen, nheads, headdim]
q = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)
k = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)

# FP8 forward (inference / forward-only): cast to float8_e4m3fn
q_fp8 = q.to(torch.float8_e4m3fn)
k_fp8 = k.to(torch.float8_e4m3fn)
v_fp8 = v.to(torch.float8_e4m3fn)

out = flash_attn_func(q_fp8, k_fp8, v_fp8, causal=True)
# FP16/BF16 forward+backward is also supported by the FA3 interface.

FP8 is forward-only (no backward pass) in FA3. FP16/BF16 forward+backward is also supported by the FA3 interface.

Parameters

Parameter	Meaning	Required
dropout_p	Dropout probability applied to attention weights	No (default 0.0)
causal	If True, applies causal masking (for autoregressive models)	No (default False)
window_size	Tuple (left, right) for sliding window attention; (-1,-1) means no window	No (default (-1,-1))
softmax_scale	Scaling factor for softmax; if None, auto-scales by 1/√(head_dim)	No (default None)
attn_mask	Optional attention mask passed to F.scaled_dot_product_attention	No

Constraints and Caveats

• Flash Attention requires a GPU; not supported on CPU
• PyTorch native (F.scaled_dot_product_attention) requires PyTorch 2.2+
• The flash-attn library (pip) requires CUDA 12.0+ (minimum 11.8)
• FlashAttention-3 (FP8) is a separate beta build from source (hopper/ directory); not included in the pip package
• flash_attn_func from pip (flash-attn 2.8.x) does not auto-use FP8; it is FA2 only
• FP8 is forward-only (no backward pass) in FA3
• Flash Attention uses float16 or bfloat16; float32 is not supported
• For sequences under 256 tokens, standard attention may be faster due to overhead
• Volta (V100) GPUs are not supported; requires Turing+ (compute capability 7.5 or higher)
• VRAM usage is the same as standard attention (does not increase memory)
• Accuracy degradation is possible if dtype is float32; cast to float16 or bfloat16

Failure Modes

• ImportError: cannot import flash_attn – Install with --no-build-isolation or install the CUDA toolkit first
• Slower than expected – Sequence length is too short (under 512 tokens gives minimal speedup)
• RuntimeError: CUDA error – GPU not supported (check compute capability is 7.5 or higher)

Check compute capability:

import torch
print(torch.cuda.get_device_capability())
# Should be ≥(7, 5) for Turing+

Integration Checklist

Use these checklists to track your integration progress.

Flash Attention Integration:

Flash Attention Integration:
- [ ] Step 1: Check PyTorch version (≥2.2)
- [ ] Step 2: Enable Flash Attention backend
- [ ] Step 3: Verify speedup with profiling
- [ ] Step 4: Test accuracy matches baseline

flash-attn Library Setup:

flash-attn Library Setup:
- [ ] Step 1: Install flash-attn library
- [ ] Step 2: Modify attention code
- [ ] Step 3: Enable advanced features
- [ ] Step 4: Benchmark performance

FP8 Setup:

FP8 Setup:
- [ ] Step 1: Verify Hopper (H100) GPU available
- [ ] Step 2: Build & install FlashAttention-3 from source (hopper/)
- [ ] Step 3: Use the FA3 interface (FP8 forward)