flash-attention

Part I: The mental model

flash_attn_func(q, k, v, causal=True)        Python API         flash_attn/flash_attn_interface.py
        |
FlashAttnFunc.apply                          autograd binding   same file, torch.autograd.Function
        |
torch op "flash_attn::_flash_attn_forward"   custom-op shim     same file
        |
flash_attn_2_cuda.fwd                        pybind11 entry     csrc/flash_attn/flash_api.cpp (mha_fwd)
        |
run_mha_fwd -> template instantiation        static dispatch    csrc/flash_attn/src/static_switch.h
        |
flash_fwd_kernel, one CTA per Q tile         CUDA kernel        csrc/flash_attn/src/flash_fwd_kernel.h
  loop over KV tiles in SRAM:                                   csrc/flash_attn/src/softmax.h
  QK^T -> mask -> online softmax -> PV
        |
O (bf16/fp16) + logsumexp LSE (fp32)         HBM write          consumed later by mha_bwd

The one-sentence identity: flash-attention is a fused CUDA kernel family that computes exact softmax attention by streaming over the key/value sequence in tiles that fit in on-chip SRAM, so the N by N attention matrix is never written to GPU main memory. Everything else in the repository, the Python wrappers, the C++ dispatch, the per-architecture rewrites, exists to deliver that one loop to as many shapes, dtypes, and GPUs as possible.

The mental model has three layers. At the top is a thin Python package, flash_attn/, whose job is to look like a normal PyTorch function: it validates tensors, binds forward and backward into autograd, and calls into a compiled extension. In the middle is one C++ translation unit per kernel family, csrc/flash_attn/flash_api.cpp for FlashAttention-2, which owns the contract: what layouts are accepted, what gets padded, which of the hundred-plus precompiled kernel instantiations to launch. At the bottom are the kernels themselves, written against NVIDIA's CUTLASS and its CuTe layout algebra, where the algorithm finally meets the memory hierarchy.

Two facts make the whole design make sense. First, attention on real hardware is limited by memory traffic, not arithmetic, so the win comes from never touching slow memory with the quadratic intermediate. Second, softmax can be computed incrementally, one block of scores at a time, by carrying a running maximum and a running sum and rescaling past work when the maximum moves. The first fact says what to optimize; the second says the optimization is possible without approximating anything.

Part II: Using it

Installing

On Linux with an NVIDIA GPU (Ampere, Ada, or Hopper for the main package; Turing has a community fork), the install is one line, and the flag matters because the build reads your already-installed torch:

pip install ninja
pip install flash-attn --no-build-isolation

If no prebuilt wheel matches your Python, torch, and CUDA combination, pip compiles the CUDA sources, which takes a few minutes on a large machine with ninja and hours without it. On machines with many cores but modest RAM the parallel compile can exhaust memory, and the README's fix is to cap the job count:

MAX_JOBS=4 pip install flash-attn --no-build-isolation

There is no macOS build and cannot be one, because the package is CUDA: on a Mac you study the code and use PyTorch's built-in scaled_dot_product_attention, which implements the same algorithm on its supported backends. The honest macOS setup for this chapter is a clone for reading plus a rented Linux GPU for the labs:

git clone https://github.com/Dao-AILab/flash-attention.git

First session

import torch
from flash_attn import flash_attn_func

q = torch.randn(2, 4096, 16, 64, dtype=torch.bfloat16, device="cuda")
k = torch.randn(2, 4096, 16, 64, dtype=torch.bfloat16, device="cuda")
v = torch.randn(2, 4096, 16, 64, dtype=torch.bfloat16, device="cuda")

out = flash_attn_func(q, k, v, causal=True)
print(out.shape, out.dtype)   # torch.Size([2, 4096, 16, 64]) torch.bfloat16

Note the layout: flash-attention takes (batch, seqlen, nheads, headdim), while PyTorch's own attention APIs take (batch, nheads, seqlen, headdim). Mixing these up is the single most common beginner failure, and it often does not crash, because both layouts have four dimensions; it silently computes attention over the wrong axis when shapes happen to be compatible. The wrong and right versions side by side:

# WRONG: PyTorch SDPA layout fed to flash_attn_func
q = torch.randn(2, 16, 4096, 64, ...)   # (B, H, S, D)
out = flash_attn_func(q, k, v)          # treats 16 as seqlen, 4096 as nheads

# RIGHT: flash-attn layout, or transpose explicitly
q = torch.randn(2, 4096, 16, 64, ...)   # (B, S, H, D)
out = flash_attn_func(q, k, v)

The second classic mistake is dtype. The kernels support fp16 and bf16 only (FlashAttention-3 adds an fp8 forward), so a float32 tensor raises a runtime error; cast at the boundary and keep your master weights in whatever precision your training recipe wants.

Grouped-query attention comes for free

If K and V have fewer heads than Q, the kernel broadcasts them: with 16 query heads and 2 KV heads, query heads 0 through 7 attend to KV head 0 and the rest to KV head 1. This is how MQA and GQA models run without any reshaping:

q = torch.randn(2, 4096, 16, 64, dtype=torch.bfloat16, device="cuda")
k = torch.randn(2, 4096, 2, 64, dtype=torch.bfloat16, device="cuda")
v = torch.randn(2, 4096, 2, 64, dtype=torch.bfloat16, device="cuda")
out = flash_attn_func(q, k, v, causal=True)   # works, no repeat_interleave

Variable-length batches: stop padding

