FlashAttention

FlashAttention (Dao et al., 2022) computes exact standard attention — the same mathematical operation as naive multi-head attention — but restructures the computation to minimize data movement between GPU high-bandwidth memory (HBM) and on-chip SRAM. It is not a new attention variant, an approximation, or an architectural change. It is a kernel-level optimization that makes dense attention faster and more memory-efficient by never materializing the N x N attention score matrix in HBM.

The memory wall in naive attention

Standard attention for a single head computes three steps:

$S = Q @ K^{T}$ — score matrix, shape (N, N)

P = softmax(S, dim=-1) — attention weights, shape (N, N)

O = P @ V — output, shape (N, d_head)

In a naive PyTorch implementation, S and P are both materialized as full N x N tensors in GPU HBM:

python

def naive_attention(Q, K, V):
    # Q, K, V: (batch, n_heads, seq_len, d_head)
    d_head = Q.shape[-1]
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_head)  # (B, H, N, N)
    attn_weights = F.softmax(scores, dim=-1)                            # (B, H, N, N)
    output = torch.matmul(attn_weights, V)                              # (B, H, N, d_head)
    return output

For sequence length N = 128K tokens in float16, the score matrix S alone is:

128, 000 \cdot 128, 000 \cdot 2 \text{bytes} = 32.8 \text{GB}

A single A100 GPU has 80 GB of HBM. The attention matrix for one head at 128K tokens consumes 41% of total GPU memory — and a model with 32 heads would need $32 \cdot 32.8 \text{GB} = 1, 049 \text{GB}$ if computed simultaneously. This doesn't fit. Even at 8K tokens, the matrix is $8, 000 \cdot 8, 000 \cdot 2 = 128 \text{MB}$ per head, 4 GB for 32 heads — manageable but wasteful, since the matrix is intermediate and discarded after the P @ V multiply.

The memory problem is O(N^2): quadratic in sequence length.

The bandwidth bottleneck

GPU SRAM (on-chip shared memory) is fast but tiny: an A100 has ~20 MB of SRAM across all streaming multiprocessors, with ~200 KB available per thread block. HBM is large (80 GB) but slow: A100 HBM bandwidth is ~2 TB/s, while SRAM bandwidth is ~19 TB/s — roughly 10x faster.

Naive attention writes the full N x N score matrix to HBM, reads it back for softmax, writes the softmax result to HBM, reads it back for the P @ V multiply. Each of these HBM round-trips takes time proportional to N^2. The actual arithmetic (the matrix multiplies and exponentials) is fast — the bottleneck is data movement. This is what makes attention memory-bound rather than compute-bound for the decode phase and for moderate sequence lengths during prefill.

Tiling and online softmax

FlashAttention's core technique is tiling: instead of computing the full N x N matrix, divide Q, K, and V into blocks and compute attention one block at a time, keeping all intermediate results in SRAM.

The challenge is softmax. Standard softmax over a row of scores requires seeing all scores in that row to compute the denominator:

\text{softmax}(s_{i}) = \frac{\exp(s_{i})}{\text{sum}_{j}(\exp(s_{j}))}

If you've only seen a block of scores, you can't compute the final softmax — the denominator is incomplete.

The online softmax trick (Milakov and Gimelshein, 2018) solves this by maintaining running statistics. For each row of the output, FlashAttention tracks:

m: the running maximum score seen so far (for numerical stability)
l: the running sum of exp(s_j - m) (the partial softmax denominator)
o: the running weighted sum of V (the partial output)

When a new block of K scores arrives, the algorithm:

Computes new block scores in SRAM
Updates m if any new score exceeds the current max
Rescales the existing partial sums by exp(m_old - m_new) to account for the changed max
Accumulates the new block's contribution to l and o

After processing all K blocks, o / l gives the exact softmax-weighted output — identical to the result of materializing the full N x N matrix.

The tiled algorithm (simplified)

for each block Q_i of Q (block size B_r):
    initialize: m_i = -inf, l_i = 0, o_i = 0
    for each block K_j, V_j of K, V (block size B_c):
        load Q_i, K_j, V_j into SRAM
        compute S_ij = Q_i @ K_j^T         (B_r x B_c block, in SRAM)
        compute m_new = max(m_i, rowmax(S_ij))
        compute P_ij = exp(S_ij - m_new)    (in SRAM)
        rescale: l_i = l_i * exp(m_i - m_new) + rowsum(P_ij)
        rescale: o_i = o_i * exp(m_i - m_new) + P_ij @ V_j
        update: m_i = m_new
    write o_i / l_i to HBM

