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. This deep dive connects that caller contract to the tiled computation inside FlashAttention.

The FlashAttention & Memory Efficiency lesson explained why attention can wait on GPU memory traffic. Here, follow the repository from Python interfaces to online softmax, backward recomputation, and tests. You'll distinguish a mathematically equivalent schedule from a compatible package, mask, and tensor layout.

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}

What stays mathematically unchanged when FlashAttention is used?

Answer

The operator is still softmax of scaled $QK^T$ multiplied by $V$. FlashAttention changes scheduling and memory movement, not the attention definition.

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 are documentation checks as of September 2, 2026, not commands executed on a GPU here:^{[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 is pinned to commit ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820. 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.

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.^{[6]Reference 6FlashAttention-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. This is an IO bound under that memory model, not a byte count or a runtime prediction for every later kernel.^{[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:

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.

Why can a fused attention kernel still be compute-heavy at long context?

Answer

It avoids storing the full score matrix, but it still evaluates the allowed query-key pairs. IO savings remove a large memory cost; they don't make the quadratic dot-product work disappear.

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:

$m_{\text{new}}=\max(m_{\text{old}},\max(s))$

$\ell_{\text{new}}=e^{m_{\text{old}}-m_{\text{new}}}\ell_{\text{old}}+\sum_k e^{s_k-m_{\text{new}}}$

$u_{\text{new}}=e^{m_{\text{old}}-m_{\text{new}}}u_{\text{old}}+\sum_k e^{s_k-m_{\text{new}}}V_{j,k}$

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.^{[7]Reference 7Online 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$:

$\ell_2=0.3679\times 1.3679+(1+e^{-3})=1.5530$

$0.3679[3.6788,10]+[5,5]+e^{-3}[9,1]=[6.8014,8.7286]$

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.

What error would appear if a later tile had a larger maximum but the old numerator wasn't rescaled?

Answer

The old contributions would be compared in the wrong exponential scale. Their relative probability would be too large, so the result wouldn't equal dense softmax. Rescaling by $e^{m_{\text{old}}-m_{\text{new}}}$ is the invariant that makes tile order exact.

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.

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[8]Reference 8torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/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. The pinned wrapper attempts that path after a HIP extension import failure, with a warning. It still needs compatible AIter dependencies and hardware; the attempt isn't a guarantee of successful fallback.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820}

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. Those are FA2-specific qualifications, not universal limits for every generation.^{[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}

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.

Causal 2.1 is the decode gotcha

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):

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.^{[8]Reference 8torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/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.

⚠️ 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 is unbounded, not a negative radius. Causal masking can further restrict that set. A smaller window changes the operator and may affect model quality; fewer allowed scores don't guarantee lower latency for every tile shape.^{[5]Reference 5FlashAttention Source Snapshot (ce088ab)https://github.com/Dao-AILab/flash-attention/tree/ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820}

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 is that API's constraint, not a universal paged-cache rule. 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}

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.

Backward recomputes tiles on purpose

Training needs gradients for $Q$, $K$, and $V$. 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

$dS=P\odot\left(dP-\operatorname{rowsum}(P\odot dP)\right),\qquad dQ=\frac{dS K}{\sqrt d},\quad dK=\frac{dS^TQ}{\sqrt d}.$

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[6]Reference 6FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691}

Recomputing costs extra FLOPs and HBM reads. It avoids storing and loading a full $P$, so long-context training can fit when a materializing kernel doesn't. Memory headroom isn't a free speedup: the extra work still has to be measured.

FlashAttention-2 partitions query tiles across thread blocks and reduces communication between warps.^{[6]Reference 6FlashAttention-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. “Forward is deterministic” in the README doesn't promise identical results across different kernels, devices, or dropout RNG states.^{[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. It does not execute or benchmark FlashAttention.^{[8]Reference 8torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html}

Why does backward recomputation often make sense for training?

Answer

Saving a quadratic probability matrix can exceed activation memory at long context. Recomputing score tiles spends extra arithmetic and reads, but keeps only linear-size output and LSE state, often enabling a larger batch or sequence length.

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 are published measurements, not results reproduced here. The same caption says a newer cuDNN version reaches similar performance, so the ratios aren't a current vendor ranking.^{[10]Reference 10FlashAttention-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.^{[10]Reference 10FlashAttention-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.^{[10]Reference 10FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scalinghttps://arxiv.org/abs/2603.05451}

🔬 Research insight: FA4's claim isn't "Blackwell is twice as fast, so attention is twice as fast." Tensor-core throughput doubled from H100 to B200 in that paper's comparison, while shared-memory bandwidth and exponential units didn't. The kernel is redesigned around those slower units.

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.^{[8]Reference 8torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/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.^{[8]Reference 8torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/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.

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.

"Flash" isn't a guarantee that a tiny or unusual tensor wins. Measure 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 isn't a measurement. Build 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.

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.^{[11]Reference 11Nsight Systems User Guidehttps://docs.nvidia.com/nsight-systems/UserGuide/[12]Reference 12Nsight 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.

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