Real batches are ragged, and the naive treatment pads every sequence to the longest one, then computes attention over the padding. flash-attention's answer is flash_attn_varlen_func: concatenate all sequences into one long packed tensor with no batch dimension, and pass cumulative-length offsets so the kernel knows where each sequence begins:

import torch
from flash_attn import flash_attn_varlen_func

lens = [512, 37, 1024]
total = sum(lens)
q = torch.randn(total, 16, 64, dtype=torch.bfloat16, device="cuda")
k, v = q.clone(), q.clone()
cu = torch.tensor([0, 512, 549, 1573], dtype=torch.int32, device="cuda")

out = flash_attn_varlen_func(
    q, k, v,
    cu_seqlens_q=cu, cu_seqlens_k=cu,
    max_seqlen_q=max(lens), max_seqlen_k=max(lens),
    causal=True,
)

Padding tokens are simply never computed, which for realistic length distributions is a large fraction of the batch's FLOPs. There is deliberately no boolean attention-mask argument anywhere in the API; if you find yourself building one, the library is telling you to pack instead.

Decoding against a KV cache

flash_attn_with_kvcache serves the inference shape: one or a few new query tokens attending to a long cache, with optional in-place cache append, rotary embedding application, and a paged cache via a block_table argument, which is exactly the shape vLLM-style paged serving needs. Most people consume all of this indirectly: PyTorch routes torch.nn.functional.scaled_dot_product_attention to a flash backend when it can, and Hugging Face models accept attn_implementation="flash_attention_2" in from_pretrained, so the repository's reach is far larger than its import count.

Part III: When it is the right tool

Use the package directly when you are writing model code and need the full contract: varlen packing, KV-cache decoding, sliding windows (window_size), soft capping (softcap), ALiBi slopes, or determinism control in the backward. Training frameworks and long-context finetuning are the core audience, and any transformer with head dimension up to 256 in fp16 or bf16 on Ampere or newer qualifies.

The main alternative for everyday work is not another repository, it is PyTorch itself: scaled_dot_product_attention ships flash-style kernels with zero extra dependencies, and if you do not need varlen or kvcache specifics it is the right default. xFormers offers memory_efficient_attention with a similar CUTLASS lineage. For inference serving, FlashInfer specializes in paged KV and decode-shaped workloads, and cuDNN ships its own fused attention that compilers like torch.compile can pick. On AMD, this same repository builds a ROCm backend from csrc/flash_attn_ck (Composable Kernel) or a Triton path.

The architecture-shaped warning is about the build, not the math. flash-attn compiles against the C++ ABI of one specific torch build, so treating it as a loosely-pinned dependency is how production images die at import time with undefined-symbol errors: a base image upgrades torch, the cached flash-attn wheel was compiled against the old one, and nothing fails until runtime. Pin torch and flash-attn as a pair, build wheels once in CI, and never let two versions of that pair drift independently.

SAFE                                      DANGEROUS
torch==X and flash-attn wheel             requirements.txt: torch (floating)
built against X, shipped together,        + pip cache with flash-attn built
upgraded together in one commit           against last month's torch
                                          -> ImportError: undefined symbol
                                             ...c10...

A second, quieter version of the same warning: padded rectangular batches. A serving or training system that pads everything to max_seqlen and masks is architecturally simple and quadratically wasteful, and because flash-attention makes attention fast, the waste hides in the profile. Packed varlen batches are the safe architecture; the padded rectangle is the NFS-mounted SQLite file of this chapter.

Part IV: The full life of one call

The canonical operation: flash_attn_func(q, k, v, causal=True) on a (2, 4096, 16, 64) bf16 tensor on an A100, forward and then backward. Every stage below names the file that owns it, all under the repository root.

Stage 1: the Python entry point

flash_attn/flash_attn_interface.py defines flash_attn_func with the full public contract in its signature: dropout_p, softmax_scale (defaulting to 1/sqrt(headdim)), causal, window_size, softcap, alibi_slopes, deterministic. The module's import block is worth reading first: it imports the compiled extension flash_attn_2_cuda under the neutral name flash_attn_gpu, falling back to a Triton implementation on ROCm, which is the whole multi-backend story in twenty lines.

Stage 2: autograd and the custom op

The function does no math; it calls FlashAttnFunc.apply, a torch.autograd.Function in the same file. Its forward invokes _flash_attn_forward, which is registered as a torch custom op (flash_attn::_flash_attn_forward) with a "fake" shape-propagation twin so torch.compile can trace through it without running the kernel. The op calls flash_attn_gpu.fwd(...) and returns four things: out, softmax_lse, an optional dropout mask tensor, and the RNG state. ctx.save_for_backward keeps q, k, v, out, softmax_lse, rng_state, and that list is the recomputation strategy stated as code: no attention matrix is saved, only a float32 logsumexp of shape (batch, nheads, seqlen_q), one number per query row.

Stage 3: the C++ contract

flash_attn_gpu.fwd is pybind11 for mha_fwd in csrc/flash_attn/flash_api.cpp (the module exposes exactly five entry points: fwd, varlen_fwd, bwd, varlen_bwd, fwd_kvcache). This function is the contract made explicit: it checks dtype and device, requires the last dimension contiguous, pads head dimension up to a multiple of 8, allocates the output and the LSE tensor, and fills a plain Flash_fwd_params struct with every pointer, stride, and scale the kernel will need. Reading it answers most "does it support X" questions faster than the README.

Stage 4: static dispatch to a precompiled kernel

