Flash Attention Explained: How IO-Aware Computing Made Long Context Practical
TL;DR
Flash Attention is an IO-aware rewrite of the standard attention algorithm that produces identical results but runs 5 to 9 times faster and uses far less GPU memory. It does not approximate or skip any computation. Instead, it moves computation to stay within the GPU's fast on-chip cache rather than shuttling data back and forth to slow video memory. The practical consequences: longer context windows at lower cost, larger batch sizes, and the ability to train models that would otherwise not fit on available hardware. Every Llama, Mistral, Claude, and GPT-4-class model you work with today uses some form of Flash Attention internally.
The AI PM Minute
One tactic to make you a sharper AI PM, twice a week. 60 seconds to read. Free.
No fluff. Unsubscribe anytime.
The Problem Standard Attention Has
To understand why Flash Attention matters, you need to know one fact about how GPUs work: not all memory is created equal. GPUs have two kinds of memory, and the difference between them is enormous.
HBM (High-Bandwidth Memory)
40 to 80 GB on an A100 or H100 — ~2 TB/s bandwidth
This is the large video memory your model weights live in. Fast by consumer standards, but the bottleneck for attention computation. Reading from and writing to HBM takes orders of magnitude longer than computing.
SRAM (Static RAM / on-chip cache)
~20 MB total on an A100 — ~19 TB/s bandwidth — roughly 10x faster
This is the tiny, blazingly fast memory right next to the compute units. Computation that stays in SRAM is an order of magnitude faster. The limitation: it can only hold a fraction of the attention matrix for realistic sequence lengths.
Standard attention computes the full N x N attention matrix (where N is sequence length) and writes it entirely to HBM before reading it back for the softmax step. For a 32K token sequence, that matrix is 32,000 x 32,000 values. Writing that to slow HBM and reading it back dominates total runtime. The GPU arithmetic units sit idle waiting for data to move.
The IO bottleneck in numbers
For a 4K sequence, standard attention spends roughly 40% of wall-clock time doing arithmetic and 60% moving data. At 32K, that ratio flips to roughly 10% arithmetic and 90% memory IO. The GPU is not compute-bound; it is memory-bandwidth-bound. Flash Attention attacks the bottleneck that actually matters.
How Flash Attention Fixes It: Tiling and Recomputation
Flash Attention, introduced by Tri Dao et al. in 2022, makes one key insight: you do not need to materialize the full N x N attention matrix. You can compute the exact same output by processing the attention in tiles that fit inside SRAM, fusing all the attention sub-operations (QK multiply, softmax, multiply by V) into a single GPU kernel pass.
Tile the input
Flash Attention loads small blocks of the Q, K, and V matrices from HBM into SRAM, one tile at a time. The tile size is chosen to fit within the on-chip cache. This is the 'tiling' step.
Fuse into one kernel
Standard attention uses multiple GPU kernels: one for QK, one for softmax, one for the V multiply, with HBM reads and writes between each. Flash Attention fuses all of these into a single kernel that never writes intermediate results to HBM.
Compute softmax incrementally
The tricky part is softmax, which needs the full row sum. Flash Attention uses a numerically stable online softmax algorithm (Milakov & Gimelshein, 2018) that updates the running max and running sum as it processes each tile, reaching the exact same final result.
Recompute during backward pass
To save memory during training, Flash Attention does not save the attention matrix for the backward pass. It recomputes it on-the-fly from the saved Q, K, V tiles. This trades a small amount of extra compute for a massive memory saving. In practice, the extra compute is negligible versus the memory savings.
The result: Flash Attention achieves 2 to 4x speedup for forward passes and 5 to 9x speedup for training (which includes the backward pass) compared to standard attention, with no loss in output quality. It is not an approximation. The math is identical. Only the memory access pattern changes.
Flash Attention 1, 2, and 3: What Changed
Flash Attention 1 (2022)
What it did: Original tiling algorithm by Tri Dao, Daniel Fu, and colleagues at Stanford. Proved exact attention was possible without materializing the full attention matrix. Achieved 2 to 4x speedup on A100s. Required writing custom CUDA kernels.
Remaining limit: Work partitioning was suboptimal; did not fully utilize all GPU parallelism dimensions.
Flash Attention 2 (2023)
What it did: Redesigned the work partitioning to reduce non-matmul FLOPs and better parallelize across the sequence dimension. Result: 50 to 75% faster than Flash Attention 1, reaching 50 to 73% of theoretical GPU peak FLOPs. Now the standard implementation in most major frameworks.
Remaining limit: Optimized for A100/H100 architecture; less tuned for non-NVIDIA hardware.
Flash Attention 3 (2024)
What it did: Purpose-built for H100 Hopper architecture. Exploits new hardware features: asynchronous tensor core operations, warp specialization, FP8 support. Achieves 1.5 to 2x speedup over FA2, approaching 75% of H100 peak FLOPs.
Remaining limit: H100-specific; benefits require the newer hardware generation.
Bridge Technical Depth and Product Decisions
The AI PM Masterclass covers how architectural choices like attention algorithms translate into real product tradeoffs around cost, latency, and context length.
What Flash Attention Means for Products You Ship
Flash Attention is invisible to you at the API level — you call an endpoint and get tokens back. But the architectural decisions made possible by Flash Attention directly shape what you can build, at what cost.
Long context windows became economically viable
Without Flash Attention, 128K token context windows would require holding a 128K x 128K attention matrix in GPU memory — roughly 64 GB per attention layer per request. With Flash Attention, memory scales linearly rather than quadratically. This is the primary reason 100K to 1M token context lengths are available today at prices that can be productized.
Your batch size affects latency differently than you might expect
Flash Attention is especially efficient at larger batch sizes because it amortizes the memory access cost across more sequences. If your product processes requests in batches (background jobs, document pipelines), you pay less per token than if you process single requests at low traffic. Size your inference workloads accordingly.
Context window cost is not linear with length
Even with Flash Attention, attention is IO-quadratic in the worst case (though memory is linear). A 128K context call still costs significantly more than two 64K calls. When designing product features that use long context, model the cost curve, not just the price per 1K tokens.
Hardware generation matters for your provider choices
Flash Attention 3 is H100-specific. Providers on older A100 infrastructure run FA2 at lower throughput. When choosing between inference providers, GPU generation is a meaningful proxy for attention efficiency, which translates to latency and cost for context-heavy workloads.
Open-weight model performance depends on whether FA is enabled
If you self-host or fine-tune open models like Llama, Mistral, or Falcon, ensure Flash Attention is enabled in your inference stack (vLLM, TGI, and SGLang all support it). A default PyTorch setup without Flash Attention is materially slower and more expensive for sequences above 2K tokens.
Flash Attention vs. Related Techniques: What PMs Should Know
Flash Attention is often confused with other attention optimizations. They address different problems and are often used together.
Flash Attention
Addresses: IO efficiency: how fast data moves between memory tiers
Result: Same output, faster compute, lower memory footprint
With Flash Attention: Yes, used alongside all others
Grouped Query Attention (GQA)
Addresses: KV cache size: how much memory is needed per token to serve inference
Result: Smaller KV cache, faster decoding, some quality tradeoff
With Flash Attention: Yes, GQA and Flash Attention are orthogonal optimizations
Sparse Attention
Addresses: Computation reduction: skipping attention between distant tokens
Result: Approximate output, significant speedup for very long sequences
With Flash Attention: Flash Attention can accelerate sparse patterns too
Sliding Window Attention
Addresses: Context locality: each token only attends to nearby tokens in a window
Result: Linear complexity, but no global context beyond the window
With Flash Attention: Used in Mistral models; Flash Attention accelerates the window ops
The practical summary for AI PMs
Flash Attention makes exact attention faster. Grouped Query Attention makes the memory cache of token representations smaller. Sparse and sliding window attention reduce which tokens talk to which. A high-performance production LLM typically uses all of these simultaneously. You do not need to choose between them.
When Flash Attention Should Come Up in Product Conversations
Most of the time you will not need to ask about Flash Attention explicitly. The major API providers handle it. But these are the situations where knowing about it sharpens a product decision:
Evaluating an inference provider for long-context workloads
Ask: What GPU generation is your inference cluster running? Do you use Flash Attention 2 or 3? What is your measured throughput for 64K versus 128K context lengths?
Debugging unexpectedly high latency on document-heavy features
Ask: Is the serving stack actually using Flash Attention? (Many teams self-host open models without enabling it.) What is the sequence length distribution of our real traffic?
Designing pricing or cost models for a context-heavy feature
Ask: Our cost per token goes up nonlinearly with context length even with Flash Attention. How is our pricing reflecting the actual compute cost curve, not just a flat per-token rate?
Deciding whether to self-host or use API for an open-weight model
Ask: If we self-host, will the team enable Flash Attention? What GPU hardware will we run on? Without FA3-class optimization, H100 cost savings over API may not materialize.
Go Deeper on Technical AI for Product Managers
The AI PM Masterclass covers how infrastructure and architecture decisions translate into product strategy. Learn from a Salesforce Sr. Director PM who built AI products at scale.
Related Articles
Before you go: get the AI PM Minute
One tactic to make you a sharper AI PM, twice a week. 60 seconds to read. Free.
No fluff. Unsubscribe anytime.