FlashAttention

Deep dive

How it works

In: Q, K, V of shape N×d (per head) · Out: attention output O and one logsumexp per row · State: none between calls; per tile m, l, acc in SRAM · Cost: 4·N²·d FLOPs, O(N·d) HBM traffic, O(N) extra memory.

1. Load one Q block into SRAM.
2. For each K/V block compute the score tile, update the running row max m and sum l, rescale the partial output.
3. After the last block divide by l and store O and L = m + log l.
4. Backward recomputes score tiles from L.

For one query block Qᵢ (Br rows), iterate over key/value blocks Kⱼ, Vⱼ (Bc rows):

m = −∞ (Br),  l = 0 (Br),  acc = 0 (Br×d)                 all in registers/SRAM
for j:
    S    = Qᵢ Kⱼᵀ · scale                                  Br×Bc, never leaves the SM
    mₙ   = max(m, rowmax(S))
    P    = exp(S − mₙ)
    α    = exp(m − mₙ)                                      rescale old partial results
    l    = α·l + rowsum(P)
    acc  = α·acc + P Vⱼ
    m    = mₙ
Oᵢ = acc / l ;   Lᵢ = m + log(l)                            store O and logsumexp only

The backward pass recomputes P = exp(QKᵀ·scale − L) per tile from the saved L, then uses D = rowsum(dO ⊙ O) and dS = P ⊙ (dP − D) to form dQ, dK and dV. FLOPs rise to about 2.5x forward, but HBM traffic stays O(N·d).

FlashAttention-2 parallelizes over the sequence as well as batch×heads, splits Q (not K) across warps to avoid shared-memory reductions, and trims non-matmul FLOPs (rescaling once at the end). FlashAttention-3 adds Hopper warp specialization (producer warps issue TMA loads, consumers run wgmma), ping-pong scheduling that overlaps softmax with GEMMs, and fp8 with incoherent processing.

Code

The algorithm on plain slices (one head, row-major N×d), checked against naive softmax attention:

/// Tiled attention with online softmax. Only a BM×BN score tile exists at any time.
fn flash_attention(q: &[f32], k: &[f32], v: &[f32], n: usize, d: usize, bm: usize, bn: usize) -> (Vec<f32>, Vec<f32>) {
    let scale = 1.0 / (d as f32).sqrt();
    let (mut out, mut lse) = (vec![0.0f32; n * d], vec![0.0f32; n]);
    let mut s = vec![0.0f32; bm * bn];                         // the "SRAM" tile
    for i0 in (0..n).step_by(bm) {
        let rows = bm.min(n - i0);
        let (mut m, mut l) = (vec![f32::NEG_INFINITY; rows], vec![0.0f32; rows]);
        let mut acc = vec![0.0f32; rows * d];
        for j0 in (0..n).step_by(bn) {
            let cols = bn.min(n - j0);
            for r in 0..rows {
                let qi = &q[(i0 + r) * d..(i0 + r + 1) * d];
                for c in 0..cols {                                // S = Q_i K_j^T · scale
                    let kj = &k[(j0 + c) * d..(j0 + c + 1) * d];
                    s[r * bn + c] = qi.iter().zip(kj).map(|(a, b)| a * b).sum::<f32>() * scale;
                }
                let row = &mut s[r * bn..r * bn + cols];
                let m_new = row.iter().fold(m[r], |a, &b| a.max(b));
                let alpha = (m[r] - m_new).exp();                 // rescale earlier partial sums
                let acc_r = &mut acc[r * d..(r + 1) * d];
                acc_r.iter_mut().for_each(|a| *a *= alpha);
                l[r] *= alpha;
                for (c, p) in row.iter_mut().enumerate() {        // P = exp(S − m), acc += P V_j
                    *p = (*p - m_new).exp();
                    l[r] += *p;
                    let vj = &v[(j0 + c) * d..(j0 + c + 1) * d];
                    acc_r.iter_mut().zip(vj).for_each(|(a, b)| *a += *p * b);
                }
                m[r] = m_new;
            }
        }
        for r in 0..rows {
            for t in 0..d { out[(i0 + r) * d + t] = acc[r * d + t] / l[r]; }
            lse[i0 + r] = m[r] + l[r].ln();                         // all the backward pass needs
        }
    }
    (out, lse)
}

fn main() {
    let (n, d) = (300, 64);
    let gen = |seed: usize| -> Vec<f32> { (0..n * d).map(|i| (((i * 7919 + seed) % 1000) as f32 / 500.0) - 1.0).collect() };
    let (q, k, v) = (gen(1), gen(2), gen(3));
    let (o, _) = flash_attention(&q, &k, &v, n, d, 64, 64);
    for i in 0..n {                                                 // naive: full softmax row
        let s: Vec<f32> = (0..n).map(|j| (0..d).map(|t| q[i * d + t] * k[j * d + t]).sum::<f32>() / (d as f32).sqrt()).collect();
        let mx = s.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let z: f32 = s.iter().map(|x| (x - mx).exp()).sum();
        for t in 0..d {
            let want: f32 = (0..n).map(|j| (s[j] - mx).exp() / z * v[j * d + t]).sum();
            assert!((o[i * d + t] - want).abs() < 1e-4);
        }
    }
    println!("tiled online-softmax attention matches naive attention");
}

The same loop as a GPU kernel in Triton: one program per (query block, batch·head), tensors contiguous as (B·H, N, D), launched on grid (cdiv(N, BM), B·H).