No N x N matrix is ever stored in HBM. The largest intermediate tensor is B_r x B_c — a block of the score matrix that fits in SRAM. Typical block sizes on A100: $B_{r} = B_{c} = 128$ .

Memory and speed improvements

FlashAttention changes the memory complexity of attention:

Naive: O(N^2) — the full score matrix
FlashAttention: O(N) — linear in sequence length. Only the output (N, d_head), plus O(N) for the running statistics m and l

The FLOPs are identical — FlashAttention computes the same matrix multiplications. The speedup comes entirely from reduced HBM traffic:

Naive attention performs O(N^2 * d_head) HBM reads/writes (loading and storing the score matrix, the softmax output, etc.)
FlashAttention performs O(N^2 * d_head^2 / M) HBM accesses, where M is the SRAM size. For typical $d_{\text{head}} = 128$ and M = 200 KB, this is a significant reduction.

Benchmarked results (Dao et al., 2022)

On A100-80GB with GPT-2 style models:

Sequence length 1K: FlashAttention is 1.5–2x faster than PyTorch naive attention
Sequence length 4K: 2–3x faster
Sequence length 16K: 3–4x faster (memory savings prevent OOM that kills naive)
Training speed: BERT-large training wall-clock time reduced by 15%. GPT-2 medium (1.5B) end-to-end training 3x faster at sequence length 8K.

The speedup increases with sequence length because the O(N^2) memory traffic of naive attention grows quadratically while FlashAttention's tiled approach grows much more slowly.

FlashAttention-2 and FlashAttention-3

FlashAttention-2 (Dao, 2023) improved on the original in two main ways:

Better work partitioning: FlashAttention-1 parallelized over batch and heads, leaving sequence length serial within each thread block. FlashAttention-2 also parallelizes across the sequence dimension, improving GPU occupancy on long sequences.
Reduced non-matmul FLOPs: Restructured the online softmax rescaling to minimize the number of non-tensor-core operations (element-wise exponentials, reductions), which are slower than matrix multiplies on modern GPUs.

Result: FlashAttention-2 achieves 50–73% of theoretical peak A100 FLOPS (measured in TFLOPS for the attention operation), compared to ~30–50% for FlashAttention-1.

FlashAttention-3 (Shah et al., 2024) targets NVIDIA Hopper GPUs (H100, H200):

Warp specialization: Hopper's architecture supports asynchronous execution of different warps. FlashAttention-3 pipelines data loading (one warp group) with computation (another warp group), hiding memory latency.
FP8 support: Hopper's FP8 tensor cores deliver 2x the FLOPS of FP16. FlashAttention-3 computes attention in FP8 with selective FP32 accumulation for numerical stability (the softmax reduction is done in higher precision).
Asynchronous block-level operations: Uses Hopper's Tensor Memory Accelerator (TMA) for bulk asynchronous data copies between HBM and SRAM, overlapping with computation.

FlashAttention-3 on H100 achieves 740 TFLOPS in FP16 (versus ~400 for FlashAttention-2 on the same hardware) and up to 1.2 PFLOPS in FP8.

Using FlashAttention in practice

PyTorch 2.0+ exposes FlashAttention through torch.nn.functional.scaled_dot_product_attention (SDPA). This function automatically dispatches to the FlashAttention kernel when the inputs are on CUDA, in float16 or bfloat16, and the sequence length is compatible:

python

import torch
import torch.nn.functional as F
import time

