Why It Matters
For agents, every serving-layer gain compounds across long contexts and many calls; FlashAttention is the kernel-level reason longer prompts and cheaper serving became practical.
Problem
- Self-attention has quadratic time and memory complexity in sequence length.
- Approximate attention methods reduce theoretical compute but often fail to improve wall-clock time or reduce quality.
- Standard attention materializes the N by N attention matrix and moves it through GPU HBM, making memory traffic a bottleneck.
- Deep learning framework interfaces often hide fine-grained memory movement decisions needed for optimal kernels.
Key gap: The paper asks how to compute exact attention with fewer HBM reads/writes by explicitly designing for the GPU memory hierarchy.
Method
Load blocks of Q, K, and V into SRAM, compute attention tile by tile with online softmax, write only final outputs and small normalization state to HBM, and recompute tiles during backward rather than storing the whole attention matrix.
- Partition Q, K, and V into blocks sized for GPU SRAM.
- Loop over K/V blocks and Q blocks, loading only the active tiles into SRAM.
- Compute local attention scores for the tile.
- Update online softmax statistics and accumulated outputs without storing the full N by N score matrix.
- Write final output blocks and small normalization values to HBM.
- During backward pass, recompute required attention tiles on chip using saved normalization data.
- Fuse attention operations into one CUDA kernel for fine-grained memory control.
- Extend the same IO-aware pattern to block-sparse attention for longer contexts.
Claims
C1 · paper-supported
Approximate attention methods often reduce compute complexity but fail to deliver wall-clock speedups because they ignore memory-access overhead.
Scope: Abstract and Introduction.
- E-approximate-attention-gap: The paper argues approximate attention methods focus on FLOP reduction and often do not achieve wall-clock speedup because they ignore IO and memory-access overheads.
C2 · paper-supported
FlashAttention is an exact attention algorithm that uses tiling to reduce reads and writes between GPU HBM and on-chip SRAM.
Scope: Abstract and Introduction.
- E-tiling: FlashAttention uses tiling to load Q, K, and V blocks into SRAM and avoid reading/writing the large attention matrix to HBM.
C3 · paper-supported
The paper analyzes FlashAttention's IO complexity and proves it is IO-optimal for a range of SRAM sizes.
Scope: Abstract and Introduction.
- E-io-complexity: The paper analyzes IO complexity, states FlashAttention requires fewer HBM accesses than standard attention, and proves optimality for a range of SRAM sizes.
C4 · paper-supported
FlashAttention reports major attention-kernel speedups by avoiding materialization of the N by N attention matrix in HBM.
Scope: Introduction and Figure 1.
- E-attention-kernel-speed: The paper reports that not materializing the N by N attention matrix gives a 7.6x speedup on GPT-2 attention computation and up to 9x fewer HBM accesses in typical settings.
C5 · paper-supported
FlashAttention reports 15% end-to-end speedup on BERT-large, 3x speedup on GPT-2, and 2.4x speedup on Long Range Arena training workloads.
Scope: Abstract and Introduction.
- E-training-speedups: FlashAttention trains BERT-large 15% faster than the MLPerf 1.1 training speed record, trains GPT-2 at sequence length 1K 3x faster than HuggingFace/Megatron baselines, and trains Long Range Arena at sequence length 1K-4K 2.4x faster than baselines.
C6 · paper-supported
FlashAttention and block-sparse FlashAttention enable longer context, including better-than-chance performance on Path-X at 16K length and Path-256 at 64K length.
Scope: Abstract and Introduction.
- E-long-context: The paper reports 0.7 better GPT-2 perplexity, 6.4 points lift on long-document classification, 61.4% accuracy on Path-X at 16K, and 63.1% on Path-256 at 64K using FlashAttention/block-sparse FlashAttention.
C7 · paper-supported
FlashAttention uses recomputation in backward pass to avoid storing the large attention matrix, trading extra FLOPs for much lower HBM traffic and memory footprint.
Scope: Introduction.
- E-recompute: The paper stores softmax normalization factors from the forward pass and recomputes attention on-chip during the backward pass rather than reading the intermediate attention matrix from HBM.
Agent Infrastructure Takeaways
- Serving improvements start at the kernel. Agent workflows become cheaper when the primitive attention operation stops wasting HBM bandwidth.
- FLOPs are not the whole cost model. For long-context agents, memory traffic and KV residency often dominate.
- FlashAttention is exact, so it can improve throughput/latency without changing model semantics the way approximate attention might.
- Longer context is an infrastructure feature. It changes what agents can carry in prompt state, tool traces, and retrieval context.
- IO-awareness recurs up the stack: FlashAttention for attention matrices, PagedAttention for KV cache, SGLang for prefix reuse.
- Better kernels do not remove the need for scheduling, paging, prompt compression, or durable state; they make those layers cheaper.
Caveats
- The paper is primarily about training and attention kernels, though the kernel later became central to inference stacks.
- Kernel speedups are hardware- and implementation-dependent.
- Block-sparse FlashAttention is approximate and should not be conflated with exact FlashAttention.
- Exact attention still has quadratic work in sequence length; FlashAttention reduces memory traffic and footprint, not the mathematical attention dependency.
- The original 2022 results predate FlashAttention-2/3 and current production kernels.
- Longer context can help agents but can also increase latency, prompt-management complexity, and exposure to irrelevant or malicious context.