What problem does FlashAttention solve, and is it an approximation?

How to think about it

FlashAttention solves the memory-traffic problem in dense attention. It computes the same softmax attention as the textbook formula, but processes it in tiles that stay in fast on-chip memory instead of materializing the full N × N attention matrix in GPU memory. It is not an approximation; its algorithm is exact apart from ordinary floating-point rounding.

Why ordinary attention becomes expensive

For one attention head, let N be the sequence length and d the head dimension. Attention computes:

S = QKᵀ / sqrt(d)

Here, Q contains queries, K contains keys, and V contains values. Each token supplies one row of each matrix. The model applies a row-wise softmax to S, then multiplies the result by V.

The problem is S. It has one score for every query-key pair, so it contains N² numbers.

A straightforward GPU implementation does roughly this:

1. Compute QKᵀ and write the score matrix to HBM, the GPU’s large high-bandwidth memory.
2. Read the scores back to apply softmax.
3. Write the probabilities.
4. Read the probabilities again to multiply by V.

The GPU’s tensor cores can perform the matrix multiplications very quickly. The bottleneck is repeatedly moving a huge intermediate matrix between HBM and the compute units. This is why the problem is called IO-bound: the time and energy spent moving data dominate the arithmetic.

Modern GPU kernels can fuse some of these operations, so the exact number of reads and writes depends on the implementation. The underlying issue remains: materializing an N × N object creates far more memory traffic than the final output requires.

A concrete 32k-token example

Take one sequence with:

• N = 32,768 tokens
• d = 128
• FP16 scores, using 2 bytes per number
• 32 attention heads

For one head, the score matrix contains:

32,768 × 32,768 = 1,073,741,824 scores

At 2 bytes per score, that is exactly 2 GiB. If the implementation separately materializes both logits and the post-softmax probability matrix, those two matrices consume about 4 GiB for one head. Across 32 heads, that becomes 64 GiB for scores alone, or about 128 GiB for both matrices.

By comparison, one head’s Q, K, and V tensors each contain only 32,768 × 128 values, about 8 MiB in FP16. The full inputs for all three are about 24 MiB per head. The attention output is another 8 MiB.

The mathematical output is small. The intermediate score matrix is the balloon tied to its ankle.

FlashAttention avoids that balloon. It loads a block of queries and a block of keys into SRAM, which is small static memory located on the GPU chip close to the arithmetic units. It computes a score tile, uses that tile immediately, and discards it before loading the next one.

For intuition, a 128 × 128 score tile in FP16 is only 32 KiB. A real kernel also needs space for query, key, value, and accumulator tiles, and the exact block sizes vary by GPU and kernel version. The important point is that the tile fits in on-chip memory; the full 32k-by-32k matrix does not.

The mechanism: tiled attention with online softmax

Tiling alone is not enough. Softmax normally needs every score in a row before it can calculate the denominator. If the row is split across many tiles, FlashAttention needs a way to normalize each row while seeing only one tile at a time.

It uses online softmax, which means carrying a small amount of state from one tile to the next.

For each query row, the kernel maintains:

• m, the largest score seen so far
• ℓ, the running softmax denominator after stabilisation
• o, the running, unnormalised weighted sum of value vectors

Suppose the next tile contains scores s_t. The update is:

m_t = max(m_{t-1}, max(s_t))

ℓ_t = exp(m_{t-1} - m_t) ℓ_{t-1} + sum(exp(s_t - m_t))

o_t = exp(m_{t-1} - m_t) o_{t-1} + sum(exp(s_t - m_t) V_t)

After the final tile, the output row is o_t / ℓ_t.

The rescaling term is the key. Imagine the first tile’s largest score is 4, so its exponentials are measured relative to 4. A later tile contains a score of 9. The running maximum becomes 9, so the contribution accumulated from the earlier tile is multiplied by exp(4 - 9). That puts old and new contributions on the same scale.

Nothing was discarded. The algorithm has simply evaluated the same sums in pieces. In exact arithmetic, the final result is the same as applying softmax to the complete row.

The kernel also applies the usual mask while processing tiles. With causal attention, for example, it ignores scores from a token to future tokens. FlashAttention is exact relative to that requested masked attention; it does not secretly turn causal attention into bidirectional attention.

What FlashAttention changes, and what it does not