run_mha_fwd selects the kernel with compile-time switches from csrc/flash_attn/src/static_switch.h: macros expand booleans and the head-dimension bucket into template arguments, landing on run_mha_fwd_<elem_type, kHeadDim, Is_causal>. Each combination lives in its own tiny .cu file, over a hundred of them, with names that are the dispatch table spelled out: flash_fwd_hdim128_bf16_causal_sm80.cu and its siblings for head dims 32 through 256, fp16 and bf16, causal and not. Splitting per file exists purely to parallelize compilation. For decode shapes with tiny query length there is a separate split-KV path (run_mha_fwd_splitkv_dispatch) that parallelizes over the KV dimension instead and merges partials.

Stage 5: launch geometry

csrc/flash_attn/src/flash_fwd_launch_template.h computes the grid, and this single line is FlashAttention-2's headline improvement over version 1:

const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
dim3 grid(num_m_block, params.b, params.h);

One thread block per (query tile, batch, head). Parallelizing over the query-sequence dimension, not just batch times heads, is what keeps a long-context, small-batch workload from leaving most of the GPU idle. Tile sizes (kBlockM, kBlockN) come from kernel_traits.h and vary with head dimension and architecture, on the order of 64 to 128 queries by 32 to 128 keys.

Stage 6: the main loop

csrc/flash_attn/src/flash_fwd_kernel.h is the heart. Each thread block loads its Q tile into SRAM once, then iterates over KV tiles from the last block backward. The loop body is four steps: copy the K tile and compute the score tile QK^T with tensor cores, apply masking (causal, local window, ALiBi) from mask.h, fold the scores into the running softmax state via softmax.template softmax_rescale_o(...) from softmax.h, then multiply the exponentiated tile against the V tile and accumulate into the output registers. The code structurally separates the iterations that need mask checks (n_masking_steps) from the clean interior iterations so the common path pays nothing for masking. The score tile lives and dies in registers and shared memory; it is never written out.

Stage 7: epilogue

After the last tile, normalize_softmax_lse divides the accumulated output by the final softmax denominator, converts to the output dtype, and writes two things to HBM: the O tile, and the per-row logsumexp into params.softmax_lse_ptr. The kernel that computed a 4096 by 4096 attention matrix per head wrote 4096 floats of evidence that it ever existed.

Stage 8: the backward, briefly

When loss.backward() reaches this node, autograd calls mha_bwd, which first launches a preprocess kernel (flash_bwd_preprocess_kernel.h, compute_dot_do_o) computing the row-wise dot product of dO and O, then the main backward kernel (flash_bwd_kernel.h) which recomputes each score tile from Q, K, and the saved LSE, and accumulates dQ, dK, dV. Part VII walks the algebra; the point at this altitude is that the backward trades recomputed FLOPs for never having stored the forward's quadratic intermediate.

Part V: Deep dive: the IO argument

A GPU is two memory systems glued to arithmetic units. Using the A100 figures from the FlashAttention paper: HBM offers 40 to 80 GB at 1.5 to 2.0 TB/s, while on-chip SRAM is 192 KB per each of 108 streaming multiprocessors with aggregate bandwidth around 19 TB/s, an order of magnitude faster and about four orders of magnitude smaller. Whether a kernel is limited by compute or by memory is decided by its arithmetic intensity, FLOPs per byte moved, against the machine's ratio. Big matrix multiplies do O(n³) work on O(n²) data and saturate the tensor cores. Elementwise ops and softmax do O(n) work on O(n) data and are pure memory traffic.

Standard attention implements a memory-bound pattern at quadratic scale. It computes S = QK^T and writes all N² scores to HBM, reads them back to compute P = softmax(S), writes P, then reads P again to compute PV. In HBM-access terms that is Θ(Nd + N²) reads and writes, and for realistic sizes the N² term towers over everything: at N = 4096, d = 64, one head's score matrix is 16.8 million entries against 262 thousand entries of Q, K, V combined, a 64-to-1 ratio of intermediate traffic to actual input. The matmul FLOPs were never the bottleneck; shuttling the score matrix through slow memory three times was.

FlashAttention's analysis counts memory accesses instead of FLOPs and shows the tiled algorithm needs Θ(N²d²/M) HBM accesses, where M is the SRAM size. With M around 100 KB and d = 64, that is roughly a 10 to 20 times reduction in traffic, which matched the observed wall-clock speedups, and the paper proves no exact attention algorithm can do asymptotically better over the whole range of SRAM sizes. Two corollaries matter in practice. Memory for the attention op drops from O(N²) to O(N), because the score matrix never materializes, which is what makes 32k and 128k contexts feasible at all. And the speedup grows with sequence length, since the eliminated term is the quadratic one; at short sequences flash and naive attention are close, which surprises people benchmarking at N = 512.

The trap to correct explicitly: FlashAttention is not an approximation. It computes the identical mathematical function, exactly, in exact arithmetic; years of sparse and low-rank attention research were aimed at FLOPs, the resource that was not scarce. Floating-point outputs differ from a naive implementation in the last bits because summation order differs, which is why the repository's own tests compare against a reference within tolerance rather than bit-for-bit.

Part VI: Deep dive: tiling and the online softmax

The obstacle to tiling attention was always the softmax denominator, which sums exp(s_j - m) over an entire row of scores using the row maximum m for numerical safety. A tile of scores seems useless until the whole row is known. The escape is a recurrence: keep a running maximum m and running sum l, and when a new tile arrives with its own maximum, rescale everything computed so far by exp(m_old - m_new). Milakov and Gimelshein worked out this recurrence for plain softmax in the online normalizer paper, and FlashAttention extends it by carrying the output accumulator through the same rescaling. I derive this fully on the softmax page; here is the algebra worked once numerically, one query row, four scores split into two tiles, [1, 3] then [2, 5]:

