Why FlashAttention Is a Kernel Story
FlashAttention achieves its speedups not by reducing FLOPs but by restructuring the attention computation into a single tiled CUDA kernel that fits working data in on-chip SRAM, eliminating the dominant cost of round-tripping through GPU HBM.
Standard scaled dot-product attention on a sequence of length N requires materialising an N x N matrix: the full attention score grid. On an A100 GPU with 80 GB of HBM running at roughly 2 TB/s, writing and re-reading that matrix for a 2048-token sequence with 64 heads costs tens of milliseconds of pure memory time, even though the arithmetic is trivial. FlashAttention (Dao et al., 2022) reports a 3x wall-clock speedup on GPT-2 style training without changing a single attention output value. The source of that speedup is not new maths. It is a kernel rewrite.
The real bottleneck: memory bandwidth, not FLOPs
Modern GPU compute has outrun its memory system. An A100 delivers around 312 TFLOPS of FP16 throughput but only ~2 TB/s of HBM bandwidth. A simple back-of-envelope: reading a 1 GB tensor takes ~0.5 ms; executing 10^12 FLOPs on it takes ~3 ms. For workloads that are compute-bound, bandwidth is irrelevant. For workloads where the ratio of arithmetic operations to bytes moved is low (called arithmetic intensity), bandwidth is everything.
Attention has terrible arithmetic intensity in its naive form. The forward pass for one head reads Q, K, V (all O(Nd)), writes S = QK^T (O(N^2)), reads S again for softmax, writes P (O(N^2)), reads P and V, writes O (O(Nd)). The N^2 matrices dominate for long sequences. Each one makes a full round-trip through HBM, even though the computation each byte participates in is a handful of multiplications and adds.
Standard attention IO (per head, sequence length N, head dim d):
Read: Q, K, V -> 3Nd floats from HBM
Write: S = QK^T -> N^2 floats to HBM
Read: S (for softmax) -> N^2 floats from HBM
Write: P = softmax(S) -> N^2 floats to HBM
Read: P, V -> N^2 + Nd floats from HBM
Write: O = PV -> Nd floats to HBM
Total reads + writes: ~4N^2 + 6Nd floats
For N=2048, d=64, that is roughly 128 M floats (256 MB in FP16) flowing through HBM per head. With 64 heads and the whole model, this adds up fast.
Tiling: fitting the work inside SRAM
The key insight is that softmax and the output accumulation can be computed in tiles without ever materialising the full N x N score matrix in HBM, as long as a numerically stable online softmax update is used.
The SRAM of a single A100 streaming multiprocessor (SM) is around 192 KB. A tile of Q with block size Br and a tile of K with block size Bc together occupy 2 * Br * Bc * d floats. For Br = Bc = 64 and d = 64, that is 2 * 64 * 64 * 64 * 2 bytes = 1 MB, still too large. In practice FlashAttention uses Br, Bc around 32-64 and fits each tile inside the shared memory budget.
The online softmax trick (Milakov and Gimelshein, 2018) allows incremental updates. For each new tile of scores s_i, maintain a running maximum m and a running denominator l:
m_new = max(m_old, rowmax(s_i))
l_new = exp(m_old - m_new) * l_old + rowsum(exp(s_i - m_new))
O_new = diag(exp(m_old - m_new)) * O_old + exp(s_i - m_new) * V_i
After all tiles, rescale: O_final = O / l. The final output is identical to full attention. No approximation. The N^2 HBM traffic is replaced by O(N * d) HBM traffic, because tiles of K and V are loaded once and discarded.
FlashAttention IO (per head):
Read: Q, K, V -> 3Nd floats from HBM (same as before)
Write: O, l, m -> Nd + 2N floats to HBM
Total: ~4Nd + 2N floats (no N^2 term)
The algorithmic complexity of FLOPs is unchanged at O(N^2 d). Only the IO complexity improves, from O(N^2) to O(Nd). For N >> d (the common regime), this is the dominant factor.
Why this only works as a fused kernel
The tiling only eliminates HBM traffic if Q, K, V tiles, score computation, softmax update, and output accumulation all happen inside a single kernel launch. If those operations are implemented as separate PyTorch ops (matmul, softmax, matmul), each op forces intermediates to HBM because PyTorch cannot hold state between kernel launches.
Fusion means: one kernel, one set of register files and shared memory, no eviction to HBM between the sub-steps. This is why FlashAttention is not expressible as a composition of existing high-level ops. It requires a hand-written (or Triton-compiled) kernel that takes full ownership of the data flow.
FlashAttention was originally written in CUDA C++. FlashAttention-2 (Dao, 2023) additionally tuned work partitioning: rather than assigning one tile per thread block and serialising across tiles for a single head, it parallelises across the sequence dimension, increasing GPU occupancy and cutting non-matmul overhead. The result on A100s is 50-73% of peak theoretical FLOPs/s, compared to roughly 25-35% for the original FlashAttention and well under 10% for the naive PyTorch implementation.
The triton path and torch.compile
Writing CUDA kernels is expensive. Triton (Tillet et al.) provides a Python-level DSL that compiles blocked operations to PTX, handling many of the low-level concerns (shared memory layout, async copy, warp synchronisation) automatically. A fused softmax tutorial in Triton achieves ~4x bandwidth improvement over an unfused implementation with around 30 lines of Python, because it eliminates the same HBM round-trips described above.
torch.compile with the Inductor backend can fuse point-wise ops automatically (e.g., bias add + dropout + activation) but cannot synthesise a tiled attention kernel from scratch. It does not know that softmax and two matmuls can be merged into one tiled loop. This is the limit of general-purpose fusion: it handles elementwise and reduction fusions, but structure-specific algorithms like FlashAttention require domain knowledge that compilers do not currently infer. F.scaled_dot_product_attention in PyTorch 2.x routes to FlashAttention or memory-efficient attention kernels when available, effectively hard-coding the domain insight.
When it falls down
Very short sequences. For N below ~128, the N^2 matrices are small enough to fit in cache regardless. FlashAttention's tile management adds overhead (synchronisation, index arithmetic) that can make it slower than a standard unfused implementation.
Non-standard attention variants. Any modification that changes the access pattern (sparse attention, relative position biases with learned kernels, custom masking that cannot be expressed as a tile-local mask) may require a new kernel. Adapting FlashAttention to sliding-window attention (Longformer/Mistral-style) or cross-attention with non-square K,V grids requires non-trivial modifications to the tiling logic.
Backward pass complexity. The forward pass does not store S or P (to save HBM). The backward pass must recompute them from the stored output and log-sum-exp values. This recomputation is correct but increases the FLOP count in the backward relative to storing intermediates. For inference-only workloads, the forward-pass savings dominate. For training, the FLOP overhead in the backward is real and worth profiling.
Head dimension limits. The tile sizes that fit in SRAM depend on d. Standard head dimensions (64, 128) work well. Very large head dimensions (e.g., 256 as used in some multi-query architectures) require smaller tiles or spill to HBM, reducing the bandwidth advantage.
Precision edge cases. The online softmax uses floating-point max and exponential operations in tiles. With FP16, the dynamic range is limited. FlashAttention uses FP32 accumulators internally for stability, but using BF16 or FP8 for the score tiles introduces additional precision considerations that naive usage can overlook.
Further reading
7 flashcards for this concept
Click a card to reveal the answer.