What is FlashAttention?

How it works: tiling, and the memory traffic it removes

Attention is a three-step recipe: multiply the query matrix by the transposed key matrix, apply a softmax to each row, then multiply by the value matrix. Written that way, the intermediate N by N matrix of scores exists in full. On a long context it is the largest thing in the program, and writing it out and reading it back dominates the running time. The arithmetic units finish their work and then wait.

FlashAttention attacks that. The queries and keys are cut into blocks that fit in the on-chip SRAM next to the compute units, and the algorithm never sends the score matrix to main memory. Instead of computing one softmax over a whole row at the end, it computes a running maximum and a running sum as it goes, rescales the partial results each time a new block arrives a larger maximum, and combines the blocks into the final output on chip. Each block is read once.

The consequence is that the answer is identical to the reference implementation. This is not an approximation, not a low-rank substitute and not a sparse pattern that drops some pairs. The name is a description of how the data moves, not a claim about the arithmetic.

The technique is generally described as IO-aware, and it is also called tiling with recomputation. The second half matters in the backward pass: instead of saving intermediate activations to memory for the backward computation, the kernel recomputes them from the inputs it already has. That deliberately spends more arithmetic in exchange for fewer memory round trips, which is the trade the whole design rests on.

What it does not change

It is worth being precise about the boundaries, because the name invites over-reading.

The result is exact, but the cost is still quadratic in the sequence length. Tiling reduces the constant factor dramatically and removes the memory bottleneck, but a context twice as long still costs roughly four times the arithmetic. FlashAttention is not an answer to long context in the algorithmic sense.

It is a kernel, not a model change. You keep the same layers, the same weights and the same outputs; you replace the implementation of the attention operation. That is why it can be adopted one call site at a time, and why it shows up inside inference servers and model implementations rather than as a separate project you install alongside them.

It does not remove the need for the other parts of attention. The softmax, the masking and the causal pattern all remain, and the kernel still has to respect the causal constraint by not reading blocks that lie in the future.

And it is not a universal speedup. The gain comes from replacing memory traffic with on-chip computation, so it shows up where that trade is favourable: long sequences, GPU hardware with enough on-chip memory to hold a useful block, and data types the kernel supports. A workload that is latency bound on something else, or a model small enough that the launch overhead dominates, will not notice.

Where it pays, and where it does not

The workloads that benefit are the ones where the attention matrix is large relative to the compute.

Long context, where the quadratic term dominates everything else. A generation pass over a long document spends most of its time moving score data rather than multiplying it, and that is exactly the traffic FlashAttention removes.

Decoder-heavy serving, where each step processes one token against a long cached context. The per-step work is small, so the ratio of memory traffic to arithmetic is at its worst, and the kernel helps most there. This is why it turns up in inference engines by default.

Training with long sequences, for the same reason, with the extra benefit that recomputation cuts the activation memory saved for the backward pass.

Where you will not feel it: short sequences, where the matrix fits in cache anyway; small models, where the launch and memory overhead of a specialised kernel can exceed the work it saves; anything not running on a supported GPU; and the prefill-versus-decode split, where prefill is compute bound and the memory-traffic argument is weaker than it is during decoding.

There is also a practical threshold most people discover late. The kernels need a matching build for your Python, CUDA and PyTorch versions, and a kernel that is not built for your configuration either fails to load or falls back to a slower path. A large share of the difficulty people report with attention performance is a build problem rather than a kernel problem.

The build and the compatibility problem

This is where the open-source story gets concrete, and the reference implementation is the place to start.

Dao-AILab/flash-attention is described as the official implementation of FlashAttention and FlashAttention-2 from the papers behind them. Our analysis of flash-attention describes it as an IO-aware exact attention kernel for PyTorch, shipped as four separate builds: v2, v3, v4 and the ROCm Triton backend. Each has its own requirements, so choosing one is a decision about your hardware and your toolchain rather than a default.

mjun0812/flash-attention-prebuild-wheels exists because compiling from source is slow. Our analysis of that repository says it publishes prebuilt flash-attn 2 and 3 wheels for Linux and Windows using GitHub Actions, so you can skip a long source build, and also that the wheel you need depends on an exact Python, CUDA and PyTorch combination. In other words, it trades compile time for a matching problem: you have to know all three numbers before you can pick a file, and a near miss is a miss.

The AMD path runs through the ROCm Triton backend rather than a CUDA build, which is a separate thing to learn and a separate set of version constraints.

So the practical checklist is short: know your GPU architecture, know your CUDA, PyTorch and Python versions, pick the matching build, and verify that the kernel actually loads rather than silently falling back. Every one of those steps is where a slow attention implementation usually turns out to be a mismatched build.

How the surrounding projects use it

Almost nobody vendors these kernels directly. They arrive through a framework that picks one for you.

flashinfer-ai/flashinfer is the clearest example. It is described as a kernel library for LLM serving that provides unified APIs for attention, GEMM and MoE operations with multiple backend implementations, including FlashAttention-2 and 3, cuDNN, CUTLASS and TensorRT-LLM. Our analysis of FlashInfer says it packages attention, GEMM, MoE and sampling kernels behind one Python API and picks a backend per GPU, that the pip install is one line, and that the architecture coverage table is where the caveats start. That is the pattern to copy: you depend on a library that has already worked out which kernel suits which card.

Some projects use the kernel as a building block inside a larger engine. Wenyueh/MinivLLM rebuilds the vLLM inference engine on top of Nano-vLLM with its own paged attention and flash attention kernels, plus prefilling and decoding benchmarks, and our analysis calls it a study artifact rather than a serving stack. Reading it is the way to see how the kernel is called from a real engine.

Some are teaching material. xlite-dev/LeetCUDA is described as a collection of more than two hundred CUDA kernels including FlashAttention, an HGEMM implementation and a four-hundred-page open source book, aimed at beginners who want to read real MMA PTX. Our analysis of LeetCUDA notes it is GPL-3.0, covers how the build works and what its benchmarks claim, and where it stops being the right tool.

The pattern across all of them is the same: the kernel is stable and old, the difficulty has moved to the packaging and the hardware coverage.

In practice

FlashAttention is worth reaching for when attention runs over long sequences on a GPU, because it removes the memory traffic that dominates the textbook implementation while returning an identical result, and it is not worth the trouble when your sequences are short or your model is small. Start from a serving library that selects a kernel per GPU rather than compiling one yourself, and if you do compile, match the build to your exact Python, CUDA and PyTorch versions before blaming the kernel for being slow. When you want to understand why a particular kernel is faster, Wenyueh/MinivLLM shows the call site inside an engine and xlite-dev/LeetCUDA shows the instructions themselves.