Tile 1: scores [1, 3]
  m1 = 3
  l1 = e^(1-3) + e^(3-3) = 0.1353 + 1 = 1.1353
  acc = 0.1353*v1 + 1.0000*v2          (unnormalized)

Tile 2: scores [2, 5]
  m2 = max(3, 5) = 5
  correction c = e^(3-5) = 0.1353      (old work was scaled by wrong max)
  l2 = c*l1 + e^(2-5) + e^(5-5)
     = 0.1536 + 0.0498 + 1 = 1.2034
  acc = c*acc + 0.0498*v3 + 1.0000*v4

Check against one pass over [1,3,2,5], m=5:
  e^-4 + e^-2 + e^-3 + e^0 = 0.0183+0.1353+0.0498+1 = 1.2034  ✓
Final: out = acc / l2,  LSE = m2 + ln(l2) = 5.1851

Each KV tile is loaded to SRAM, used, and discarded; the state carried between tiles is one maximum, one sum, and the unnormalized output accumulator per query row. That is the entire replacement for N² stored scores. The kernel encodes this in softmax.h as softmax_rescale_o, with two production details worth knowing. First, it computes exponentials as exp2f(x * log2(e) * scale - m * log2(e) * scale) rather than expf, folding the softmax scale and the base conversion into one fused multiply so the hardware's fast base-2 exponential unit does the work; that is why the params struct carries scale_softmax_log2. Second, a row whose running maximum is negative infinity (fully masked, which happens legitimately in causal cross-attention shapes) is detected and produces zeros instead of NaNs.

One more identity closes the loop: instead of saving m and l separately for the backward, the kernel saves the single value LSE = m + log(l), the log-sum-exp itself, in float32. Any probability can then be reconstructed later as p_ij = exp(s_ij - LSE_i), no second maximum needed. That is the (batch, nheads, seqlen_q) float32 tensor from Stage 2, and it is also what split-KV decoding and ring-attention style distributed tricks use to merge partial attentions computed over disjoint key ranges: two partial (out, LSE) pairs combine into one by exactly the tile 2 arithmetic above.

Misconception to retire: the rescaling does not accumulate error across hundreds of tiles in any way that matters. The corrections are exact identities in real arithmetic, accumulation happens in float32 regardless of input dtype, and the repository's test suite (tests/test_flash_attn.py) asserts the kernel's error versus an fp64-ish reference stays within a small multiple of the baseline fp16 implementation's own error.

Part VII: Deep dive: the backward pass and recomputation

Backpropagating through attention needs the probability matrix P twice: dV = P^T dO, and dS = P ∘ (dP - D) where dP = dO V^T. The forward refused to store P, so the backward recomputes it, tile by tile, from Q, K, and the saved LSE: p_ij = exp(q_i·k_j·scale - LSE_i). This is the second act of the compute-for-memory trade, the same exchange as activation checkpointing but inside a single operator, and on memory-bound hardware the recomputation is the cheap side of the ledger.

The term D deserves its own sentence because it is the trick that makes the tiled backward local. The softmax Jacobian couples a whole row: dS_ij = P_ij (dP_ij - Σ_k P_ik dP_ik). That row sum looks like it needs the full row of P again, but it equals rowsum(dO ∘ O), computable from tensors we already have. So the backward begins with a small preprocess kernel, compute_dot_do_o in csrc/flash_attn/src/flash_bwd_preprocess_kernel.h, writing D once per query row, and every tile of the main backward (flash_bwd_kernel.h) is then independent:

preprocess:  D_i = rowsum(dO_i ∘ O_i)                    one pass, linear
main loop over (KV tile j):        (parallelized over KV blocks)
  recompute S_j  = Q K_j^T
  recompute P_j  = exp(S_j*scale - LSE)                   uses saved fp32 LSE
  dV_j += P_j^T dO
  dP_j  = dO V_j^T
  dS_j  = P_j ∘ (dP_j - D)
  dK_j += dS_j^T Q
  dQ   += dS_j K_j        <- crosses tiles: atomicAdd

The dQ line explains the API's deterministic flag: the backward parallelizes over KV blocks, so multiple thread blocks contribute to the same dQ rows, and by default they combine with atomic adds, whose ordering varies run to run. That is why gradients from flash-attention are correct but not bit-reproducible by default; passing deterministic=True switches to a slower deterministic accumulation. Teams chasing exact loss-curve reproducibility hit this constantly and suspect a bug; it is summation order, the same phenomenon as the forward's last-bit differences, amplified because gradients feed back into training.

Cost accounting, so the trade is concrete: the backward performs roughly 2.5 times the forward's matmul FLOPs (five tile matmuls versus two) plus the recomputation of S, yet runs at the same order of speed as the forward because it, too, never touches HBM with a quadratic tensor. Storing P instead would cost N² times heads times layers of memory and the bandwidth to write and read it; recomputing costs FLOPs the tensor cores had to spare. The saved-tensors list in FlashAttnFunc is this paragraph in one line of Python.

Part VIII: Deep dive: one algorithm, four generations of kernels

The repository is organized by hardware generation, and reading its directory layout is reading a history of NVIDIA architectures. Everything below csrc/ is compiled extension code: csrc/flash_attn/ is FlashAttention-2 proper (flash_api.cpp plus src/ with the kernel headers and the wall of per-shape .cu instantiation files), csrc/cutlass/ is NVIDIA's CUTLASS vendored as a submodule, csrc/flash_attn_ck/ binds AMD's Composable Kernel implementation, and csrc/fused_dense_lib/ and csrc/layer_norm/ are auxiliary fused ops for the training stack under flash_attn/models/ and flash_attn/modules/.

