What problem does FlashAttention actually solve, and wh...

Standard attention computes the full query-key similarity matrix, applies softmax, and multiplies by values — and the naive implementation materializes that full N×N matrix in GPU high-bandwidth memory (HBM) at each step. HBM is large but comparatively slow; on-chip SRAM is tiny but far faster. For long sequences, the dominant cost of attention isn’t the arithmetic at all — it’s shuttling that large intermediate matrix in and out of slow memory.

FlashAttention restructures the same computation to process it in tiles small enough to fit in fast on-chip SRAM, computing partial softmax statistics incrementally (an online softmax) so the full N×N matrix is never fully materialized in HBM at once. Critically, this produces the mathematically exact same output as standard attention — it’s not an approximation. The speedup comes entirely from moving less data through slow memory, which is exactly the kind of win the roofline model predicts matters most for a memory-bound operation like attention over long sequences.