FlashAttention changes the execution order and memory access pattern:

• It computes the same query-key scores.
• It applies the same scaling, masking, softmax, and value aggregation.
• It keeps intermediate tiles on chip.
• It writes the output rather than the full attention matrix.
• During training, it usually saves compact row statistics and recomputes tiles during the backward pass instead of storing every score.

That last choice trades some extra arithmetic for much lower memory use. Recomputing a score tile costs work, but saving and rereading a massive matrix costs memory capacity and bandwidth. On modern GPUs, that trade is often favorable.

FlashAttention therefore reduces the attention activation memory from quadratic in N to roughly O(Nd), or simply linear in N when the head dimension is treated as fixed. It does not reduce the dense attention computation from quadratic time. The kernel still considers every allowed query-key pair, so its arithmetic cost remains roughly O(N²d).

This distinction matters. Sparse attention, low-rank attention, and linear attention change the computation itself to avoid considering every pair. FlashAttention does not. It is an exact implementation of dense attention with better IO behaviour.

The production pattern

In PyTorch, the usual production-level interface is the semantic operation rather than a hand-written FlashAttention call:

import torch
import torch.nn.functional as F

q = torch.randn(
    1, 32, 32768, 128,
    device="cuda",
    dtype=torch.float16,
)
k = torch.randn_like(q)
v = torch.randn_like(q)

output = F.scaled_dot_product_attention(
    q, k, v,
    is_causal=True,
)

The tensor shape here is batch, heads, sequence, head dimension. The call requests scaled dot-product attention with a causal mask. Depending on the GPU, dtype, tensor shapes, mask, and installed framework, PyTorch may dispatch to a FlashAttention backend, another memory-efficient backend, or a conventional math implementation. The API expresses the operation; it does not guarantee one particular kernel.

That is a useful production lesson: do not assume that importing a library or calling a similarly named function means FlashAttention is active. Profile the workload and inspect the framework’s current backend-selection documentation. A strange mask or unsupported dtype can quietly send execution down a fallback path.

The senior-level nuance

“Exact” does not mean bit-for-bit identical in every run. Different tiling orders change the order of floating-point additions. FP16 and BF16 also have limited precision. A standard implementation and FlashAttention can therefore differ by tiny numerical amounts while implementing the same mathematical function. Those differences should be checked against tolerances, not treated as evidence of approximation.

FlashAttention also does not solve every long-context memory problem.

During autoregressive generation, the model usually stores a KV cache, which is the collection of past keys and values reused for each newly generated token. That cache grows linearly with context length and can itself consume many gigabytes. FlashAttention avoids the quadratic attention matrix, but it does not eliminate the KV cache. Decode-specific kernels, grouped-query attention, multi-query attention, quantised caches, and paged cache layouts address different parts of that problem.

It may also fail to be faster for very short sequences or unusual shapes. Kernel launch overhead, unsupported masks, dtype conversions, and fallback implementations can erase the benefit. If an application needs the full attention probability matrix for inspection, it must materialize something FlashAttention is designed specifically not to materialize, giving back much of the memory advantage.

So the accurate claim is not “FlashAttention makes all long-context models cheap.” It makes dense attention substantially more practical by removing an unnecessary memory bottleneck. The quadratic compute cost, model quality, KV-cache cost, and hardware limits still exist.

What they’ll ask next

Does FlashAttention reduce the time complexity of attention?

No. Dense attention still performs roughly N² query-key interactions. FlashAttention reduces HBM traffic and peak activation memory. It is an IO optimisation, not a sparse or linear-time attention algorithm.

If it is exact, why might its output differ slightly from standard attention?

Because floating-point arithmetic is not associative. Tiling changes the order of accumulation, and reduced-precision inputs introduce rounding. The formula is the same, but bitwise equality is not guaranteed. In practice, outputs should be compared with an appropriate numerical tolerance.

Is scaled_dot_product_attention always using FlashAttention?

No. Frameworks choose among several backends based on hardware, dtype, shapes, masks, and other constraints. I would profile the real workload and verify the selected backend rather than relying on the function name.

Say this in the interview

“FlashAttention is exact dense softmax attention implemented as a tiled, IO-aware kernel: it keeps tiles and online-softmax state on chip, avoids materializing the N × N matrix in HBM, and cuts memory without changing the attention computation.”