FlashAttention Primer: IO-Aware Attention, from FA1 to FA4
On this page
FlashAttention Primer: IO-Aware Attention, from FA1 to FA4
The attention layer is the main bottleneck when scaling Transformers to long sequences: its time and memory cost grow quadratically with sequence length. FlashAttention is a family of algorithms that keeps attention exact while making it fast and memory-efficient, by treating GPU memory traffic — not just FLOPs — as the resource to optimize. This primer walks through the core ideas with original diagrams and derivations: the GPU memory hierarchy, tiling, online softmax, backward recomputation, then the FA1 → FA2 → FA3 → FA4 timeline, the difference between the standalone flash-attn package and PyTorch’s scaled_dot_product_attention, and the hardware caveats that matter in practice.
Two ground rules for this article. First, every performance figure is cited from the primary papers or the official Dao-AILab repository and PyTorch documentation, and each number is tied to the hardware it was measured on; nothing below was benchmarked locally. Second, the article was written on a workstation whose GPU is a TITAN RTX (Turing, compute capability 7.5) — that is not a coincidence, because the official FlashAttention-2/3/4 CUDA kernels cannot run on Turing, and the last section turns that limitation into an honest verification plan.
All “as of” statements refer to 2026-08-09.
Why attention is quadratic#
For a single head, with query, key, value matrices $Q, K, V \in \mathbb{R}^{N \times d}$ ($N$ tokens, head dimension $d$), scaled dot-product attention is:
$$S = \frac{Q K^{\top}}{\sqrt{d}}, \qquad P = \operatorname{softmax}(S) \text{ (row-wise)}, \qquad O = P V,$$
where $O \in \mathbb{R}^{N \times d}$ is the output. The arithmetic cost is $O(N^2 d)$ FLOPs: two matrix multiplications of $N \times N \times d$ each. That part is inherent — the memory cost is what makes naive attention painful. The score matrix $S$ and the probability matrix $P$ both have $N^2$ entries, and a standard implementation materializes both in HBM (GPU main memory).
A worked example, pure arithmetic: with $N = 8192$ tokens, head dimension $d = 128$, in fp16 (2 bytes), one head needs $2 \times 8192^2 \times 2$ bytes ≈ 268 MB just to hold $S$ and $P$. With 32 heads that is ≈ 8.6 GB — before the model weights, KV cache, or activations. At $N = 64,K$ (a long-context setting), one head alone needs ≈ 17 GB, and 32 heads ≈ 550 GB, which no current GPU holds. This is why approximate attention (sparse, low-rank, hashing) became popular: it tried to cut the $N^2$ cost. FlashAttention takes a different route — it computes the same matrices, but never lets them live in HBM.
The GPU memory hierarchy and IO-awareness#
A GPU has two relevant memory levels. HBM is large (tens of GB) and comparatively slow: roughly 1.5 TB/s on an A100 and 3.35 TB/s on an H100 (NVIDIA datasheets), and ~672 GB/s on the TITAN RTX. SRAM is the on-chip memory right next to the compute units: a few tens to a few hundred KB per streaming multiprocessor (the FlashAttention v1 README cites 64 KB on a T4; NVIDIA datasheets list up to 164 KB usable shared memory per block on A100 and 228 KB per SM on H100). SRAM is an order of magnitude faster than HBM but far too small to hold an $N \times N$ matrix.
Figure 1. GPU memory hierarchy and the FlashAttention dataflow.
+------------------------------------------------------------------+
| HBM - main memory, tens of GB |
| Q, K, V, O (A100 ~1.5 TB/s, H100 ~3.35 TB/s, |
| TITAN RTX ~672 GB/s) |
| |
| | load tiles (Q_i, K_j, V_j) ^ write O blocks |
| v | |
| +----------------------------------------------+ |
| | SRAM - on-chip, per SM, ~64-228 KB | |
| | S_ij, P_ij, running m, l, O_i live here | |
| | attention computed locally, block by block | |
| +----------------------------------------------+ |
+------------------------------------------------------------------+
The FLOP count is the same as standard attention; what changes is where the intermediate matrices are stored. The FA1 paper’s key observation is that many operations are memory-bound: their runtime is determined by the number of memory accesses, not by arithmetic. A standard attention pass is exactly such a case:
Figure 2. HBM traffic, standard attention vs FlashAttention (forward).
Standard attention: FlashAttention:
Q, K --HBM--> S = QK^T / sqrt(d) Q, K, V stream in as blocks;
S --HBM--> rowmax/rowsum (softmax) S_ij and P_ij are computed and
P --HBM--> O = P V consumed inside SRAM, never
HBM traffic: Theta(Nd + N^2) written out
extra memory: O(N^2) HBM traffic: Theta(N^2 d^2 / M)
extra memory: O(N)
Standard attention writes $S$ to HBM, reads it back for softmax, writes $P$, reads it back for the multiply with $V$ — four quadratic round trips. FlashAttention’s IO complexity analysis (Theorem 2 of 2205.14135 ) shows standard attention needs $\Theta(Nd + N^2)$ HBM accesses, while FlashAttention needs $\Theta(N^2 d^2 / M)$, where $M$ is the SRAM size. For typical $d = 64$–$128$ and $M$ in the tens-to-hundreds of KB, $M \gg d^2$, so FlashAttention performs many times fewer HBM accesses. The paper also proves a matching lower bound: no exact attention algorithm can asymptotically beat this over all SRAM sizes.
FlashAttention-1: tiling plus online softmax#
The obstacle to fusing everything into one kernel is softmax: row $i$ of $P$ depends on the maximum and the sum over all $N$ scores in that row, so a naive implementation must see the full row before emitting any output. FlashAttention (Algorithm 1 of 2205.14135 ) solves this by combining two classic techniques — tiling and an online softmax — plus recomputation for the backward pass (next section).
Tiling. Split $Q$ into row blocks $Q_i$ of size $B_r \times d$, and $K, V$ into column blocks $K_j, V_j$ of size $B_c \times d$. The paper sets $B_c = \lceil M / 4d \rceil$ and $B_r = \min(\lceil M / 4d \rceil, d)$, so that the four matrices $Q_i, K_j, V_j, S_{ij}$ (each $B_r \times d$, $B_c \times d$, $B_c \times d$, $B_r \times B_c$) fit in SRAM together. The kernel loops over all $(i, j)$ tiles, computing $S_{ij} = Q_i K_j^{\top} / \sqrt{d}$ and folding the result into a running output — no $N \times N$ matrix ever exists.
Figure 3. Tiling of the attention computation. Each tile S_ij is
computed, softmaxed, and consumed inside SRAM.
K blocks (each B_c x d) ->
+----------------------------------+
| S_11 S_12 S_13 ... S_1,nc | inside one tile:
| S_21 S_22 S_23 ... S_2,nc | K_j (B_c x d) V_j (B_c x d)
| ... | \ /
| S_nr1 S_nr2 ... S_nr,nc | S_ij = Q_i K_j^T / sqrt(d)
+----------------------------------+ P_ij = softmax(S_ij)
Q blocks (each B_r x d) O_i += P_ij V_j
Online softmax. The standard trick is to write softmax in “max-shifted” form and keep only a running maximum and normalization. Let $m^{(j)}$ be the running row-max after processing key blocks $1 \dots j$, $\ell^{(j)}$ the running sum of $\exp(s_k - m^{(j)})$, and $O^{(j)}$ the running unnormalized output $\sum_{k \le j} \exp(s_k - m^{(j)}) v_k$. The invariant is that at every step the true output would be $O^{(j)} / \ell^{(j)}$; the final step divides once. When a new block with scores $s_{j+1}$ arrives, we first raise the new max, $m’ = \max(m^{(j)}, \max s_{j+1})$, then rescale the old state by $\alpha = e^{m^{(j)} - m’}$ and the new block by $\beta = e^{\max s_{j+1} - m’}$:
$$\ell^{(j+1)} = \alpha, \ell^{(j)} + \beta, \ell_{j+1}, \qquad O^{(j+1)} = \alpha, O^{(j)} + \beta, \tilde{P}{j+1} V{j+1},$$
where $\tilde{P}{j+1} = \exp(S{j+1} - \max s_{j+1})$ and $\ell_{j+1} = \sum \tilde{P}{j+1}$ row-wise. Why this is correct: the old rows’ exponentials were all scaled by $e^{m^{(j)}}$ relative to the old max; changing the reference max from $m^{(j)}$ to $m’$ multiplies every old exponential by exactly $\alpha$, and every new-block exponential by $\beta$. The output $O$ is a weighted sum of value vectors with those same exponential weights, so it must be rescaled identically. A short induction on $j$ confirms the invariant: after any prefix of blocks, $O^{(j)} = \sum{k \le j} \exp(s_k - m^{(j)}) v_k$, hence the final $O / \ell$ is exactly $\operatorname{softmax}(S) V$.
Figure 4. Online-softmax state update when a new key block arrives.
state before: new block j+1: state after:
m = running max m' = max(m, max s_j+1) m <- m'
l = sum exp(s_k - m) a = exp(m - m') l <- a*l + b*l_j+1
O = sum exp(s_k - m) v_k b = exp(max s_j+1 - m') O <- a*O + b*P~ V_j+1
Causal masking is nearly free: for causal attention, tiles with $j > i$ are skipped (or masked with $-\infty$), halving the work. The FA1 paper also extends the scheme to block-sparse attention, which the paper reports is faster than existing approximate attention methods while remaining exact on the kept blocks.
The forward pass therefore uses linear extra memory — Theorem 1 of the paper: Algorithm 1 returns $\operatorname{softmax}(QK^{\top})V$ using $O(N)$ extra memory beyond the inputs and output — and $\Theta(N^2 d^2 / M)$ HBM accesses instead of $\Theta(Nd + N^2)$.
Backward pass: recomputation instead of storage#
Training also needs gradients. The naive backward needs $P$ (or $S$) again; storing it costs $O(N^2)$ memory. FlashAttention instead stores only the small per-row statistics — $O$ and the log-sum-exp of each row, $m_i$ — and recomputes each $S_{ij}$ block inside the backward kernel. For an output gradient $dO$, with $\tilde{P}{ij} = \exp(S{ij} - m_i)$:
$$dV_j = \tilde{P}{ij}^{\top} dO_i, \qquad dS{ij} = \tilde{P}{ij} \odot \left( dO_i V_j^{\top} \right), \qquad dQ_i \mathrel{+}= dS{ij} K_j, \qquad dK_j \mathrel{+}= dS_{ij}^{\top} Q_i.$$
Figure 5. Backward pass: recompute, don't store.
forward stores: O, m (per-row log-sum-exp) -- NOT the N x N matrix
backward: reload Q, K, V blocks
S_ij = Q_i K_j^T / sqrt(d) (recomputed)
P_ij = exp(S_ij - m_i) (recomputed)
dV_j = P_ij^T dO_i
dS_ij = P_ij o (dO_i V_j^T)
dQ_i += dS_ij K_j, dK_j += dS_ij^T Q_i
This trades a bit of extra arithmetic for dramatically fewer HBM accesses; the paper notes the recomputed backward pass is faster, not slower, precisely because HBM traffic dominates. The official repository’s flash_attn_func documents the determinism consequence: the forward pass is always deterministic, while the default backward pass accumulates $dQ$/$dK$ across thread blocks with atomic adds, which is bitwise nondeterministic run to run; a deterministic=True flag exists and is documented as “slightly slower and uses more memory.“ FlashAttention-4’s paper later lists “reduce atomic adds in the backward pass” as one of its explicit design goals.
FlashAttention-2: better work partitioning#
The FA2 paper (2307.08691 ) starts from a sobering measurement: FA1 reached only 25–40% of the GPU’s theoretical FLOPs/s — far below what optimized GEMMs achieve. Its diagnosis: suboptimal partitioning of work between thread blocks and warps, causing low occupancy and unnecessary shared-memory traffic. Three fixes follow directly:
- Fewer non-matmul FLOPs: rescale the running statistics with cheaper math (e.g., dividing instead of re-exponentiating where possible), so a larger fraction of cycles go to tensor-core matmuls.
- Parallelize over the sequence dimension, even for one head: in FA1 a single head’s output was computed by a limited number of blocks; FA2 lets the grid cover $(batch, heads, \text{seqlen}_q \text{ blocks})$, so long sequences get many more thread blocks (higher occupancy) — and the forward pass needs no atomics because each output tile belongs to exactly one block.
- Better intra-block warp partition: instead of splitting the rows of the output among warps (which forces shared-memory round trips), split the columns of $K, V$ across warps so each warp’s partial results can be combined with less communication.
Figure 6. FA2 forward: a CTA grid over (batch, heads, seqlen_q blocks).
grid: batch x heads x (N / B_r) thread blocks
seqlen_q blocks ->
+----+----+----+----+----+
| C | C | C | C | C | each C owns one B_r x d output tile,
| C | C | C | C | C | streams K/V blocks through SRAM,
+----+----+----+----+----+ no atomics in the forward pass
The abstract reports “around 2× speedup compared to FlashAttention, reaching 50–73% of the theoretical maximum FLOPs/s on A100,“ and end-to-end GPT-style training at up to 225 TFLOPs/s per A100 (72% model-FLOPs utilization). The official README’s changelog shows the release line then grew into a full inference toolkit: varlen (unpadded) entry points, sliding-window attention, ALiBi, a deterministic backward, paged KV caches, softcapping, and torch.compile compatibility. On the hardware side the README is explicit: FA2 CUDA kernels support Ampere, Ada, and Hopper (sm80+); for Turing it points to a separate community repository, ssiu/flash-attention-turing , “which supports a core subset of FlashAttention features on Turing."
FlashAttention-3: Hopper asynchrony and FP8#
FA3 (2407.08608 ) targets the H100, where FA2 achieves only ~35% utilization. The paper identifies the cause as FA2’s failure to exploit Hopper’s new hardware capabilities, and contributes three techniques:
- Warp specialization + TMA: dedicate some warps to data movement via the Tensor Memory Accelerator (asynchronous bulk copies from HBM into shared memory) and other warps to tensor-core math, so memory transfer and compute overlap instead of alternating.
- Interleaving matmul and softmax: software-pipeline the two GEMMs of one iteration with the softmax of the previous one, so the tensor cores never wait on the exponential/softmax units.
- FP8 with block quantization and incoherent processing: quantize scores in blocks and apply a random (incoherent) rotation before quantization, which the paper reports reduces FP8 numerical error; FA3’s FP8 is validated at 2.6× lower error than a baseline FP8 attention.
Figure 7. FA3 warp-specialized pipeline on Hopper.
producer warps: TMA: HBM -> SRAM (fetch K_{j+1}, V_{j+1} while ...)
consumer warps: MMA: S_ij = Q_i K_j^T | softmax(S_ij)
MMA: O_i += P_ij V_j | (interleaved, pipelined)
time -> overlap: compute tile j while fetching tile j+1
Reported numbers (H100, from the abstract): 1.5–2.0× speedup over FA2, up to 740 TFLOPs/s with FP16 (75% utilization), and close to 1.2 PFLOPs/s with FP8. The official repository still labels FA3 a beta release as of 2026-08-09: it requires an H100/H800, CUDA ≥ 12.3 (12.8 recommended), and currently ships FP16/BF16 forward+backward with FP8 forward only.
FlashAttention-4: Blackwell co-design#
FA4 (2603.05451 , March 2026) is the response to a hardware curveball: on Blackwell (B200/GB200), tensor-core throughput roughly doubled while shared-memory bandwidth and the exponential units scaled much less — so FA3’s Hopper-tuned pipeline leaves the new bottlenecks on the table. FA4’s techniques are a direct mapping onto that asymmetry:
- Redesigned pipelines with fully asynchronous MMA operations and larger tiles, so the doubled tensor cores stay saturated without waiting on shared memory.
- Software-emulated exponential and conditional softmax rescaling, moving work off the (barely-scaled) special-function units.
- Tensor memory and 2-CTA MMA mode, cutting shared-memory traffic and — explicitly — the atomic adds in the backward pass.
The paper reports up to 1.3× over cuDNN 9.13 and 2.7× over a Triton baseline on B200 with BF16, reaching 1613 TFLOPs/s (71% utilization). A notable engineering shift: FA4 is implemented entirely in CuTeDSL (a Python-embedded DSL), with the paper reporting 20–30× faster compile times than C++ template kernels, and the official README installs it as pip install flash-attn-4, targeting Hopper and Blackwell (an optional cu13 extra for CUDA 13).
Table 1 summarizes the timeline. Every number is as cited in the respective paper; none was reproduced here.
Table 1. The FlashAttention timeline (all figures cited, not reproduced).
| Generation | Paper / release | Target hardware | Key ideas | Reported performance |
|---|---|---|---|---|
| FA1 | 2205.14135 , NeurIPS 2022 | Ampere; Turing also supported in the v1 release | tiling, online softmax, backward recomputation, block-sparse extension | 2–4× vs standard attention on A100 (seq 128–4K); 15% BERT-large wall-clock gain vs the MLPerf 1.1 record; 189 TFLOPs/s per A100 (60.6% MFU) |
| FA2 | 2307.08691 , ICLR 2024 | Ampere, Ada, Hopper (sm80+) | work partitioning, fewer non-matmul FLOPs, seqlen parallelism | ~2× vs FA1; 50–73% of A100 peak FLOPs/s; 225 TFLOPs/s per A100 (72% MFU) |
| FA3 | 2407.08608 , 2024 (beta) | Hopper H100/H800 | TMA, warp specialization, matmul/softmax interleaving, FP8 | 1.5–2.0× vs FA2 on H100; 740 TFLOPs/s FP16 (75%); ~1.2 PFLOPs/s FP8 |
| FA4 | 2603.05451 , 2026 | Hopper + Blackwell (B200/GB200) | CuTeDSL, async MMA pipelines, emulated exp/rescaling, tensor memory, 2-CTA MMA | 1.3× vs cuDNN 9.13, 2.7× vs Triton on B200 BF16; 1613 TFLOPs/s (71%) |
Standalone flash-attn versus PyTorch SDPA#
Two very different ways to reach FlashAttention-style kernels, frequently confused:
- The standalone package (Dao-AILab/flash-attention ) is the kernel library from the papers. It exposes
flash_attn_func,flash_attn_qkvpacked_func,flash_attn_varlen_func,flash_attn_with_kvcache, and friends, and its README documents a long feature list: dropout, causal and sliding-window masks, ALiBi, MQA/GQA, paged KV caches, fused rotary embeddings, and thedeterministicbackward flag. It installs by compiling CUDA kernels (pip install flash-attn --no-build-isolation; needs a CUDA/ROCm toolkit, PyTorch ≥ 2.2, Linux); FA3 and FA4 ship as separate packages (flash-attn-3from thehopper/subdirectory,flash-attn-4). - PyTorch SDPA (
torch.nn.functional.scaled_dot_product_attention) is a framework-level API with an automatic dispatcher. The PyTorch 2.13 docs list three backends — a FlashAttention-2-style kernel, a memory-efficient (xFormers-style) kernel, and a C++ math implementation — plus a cuDNN backend in recent versions. Backends are selected automatically and can be forced or disabled withtorch.nn.attention.sdpa_kernel()ortorch.backends.cuda.enable_flash_sdp()/enable_mem_efficient_sdp()/enable_math_sdp(). The flash backend is PyTorch’s own implementation (the source even documents “FlashAttentionV2 requires that head dimension be a multiple of 8”), not the Dao-AILab package, and it has its own constraint list: no explicitattn_mask(only theis_causalflag), causal + non-square sequence lengths rejected, head dimension ≤ 256.
The practical distinction: SDPA gives you one call that picks a safe kernel for your inputs on whatever GPU you have; the standalone package gives you the full kernel feature set, and the two are not interchangeable drop-ins — the docs warn that “the output of this function may be different depending on what backend kernel is chosen,“ and that the math backend keeps intermediates in fp32 for fp16/bf16 inputs, which is why SDPA results can differ slightly across backends and across GPUs.
Hardware, dtype, determinism, and padding caveats#
Table 2 collects what actually supports what, from the official READMEs and the PyTorch v2.13 dispatch code (sdp_utils.cpp , which gates the flash backend to GPU architectures in the range sm80–sm121, and the mem-efficient backend to sm50–sm121).
Table 2. Support matrix for FlashAttention-family kernels (as of 2026-08-09).
| Implementation | Min GPU (CC) | fp16 | bf16 | fp8 | Head dim | Notes |
|---|---|---|---|---|---|---|
| FA1 v1 (official) | Turing sm75+ | yes | yes, sm80+ | no | ≤ 128, multiple of 8; backward > 64 needs A100/H100 | varlen API; 128/256-token blocks |
| FA2 (official CUDA) | Ampere sm80+ | yes | yes | no | up to 256 | Turing users directed to the community Turing repo |
| FA3 (official, beta) | Hopper H100/H800 | yes | yes | forward only | per kernel | CUDA ≥ 12.3, 12.8 recommended |
| FA4 (official) | Hopper + Blackwell | — | yes (benchmarked) | — | per kernel | pip install flash-attn-4; CuTeDSL |
| PyTorch SDPA flash | sm80–sm121 | yes | yes | sm90+, if FA3 enabled | ≤ 256 | no attn_mask; non-square causal rejected |
| PyTorch SDPA mem-efficient | sm50–sm121 | yes | sm80+ | no | alignment 8 (fp16) / 4 (fp32) | on pre-sm80 GPUs restricted to fp16/fp32 |
| PyTorch SDPA math | any (CPU included) | yes | yes | no | any | fp32 accumulation; fp64 supported |
| ssiu/flash-attention-turing | sm75 | yes | see repo | no | core subset | community port named by the official README |
Caveats that bite in practice:
- Dtype. The fused kernels are fp16/bf16 affairs. bf16 needs sm80+ (both FA1’s README and PyTorch’s dtype gate agree); below sm80, PyTorch’s fused kernels accept fp16 only. fp32 runs only through the math or mem-efficient backends. FP8 is FA3-era territory (e4m3) and forward-only in the official beta; FA4’s headline number is BF16.
- Determinism. Standalone FA forward is always deterministic; the default backward is bitwise nondeterministic (atomic accumulation), with
deterministic=Trueas the opt-in. PyTorch’s docs note backend-dependent outputs and that the cuDNN path may select nondeterministic algorithms. If you need reproducible runs, force the math backend or use the deterministic flag — at a measured-in-your-workload performance cost. - Padding and shapes. Head dimension must be a multiple of 8 for the fused kernels (FA1 asserts it; PyTorch pads head dims to a multiple of 8 in the composite layer). FA1 caps head dim at 128 (and at 64 for backward on non-A100/H100); FA2 and PyTorch’s flash backend allow up to 256. Causal +
seqlen_q != seqlen_kis rejected by PyTorch’s flash backend. The standalone package handles variable-length sequences without padding viaflash_attn_varlen_func(and the FA1-eracu_seqlensAPI), and FA4’s headline is explicitly un-padded attention; the cuDNN SDPA backend historically requiredseq_kvmultiples of 64 (before cuDNN 8.9.6) and head-dim multiples of 8, so padded shapes that work in one backend can silently fall back in another.
What can run on this machine (Turing), and honest checks we propose#
This article’s workstation has a TITAN RTX (TU102, compute capability 7.5, 24 GB HBM2, ~672 GB/s per NVIDIA’s spec sheet). Folding that into Table 2:
- Runs: the official FA1 v1 kernels (fp16; head dim ≤ 128, and ≤ 64 for backward) — FA1 is the only generation whose official CUDA kernels list Turing support; PyTorch SDPA’s mem-efficient (fp16/fp32) and math backends; the community
ssiu/flash-attention-turingport. - Cannot run: the official FA2, FA3, and FA4 CUDA kernels — FA2 requires Ampere or newer (sm80+), FA3 requires H100/H800, FA4 requires Hopper/Blackwell. PyTorch’s SDPA flash backend is likewise gated to sm80–sm121 and will not be selected on this GPU, and bf16 fused kernels are unavailable (bf16 needs sm80+).
Rather than report anything we did not run, here are four proposed verification checks for Turing, with the outcome each should produce according to the cited sources. None of them was executed while writing this article; treat them as a concrete plan, not results.
- Backend probe (PyTorch). Call
F.scaled_dot_product_attentioninsidesdpa_kernel([SDPBackend.FLASH_ATTENTION])on this GPU. Expected: the dispatcher refuses, with the warning text from the v2.13 source — “Flash attention only supports gpu architectures in the range [sm80, sm121]“ — then repeat with[SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]and verify outputs agree with the math backend within fp16 tolerance. - Algorithm reference (eager PyTorch). Implement Algorithm 1 (tiling + online softmax + final rescale, per Figure 4 above) in pure torch ops, fp16, head dim 64/128, seq 512/2048, causal and not, and compare against the SDPA math backend on the same tensors; mirror the official repo’s test criterion that the max error is at most ~2× the baseline implementation’s error. This validates the derivation in this article on Turing without any fused kernel.
- Triton on Turing (pinned release). Adapt the OpenAI 06-fused-attention.py tutorial kernel, which the official FA README points to as the readable reference implementation, and run it on the TITAN RTX with a Triton 2.1.x release — its README lists “NVIDIA GPUs (Compute Capability 7.0+)“. Note the dated caveat: current Triton (main, 2026) requires compute capability 8.0+, so this check specifically needs the older release; expected outcome is a working fp16 kernel whose output matches the math backend.
- Community Turing kernel. Install
ssiu/flash-attention-turing(the repository the official FA2 README names for Turing) and benchmark its forward/backward against SDPA mem-efficient on the TITAN RTX, verifying correctness first and only then comparing timings.
If you run any of these on a Turing GPU, the interesting questions are exactly the ones the papers’ IO analysis predicts: how much of the speedup survives on a GPU with 64 KB-class SRAM and no async-copy hardware (FA1’s T4 measurements suggest 2.5–4.5× on the forward pass is plausible territory), and how the mem-efficient cutlass kernels compare to the tiled online-softmax formulation at equal precision.
Summary#
- Attention’s quadratic memory cost, not its FLOPs, is the original bottleneck; FlashAttention keeps attention exact by being IO-aware: it tiles $Q, K, V$, computes $S$ and $P$ in SRAM, and never materializes the $N \times N$ matrices (HBM traffic $\Theta(N^2 d^2 / M)$ vs $\Theta(Nd + N^2)$, per Theorem 2 of the FA1 paper).
- Online softmax (running max, running sum, rescale) makes fusion possible; recomputation makes the backward pass memory-lean and, despite extra FLOPs, faster.
- FA2 fixed work partitioning (~2× over FA1; 50–73% of A100 peak); FA3 exploited Hopper’s TMA/warp-specialization/FP8 (740 TFLOPs/s FP16, ~1.2 PFLOPs/s FP8 on H100); FA4 re-co-designed for Blackwell’s asymmetric scaling (1613 TFLOPs/s BF16 on B200, 1.3× over cuDNN 9.13, in CuTeDSL).
- Standalone
flash-attnis the kernel library; PyTorch SDPA is the auto-dispatching API with its own FA2-style backend and stricter constraints. They are not drop-in equivalents. - Hardware reality: official FA2/FA3/FA4 kernels need Ampere/Hopper/Blackwell. On this workstation’s TITAN RTX (Turing, sm75), FA1 v1 and PyTorch’s mem-efficient/math backends are the honest options; the four checks above are proposed, not yet executed — no local benchmark numbers are claimed in this article.
References#
- Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness . arXiv:2205.14135, NeurIPS 2022.
- Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning . arXiv:2307.08691, ICLR 2024.
- Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision . arXiv:2407.08608, 2024.
- Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, Tri Dao. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling . arXiv:2603.05451, 2026.
- Dao-AILab. flash-attention: official implementation of FlashAttention and FlashAttention-2 (with FA3/FA4 releases) . GitHub repository (README and v1.0.9 tag cited).
- PyTorch. torch.nn.functional.scaled_dot_product_attention and torch.nn.attention.sdpa_kernel . PyTorch 2.13 documentation.
- PyTorch. aten/src/ATen/native/transformers/cuda/sdp_utils.cpp (v2.13.0) . Backend capability gates cited.
- Shengqi Chen (ssiu). flash-attention-turing: FlashAttention for Turing GPUs . Community repository named by the official FA2 README.
- Triton. triton-lang/triton README compatibility notes (main: NVIDIA compute capability 8.0+; v2.1.0 tag: 7.0+) and 06-fused-attention.py tutorial .