The Wall That Stopped Long Context
A few years ago, if you asked a large language model to read an entire book, it simply couldn't. The "context window" — the amount of text a model can consider at once — topped out at a few thousand words. Feed it more, and the computation became too expensive to run. This wasn't a lack of raw processing power. It was a memory problem.
At the center of every modern language model is a step called "attention," which lets the model work out how much each word should influence its understanding of every other word. The trouble is that the memory this step needs grows far faster than the text does. Double the length of the input, and instead of needing twice the memory, the model needs four times as much. Feed it ten times more text, and its memory needs balloon by a hundred times. That runaway growth — called quadratic scaling — created a hard wall. Past a certain length, there simply wasn't enough memory available to hold everything attention needed, no matter how fast the chip could calculate.
Then, in 2022, a paper titled FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness proposed a way through. The researchers didn't invent a new kind of attention — the model still weighs words against each other exactly as before. Instead, they found a smarter way to organize the arithmetic, paying close attention to how data moves between a GPU's small, extremely fast working memory and its much larger, slower main memory. By keeping the heaviest intermediate data out of the slow memory almost entirely, they eliminated the memory wall. The technique, now known as FlashAttention, produces exactly the same output as the original method, using a fraction of the memory and running much faster. It turned attention from a problem limited by memory traffic into one limited by raw computing power — which is exactly what GPUs are built to provide. The long-context models we have today, able to process entire books or codebases in one pass, exist because of this piece of engineering.
The Quadratic Memory Bottleneck
To understand why FlashAttention was so revolutionary, we first need to understand the problem it solved. At the heart of the Transformer architecture, which powers nearly all modern LLMs, is the self-attention mechanism. This is how a model learns the relationships between words in a sequence. For every word, it generates three vectors: a Query (Q), a Key (K), and a Value (V).
The Query vector is like a question: "What context is relevant to me, the word at this position?"
The Key vector is like a label: "I am the word at this position, and this is what I'm about."
The Value vector contains the actual information or meaning of the word.
To figure out how much attention the word at position i should pay to the word at position j, the model computes the dot product of Query i and Key j. This is done for all pairs of words, producing a large matrix of attention scores. If your input sequence has N tokens, this calculation creates an N×N matrix.
The standard attention algorithm, written in a framework like PyTorch, looks something like this:
Calculate the raw scores:
S = Q @ K.TScale and normalize the scores into probabilities using the softmax function:
P = softmax(S)Use the probability matrix to create a weighted sum of the Value vectors:
O = P @ V
The problem lies in step 1. The matrix S, often called the attention matrix, is huge. For a sequence length N of just 4,096 (4k), an N×N matrix of 32-bit floats requires 4096 * 4096 * 4 bytes, which is about 67 MB. For a 64k context window, that same matrix explodes to over 16 GB. For a 1 million token context window, it would require over 3.7 terabytes of memory, just for this one intermediate step.
This is the quadratic bottleneck. The memory required to explicitly create this matrix scales with N², and it must be read from and written to the GPU's main memory.
Memory-Bound vs. Compute-Bound
This enormous memory requirement made attention a memory-bound operation. A processor, like a GPU, has a small amount of extremely fast on-chip memory (SRAM) and a much larger pool of slower off-chip memory (High-Bandwidth Memory, or HBM).
A compute-bound task is one where the GPU spends most of its time performing calculations. The data fits in fast memory, and the processor is the bottleneck. This is what GPUs are designed for.
A memory-bound task is one where the GPU spends most of its time waiting for data to be moved back and forth from the slower HBM. The data transfer (the I/O operation) is the bottleneck, not the calculation itself.
Standard attention was severely memory-bound. The GPU would compute a small part of the giant S matrix, write it out to HBM, read another part, compute some more, and so on. The actual math is composed of matrix multiplications (GEMMs), which GPUs excel at. But the sheer size of the intermediate S and P matrices forced a constant, slow conversation with HBM, leaving the powerful processing cores idle for much of the time. The problem wasn't the number of calculations (FLOPs), but the number of memory accesses.
IO-Awareness: The FlashAttention Solution
The authors of FlashAttention recognized that the core problem was the materialization of the N×N attention matrix in HBM. Their solution was simple in concept but complex in execution: just don't write it.
This is the core of IO-awareness. An IO-aware algorithm is designed with a deep understanding of the memory hierarchy it's running on: transfers between the GPU cores and SRAM are lightning-fast, while transfers to and from HBM are a performance killer. The goal is to minimize those slow HBM accesses by doing as much work as possible in the fast SRAM.
FlashAttention achieves this through two key techniques: tiling and computing the softmax online.
Tiling: Breaking the Problem into Pieces
Instead of trying to compute the entire N×N matrix at once, FlashAttention breaks the input Q, K, and V matrices into smaller blocks, or tiles. It then loads these smaller blocks into the GPU's fast SRAM and performs the full attention calculation for just that small part of the output.
The process looks roughly like this:
Divide the Q matrix into blocks along the sequence length dimension.
For each block of Q (
Q_i), iterate through all the blocks of K and V (K_j,V_j).In the inner loop, load
Q_i,K_j, andV_jfrom HBM into the fast SRAM.Once in SRAM, perform the attention calculation for just these blocks: compute the score sub-matrix
S_ij = Q_i @ K_j.T, apply the softmax, and multiply byV_j.Crucially, the result of this is immediately used to update the final output block
O_i, which is also kept in SRAM. The intermediateS_ijmatrix never leaves the fast on-chip memory.Repeat for all blocks
K_jandV_j, accumulating the results intoO_i.Once all inner loops for
Q_iare done, write the final output blockO_iback to HBM.
By fusing these operations into a single GPU kernel, FlashAttention avoids the expensive round trips to HBM for the giant intermediate matrices. The memory usage is no longer quadratic, because the N×N matrix is never fully formed. Instead, memory scales linearly with sequence length, dominated by the storage for Q, K, and V themselves.
The Trick: Online Softmax
This sounds simple, but there's a mathematical catch. The softmax function normalizes a vector of numbers by dividing each by the sum of all of them.
softmax(x_i) = exp(x_i) / sum(exp(x_j) for all j)
To compute the denominator, you need to see all the elements of the vector. But with tiling, we only have a small block of the attention scores at a time. How can we compute a global function like softmax one piece at a time?
This is the second clever insight of FlashAttention. It uses an "online" algorithm to compute the softmax correctly without having the full input vector. As it iterates through the blocks of K, it keeps track of the running statistics needed for the softmax normalization. For each row of the attention matrix (corresponding to a single query token), it maintains two values:
The maximum value seen so far in the row.
The sum of the exponentials of the values seen so far (the normalization factor).
When a new block S_ij is computed in SRAM, the algorithm updates these running statistics and correctly rescales the contribution from the previous blocks. This ensures that by the time the final block of K has been processed, the resulting output is numerically identical to what standard attention would have produced. It's a re-association of the math that makes it possible to compute a global property in a streaming, block-by-block fashion.
A Worked Example in Pseudocode
To make this concrete, let's compare the standard approach to the tiled approach in simplified pseudocode.
Standard Attention
This is simple and readable, but has a huge memory footprint for the S matrix.
# Q, K, V are N x d matrices, stored in HBM
# N = sequence length, d = head dimension
def standard_attention(Q, K, V):
# Step 1: Compute the full N x N matrix.
# This requires a massive write to HBM.
S = Q @ K.T # S is N x N
# Step 2: Compute the softmax.
# This reads the N x N matrix S from HBM
# and writes a new N x N matrix P to HBM.
P = softmax(S) # P is N x N
# Step 3: Compute the final output.
# This reads the N x N matrix P and the N x d matrix V,
# and writes the N x d output O to HBM.
O = P @ V # O is N x d
return OFlashAttention (Conceptual)
This is far more complex, but its memory accesses are much more efficient. The operations inside the loop all happen in fast SRAM.
# Q, K, V are N x d matrices, stored in HBM
# O is an N x d zero matrix, stored in HBM
def flash_attention(Q, K, V):
# Define block sizes for tiling
BLOCK_SIZE_Q = 128
BLOCK_SIZE_K = 128
# Outer loop: iterate over blocks of the output
for i in 0..num_blocks(N, BLOCK_SIZE_Q):
# Load a block of Q into fast SRAM
Q_block = load_from_hbm(Q, i, BLOCK_SIZE_Q)
# Initialize output block and softmax stats in SRAM
O_block = zeros(BLOCK_SIZE_Q, d)
running_max = -infinity
running_normalizer = 0
# Inner loop: iterate over blocks of the inputs
for j in 0..num_blocks(N, BLOCK_SIZE_K):
# Load blocks of K and V into fast SRAM
K_block = load_from_hbm(K, j, BLOCK_SIZE_K)
V_block = load_from_hbm(V, j, BLOCK_SIZE_K)
# --- All operations below happen in fast SRAM ---
# 1. Compute score sub-matrix
S_block = Q_block @ K_block.T
# 2. Find the new max for the online softmax
local_max = row_max(S_block)
new_max = max(running_max, local_max)
# 3. Rescale previous output and normalizer with new max
O_block *= exp(running_max - new_max)
running_normalizer *= exp(running_max - new_max)
# 4. Compute softmax for the current block
P_block = exp(S_block - new_max)
local_normalizer = row_sum(P_block)
# 5. Update the running normalizer
running_normalizer += local_normalizer
# 6. Update the output block with the contribution from this V_block
O_block += (P_block @ V_block)
# 7. Update the running max for the next iteration
running_max = new_max
# --- End of SRAM-only operations ---
# Rescale the final output block and write back to HBM
O_block /= running_normalizer
write_to_hbm(O, i, O_block)
return OThis pseudocode simplifies the online softmax update, but it illustrates the core loop structure. The key takeaway is that the N×N matrices S and P are never created. Instead, small sub-matrices are computed on the fly in fast memory and immediately used to update the final output, dramatically reducing the required memory bandwidth.
The Payoff and The Price
The benefits of FlashAttention are clear, but it's not a free lunch.
The Price: Complexity. Implementing FlashAttention requires writing low-level CUDA code. This is a highly specialized skill, far removed from the high-level Python of typical machine learning engineering. Standard attention is a single line of code in most frameworks; FlashAttention is a complex, hand-tuned kernel that must be carefully managed. This is why it's a library you install, not an algorithm you re-implement for every project.
The Payoff: Speed, Memory, and the Dawn of Long Context. The reward for this complexity is enormous.
Speed: By eliminating the HBM bottleneck, FlashAttention makes attention compute-bound again. It can be up to an order of magnitude faster than standard attention, especially for long sequences.
Memory: Memory usage scales linearly (O(N)) with sequence length, not quadratically (O(N²)). This is the game-changer. It makes context lengths of 64k, 128k, or even a million tokens computationally feasible on current hardware.
Exactness: Unlike other methods like sparse attention or low-rank factorization, FlashAttention is an exact algorithm. It doesn't approximate the attention matrix; it computes the exact same output as the standard implementation. The result is numerically identical, just arrived at much more efficiently.
This combination of speed and memory efficiency directly enabled the rapid expansion of LLM context windows that started in 2023. Before FlashAttention, long context was a theoretical research area. After, it became a core product feature. The ability to feed entire documents, codebases, or transcripts into a model is a direct consequence of solving this fundamental systems problem. You can find the original research on arXiv: https://arxiv.org/abs/2205.14135 ↗.
What This Means for You
For the AI practitioner, the good news is that you are almost certainly already benefiting from this work without needing to know it happened. Optimized attention kernels based on FlashAttention are now built into the core AI software stack: PyTorch's scaled_dot_product_attention function automatically switches to a FlashAttention-based implementation when it detects a compatible GPU and input shape, and libraries from Hugging Face and others turn it on by default. When you pick a model advertising a 200k-token context window, you are relying on this specific piece of engineering, whether or not the vendor mentions it.
What's still unresolved is the complexity tax described above: writing and maintaining these kernels requires low-level CUDA skill that is scarce, and each new generation of GPU hardware means the kernels have to be rewritten and re-tuned, not just recompiled. That's part of why support for FlashAttention-style optimizations lags on some hardware and some model architectures, and why the "default fast path" doesn't automatically cover every combination of model, GPU, and sequence length you might want to use.
For the technically curious, FlashAttention is a reminder that progress in AI isn't only about bigger models and novel architectures. It's also about engineers taking the physical realities of the hardware seriously — how far data has to travel, how fast it can move — and redesigning the arithmetic around those realities. FlashAttention didn't change what attention computes; it changed how the computation is scheduled on real chips. That distinction is why context windows went from a few thousand tokens to hundreds of thousands in the space of about a year.
References
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135 ↗.