The FA2 kernels are written against CuTe, CUTLASS's layout algebra, where tensor shapes, strides, and the partitioning of tiles across warps and threads are types manipulated algebraically instead of hand-written index arithmetic. That is what keeps one algorithm maintainable across head dimensions 32 to 256, two dtypes, causal and local masking, dropout, and forward plus backward: the algorithm is written once over abstract layouts, and C++ templates instantiate concrete variants, one .cu file each so ninja can compile them in parallel. Setup.py targets sm80 and sm90, with sm100 and sm120 added on CUDA 12.8 or newer toolkits.

FlashAttention-3 lives in hopper/ as its own package with its own setup.py (H100/H800, CUDA 12.3 or newer required). Same mathematics, new choreography: Hopper's TMA engine copies tiles asynchronously, warpgroup-wide WGMMA instructions run matmuls asynchronously, and FA3's kernels use warp specialization, in mainloop_fwd_sm90_tma_gmma_ws.hpp, so producer warps move data while consumer warps compute, with softmax of one tile overlapped against the matmuls of the next; a ping-pong schedule across warpgroups hides the softmax latency further, and an fp8 forward path with block quantization recovers accuracy that naive fp8 would lose. The FA3 paper reports 1.5 to 2 times FA2's speed on H100, up to about 740 TFLOPs/s in fp16, near 75 percent utilization. The hopper/instantiations/ directory holds 451 generated .cu files, produced by generate_kernels.py: at this variant count, the dispatch table is itself generated code. The fourth generation is visible in flash_attn/cute/, a Python CuTe-DSL implementation targeting Hopper and Blackwell, the README's "FlashAttention-4" direction, which trades C++ templates for Python-generated kernels.

The misconception to retire here: "FlashAttention" is not one kernel you install but a family indexed by (architecture, dtype, head dim, mask type, forward/backward), plus independent reimplementations, in PyTorch SDPA, Triton, cuDNN, FlashInfer, and every serious inference engine. When someone reports "flash attention is slow" the first question is which implementation on which chip at which shape, because a decode-shaped call (seqlen_q of 1) that misses the split-KV path, for instance, wastes almost the whole GPU regardless of how good the kernel is.

Part IX: Reading the repository

A staged plan, top altitude first. Every path below exists in the repository as of v2.8.4.

Stage 0: the contract from outside

Read README.md (usage, requirements, the FA3 section) and the docstrings of flash_attn_func and flash_attn_varlen_func in flash_attn/flash_attn_interface.py, including the bottom-right causal-mask alignment diagrams. You should be able to answer: what layouts and dtypes are accepted; how GQA is expressed; what causal means when seqlen_q differs from seqlen_k; what varlen's cu_seqlens encode.

Stage 1: the Python layer completely

Read all of flash_attn/flash_attn_interface.py: the import of flash_attn_2_cuda, the custom-op wrappers, the FlashAttnFunc family of autograd Functions and what each saves for backward. Then skim flash_attn/modules/mha.py to see the kernels embedded in a full attention module, and flash_attn/bert_padding.py for the pad/unpad utilities that feed varlen. Questions: what exactly is saved between forward and backward, and what is recomputed? Why is the LSE float32? Where does dropout's RNG state live?

Stage 2: the C++ contract

Read csrc/flash_attn/flash_api.cpp: just mha_fwd and mha_bwd, plus the pybind block at the bottom. Note every shape check, the head-dim padding to multiples of 8, and run_mha_fwd's dispatch. Questions: which five functions does the extension export; what triggers the split-KV path; where do kernel launch parameters come from?

Stage 3: the forward kernel with the derivation beside you

Read csrc/flash_attn/src/flash_fwd_kernel.h main loop with softmax.h open in a second pane and the recurrence from Part VI on paper. Match each stage: tile load, QK^T gemm, mask, softmax_rescale_o, PV accumulate, epilogue with normalize_softmax_lse. Then flash_fwd_launch_template.h for the grid and kernel_traits.h for the tile shapes. Questions: why two loops (masking steps versus clean steps); what state persists across KV iterations; why exp2?

Stage 4: backward, then generations

Read flash_bwd_preprocess_kernel.h (small) and as much of flash_bwd_kernel.h as you can stomach, then switch to hopper/flash_fwd_kernel_sm90.h and hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp to see the same loop re-choreographed with producer/consumer warps. tests/test_flash_attn.py is ground truth throughout: every variant is checked against a plain PyTorch reference, which also makes it the best catalog of what is actually supported.

Where not to start

Do not start in hopper/, whose asynchronous pipeline obscures the algorithm under Hopper-specific machinery; do not start in the CUTLASS submodule, which is a semester of its own; and do not start by reading the hundred instantiation .cu files, which contain one line each. The readable spine is interface, api, forward kernel, in that order.

Part X: Hands-on labs

Labs 1, 2, 3, and 5 need a CUDA GPU (Ampere or newer); labs 4 and 6 run anywhere. Numbers shown vary with hardware and versions; shapes and trends should reproduce.

Lab 1: see the memory argument (concept: Part V)

Compare peak memory of the math backend versus the flash backend through PyTorch SDPA, no flash-attn install needed:

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