@triton.jit
def attn_fwd(Q, K, V, O, LSE, N, scale, D: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr):
    pid_m, bh = tl.program_id(0), tl.program_id(1)
    base = bh.to(tl.int64) * N * D                              # int64: avoids 2^31 overflow
    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = tl.arange(0, BN)
    offs_d = tl.arange(0, D)
    q = tl.load(Q + base + offs_m[:, None] * D + offs_d[None, :], mask=offs_m[:, None] < N, other=0.0)
    m_i = tl.full([BM], float("-inf"), tl.float32)
    l_i = tl.zeros([BM], tl.float32)
    acc = tl.zeros([BM, D], tl.float32)
    qk_scale = scale * 1.4426950408889634                       # fold log2(e): exp2 is one MUFU op
    for start in range(0, N, BN):
        kn = start + offs_n
        kt = tl.load(K + base + kn[None, :] * D + offs_d[:, None], mask=kn[None, :] < N, other=0.0)
        s = tl.dot(q, kt) * qk_scale                            # tensor cores, fp32 accumulate
        s = tl.where(kn[None, :] < N, s, float("-inf"))
        m_new = tl.maximum(m_i, tl.max(s, 1))
        p = tl.exp2(s - m_new[:, None])
        alpha = tl.exp2(m_i - m_new)
        l_i = l_i * alpha + tl.sum(p, 1)
        v = tl.load(V + base + kn[:, None] * D + offs_d[None, :], mask=kn[:, None] < N, other=0.0)
        acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
        m_i = m_new
    acc = acc / l_i[:, None]
    tl.store(O + base + offs_m[:, None] * D + offs_d[None, :], acc.to(O.dtype.element_ty), mask=offs_m[:, None] < N)
    tl.store(LSE + bh.to(tl.int64) * N + offs_m, m_i + tl.log2(l_i), mask=offs_m < N)  # base-2 LSE

Complexity and constants

• FLOPs: 4·N²·d per head forward (two GEMMs). Causal masking halves it if fully masked tiles are skipped.
• HBM bytes: Q, K, V, O ≈ 4·N·d·2 B, plus K/V re-reads per query block (N/Br times). With Br = 128 and d = 128, the kernel is compute-bound from N of about 1k.
• SRAM budget per CTA: Q tile Br×d plus double-buffered K and V tiles of Bc×d each. At d = 128 with Br = Bc = 128 in bf16 that is about 32 KB + 64 KB + 64 KB. This is why head dim 256 forces smaller tiles.
• End to end: about 15% faster BERT-large and 3x faster GPT-2 training in the original paper. Long-context runs that OOM without it become feasible.

Efficiency

• Cost: 4·N²·d FLOPs per head forward. Standard attention moves the N×N matrix through HBM (128 MB per head at N = 8k in fp16); the tiled kernel keeps it in SRAM (bandwidth about 19 TB/s versus 1.5-2.0 TB/s for A100 HBM, FlashAttention paper). Bound: HBM bandwidth for standard attention, tensor-core compute for FlashAttention.
• Levers:
  • tiling with online softmax — 15% faster BERT-large and 3x faster GPT-2 training (FlashAttention paper) — any sequence beyond a few hundred tokens
  • better work partitioning — about 2x over v1, 50-73% of A100 peak (FlashAttention-2) — A100 and newer
  • Hopper async warp specialization and fp8 — up to 740 TFLOPs/s fp16 on H100, 75% utilization (FlashAttention-3) — H100
  • skip fully masked tiles — about half the FLOPs — causal masks
  • split-KV for single-query decode (Flash-Decoding) — parallelism over a long cache — decode
• Naive → standard → state of the art:
  • materialized N×N scores · O(N²) memory, HBM-bound
  • FlashAttention-2 · exact, O(N) memory, 50-73% of A100 peak
  • FlashAttention-3 on Hopper · up to 740 TFLOPs/s fp16

Implementation notes

• Work in base 2 (exp2 with log2(e) folded into the scale). It is one MUFU instruction.
• Accumulate S, l and acc in fp32. Cast P to bf16/fp16 only for the PV matmul.
• Causal masks: skip tiles entirely above the diagonal, and mask only the diagonal tiles. Process the masked diagonal tile last to keep the main loop branch-free.
• Dropout: regenerate the mask in backward from a Philox seed and offset rather than storing it.
• GQA/MQA: map several query heads to one KV head via the program id. The loads are shared, and nothing else changes.
• Decode (Flash-Decoding): with 1 query row, parallelize over KV chunks and merge partial (m, l, acc) triples with the same rescaling rule.

Pitfalls

• Fully masked rows (padding, some causal block layouts): m stays −∞, and exp(−∞ − (−∞)) is NaN. Guard with m_safe = where(m == −inf, 0, m) and output 0.
• int32 pointer offsets overflow beyond 2³¹ elements (e.g. B·H·N·D at long context). Promote to int64, as the kernel does.
• Some backward implementations accumulate dQ with atomics, which is non-deterministic run to run. Use the deterministic flag when reproducibility matters.
• Different tile sizes or split-KV change the summation order. Expect differences at the 1e-3 level in bf16, not bitwise equality.
• Head dims that are not multiples of 8/16 need padding to satisfy tensor-core shapes.

Variants and alternatives

In production

PyTorch scaled_dot_product_attention (flash and memory-efficient backends), the flash-attn package (Dao-AILab), xFormers, cuDNN fused attention, vLLM and TensorRT-LLM, JAX Pallas kernels, Triton's fused-attention tutorial.

Implementations

Sources

• Dao, Fu, Ermon, Rudra, Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022. paper · code
• Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023. paper
• Shah et al., FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, 2024. paper
• Milakov and Gimelshein, Online normalizer calculation for softmax, 2018. paper
• Rabe and Staats, Self-attention Does Not Need O(n²) Memory, 2021. paper