Deep Dive - FlashAttention

A large language model (LLM) can produce the wrong next-token attention result even when its kernel computes softmax perfectly. Two new query tokens attending to five cached keys expose one such bug: aligning a causal mask to the wrong corner hides valid keys. The tiled kernel inside FlashAttention has to keep the caller contract that chooses that corner.

The FlashAttention & Memory Efficiency lesson explained why attention can wait on GPU memory traffic. The two-token mask is a different failure. The math can be exact and the installed package can still hide legal keys. A compatible mask, layout, and return contract have to travel with the kernel.

Hold the equation steady before changing the implementation. Scaled dot-product attention is still softmax of scaled $QK^T$ times $V$. FlashAttention changes the schedule and the bytes that cross HBM (high-bandwidth memory), not the definition.^{[1]Reference 1FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135}

Project identity

The papers name the algorithms; the Dao-AILab repository tells you what a process can actually import, test, and run. Kernels, wrappers, tests, and install paths don't all ship as one wheel.^{[2]Reference 2FlashAttentionhttps://github.com/Dao-AILab/flash-attention}

Three install surfaces matter more than the paper titles. These documentation requirements were rechecked September 23, 2026; no GPU installation is claimed:^{[2]Reference 2FlashAttentionhttps://github.com/Dao-AILab/flash-attention}

For CUDA 13, the current README recommends pip install "flash-attn-4[cu13]". The FA2 build instructions assume Linux and an installed CUDA/ROCm toolkit, plus PyTorch, packaging, psutil, and working ninja. A compatible GPU driver alone isn't the source-build toolchain; Windows support remains qualified.^{[2]Reference 2FlashAttentionhttps://github.com/Dao-AILab/flash-attention}

A paper name in a citation doesn't tell you which code a wheel will import. FA3 is a separate Hopper path, FA4 is a separate CuTe-DSL package, and FA2 remains the default flash-attn import.^{[2]Reference 2FlashAttentionhttps://github.com/Dao-AILab/flash-attention} Before benchmarking, print the package and function that actually loaded; otherwise a result can wear the wrong generation's label.

The source walkthrough stays pinned to commit ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820. Its selected interfaces, architecture guards, save-state and tolerance assertions were rechecked September 23, 2026. A wheel release needn't match that commit. Read these files in order:^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820}

The table's hardware targets summarize the README, not an exhaustive runtime selector. In particular, the pinned FA4 source has architecture-specific guards beyond the README's Hopper/Blackwell summary, including selected value dimensions of 512. Don't copy a universal head_dim <= 256 gate from FA2 into a FA4 caller.

The October 10, 2026 check uses commit 94e22c906678e5483fa0e9e24d8e787bc2c0ed4c for current FA4 behavior while preserving the older walkthrough. Its public call still returns (out, lse). Its source also records constraints that depend on the operation, such as forward-only FP8 and the paged-load boundary below. Test your training and cache configuration against those constraints.^{[6]Reference 6FlashAttention Source Snapshot, October 10, 2026https://github.com/Dao-AILab/flash-attention/tree/94e22c906678e5483fa0e9e24d8e787bc2c0ed4c}

Exact attention without an $n \times n$ workspace

For batch size 8, 32 heads, and 8,192 tokens, one 16-bit score tensor occupies $8\times32\times8192^2\times2=32$ GiB. That's one intermediate, not the whole model. A materializing baseline can write scores, reread them for softmax, write probabilities, and reread them for the value multiply. Each round trip crosses HBM, the large GPU memory pool. On-chip static random-access memory (SRAM) is much smaller, so a tile that stays there can be revisited cheaply.

A query-owned forward loop, as used in FA2's work partitioning, keeps a query tile $Q_i$ while walking key/value tiles $(K_j, V_j)$. For each pair, it computes $S_{ij}=Q_i K_j^T / \sqrt{d}$ on chip, updates a row-wise softmax state, then discards $S_{ij}$. The dot products are still dense, so arithmetic remains $O(n^2 d)$ FLOPs. The live score workspace drops from $n \times n$ to one tile.^{[7]Reference 7FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691}

The original FA1 analysis gives $O(n^2 d^2 / M)$ HBM accesses for $d\le M\le nd$, with $M$ measured in scalar elements, versus $\Omega(nd+n^2)$ for its materializing baseline. Under this memory model, the IO bound counts HBM accesses in scalar elements; applying it to bytes or predicting runtime for later kernels requires additional assumptions.^{[1]Reference 1FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135}

After one tile, each query row retains $m_i$ (largest score so far), $\ell_i$ (sum of shifted exponentials), and a vector $u_i$ (the unnormalized, value-weighted numerator). Only at the end do we divide: $O_i=u_i/\ell_i$. For a nonempty row, $\operatorname{LSE}_i=m_i+\log\ell_i$ is the log-sum-exp summary. At the example's shape, FP32 LSE needs $8\times32\times8192\times4=8$ MiB, not 32 GiB.

The dependency is easier to see as a loop:

View larger

The “yes” edge advances to a new key/value tile. The query tile and softmax state stay live. This conceptual forward loop omits parallel split reductions and hardware-specific pipelines.

Online softmax is what makes the tiles exact

Ordinary softmax looks like it needs every score before it can form a denominator. Suppose tile two has a larger row maximum than tile one. Throwing away tile one's state would lose its contribution, but keeping it on the old scale would misweight it. Online softmax merges partial summaries. Let the retained state be $(m_{\text{old}},\ell_{\text{old}},u_{\text{old}})$ and the new tile's row scores be $s$. Then:

When a later tile contains a larger score, the old accumulator is rescaled instead of thrown away. Milakov and Gimelshein describe this streaming normalizer, which FlashAttention uses inside its tiled kernel.^{[8]Reference 8Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

Initialize $m=-\infty$, $\ell=0$, and $u=0$. Skip a tile with no allowed key before evaluating exponentials: $-\infty-(-\infty)$ isn't a usable shift. If the entire row is masked, return zero by the kernel's documented convention; ordinary softmax on an all-$-\infty$ row is undefined. Exactness here means the same operator in real arithmetic, not bitwise equality after floating-point reordering.

Two tiles by hand

Use one query row and four keys split into two tiles. Tile one has scores $[1,2]$ and values $[[10,0],[0,10]]$. Tile two has scores $[3,0]$ and values $[[5,5],[9,1]]$. The common $1/\sqrt{d}$ scale is omitted so the state update stays readable. Before calculating, predict what a new maximum of 3 should do to the state built from tile one.

For tile one, $m_1=2$ and $\ell_1=e^{-1}+1=1.3679$. The numerator is $u_1=e^{-1}[10,0]+[0,10]=[3.6788,10]$. Keep that numerator, rather than dividing it after every tile.

Tile two contains a larger score, so $m_2=3$. Old state is rescaled by $e^{2-3}=0.3679$:

The final result is $O_2=u_2/\ell_2=[4.380,5.620]$, rounded to three decimals. Ordinary softmax over all four scores matches. The tile boundary changed storage, not the answer: no $1\times4$ score row had to survive between tiles.

View larger

This standard-library CPU reference accepts finite values and scores that are finite or -inf (masked). It rejects malformed shapes instead of silently truncating zip. Moving tile boundaries or adding a large constant to every score should preserve the result. It tests the recurrence, not a GPU kernel.

import math

scores = [1.0, 2.0, 3.0, 0.0]
values = [[10.0, 0.0], [0.0, 10.0], [5.0, 5.0], [9.0, 1.0]]

def validate(scores, values):
    if not scores or len(scores) != len(values) or not values[0]:
        raise ValueError("scores and nonempty value rows must match")
    dim = len(values[0])
    if any(len(row) != dim for row in values):
        raise ValueError("ragged values")
    if any(math.isnan(x) or x == math.inf for x in scores):
        raise ValueError("scores must be finite or -inf")
    if any(not math.isfinite(x) for row in values for x in row):
        raise ValueError("values must be finite")
    return dim

def full(scores, values):
    dim = validate(scores, values)
    m = max(scores)
    if m == -math.inf:
        return [0.0] * dim
    weights = [math.exp(x - m) for x in scores]
    return [sum(w * v[d] for w, v in zip(weights, values)) / sum(weights)
            for d in range(dim)]

def tiled(scores, values, tile_size):
    dim = validate(scores, values)
    if type(tile_size) is not int or tile_size < 1:
        raise ValueError("tile_size must be a positive integer")
    m = float("-inf")
    total = 0.0
    numerator = [0.0] * dim
    for start in range(0, len(scores), tile_size):
        current_scores = scores[start:start + tile_size]
        current_values = values[start:start + tile_size]
        if max(current_scores) == -math.inf:
            continue
        new_m = max(m, max(current_scores))
        old_scale = 0.0 if m == float("-inf") else math.exp(m - new_m)
        tile_weights = [math.exp(x - new_m) for x in current_scores]
        total = old_scale * total + sum(tile_weights)
        for d in range(dim):
            numerator[d] = old_scale * numerator[d] + sum(
                w * v[d] for w, v in zip(tile_weights, current_values)
            )
        m = new_m
    return [x / total for x in numerator] if total else [0.0] * dim

dense = full(scores, values)
split = tiled(scores, values, 2)
print("dense:", [round(x, 6) for x in dense])
print("tiled:", [round(x, 6) for x in split])
assert max(abs(a - b) for a, b in zip(dense, split)) < 1e-12
for size in [1, 2, 3, 4, 10]:
    for shift in [-1000.0, 0.0, 1000.0]:
        result = tiled([s + shift for s in scores], values, size)
        assert max(abs(a - b) for a, b in zip(dense, result)) < 1e-12
assert tiled([-math.inf] * 4, values, 2) == [0.0, 0.0]
assert tiled([-math.inf, -math.inf, 0.0, -math.inf], values, 2) == values[2]

dense: [4.379542, 5.620458]
tiled: [4.379542, 5.620458]

The public interface is several kernels

The FA2 flash_attn_func accepts tensors shaped (batch, sequence, heads, head_dim) for $Q$, and the same layout for $K$ and $V$ except sequence length and head count may differ from $Q$. PyTorch SDPA puts heads before sequence: (batch, heads, sequence, head_dim). A layout mismatch can run without a shape error when those dimensions happen to fit.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820[9]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

Separate public and private returns. The public FA2 function normally returns only the output tensor. Its private CUDA forward helper returns output, FP32 row LSE shaped (batch, heads, seqlen_q), a probability/dropout buffer, and random-number-generator (RNG) state for autograd. Public return_attn_probs=True instead returns a three-tuple (out, lse, S_dmask) for testing. The source warns that this debug probability buffer may have incorrect scaling; requesting it can also allocate a quadratic-size buffer. Don't enable it in a memory benchmark.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820}

Even identical function names don't imply identical returns: the pinned FA4 flash_attn.cute.flash_attn_func returns (out, lse), with lse potentially None when it isn't needed or requested. Check the installed signature and source before substituting it for FA2's tensor-returning call.

ROCm has two backends. Composable Kernel is the default, while FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE selects the AIter Triton path. When the HIP extension import fails, the pinned wrapper tries the AIter path and emits a warning. Fallback succeeds only with compatible AIter dependencies and hardware, so verify it in the target environment.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820}

Current ROCm installation instructions add an environment hazard: bundled AIter installation can replace Triton. AITER_USE_SYSTEM_TRITON=1 retains the environment's Triton; FLASH_ATTENTION_USE_SYSTEM_AITER=TRUE selects an already-installed AIter. These control different dependencies. Record the chosen versions and verify compatibility before treating either setting as a fix.^{[2]Reference 2FlashAttentionhttps://github.com/Dao-AILab/flash-attention}

FA2 CUDA documents 16-bit floating point (FP16) or brain floating point (BF16) and head dimensions up to 256 on its supported GPUs. Consumer-GPU backward at dimension 256 without dropout is documented from flash-attn 2.5.5. These qualifications apply to the stated FA2 releases and supported GPUs; check the documentation for the generation you use.^{[2]Reference 2FlashAttentionhttps://github.com/Dao-AILab/flash-attention}

Multi-query attention (MQA) uses one K/V head; grouped-query attention (GQA) uses several, but fewer than Q. FA2 maps consecutive groups of Q heads onto K/V heads. Six Q heads and two K/V heads map as 0,0,0,1,1,1, not alternating 0,1,0,1,0,1. This CPU check validates that caller-side contract. It doesn't predict which GPU kernel will run.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820}

def kv_head_for_queries(query_heads, key_heads, value_heads):
    if any(type(h) is not int or h <= 0
           for h in (query_heads, key_heads, value_heads)):
        raise ValueError("head counts must be positive integers")
    if key_heads != value_heads or query_heads % key_heads:
        raise ValueError("K/V heads must match and divide Q heads")
    group_size = query_heads // key_heads
    return [q_head // group_size for q_head in range(query_heads)]

print("GQA:", kv_head_for_queries(6, 2, 2))
print("MQA:", kv_head_for_queries(6, 1, 1))
assert kv_head_for_queries(3, 3, 3) == [0, 1, 2]
for invalid in [(0, 1, 1), (5, 2, 2), (6, 2, 1), (True, 1, 1)]:
    try:
        kv_head_for_queries(*invalid)
    except ValueError:
        pass
    else:
        raise AssertionError(f"accepted invalid head counts: {invalid}")

GQA: [0, 0, 0, 1, 1, 1]
MQA: [0, 0, 0, 0, 0, 0]

The equivalent reference repeats each K/V head consecutively. That checks kernel semantics for fixed Q/K/V. Whether a model trained with fewer K/V heads preserves task quality is a separate architectural evaluation. Head mapping can still be right while the decode result is wrong. The mask corner decides which cached keys the first new query is allowed to see.

Causal 2.1 is the decode gotcha

The corner is the bug. Autoregressive decoders can't read future tokens. A causal tile treats those scores as $-\infty$ before softmax. Before FlashAttention 2.1, causal=True aligned the triangle to the top-left. From 2.1 onward, it aligns to the bottom-right. This matches incremental decode when the query rows correspond to the final positions of the key sequence. It isn't the right alignment for every arbitrary short-query layout.^{[2]Reference 2FlashAttentionhttps://github.com/Dao-AILab/flash-attention}

For seqlen_q = 2 and seqlen_k = 5, ask whether the first query sees one old key or four. Keep-cells look like this (1 means the key participates):

View larger

If every cell in a query row is masked, the documented output for that row is zero. That shows up when seqlen_q > seqlen_k under bottom-right alignment: the first rows have no legal keys.

Equal query and key lengths hide the distinction: top-left and bottom-right masks are the same lower triangle. A short query over a long cache exposes it. PyTorch's scaled_dot_product_attention(..., is_causal=True) uses a lower-triangular mask on square inputs and documents an upper-left causal bias for non-square inputs. Matching FlashAttention 2.1 from SDPA means passing an explicit lower-right causal bias, not assuming is_causal=True did that for you.^{[9]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

Before running the helper, predict the seqlen_q > seqlen_k case: early rows should have no legal keys and therefore be all zero. It prints both corners.

def causal_keep(seqlen_q: int, seqlen_k: int, align: str) -> list[list[int]]:
    if any(type(n) is not int or n <= 0 for n in (seqlen_q, seqlen_k)):
        raise ValueError("sequence lengths must be positive integers")
    if align not in {"top-left", "bottom-right"}:
        raise ValueError("unknown causal alignment")
    rows = []
    for i in range(seqlen_q):
        row = []
        for j in range(seqlen_k):
            if align == "top-left":
                keep = j <= i
            else:
                keep = j <= i + seqlen_k - seqlen_q
            row.append(1 if keep else 0)
        rows.append(row)
    return rows

print("FA 2.0 q=2,k=5:", causal_keep(2, 5, "top-left"))
print("FA 2.1 q=2,k=5:", causal_keep(2, 5, "bottom-right"))
print("FA 2.1 q=5,k=2:", causal_keep(5, 2, "bottom-right"))
assert causal_keep(2, 5, "top-left") == [[1, 0, 0, 0, 0], [1, 1, 0, 0, 0]]
assert causal_keep(2, 5, "bottom-right") == [[1, 1, 1, 1, 0], [1, 1, 1, 1, 1]]
assert causal_keep(5, 2, "bottom-right") == [
    [0, 0],
    [0, 0],
    [0, 0],
    [1, 0],
    [1, 1],
]

FA 2.0 q=2,k=5: [[1, 0, 0, 0, 0], [1, 1, 0, 0, 0]]
FA 2.1 q=2,k=5: [[1, 1, 1, 1, 0], [1, 1, 1, 1, 1]]
FA 2.1 q=5,k=2: [[0, 0], [0, 0], [0, 0], [1, 0], [1, 1]]

⚠️ Common mistake: Comparing a FlashAttention 2.1 decode call against an SDPA is_causal=True reference on unequal query and key lengths, then treating the mismatch as a kernel bug. Check the causal corner before you check numerics.

Local attention is related, but not identical. In FA2, finite window_size=(left, right) keeps key indices from $i+S-L-\text{left}$ through $i+S-L+\text{right}$, inclusive, clipped to valid keys, where $L$ and $S$ are query and key lengths. A side set to -1 means unbounded attention on that side. Causal masking can further restrict that set. A smaller window changes the operator and may affect model quality. Tile shape and memory access determine latency, so fewer allowed scores can still run slower.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820} The mask says which keys inside one sequence are legal. A packed batch still has to say where one example ends and the next begins, or attention leaks across examples.

Paged KV is a lookup, not a cache manager

Training batches can pack unequal lengths through flash_attn_varlen_func. Its int32 cumulative-length arrays (cu_seqlens_q, cu_seqlens_k) start at zero, are nondecreasing, and end at the corresponding packed token count. They have batch + 1 entries. An incorrect offset can leak attention across examples even if every dot product is right; maximum lengths must also agree with the sequences.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820}

Decode has a different boundary. Serving engines store K/V in fixed-size physical pages and map each request through a block table. FA2's flash_attn_with_kvcache accepts that table and documents a page block size divisible by 256. This requirement belongs to that API; other paged-cache implementations can set different block sizes. The call can update the cache in place and apply rotary embeddings; it doesn't support backward. The caller must reserve space before updates. Repeated cache_batch_idx entries during writes can produce ambiguous competing updates, so don't use them as a cache-sharing protocol.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820}

Current FA4 documents a subtler payload constraint when its paged path loads whole key tiles through the Tensor Memory Accelerator (TMA). The final, partially used page can load rows beyond seqused_k. Masking their scores doesn't make NaN values safe: a zero probability multiplied by a NaN value is still NaN. Keep those loaded padding rows finite. Entire pages that the sequence doesn't use may contain anything, including NaNs. This finite-value requirement applies to rows loaded from the partially used final page; unused pages in the serving pool need no such cleanup.^{[6]Reference 6FlashAttention Source Snapshot, October 10, 2026https://github.com/Dao-AILab/flash-attention/tree/94e22c906678e5483fa0e9e24d8e787bc2c0ed4c}

FlashAttention handles the attention math and page lookup. A serving engine still owns allocation, eviction, prefix sharing, and request scheduling. Don't treat the kernel's paged path as a complete KV-cache manager. vLLM made that seam explicit, and FlashInfer widens it for serving-shaped work.

Forward can discard the score matrix. Training still needs gradients for $Q$, $K$, and $V$. If those probabilities were never stored, where does backward get them?

Backward recomputes tiles on purpose

Backward rebuilds the probability tiles instead of loading them. Saving the full probability matrix would recreate the quadratic memory problem. FlashAttention saves $Q$, $K$, $V$, output $O$, row-wise LSE, and any required RNG state, then rebuilds score and probability tiles during backward.

For a nonempty row without dropout, let $S=QK^T/\sqrt d$, $P=\operatorname{softmax}(S)$, and $G=\partial\mathcal L/\partial O$. Then $dV=P^TG$, $dP=GV^T$, and

The row sum broadcasts over keys. Since $O=PV$, its value is also $\operatorname{rowsum}(G\odot O)$, computable from saved output without retaining $P$. Rebuild a probability tile as $\exp(S_{ij}-\operatorname{LSE}_i)$, apply the same mask, and accumulate gradients. Masked probabilities and their gradients are zero. Dropout requires reproducing the forward random mask, not drawing a new one.^{[1]Reference 1FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135[7]Reference 7FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691}

View larger

Recomputing costs extra FLOPs and HBM reads. It avoids storing and loading a full $P$, which can let long-context training fit within memory. Measure the extra work on the target shape and hardware to determine its effect on step time.

FlashAttention-2 partitions query tiles across thread blocks and reduces communication between warps.^{[7]Reference 7FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691} The FA2 deterministic flag applies to backward and can cost memory and time. Test kernel, device, and dropout RNG changes separately when comparing reproducibility; the README's forward-determinism statement covers a narrower configuration.^{[2]Reference 2FlashAttentionhttps://github.com/Dao-AILab/flash-attention}

The test contract is numerical, not bitwise. Several pinned FA2 tests bound error by twice a lower-precision PyTorch reference's error against a higher-precision path; some assertions include an extra absolute allowance. Read the specific output and gradient assertions instead of adopting “2× error” as a universal tolerance.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820}

This small example, checked on PyTorch 2.13.0 CPU, compares the derivative above with SDPA's math backend in float64. An explicit keep-mask handles both unequal-length directions, including fully masked rows. True means allowed in SDPA, is_causal=False avoids applying another causal rule, and dropout_p=0.0 removes randomness. The reference materializes scores only because these tensors are tiny. Measure a separate GPU run to assess FlashAttention performance.^{[9]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

import math
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

torch.manual_seed(7)
for query_len, key_len in [(2, 5), (5, 2), (4, 4)]:
    q = torch.randn(1, 2, query_len, 3, dtype=torch.float64, requires_grad=True)
    k = torch.randn(1, 2, key_len, 3, dtype=torch.float64, requires_grad=True)
    v = torch.randn(1, 2, key_len, 4, dtype=torch.float64, requires_grad=True)
    keep = (torch.arange(key_len)[None, :] <=
            torch.arange(query_len)[:, None] + key_len - query_len)
    with sdpa_kernel(SDPBackend.MATH):
        out = F.scaled_dot_product_attention(
            q, k, v, attn_mask=keep, is_causal=False, dropout_p=0.0)
    upstream = torch.randn_like(out)
    actual_grads = torch.autograd.grad(out, (q, k, v), upstream)

    with torch.no_grad():
        scores = (q @ k.transpose(-2, -1) / math.sqrt(3)).masked_fill(~keep, -torch.inf)
        has_key = keep.any(dim=-1, keepdim=True)
        # Don't evaluate softmax(-inf, -inf, ...) on empty rows.
        safe_scores = torch.where(has_key, scores, torch.zeros_like(scores))
        p = torch.softmax(safe_scores, dim=-1).masked_fill(~keep, 0.0)
        expected = p @ v
        dp = upstream @ v.transpose(-2, -1)
        ds = p * (dp - (p * dp).sum(dim=-1, keepdim=True))
        expected_grads = (
            ds @ k / math.sqrt(3),
            ds.transpose(-2, -1) @ q / math.sqrt(3),
            p.transpose(-2, -1) @ upstream,
        )
        torch.testing.assert_close(out, expected, atol=1e-12, rtol=1e-12)
        for actual, reference in zip(actual_grads, expected_grads):
            torch.testing.assert_close(actual, reference, atol=1e-12, rtol=1e-12)
        assert torch.all(ds.masked_select(~keep) == 0)
        assert torch.all(out[..., ~keep.any(dim=-1), :] == 0)
    print(f"Q={query_len}, K={key_len}: output and Q/K/V gradients match")

Q=2, K=5: output and Q/K/V gradients match
Q=5, K=2: output and Q/K/V gradients match
Q=4, K=4: output and Q/K/V gradients match

FA1 through FA4 retune the same IO idea

Each release pairs the tiled operator with a hardware bottleneck. Paper numbers are snapshots from a GPU, dtype, sequence, and baseline. They shouldn't be copied into a service SLO without reproducing those conditions.

The FA4 paper's Figure 4 reports forward BF16 benchmarks at head dimension 128, sequence lengths 1k through 32k, and 32k total tokens per batch. It labels the GPU B200 and compares against cuDNN 9.13 and a Triton implementation, reporting up to 1.3× and 2.7× respectively, with a peak of 1,613 TFLOPs/s. These published measurements come from the paper; reproduce the comparison on the target hardware before making a current vendor ranking. The same caption says a newer cuDNN version reaches similar performance.^{[11]Reference 11FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scalinghttps://arxiv.org/abs/2603.05451}

There is also a reproducibility caveat in the reviewed paper: its main benchmark text names B200, while Appendix A.1 lists B100. Record that inconsistency rather than inventing one clean hardware manifest. A local comparison needs its own exact GPU, software versions, shape, mask, and baseline.

FA4 is implemented in CuTe-DSL embedded in Python. The paper also reports shorter single-kernel compile times than its C++ template baseline, but compile time and kernel runtime answer different questions.^{[11]Reference 11FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scalinghttps://arxiv.org/abs/2603.05451}

That implementation changes more than syntax. It includes software-emulated exponentials, conditional softmax rescaling, tensor memory, and backward matrix-multiply-accumulate (MMA) operations spanning two cooperative thread arrays (CTAs). These are hardware-specific scheduling techniques, not changes to which keys a query may attend.^{[11]Reference 11FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scalinghttps://arxiv.org/abs/2603.05451}

🔬 Research insight: The paper reports doubled tensor-core throughput from H100 to B200, while shared-memory bandwidth and exponential units remain bottlenecks. FA4 redesigns the kernel around those slower units, so the system-level speedup depends on the full workload.

Call SDPA unless you need a repo feature

PyTorch scaled dot-product attention (SDPA) can select among FlashAttention, memory-efficient, and math implementations from the input and device.^{[9]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html} A model library should call that API when it wants portability and let the dispatcher pick a valid backend. It should call the Dao-AILab package when it needs a generation, mask, or layout the framework path doesn't provide.

The fused kernels have input limitations. Forcing only FlashAttention through sdpa_kernel can leave no usable backend, producing diagnostics and a runtime error rather than an automatic math fallback. SDPA's backend name also doesn't identify a separately installed Dao-AILab wheel. Outputs can differ because floating-point operations aren't associative. The math backend supports float64 and keeps intermediates in FP32 for FP16/BF16 inputs.^{[9]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

Set dropout_p=0.0 explicitly for evaluation: SDPA applies the passed probability even when the surrounding model is in evaluation mode. Match scaling, masks, head mapping, and dropout before comparing outputs.

A mask can be valid and mean the opposite

SDPA's boolean attn_mask uses True to allow a key. MultiheadAttention.key_padding_mask uses True to exclude padding. Translate the meaning and broadcast shape when moving between them. Inverting a boolean padding mask makes an SDPA keep-mask; inverting a floating-point additive bias doesn't.^{[9]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

Predict the result below before running it. One real key has value 10; the padded key has value 1000. Both scores are zero, so the single allowed value decides the output. Reusing the padding mask selects the padded value without a shape error. This PyTorch 2.13.0 CPU experiment tests mask interpretation, not a FlashAttention kernel.

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

q = torch.zeros(1, 1, 1, 1, dtype=torch.float64)
k = torch.zeros(1, 1, 2, 1, dtype=torch.float64)
v = torch.tensor([[[[10.0], [1000.0]]]], dtype=torch.float64)
padding = torch.tensor([[False, True]])  # MHA-style: True excludes a key.
with sdpa_kernel(SDPBackend.MATH):
    wrong = F.scaled_dot_product_attention(
        q, k, v, attn_mask=padding[:, None, None, :], dropout_p=0.0)
    correct = F.scaled_dot_product_attention(
        q, k, v, attn_mask=(~padding)[:, None, None, :], dropout_p=0.0)
print("reused padding mask:", wrong.item())
print("SDPA keep mask:", correct.item())
assert wrong.item() == 1000.0
assert correct.item() == 10.0

reused padding mask: 1000.0
SDPA keep mask: 10.0

The repository usage page lists integrations in PyTorch, Transformers, DeepSpeed, Megatron-LM, diffusion systems, and protein-structure models.^{[2]Reference 2FlashAttentionhttps://github.com/Dao-AILab/flash-attention} Those examples cover two workloads:

1. Training: cut activation memory and raise attention throughput so longer sequences or larger batches fit.
2. Inference: accelerate prompt processing, and some decode attention, when the shape reaches an efficient regime. KV-cache layout and the scheduler still dominate many decode traces.

Tiny or unusual tensors can run slower than the default backend. Measure each shape against the backend production would otherwise choose, using the same shape, mask, dtype, and correctness tolerance.

Benchmark the path you will ship

A backend label identifies a code path. Benchmark a shape matrix from real traffic: long-query prefill, one-token decode over a long cache, unequal query and key lengths, MQA or GQA head counts, head dimensions, causal and local masks, and any variable-length packing. A kernel can win for one row and lose for another because tile occupancy and KV reads change.

Warm each implementation with identical inputs before timing. Discard compilation and first-allocation runs, then synchronize the GPU before and after the measured region, or use CUDA events. CUDA launches are asynchronous, so a CPU timer around an unsynchronized call can measure launch time instead of the kernel. Report a distribution such as median and p95, plus peak memory, rather than one favorable run.

Correctness is a gate, not a footnote. Compare the same outputs against a stable FP32 or math-backend reference with recorded absolute and relative tolerances. Exercise equal and unequal lengths, causal corners, local windows, varlen offsets, and paged indices before accepting a speed result. Raising tolerance until a mismatch disappears turns a fast wrong answer into a misleading benchmark.

Test mixed-shape warmup in one process too. FA4's October 7 MLA fix adds the KV head count to its forward compilation key. Previously, 128 query heads with two KV heads and 64 query heads with one KV head shared a GQA ratio of 64, but could reuse a kernel specialized for the wrong packed layout. Testing each shape only in a fresh process could miss that cache interaction. Check A then B and B then A against the reference, alongside isolated runs. The upstream fix establishes this failure mechanism; no GPU reproduction is claimed here.^{[12]Reference 12Key MLA forward compilation on the KV head count, PR 2978https://github.com/Dao-AILab/flash-attention/pull/2978}

When timing changes, profile a representative shape instead of guessing from the kernel name. Nsight Systems can expose CUDA API calls, kernel launches, memory operations, and stream gaps; Nsight Compute can separate memory-workload pressure from compute throughput with metrics and roofline views.^{[13]Reference 13Nsight Systems User Guidehttps://docs.nvidia.com/nsight-systems/UserGuide/[14]Reference 14Nsight Compute Profiling Guidehttps://docs.nvidia.com/nsight-compute/ProfilingGuide/} Use that evidence to decide whether HBM or shared-memory traffic, tensor-core work, launch overhead, occupancy, or register pressure is limiting the path.

Production checklist

Before enabling a FlashAttention backend, record the workload and the hardware:

• Shape: batch, query length, key length, head count, K/V head count, and head dimension.
• Numerics: FP16 or BF16 input, FP32 accumulation where supported, scale, dropout, and error versus a stable reference.
• Mask: causal corner, local window, padding, varlen offsets, and page-table indexing.
• Device: GPU family, CUDA or ROCm version, driver, PyTorch version, and which package actually imported.
• Dispatch: which backend ran, whether a fallback occurred, and why.
• Evidence: p50 and tail latency, memory peak, throughput, and correctness error on representative prompts.

If CUDA, PyTorch, GPU generation, or FlashAttention changes, rerun correctness and performance checks. Test an observable fallback within an explicit memory budget. A math fallback may materialize the quadratic intermediates you removed; at long context it can cause an out-of-memory failure. Reject or reroute shapes that lack a safe path rather than blindly retrying them.

⚠️ Common mistake: Treating "no quadratic allocation" as "constant-time attention." FlashAttention still computes dense query-key pairs, and long-context decode can remain dominated by KV-cache reads and scheduler effects.

Derive the two-tile softmax by hand, name which package a wheel actually imported, and catch the bottom-right mask an SDPA is_causal=True reference would miss. The next chapter takes the batch whose histories aren't one rectangle.

Architectural summary

• FlashAttention computes exact scaled dot-product attention while controlling HBM traffic. The live score workspace is a tile, not $n \times n$.
• Online softmax rescales the old normalizer $\ell$ and numerator $u$ when a later tile raises $m$. Divide only after merging.
• FA2, FA3, and FA4 are different install surfaces. Import paths don't upgrade themselves.
• FA2's causal mask is bottom-right from 2.1, matching queries at the end of a cached sequence. SDPA is_causal=True documents upper-left alignment for non-square inputs.
• Paged KV in the kernel is a page-table lookup. Allocation and eviction stay with the serving engine.
• Backward saves LSE and recomputes tiles. Check both output and gradient tolerances; bitwise equality isn't the contract.
• Production confidence needs a recorded backend, representative measurements, and a memory-safe fallback or rejection path.

What a kernel integration review should defend

• Derive the two-tile numerator and normalizer, including rescaling and an entirely masked row, without materializing the complete score row.
• Identify the installed package and its public return contract; translate between FA2 and SDPA layouts and justify the causal corner and GQA head mapping.
• Separate CPU reference agreement from GPU kernel correctness, and kernel timings from end-to-end service latency; include a memory budget for unsupported shapes.

Review questions

Why can a one-token decode test miss the causal-alignment bug?

If it tests only one key too, both corners produce the same mask. Use one query over several cached keys: bottom-right allows all keys, while upper-left allows only the first. Also test longer queries than keys to expose fully masked rows.

Can a faster FA4 paper result justify replacing FA2 in a service?

No. First match the installed API, return values, device support, shapes, masks, and numerical tolerances. Then benchmark the actual fallback baseline, including compile warmup and cache behavior. A paper's peak throughput neither validates that integration nor predicts service tail latency.

Mask, state, and package checks