def peak(backend, S):
    q = torch.randn(1, 16, S, 64, dtype=torch.float16, device="cuda")
    torch.cuda.reset_peak_memory_stats()
    with sdpa_kernel(backend):
        F.scaled_dot_product_attention(q, q, q, is_causal=True)
    torch.cuda.synchronize()
    return torch.cuda.max_memory_allocated() / 2**20

for S in [1024, 2048, 4096, 8192]:
    print(S, round(peak([SDPBackend.MATH], S)),
             round(peak([SDPBackend.FLASH_ATTENTION], S)))

Observe: math-backend peak memory grows roughly 4x per doubling of S (the N² matrix, times heads), flash grows roughly linearly. At 8192 the math backend allocates gigabytes for one layer's attention; flash allocates megabytes.

Lab 2: verify exactness and its limits (concept: Parts V, VI)

import torch
from flash_attn import flash_attn_func

torch.manual_seed(0)
q = torch.randn(1, 2048, 8, 64, dtype=torch.bfloat16, device="cuda")
k, v = torch.randn_like(q), torch.randn_like(q)

out = flash_attn_func(q, k, v, causal=True)

qf, kf, vf = [t.transpose(1, 2).float() for t in (q, k, v)]
s = qf @ kf.transpose(-1, -2) / 64**0.5
s = s.masked_fill(torch.ones(2048, 2048, device="cuda").triu(1).bool(), float("-inf"))
ref = (s.softmax(-1) @ vf).transpose(1, 2)

print((out.float() - ref).abs().max())   # ~1e-2 scale for bf16 inputs
print(torch.equal(out.float(), ref))     # False, and that is expected

Observe: max error is at the level bf16 rounding forces on any implementation, and exact equality is false. The right mental model: same function, different summation order.

Lab 3: touch the logsumexp (concept: Parts IV, VI)

import torch
from flash_attn.flash_attn_interface import _flash_attn_forward

q = torch.randn(2, 1024, 8, 64, dtype=torch.bfloat16, device="cuda")
k, v = torch.randn_like(q), torch.randn_like(q)
out, lse, _, _ = _flash_attn_forward(q, k, v, 0.0, 64**-0.5, causal=False,
                                     window_size_left=-1, window_size_right=-1,
                                     softcap=0.0, alibi_slopes=None,
                                     return_softmax=False)
print(lse.shape, lse.dtype)   # torch.Size([2, 8, 1024]) torch.float32

# reconstruct row 0's probabilities from LSE and check they sum to 1
s = (q[0, :, 0].float() @ k[0, :, 0].float().T) * 64**-0.5
p = torch.exp(s[0] - lse[0, 0, 0])
print(p.sum())                # ~1.0

Observe: the entire memory of the forward pass, per row, is one float, and it suffices to rebuild any attention probability. (Internal function; argument order may shift between versions, check the signature in your installed flash_attn_interface.py.)

Lab 4: online softmax by hand (concept: Part VI, no GPU)

import numpy as np

def online_softmax_matvec(s, v, block=4):
    m, l, acc = -np.inf, 0.0, np.zeros_like(v[0])
    for i in range(0, len(s), block):
        sb, vb = s[i:i+block], v[i:i+block]
        m_new = max(m, sb.max())
        c = np.exp(m - m_new)
        p = np.exp(sb - m_new)
        l = c * l + p.sum()
        acc = c * acc + p @ vb
        m = m_new
    return acc / l, m + np.log(l)

rng = np.random.default_rng(0)
s, v = rng.normal(size=64), rng.normal(size=(64, 8))
ref = np.exp(s - s.max()); ref /= ref.sum()
out, lse = online_softmax_matvec(s, v)
print(np.abs(out - ref @ v).max())   # ~1e-16

Observe: a dozen lines reproduce the kernel's core recurrence to machine precision. Change block and confirm the answer does not.

Lab 5: the cost of padding (concept: Parts II, III)

Build a ragged batch (lengths drawn from, say, a lognormal capped at 4096), time flash_attn_func on the padded rectangle versus flash_attn_varlen_func on the packed tensor using torch.cuda.Event timing. Observe: speedup approximately equals total padded tokens over total real tokens (squared, for the attention part); with realistic length skew, 2 to 4 times is common. This is Lab-sized proof that the missing mask argument is a feature.

Lab 6: census of the kernel zoo (concept: Part VIII, no GPU)

git clone --depth 1 https://github.com/Dao-AILab/flash-attention.git
cd flash-attention
ls csrc/flash_attn/src/ | grep -c '\.cu$'        # ~100 instantiation files
ls csrc/flash_attn/src/flash_fwd_hdim128*        # read the name grammar
ls hopper/instantiations | wc -l                 # ~451, and generated:
head -30 hopper/generate_kernels.py

Observe: the dispatch dimensions (direction, head dim, dtype, causal, arch) are literally the file-name grammar, and FA3's variant count forced the jump from hand-written files to a generator script.

Part XI: Questions and model answers

Q1. What problem does FlashAttention solve, in one sentence?

Standard attention is bottlenecked by reading and writing the N by N score matrix through GPU main memory, and FlashAttention computes the exact same function while keeping that matrix entirely in on-chip SRAM tiles, cutting HBM traffic from quadratic to roughly linear-in-N terms.

Q2. Is it an approximation?

No. In exact arithmetic the output is identical to naive attention. Floating-point results differ in the last bits because summation order differs, which is also true between any two BLAS libraries. The repository's tests compare against a PyTorch reference within tolerance for exactly this reason.

Q3. Why can softmax be computed without seeing the whole row?