def benchmark_attention(seq_len, d_model=4096, n_heads=32, batch=4, device="cuda"):
    d_head = d_model // n_heads
    q = torch.randn(batch, n_heads, seq_len, d_head, device=device, dtype=torch.float16)
    k = torch.randn(batch, n_heads, seq_len, d_head, device=device, dtype=torch.float16)
    v = torch.randn(batch, n_heads, seq_len, d_head, device=device, dtype=torch.float16)

    # Warm up
    for _ in range(3):
        _ = F.scaled_dot_product_attention(q, k, v)
    torch.cuda.synchronize()

    # Benchmark SDPA (uses FlashAttention backend)
    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    for _ in range(10):
        out_flash = F.scaled_dot_product_attention(q, k, v)
    torch.cuda.synchronize()
    flash_time = (time.perf_counter() - start) / 10
    flash_mem = torch.cuda.max_memory_allocated() / 1e9

    # Benchmark naive (force materialization of N x N matrix)
    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    for _ in range(10):
        scores = torch.matmul(q, k.transpose(-2, -1)) / (d_head ** 0.5)
        attn = F.softmax(scores, dim=-1)
        out_naive = torch.matmul(attn, v)
    torch.cuda.synchronize()
    naive_time = (time.perf_counter() - start) / 10
    naive_mem = torch.cuda.max_memory_allocated() / 1e9

    # Verify outputs match
    max_diff = (out_flash - out_naive).abs().max().item()

    print(f"seq_len={seq_len}")
    print(f"  Naive:  {naive_time*1000:.1f} ms, {naive_mem:.2f} GB peak")
    print(f"  Flash:  {flash_time*1000:.1f} ms, {flash_mem:.2f} GB peak")
    print(f"  Speedup: {naive_time/flash_time:.1f}x")
    print(f"  Max diff: {max_diff:.2e} (exact up to float16 precision)")

Typical results on A100-80GB:

seq_len=2048
  Naive:  3.2 ms, 2.15 GB peak
  Flash:  1.4 ms, 0.54 GB peak
  Speedup: 2.3x

seq_len=8192
  Naive:  48.6 ms, 34.2 GB peak
  Flash:  15.1 ms, 2.15 GB peak
  Speedup: 3.2x

seq_len=16384
  Naive:  OOM
  Flash:  58.4 ms, 8.59 GB peak

At 16K tokens, naive attention requires ~128 GB for the score matrix (4 heads 16K 16K 2 bytes batch 4) and crashes. FlashAttention completes in 58 ms using 8.6 GB.

Checking which backend SDPA uses

python

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)
    # Raises an error if FlashAttention can't be used
    # (e.g., wrong dtype, CPU tensor, unsupported head dim)

SDPA falls back to a "math" (naive) or "memory efficient" (xFormers-style) backend when FlashAttention's constraints aren't met. The FlashAttention backend requires: CUDA, float16 or bfloat16, $d_{\text{head}} \leq 256$ , and no custom attention bias in early PyTorch versions (2.0–2.1).

Orthogonality to attention variants

FlashAttention is orthogonal to the choice of MHA, GQA, MQA, or MLA. It accelerates the core softmax(Q @ K^T / sqrt(d)) @ V computation regardless of how the Q, K, V tensors were produced:

MHA + FlashAttention: Full KV cache, fast attention kernel. Used in GPT-4-era models before GQA adoption.
GQA + FlashAttention: Reduced KV cache via grouped heads, fast attention kernel. This is the standard configuration for Llama 3, Mistral, and Qwen 2.
MQA + FlashAttention: Minimal KV cache, fast attention kernel. Used in PaLM and Falcon inference.
MLA + FlashAttention: Latent-compressed KV cache, but requires custom kernel modifications for the split content/RoPE score computation. DeepSeek uses custom attention kernels based on FlashAttention principles.

The architectural choice (MHA/GQA/MQA/MLA) determines what gets cached and how much memory the KV cache consumes. FlashAttention determines how fast the attention computation itself runs and how much memory the attention operation (not the cache) uses. Both optimizations stack.

When FlashAttention doesn't help

FlashAttention's speedup is largest when attention is memory-bound — when the time spent moving data between HBM and SRAM dominates the time spent on arithmetic. This is the common case for:

Long sequences (N > 1K) where the N x N matrix is large
The decode phase of autoregressive generation (small batch of queries attending to a large KV cache)
Training with standard dense attention

FlashAttention helps less in these scenarios:

Very short sequences (N < 256): The N x N matrix fits in SRAM even without tiling. The kernel launch overhead of FlashAttention may negate the bandwidth savings.
Sparse attention patterns: If you're already using a sparse attention mechanism (local windows, sliding window, block-sparse), you're not materializing the full N x N matrix anyway. FlashAttention for sparse patterns exists (block-sparse FlashAttention) but the savings over a well-implemented sparse kernel are smaller.
CPU inference: FlashAttention is a GPU kernel optimization. On CPU, attention is compute-bound (no HBM/SRAM hierarchy to exploit), and the naive implementation with BLAS is often the best option.

← Previous