FlashAttention — how 100K-token attention fits in memory

For one head with 100K tokens:

\[ Q,K\in\mathbb{R}^{100000\times128} \]

so:

\[ QK^T\in\mathbb{R}^{100000\times100000} \]

That is \(10^{10}\) scores.

At 2 bytes each:

\[ 20GB \]

for one head if the full score matrix were materialized.

Naïve execution¶

A naïve implementation might: 1. compute QKᵀ, 2. write the giant score matrix to HBM, 3. read it, 4. softmax, 5. write probabilities, 6. read them again, 7. multiply by V.

The bottleneck is huge memory traffic.

FlashAttention idea¶

Process small tiles:

small Q block
   ↓
small K block
   ↓
partial scores
   ↓
update online softmax
   ↓
small V block
   ↓
accumulate final output

Temporary tiles stay in fast on-chip memory and are discarded.

But softmax needs the whole row — right?¶

Softmax is:

\[ p_i=\frac{e^{s_i}}{\sum_j e^{s_j}} \]

FlashAttention maintains running: - maximum, - normalization sum, - weighted V accumulator.

If a later block contains a larger maximum, previous partial sums are rescaled appropriately.

Thus the final result is exact dense attention.

What FlashAttention does not do¶

It does not change dense attention from \(O(T^2)\) to something sub-quadratic.

It changes: - memory complexity of intermediates, - HBM traffic, - practical runtime.

Why more compute can be faster¶

Modern GPUs often have abundant arithmetic relative to memory bandwidth.

So a method doing somewhat more arithmetic but dramatically fewer HBM accesses can be faster.

This is the same trade behind: - activation checkpointing, - quantization, - batching.

Training backward¶

The backward pass can recompute attention tiles rather than storing huge forward intermediates.

Again:

\[ \boxed{\text{extra compute}\leftrightarrow\text{less memory movement}} \]

IO complexity is the point¶

On modern accelerators, arithmetic units can often multiply numbers much faster than HBM can move huge intermediate tensors.

The naïve attention algorithm repeatedly spills enormous matrices to HBM.

FlashAttention reorganizes the algorithm around the memory hierarchy: - HBM holds large source tensors, - SRAM/registers hold small tiles, - intermediate scores live only long enough to update running statistics.

The algorithm is “hardware-aware” without changing the mathematical attention result.

Online softmax intuition¶

Suppose the first block has scores:

\[ [1,2,3] \]

Its max is 3, and a stable denominator can be accumulated relative to that max.

Later a new block contains score 4. The global max changes from 3 to 4, so the previous accumulator is rescaled by:

\[ e^{3-4}=e^{-1} \]

before adding the new block.

This is how blockwise processing can still produce the exact global softmax.

Training payoff¶

Naïve attention would tempt training to save giant score/probability matrices for backward.

FlashAttention can recompute local quantities from Q/K/V instead.

Again, computation is intentionally repeated to avoid expensive memory storage and traffic.

The broader lesson is that algorithmic complexity alone is not enough: where bytes move can determine real performance.

Why “exact” matters¶

Some efficient-attention methods approximate dense attention by dropping connections or compressing the score matrix.

FlashAttention does not.

For the same Q, K, V and numerical precision, it is designed to compute the dense softmax-attention result while changing the execution schedule.

That makes it easier to adopt because model semantics need not change.

Arithmetic intensity¶

A useful performance ratio is:

\[ \text{arithmetic intensity} = \frac{\text{operations}}{\text{bytes moved}}. \]

A kernel with low arithmetic intensity tends to be bandwidth-bound.

FlashAttention raises effective arithmetic intensity by reusing data in SRAM and avoiding repeated HBM traffic.

This is the same lens we used to understand why prefill and decode behave differently on GPUs.

The math of the model and the physical memory hierarchy of the accelerator meet at the kernel.

FlashAttention and long-context feasibility¶

Without memory-efficient attention, a model may theoretically support a long context but fail simply because intermediate attention tensors do not fit.

FlashAttention can turn such a workload from impossible to feasible without changing parameter count or model quality.

That is an important category of optimization: it does not make the model smarter or smaller. It makes the same mathematical model executable at a scale that naïve kernels cannot handle.

Kernel engineering can therefore unlock model capabilities that already exist architecturally.