Because max and sum-of-exponentials are both associative once you carry the maximum: given partial (m, l) statistics for a prefix and a new block, the merged statistics are m' = max, and the old l and output accumulator are multiplied by exp(m - m') before the new block's terms are added. Applied per tile, this is the online softmax; the full derivation is on my softmax page.

Q4. What does the forward save for the backward, and why is it enough?

Q, K, V, the output O, the RNG state for dropout, and a float32 logsumexp per query row. P is rebuilt tile by tile as exp(QK^T·scale - LSE), and the softmax-Jacobian row sum is recovered as rowsum(dO ∘ O), computed by a preprocess kernel, so no quadratic tensor is ever stored.

Q5. Why are training runs with flash-attention not bit-reproducible by default?

The backward parallelizes over KV blocks and accumulates dQ with atomic adds, whose ordering is nondeterministic. The deterministic=True flag buys reproducibility with a slower accumulation strategy. The gradients are correct either way.

Q6. What changed between FlashAttention-1 and 2?

The algorithm stayed; the parallelization and work partitioning changed. FA2 launches a grid over (query blocks, batch, heads) instead of just (batch, heads), so long sequences at small batch fill the GPU, rebalances work between warps to cut shared-memory traffic, and trims non-matmul FLOPs. You can see the headline change as one grid-dimension line in flash_fwd_launch_template.h.

Q7. What is FlashAttention-3, at concept level?

A Hopper-specific rewrite in hopper/: warp specialization with TMA-driven asynchronous copies so data movement overlaps compute, softmax overlapped with the next tile's WGMMA matmuls, and an fp8 forward with block quantization. Per the FA3 paper, 1.5 to 2 times FA2 on H100, up to roughly 75 percent of peak in fp16.

Q8. When does flash-attention not help much?

Short sequences, where the quadratic term it eliminates is small; decode steps with one query token, unless the split-KV path is engaged, since a single query row cannot fill a grid dimension; and workloads already dominated by MLP or communication time, where attention is a minor slice of the step.

Q9. Why does the API have no attention-mask tensor?

Arbitrary dense masks would force reading an N² mask, reintroducing the traffic the kernel exists to avoid. The supported structures, causal, sliding window, ALiBi, and variable-length packing via cu_seqlens, are all expressible as index arithmetic inside the kernel. If you want padding masks, pack with the varlen API instead.

Q10. What does the causal flag do when seqlen_q and seqlen_k differ?

The mask is aligned to the bottom-right of the score matrix: the last query attends to all keys, which is the correct convention for incremental decoding where queries are the newest tokens. The docstring in flash_attn_interface.py draws the exact picture, and rows that end up fully masked produce zeros.

Q11. How does one repository serve so many head dims, dtypes, and GPUs?

The algorithm is written once over CuTe layouts; C++ templates instantiate each (head dim, dtype, causal) variant into its own .cu file for parallel compilation, boolean flags become template parameters through static_switch.h macros, and each hardware generation gets its own kernel family (sm80 in csrc/flash_attn, sm90 in hopper/, AMD via Composable Kernel or Triton).

Q12. What is the logsumexp trick for combining partial attentions?

Attention over a union of key sets is a weighted combination of attentions over the parts, with weights derivable from the parts' LSEs; carrying (out, LSE) makes attention mergeable. Split-KV decoding uses this to parallelize one query over KV chunks and combine, and distributed ring-attention schemes use the same algebra across devices.

Q13. Your import fails with an undefined C++ symbol. What happened?

The flash-attn binary was compiled against a different torch build than the one installed; the extension links torch's C++ ABI, so the pair must be built and upgraded together. Rebuild with pip install flash-attn --no-build-isolation against the current torch, and in production pin the pair and ship CI-built wheels.

Q14. Compare using this repo directly versus PyTorch SDPA.

SDPA is zero-dependency, compiler-friendly, and picks a good backend automatically, the right default for standard shapes. The repo's own API adds the production edges: varlen packing, KV-cache with paging and in-place append, sliding window, softcap, ALiBi, determinism control, and earlier access to new-generation kernels. Frameworks tend to start with SDPA and graduate per feature.

Q15. What would you measure first if attention seems slow?

Shape and backend: confirm which implementation actually ran (SDPA can silently fall back to the math backend when a constraint fails, for instance an unsupported mask or dtype), then whether the shape is decode-like (seqlen_q of 1, needs split-KV), whether batching is padded rather than packed, and only then kernel-level metrics like achieved occupancy and DRAM throughput against the roofline.

Part XII: Design lessons

Count the scarce resource, not the obvious one. Years of attention research optimized FLOPs while the wall clock was set by memory traffic. The roofline discipline, classify the kernel, then optimize the binding resource, is the same reasoning behind external-memory algorithms in databases and cache-oblivious data structures; the B-tree and FlashAttention are answers to the same question asked of different hierarchies.

Recomputation is a first-class storage strategy. Saving one float per row and recomputing tiles beats storing the quadratic intermediate. The same trade appears in activation checkpointing, in databases replaying a WAL instead of storing every page version, and in build systems that cache hashes rather than artifacts. The general form: when bandwidth is the bottleneck and compute is idle, store the seed, not the crop.

Concentrate cleverness behind a boring contract. The public surface is one function over plain tensors; the complexity lives beneath a stable seam, which is why PyTorch, Hugging Face, and vLLM could all adopt it without redesign. This is the syscall pattern, and it is also the reason model definition code in projects like vLLM stays readably plain: the cleverness was pushed into the primitive.

