The wall you hit is not the one you expect
Suppose an Indian legal-tech startup is fine-tuning a transformer to read entire loan agreements in a single pass — a 45-page contract runs to roughly 18,000 tokens, and the model needs the whole document in context to catch a clause on page 40 that contradicts a definition on page 3. The engineering team estimates the compute: sequence length N = 18,000, hidden dimension d = 128, so a single attention layer needs on the order of N²d ≈ 4×10¹⁰ floating-point operations. An A100 GPU claims 312 teraFLOPs/s of tensor-core throughput, so that layer should finish in well under a millisecond. It does not. Profiling shows the attention layer spending over 80% of its wall-clock time not computing anything, and the training run runs out of the GPU's 40GB of memory long before the FLOP budget is anywhere near exhausted.
This is the pattern that Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré diagnosed precisely in their 2022 NeurIPS paper "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." Their claim, which sounds almost too simple, is that standard attention is not slow because there is too much arithmetic — it is slow because of where the intermediate results have to live while that arithmetic happens. A companion chapter in this course covers linear attention, which trades exact softmax attention for an approximate O(N) alternative. FlashAttention takes the opposite strategy: it changes nothing about what is computed, only how the GPU's memory is used while computing it. That distinction is the spine of this chapter, and by the end you should be able to say precisely why the two techniques solve different problems for different reasons.
Two kinds of memory live on every GPU
An A100 has two memory systems that differ by more than an order of magnitude in every dimension that matters. High-bandwidth memory (HBM) is the GPU's "main memory" — 40 or 80GB, large enough to hold an entire model's activations, but reachable at only about 1.5–2 terabytes per second. On-chip SRAM — the register files and shared memory sitting directly on each of the 108 streaming multiprocessors — totals only around 20MB across the whole chip, but its bandwidth is roughly 19 terabytes per second, close to ten times faster. Every tensor-core matrix multiply an A100 performs first has to pull its operands from HBM into SRAM, compute, and write results back out to HBM. The tensor cores themselves are almost never the bottleneck; the round trip to HBM is.
A standard, unfused implementation of attention makes this round trip badly. To compute Attention(Q,K,V) = softmax(QKᵀ/√d)V, a naive kernel: (1) reads Q and K from HBM, computes the full N×N score matrix S = QKᵀ/√d, and writes all of S back to HBM; (2) reads S back from HBM to apply softmax row-wise, and writes the full N×N probability matrix P back to HBM; (3) reads P and V back from HBM to compute O = PV, and writes O out. For the 18,000-token contract, N×N alone is 324 million entries — over a gigabyte just for one intermediate matrix, in one layer, for one attention head, materialized and re-read three separate times. This is why memory runs out before FLOPs do: the actual arithmetic is small and fast, but each of those three stages pays the full 1.5 TB/s HBM tax instead of the 19 TB/s SRAM rate, and the intermediate matrices scale as N² while the useful inputs and outputs only scale as N.
Tiling and the running (online) softmax
FlashAttention's fix is to never materialize the full N×N matrix in HBM at all. It partitions Q into row blocks of size Bᵣ and K, V into column blocks of size B꜀, chosen so that one query block plus one key/value block plus the resulting score block all fit inside the ~20MB of SRAM at once. The kernel then loops: for each query tile Qᵢ (the outer loop, i = 1…Tᵣ where Tᵣ = ⌈N/Bᵣ⌉), it loads successive key/value tiles Kⱼ, Vⱼ (the inner loop, j = 1…T꜀) into SRAM, computes only the small score tile Sᵢⱼ = QᵢKⱼᵀ/√d, and immediately folds it into a running output — never writing any block of S or P back to HBM.
The obstacle to doing this is the softmax's normalizer: softmax needs the sum of exponentials over the entire row before it can produce a single correctly normalized probability, but tiling only ever shows the kernel one column-block of a row at a time. The trick that resolves this, borrowed from earlier work on numerically stable streaming softmax, is to track three running statistics per query row as tiles arrive: a running max m, a running sum of exponentials ℓ, and a running (unnormalized) weighted output O. Every time a new tile's local maximum exceeds the old running max, the previously accumulated O and ℓ are rescaled by a correction factor exp(m_old − m_new) before adding the new tile's contribution — algebraically equivalent to redoing the whole softmax with the new, larger max, but done incrementally in place. Only after the last key/value tile has been folded in is O divided by the final ℓ to produce the true, exact softmax-weighted output. Nothing here is approximate: it is the same softmax formula, computed in a different order.
Worked example: tracing the running statistics by hand
Take one query vector q = [1, 0] attending over four keys and values (d = 2, so we omit the 1/√d scaling for arithmetic clarity — in practice it is applied to the raw scores before this point):
K = [[ 1, 0],
[ 0, 1],
[ 3, 1],
[-1, 1]]
V = [[1, 2],
[3, 4],
[5, 6],
[7, 8]]The raw scores are s = K·q = [1, 0, 3, -1]. Process these as two tiles of two keys each — tile 1 = keys {1,2}, tile 2 = keys {3,4} — exactly as a block-size-2 FlashAttention kernel would.
Initialize: m = -∞, ℓ = 0, O = [0, 0].
Tile 1 (scores [1, 0]): the tile's local max is 1, so m_new = max(-∞, 1) = 1. Unnormalized probabilities: p = exp([1,0] - 1) = [1, 0.3679]. Since the old m was -∞, the correction factor exp(m_old - m_new) = exp(-∞) = 0, which correctly zeroes out the (empty) prior accumulator. Update: ℓ = 0×0 + (1 + 0.3679) = 1.3679. O = 0×[0,0] + (1×[1,2] + 0.3679×[3,4]) = [2.1036, 3.4715]. Set m = 1.
Tile 2 (scores [3, -1]): local max is 3, so m_new = max(1, 3) = 3 — the running max just changed, which is the case that actually exercises the rescaling logic. Correction factor: c = exp(1 - 3) = exp(-2) ≈ 0.13534. New tile probabilities: p = exp([3,-1] - 3) = [1, exp(-4)] ≈ [1, 0.01832]. Update ℓ: ℓ = 0.13534 × 1.3679 + (1 + 0.01832) ≈ 0.18512 + 1.01832 = 1.20344. Update O: rescale the old accumulator first, 0.13534 × [2.1036, 3.4715] ≈ [0.2847, 0.4698], then add the new tile's contribution, 1×[5,6] + 0.01832×[7,8] ≈ [5.1282, 6.1465], giving O ≈ [5.4129, 6.6164]. Set m = 3.
Finalize: O_final = O / ℓ ≈ [5.4129, 6.6164] / 1.20344 ≈ [4.4979, 5.4979].
Now check this against computing softmax directly on the full row: softmax([1,0,3,-1]) gives probabilities [0.1125, 0.0414, 0.8310, 0.0152] (they sum to 1, and note ℓ = 1.20344 is exactly the softmax denominator computed above from the tiled trace), and Σ p·V ≈ [4.4979, 5.4979] — identical to the tiled result. The following code runs both versions and confirms the match numerically (verified by execution, not just derivation):
import numpy as np
def flash_attention_1d(q, K, V, block_size):
"""Online-softmax attention for one query vector q
against keys K (N x d) and values V (N x d), tiled."""
N, d = K.shape
m = -np.inf # running row max
l = 0.0 # running sum of exponentials
O = np.zeros(d) # running (unnormalized) output
for start in range(0, N, block_size):
end = min(start + block_size, N)
K_j = K[start:end]
V_j = V[start:end]
s_j = K_j @ q
m_new = max(m, s_j.max())
p_j = np.exp(s_j - m_new)
correction = np.exp(m - m_new) if m != -np.inf else 0.0
l = correction * l + p_j.sum()
O = correction * O + p_j @ V_j
m = m_new
return O / l
q = np.array([1.0, 0.0])
K = np.array([[1.0, 0.0], [0.0, 1.0], [3.0, 1.0], [-1.0, 1.0]])
V = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
tiled = flash_attention_1d(q, K, V, block_size=2)
scores = K @ q
probs = np.exp(scores - scores.max())
probs /= probs.sum()
direct = probs @ V
print(tiled) # [4.49786861 5.49786861]
print(direct) # [4.49786861 5.49786861]Both print [4.49786861 5.49786861] — the tiled, memory-frugal computation and the direct, memory-hungry one agree to full floating-point precision. That is the entire promise of FlashAttention: identical numbers, radically different memory traffic.
Why fewer HBM round trips actually follows from tiling
The paper proves a concrete bound: standard attention requires Θ(Nd + N²) HBM accesses (the N² term from reading and writing the full score and probability matrices), while FlashAttention requires only Θ(N²d²/M), where M is the SRAM size. Plugging in realistic numbers — GPT-2-scale N = 1024, d = 64, and an SRAM budget of roughly 25,000 32-bit values (~100KB) — standard attention touches about Nd + N² ≈ 65,536 + 1,048,576 ≈ 1.11 million elements' worth of HBM traffic, while FlashAttention's bound works out to N²d²/M ≈ 1024²×64²/25,000 ≈ 172,000 — roughly 6.5× fewer HBM accesses for this configuration, a gap that widens as N grows because the standard bound's dominant N² term has no compensating 1/M. The published results bear this out end to end: the paper reports training BERT-large 15% faster than the MLPerf 1.1 training speed record, and up to a 3× training speedup on GPT-2, purely from restructuring memory traffic — no change to model architecture, hyperparameters, or the attention formula itself.
What FlashAttention-2 changed
Tri Dao's 2023 follow-up, "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning," kept the tiling and online-softmax core but attacked three sources of remaining waste. First, the original kernel rescaled the output accumulator O at every single inner-loop step even when the running max had not changed enough to matter much; FlashAttention-2 defers most rescaling and reduces the count of non-matmul operations (exponentials, divisions) relative to matmul operations, since non-matmul FLOPs are far more expensive per operation on tensor-core hardware. Second, the original parallelized GPU work only across the batch and attention-head dimensions — fine when batch size is large, but for a single 18,000-token document processed one sequence at a time, that grid may contain far fewer blocks than the GPU has streaming multiprocessors, leaving hardware idle no matter how efficient each block's inner loop is. FlashAttention-2 additionally splits work along the sequence-length dimension itself, so long-context, small-batch workloads — exactly the contract-reading scenario that opened this chapter — can still saturate the chip. Third, it changed how work is split across the warps inside one thread block, partitioning by the key/value dimension rather than the query dimension so that warps in the forward pass avoid needing to synchronize through shared memory to combine partial results. Together these changes roughly doubled measured throughput over the original kernel, with the paper reporting up to about 72% of the A100's theoretical peak FLOPs on the attention computation, compared to roughly 25–40% for the first version.
How the tiles move: HBM and SRAM during one query tile's pass
A different bottleneck: the KV cache during inference
Training and prefill are not the only place attention's memory footprint bites. Once a model is deployed — say, an Indian e-commerce platform's chat support assistant fielding a festival-sale traffic spike, with thousands of concurrent conversations open at once — each request runs autoregressive decoding: one new token generated at a time, each attending back over every previous token's key and value vectors. Recomputing those key/value projections from scratch at every step would be wasteful, so serving systems cache them: the KV cache. But the KV cache is not a training-time score matrix that disappears after one layer's forward pass — it is per-request state that must persist in GPU memory for the entire length of a conversation, competing directly with every other concurrent request for the same fixed pool of HBM.
The size adds up fast. For a 13B-parameter model shaped like OPT-13B — 40 layers, hidden dimension 5,120, fp16 weights — each single token's KV cache entry costs 2 (for K and V) × 40 layers × 5,120 hidden dim × 2 bytes = 819,200 bytes ≈ 800 KB, a figure used directly in Woosuk Kwon and colleagues' 2023 SOSP paper "Efficient Memory Management for Large Language Model Serving with PagedAttention." At a maximum context window of 2,048 tokens, one fully-utilized request's cache alone would need 2048 × 800 KB ≈ 1,600 MB (1.5625 GB) — and a serving system that naively reserves a full 2,048-token contiguous block up front, for every incoming request, regardless of how long the actual reply turns out to be, will run out of room for concurrent users almost immediately, even on an 80GB GPU.
PagedAttention: memory management borrowed from an operating system
Kwon et al.'s fix treats the KV cache the way an operating system treats a process's virtual memory. Instead of one contiguous reservation per request, the cache is divided into fixed-size physical blocks (say, 16 tokens each), stored anywhere in GPU memory. Each request keeps a small block table — directly analogous to a page table — mapping its logical token positions to whichever physical blocks currently hold them. New blocks are allocated only as generation actually produces enough new tokens to fill one, and the attention kernel, instead of assuming its KV cache is one contiguous span, walks the block table and gathers whichever physical blocks it's pointed to. Because blocks never need to be contiguous, this also enables copy-on-write sharing: several sequences generated from the same prompt (parallel sampling, beam search) can point their block tables at the very same physical blocks for the shared prefix, only forking off a private copy once one of them starts to diverge.
Return to the median chat reply: 90 tokens generated, against a reserved maximum of 2,048. A naive contiguous allocator reserves the full 1,600 MB (1.5625 GB) and uses only 90 × 800 KB ≈ 73.7 MB of it — a utilization of just 4.4%, verified directly:
import math
block_size = 16
kv_per_token_bytes = 800 * 1024
max_context = 2048
actual_tokens = 90
naive_reserved = max_context * kv_per_token_bytes
naive_used = actual_tokens * kv_per_token_bytes
naive_util = naive_used / naive_reserved
blocks_needed = math.ceil(actual_tokens / block_size)
paged_reserved = blocks_needed * block_size * kv_per_token_bytes
paged_util = naive_used / paged_reserved
print(round(naive_util * 100, 1)) # 4.4
print(round(paged_util * 100, 1)) # 93.8With block_size = 16, the request needs ⌈90/16⌉ = 6 blocks — 96 token-slots, about 75 MB reserved instead of 1,600 MB, for a utilization of 93.8%. The remaining 6.2% waste is internal fragmentation confined entirely to the one partially-filled last block — there is no external fragmentation at all, because the block table can point to physical blocks scattered anywhere in memory, exactly as an OS page table needn't map contiguous virtual pages to contiguous physical frames. The vLLM system built on this idea reports keeping memory waste below roughly 4% in practice, against 60–80% for the contiguous-reservation systems that preceded it, translating directly into more concurrent requests served per GPU.
It's worth being precise about the boundary between the two techniques in this chapter. FlashAttention reduces HBM traffic within one attention computation, for a single sequence, mostly relevant during training and the prefill phase of inference. PagedAttention manages HBM capacity across many concurrent decoding requests, mostly relevant during the autoregressive generation phase after prefill. Modern serving stacks use both: FlashAttention-style fused kernels to compute attention efficiently once the KV blocks are in hand, and PagedAttention-style block tables to decide where those KV blocks live. They are complementary answers to the same underlying constraint — HBM is scarce and slow to reach — applied at different points in the pipeline.
Common misconception
The mistake students make most often, especially right after learning linear attention, is assuming FlashAttention is "just another approximation, like linear attention but with a faster implementation." It is not. Linear attention replaces the softmax kernel with a different, cheaper function of Q and K — it changes the mathematical operation being computed, trading exactness for an O(N) instead of O(N²) attention pattern. FlashAttention computes the exact same quantity, softmax(QKᵀ/√d)V, to the same numerical precision (up to the ordinary floating-point rounding differences that come from summing the same numbers in a different order) — the worked example above showed the tiled and direct computations agreeing to eight decimal places. What FlashAttention changes is not the formula but the memory-access pattern: it never lets the full N×N score matrix touch HBM, fusing the three separate passes of a standard implementation into one kernel that keeps every intermediate value in on-chip SRAM. The compute cost remains O(N²d), exactly as with standard attention — full, dense, unapproximated attention over every pair of tokens. If a system needs sub-quadratic compute, it needs an actual approximation like linear or sparse attention; if it needs the existing quadratic compute to run faster and fit in memory without changing the answer, FlashAttention is the right tool, and PagedAttention solves a third, separate problem — reusing that same exact KV cache efficiently across many simultaneous requests.
Active recall
Q1. In one sentence, why is standard (unfused) attention IO-bound rather than compute-bound on a modern GPU, even though the total arithmetic is the same either way?
Q2. A classmate says FlashAttention "approximates attention the same way linear attention does, just implemented more cleverly." What is wrong with this claim?
Q3. Using the same q = [1,0], K, V from the worked example, process all four keys in a single block (block_size = 4) instead of two blocks of two. Trace m and ℓ and confirm the final output. What does this tell you about what block size does, and doesn't, change?
Q4. If the SRAM budget M available to a FlashAttention kernel were cut in half, name one concrete consequence for the tiling algorithm's block sizes, and one consequence for the Θ(N²d²/M) HBM-access bound.
Q5. A serving system switches from reserving a contiguous 2,048-token slot per request to PagedAttention with block_size = 16, for a request whose actual reply is 90 tokens on an OPT-13B-shaped model (800 KB of KV cache per token). Roughly how much memory is now reserved for that request, and precisely where does the remaining waste come from?
Q6. FlashAttention already parallelizes across batch size and attention heads. Why did FlashAttention-2 add parallelization across the sequence-length dimension as well?
A1. Because the intermediate score and probability matrices are N×N and, in a naive kernel, get written to and read back from HBM (at ~1.5–2 TB/s) three separate times, while the tensor cores that would consume them — fed at up to ~19 TB/s from SRAM — sit idle waiting on those transfers; wall-clock time tracks memory bandwidth, not the FLOP count.
A2. Linear attention changes the mathematical function being computed — it swaps the softmax kernel for a cheaper approximate one to get sub-quadratic compute. FlashAttention computes the identical softmax(QKᵀ/√d)V — same compute complexity, same numerical result up to floating-point rounding order — and only changes where intermediate values are staged in the memory hierarchy. It is an exact, IO-efficient reimplementation, not an approximation.
A3. With one block covering all four keys, there is a single iteration: m_new = max(-∞, max([1,0,3,-1])) = 3; correction is exp(-∞) = 0 so the (empty) prior state contributes nothing; p = exp([1,0,3,-1] − 3) ≈ [0.1353, 0.0498, 1, 0.0183]; ℓ = 1.2034 — identical to the two-block trace's final ℓ. O = p·V ≈ [4.4979, 5.4979], matching both the two-block trace and the direct computation exactly. Block size therefore never changes the mathematical result — only how many HBM↔SRAM round trips occur and how much on-chip memory must be available at once to hold one block's Qᵢ, Kⱼ, Vⱼ, Sᵢⱼ simultaneously. Larger blocks mean fewer, bigger transfers; they must still fit within M.
A4. Tiling: the maximum Br, Bc that let Qᵢ, Kⱼ, Vⱼ, Sᵢⱼ fit in SRAM together must shrink, since their combined footprint scales with Br·d + Bc·d + Br·Bc; smaller blocks mean more inner-loop iterations (Tc grows) for the same sequence length, though correctness is unaffected per Q3. Formula: since Θ(N²d²/M) is inversely proportional to M, halving M roughly doubles the required HBM accesses, shrinking (but not eliminating) FlashAttention's advantage over standard attention's Θ(Nd+N²).
A5. ⌈90/16⌉ = 6 blocks are needed, reserving 6 × 16 = 96 token-slots — about 96 × 800 KB ≈ 75 MB, versus 2048 × 800 KB ≈ 1,600 MB (1.5625 GB) under contiguous reservation, a roughly 21× reduction for this one request. The only waste is the 6 unused slots in the last, partially-filled block (96 − 90 = 6, about 6.2%) — pure internal fragmentation. There is no external fragmentation, because the block table lets those 6 blocks live at any scattered physical addresses; nothing needs to be contiguous.
A6. For long-context, small-batch workloads — one very long document processed with batch size 1 is the extreme case — the batch×heads grid alone may contain fewer independent blocks than the GPU has streaming multiprocessors, leaving hardware idle no matter how efficient each block's own tiling is. Splitting the query-sequence dimension into additional independent units of work gives the scheduler enough parallelism to fill every SM even when batch size can't, which is exactly the situation a single 18,000-token contract creates.
Think About It
Think about this: How would you explain flash attention and memory-efficient attention variants to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind flash attention and memory-efficient attention variants, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.