When variants multiply, generate the code. The repository's dispatch is file names; FA2 maintains about a hundred hand-listed instantiations, and FA3's 451 come from generate_kernels.py. Knowing when to switch from writing instances to writing the generator is the same judgment that produces SQL query generators, protobuf compilers, and BLAS autotuners.

Test against a slow oracle. Every kernel variant is validated against a naive PyTorch reference within a tolerance that is itself calibrated to the reference's own floating-point error. Differential testing against an obviously-correct implementation is the only sane way to test fast code, the same method SQLite applies with its fuzzed query oracle and CPUs apply against architectural simulators.

Expect the fastest kernel to have a shelf life. The algorithm survived three hardware generations; the kernels did not, and the repository is structured to absorb that: per-arch directories, a shared algorithm story, tests as the invariant. Designing code so the durable part (the math, the contract) is separated from the perishable part (the choreography) is the transferable move.

Part XIII: Memorization framework

One sentence: flash-attention computes exact attention by streaming KV tiles through SRAM with an online softmax, saving only a per-row logsumexp and recomputing the rest in the backward, reimplemented per GPU generation.

func -> autograd -> api -> dispatch -> tile loop -> O + LSE -> (bwd: D, recompute, dQKV)

func       flash_attn/flash_attn_interface.py   flash_attn_func
autograd   same file                            FlashAttnFunc, saves q,k,v,out,LSE,rng
api        csrc/flash_attn/flash_api.cpp        mha_fwd: checks, pads, params struct
dispatch   csrc/flash_attn/src/static_switch.h  + ~100 per-shape .cu files
tile loop  csrc/flash_attn/src/flash_fwd_kernel.h + softmax.h  QK^T, mask, rescale, PV
bwd        flash_bwd_preprocess_kernel.h        D = rowsum(dO∘O), then flash_bwd_kernel.h

Memorize these blocks:

LAYOUT: (batch, seqlen, nheads, headdim); fp16/bf16 only; head dim <= 256;
GQA = fewer KV heads; causal mask aligns BOTTOM-RIGHT when seqlen_q != seqlen_k.

ONLINE SOFTMAX STATE per query row: running max m, running sum l, unnormalized acc.
New tile: c = exp(m_old - m_new); l = c*l + sum(exp(s - m_new)); acc = c*acc + P·V.
Saved for backward: LSE = m + ln(l), float32, shape (B, H, seqlen_q).

IO COUNTS (paper): naive Θ(Nd + N²) HBM accesses; flash Θ(N²d²/M), M = SRAM size;
~10-20x less traffic at typical d, M. Memory O(N²) -> O(N). FLOPs slightly UP.

GRID (FA2): dim3 grid(num_m_block, batch, heads); seqlen parallelism is the FA2 story.
FA3 story: hopper/, warp specialization, TMA + WGMMA overlap, FP8 forward.

BACKWARD: D = rowsum(dO∘O); P = exp(S·scale - LSE); dV = PᵀdO; dS = P∘(dP - D);
dQ via atomicAdd => nondeterministic by default; deterministic=True to fix.

Part XIV: Papers and further reading

The kernels in this repository compress a decade of attention research into one loop, and each paper below rewards a direct read. Where this site derives the same idea in depth, the companion link points there.

1. Dao et al., FlashAttention, Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022. The founding paper, with the HBM-access analysis of Part V and the tiled forward and backward this repository ships. The efficient attention note on this site places it among the approximate methods it displaced.
2. Dao, FlashAttention-2, Faster Attention with Better Parallelism and Work Partitioning, 2023. The generation in csrc/flash_attn/, the grid over query blocks and the warp rebalancing that Stage 5 and Q6 describe.
3. Shah et al., FlashAttention-3, Fast and Accurate Attention with Asynchrony and Low-precision, 2024. The Hopper generation in hopper/, warp specialization over TMA and WGMMA plus the fp8 forward. The CUTLASS walkthrough covers the CuTe machinery it is built on, and the mixed precision note covers the number formats.
4. Milakov and Gimelshein, Online Normalizer Calculation for Softmax, 2018. The single-pass softmax recurrence that Part VI extends with an output accumulator. The softmax note derives it in full.
5. Rabe and Staats, Self-attention Does Not Need O(n²) Memory, 2021. Showed shortly before FlashAttention that exact attention needs far less than quadratic memory by processing keys in chunks, without the IO argument that made the fused kernel fast.
6. Vaswani et al., Attention Is All You Need, 2017. The function every kernel here computes, worked through in the attention note.
7. Shazeer, Fast Transformer Decoding, One Write-Head is All You Need, 2019. Multi-query attention, the reason the kernels broadcast a small set of KV heads across many query heads.
8. Ainslie et al., GQA, Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, 2023. Grouped-query attention, the head-ratio contract of Part II. Both variants are compared in the attention variants note.
9. Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023. The paged KV cache that flash_attn_with_kvcache serves through its block_table argument, covered in the vLLM walkthrough.
10. Liu et al., Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023. Runs the LSE merging algebra of Q12 across devices, combining partial attentions computed over disjoint key ranges.
11. Tillet et al., Triton, an Intermediate Language and Compiler for Tiled Neural Network Computations, MAPL 2019. The tile-level compiler behind the repository's ROCm fallback and many independent FlashAttention reimplementations, covered in the Triton walkthrough.
12. Williams et al., Roofline, an Insightful Visual Performance Model for Multicore Architectures, CACM 2009. The arithmetic intensity framing Part V uses to call attention memory bound, applied throughout the parallel computing class.

Part XV: Final takeaway