How Flash Attention Works: Tiling, Softmax, GPU I/O

I still remember the exact moment a 32k-context fine-tuning job OOM'd on an 80GB A100, three hours into a run, on a batch size that had worked fine at 8k. The traceback pointed at the attention layer, and the culprit wasn't the model weights - it was the intermediate attention matrix, sitting there at seq_len squared, eating memory nobody had budgeted for. That night sent me down the rabbit hole of understanding how flash attention works, and it rewired how I think about GPU kernels entirely.

Here's the part that surprised me: FlashAttention doesn't cheat. It's not a sparse approximation, not a low-rank trick, not a "good enough" shortcut like some of the linear-attention variants floating around. It computes the exact same softmax(QK^T/√d)V that vanilla attention computes - bit-for-bit equivalent output - but it never materializes the full score matrix in slow GPU memory.

That distinction matters enormously once you're debugging memory blowups on a production box at 2am. This piece walks through the actual mechanics: tiling, online softmax, kernel fusion, recomputation, and what changes across FlashAttention 1, 2, and 3.

How Flash Attention Works: The Exact-Attention Idea

So what's flash attention, exactly? It's an I/O-aware algorithm for computing standard attention, not a new attention architecture and not an approximation. If you've read up on how attention works in transformer architecture, the formula is unchanged: softmax(QKᵀ/√d)V, with causal masking applied when you need it for autoregressive decoding.

What changes is how that formula gets executed on the GPU. The mental model I use: chop Q, K, and V into small tiles, pull each tile into fast on-chip memory, do the math there, and never write the full n×n score matrix out to slow memory. Standard attention builds that whole matrix, writes it, reads it back for softmax, writes it again - FlashAttention just skips that round trip entirely.

The output is mathematically identical to vanilla attention, down to floating-point rounding. That's the detail that made me trust it enough to ship it - you're not trading accuracy for speed, you're trading memory traffic for speed.

Why Standard Attention Becomes a GPU Memory I/O Bottleneck

Trace the standard pipeline and you'll count at least four separate kernel launches. Compute QKᵀ, apply the causal mask, run softmax, apply dropout, then multiply by V. Each of those steps is its own CUDA kernel, and each one reads its input from HBM and writes its output back to HBM before the next kernel even starts.

That's the bottleneck. GPU high-bandwidth memory is large - 80GB on an 80GB A100, the same card that OOM'd on me - but it's slow relative to on-chip SRAM and registers, which are tiny by comparison and sit right next to the compute cores. For a 32k-token sequence, the score matrix alone holds roughly 32,000² entries, about 1 billion values, and standard attention writes a matrix that size to HBM, reads it back for softmax, and writes it again. That's pure bookkeeping traffic, not useful computation.

Here's the subtlety people miss: both approaches do the same quadratic amount of arithmetic. FlashAttention doesn't reduce FLOPs - it reduces memory traffic [1]. Wall-clock time tracks bytes moved, not operations counted, and that's exactly where standard attention stalls.

How FlashAttention Computes Exact Attention Tile by Tile

Instead of building the whole matrix, FlashAttention loops over blocks, keeping everything one query tile needs on-chip until it's done.

Tiling Q, K, and V for Fast On-Chip Memory

The kernel loads one block of Q rows into SRAM, then streams K and V blocks past it one at a time, computing partial scores as it goes. Tile size isn't arbitrary - it's picked around shared-memory capacity, head dimension, and dtype, so an A100 running fp16 with a 128-dim head lands on different block sizes than an H100 running bf16. Causal masking and padding get applied per-tile too, with boundary tiles handling the ragged edge where the mask cuts a block in half.

Online Softmax Without the Full Score Matrix

Each K/V tile updates a running row-max and running normalization sum instead of waiting for the full row. When a later tile produces a bigger max, the kernel rescales the previous partial output and sum by the ratio of old to new max before adding the new contribution. This recurrence is exact - it converges to the identical value you'd get running softmax over the complete row.

Kernel Fusion and Incremental Output Accumulation

Score computation, masking, softmax stats, dropout, and the V-weighted accumulation all stay inside one fused kernel - no handoff between separate launches. That kills most of the intermediate HBM writes and synchronization stalls standard attention pays for. Only the final output block and a couple of small per-row statistics ever get written back to HBM.

The Backward Pass: Recomputation Instead of Stored Activations

Standard attention's backward pass leans on the saved probability matrix from the forward pass - you need those softmax outputs to compute gradients with respect to Q, K, and V. Storing that matrix for every layer, every head, across a long sequence is exactly the quadratic memory cost that got me into trouble that night.

FlashAttention refuses to keep it. It saves only the output and the compact per-row statistics - the running max and normalization sum - then recomputes the score tiles on the fly during backprop, tile by tile, in the same fused fashion as the forward pass.

That's more arithmetic, not less. But extra matmuls on-chip are cheap next to hauling a full quadratic matrix through HBM twice, so the trade nets out firmly in FlashAttention's favor.

Don't confuse this with layer-level gradient checkpointing. Checkpointing drops whole activations and recomputes entire forward passes; FlashAttention's recomputation is scoped to one kernel's tiles, invisible outside the attention op itself.

FlashAttention 1 vs 2 vs 3: Algorithm Generations and GPU Targets

FlashAttention 1 established the whole idea - tiling, online softmax, fused kernels, recomputed backward pass. Everything above is version 1's contribution, and it's still the algorithm the later versions build on rather than replace.

FlashAttention 2 didn't change the math; it rewrote the work partitioning. It splits work across sequence length as well as batch and heads, keeping streaming multiprocessors busier on long sequences with small batch counts, and it trims the non-matmul overhead - masking, rescaling, bookkeeping - that was quietly stealing GPU cycles from the actual matrix multiplies.

FlashAttention 3 targets Hopper specifically. It leans on asynchronous execution and warp specialization to overlap matmul and softmax work on the same streaming multiprocessor, and adds lower-precision paths on GPUs that support them.

These are algorithm generations, not interchangeable package flags:

• A Hopper-tuned FA3 kernel doesn't fall back gracefully to an Ampere card - it just doesn't apply.
• FA2's occupancy gains matter most on long sequences with small batches, less so on short, batch-heavy workloads.
• Framework backends (PyTorch SDPA, xFormers) may lag behind the standalone package's latest release.
• Support matrices shift often enough that I'd check the current repository docs before assuming your GPU and CUDA version qualify.

Is FlashAttention Faster? Memory Use and Where Gains Vary

Short answer: usually, but "faster" depends heavily on where you're measuring. End-to-end training on GPT-2-scale models runs faster compared to a standard PyTorch attention implementation, and speedups tend to grow with sequence length, because that's when standard attention's memory traffic dominates.

The arithmetic stays quadratic in sequence length either way - FlashAttention doesn't change that. What it changes is auxiliary memory: instead of storing an N² score matrix, you keep roughly linear working memory per tile, which is the real unlock for long-context LLM training.

Benchmarks are noisy because so many variables move at once: GPU generation, sequence length, head dimension, batch size, causal versus non-causal masking, dtype, and whether you're timing training or inference. Competing fused backends sometimes close the gap for specific shapes.

Gains can shrink or vanish on short sequences, odd head dimensions, unsupported shapes that fall back to a slower path, or when conversion and JIT compilation overhead eats the savings. Even then, the lower memory footprint often lets you run longer contexts or bigger batches than raw latency numbers would suggest.

Flash Attention Without CUDA: Hardware and Fallback Paths

People ask me for "flash attention without CUDA" more than any other variant of this question, and the honest answer is: that's not really a thing. The mainline package on GitHub is built around hand-written GPU kernels, tuned for specific NVIDIA architectures and CUDA versions.

I learned this the hard way trying to get a side project running on an old CI box with no GPU attached, just to sanity-check the model's forward pass logic. pip install flash-attn sat there compiling for a long stretch and then died with a nvcc: command not found error, because the wheel builder assumed a CUDA toolkit that simply wasn't there. It's not a drop-in CPU accelerator, and a CPU-only box just gets you a build failure, not a slower fallback.

AMD and other accelerator vendors have their own ports with different support windows and different kernel coverage, so don't assume an NVIDIA-flavored install guide applies to a ROCm box. Always check current repo docs rather than trusting last year's blog post.

Where this actually gets usable cross-platform is through framework-level APIs. PyTorch's scaled_dot_product_attention picks a FlashAttention-style backend when your shapes, dtype, and hardware qualify, and silently falls back to a math or memory-efficient kernel otherwise. A successful call proves nothing - check torch.backends.cuda.sdp_kernel context or profiler traces to confirm which backend actually ran.

How to Install FlashAttention 2 or 3 from GitHub

Treat the project's GitHub repository as ground truth for install commands, supported CUDA versions, and wheel availability - they shift often. Before touching pip, check whether your stack already ships a fused kernel through PyTorch's SDPA or another framework, since you may not need the standalone package at all.

FlashAttention 2 Installation Checklist

Prebuilt wheels save you a compile step when your CUDA and PyTorch versions match an available release; otherwise you're building from source, which needs a working compiler toolchain and real patience. Common failures: CUDA/PyTorch version mismatches, missing ninja or build tools, unsupported compute capability on older cards, and pip's build-isolation grabbing the wrong toolchain. In production, pin exact versions of flash-attn, torch, and CUDA together rather than trusting "latest" to stay compatible - a lesson that ties into the broader production stability guide I wrote after chasing one such mismatch across two environments.

FlashAttention 3 Installation Checklist

FA3's headline optimizations target Hopper-class GPUs specifically, so check the repository's current FA3 directory for release status and hardware requirements before assuming your card qualifies. If your hardware or CUDA version doesn't meet those requirements, fall back to FlashAttention 2 or let the framework pick a fused backend automatically.

FAQ

How does FlashAttention work?

It tiles the query, key, and value matrices into blocks that fit in on-chip SRAM, then computes attention block by block using an online softmax that tracks running max and normalization statistics instead of materializing the full score matrix. Everything runs inside one fused GPU kernel, so intermediate results never round-trip through slow HBM. The output matches standard attention exactly - it's a memory-I/O optimization, not an approximation.

What's FlashAttention?

FlashAttention is an exact, I/O-aware attention algorithm for transformers that avoids building the full N×N attention matrix in GPU memory. It reduces memory traffic rather than FLOPs, which speeds up training and inference on long sequences and cuts peak memory use. It's distributed as an open-source package on GitHub with CUDA kernels tuned for specific NVIDIA architectures.

How do I install FlashAttention?

Check the project's GitHub repo for current supported CUDA, PyTorch, and Python versions before doing anything else - this changes across releases. Install via a prebuilt wheel if one matches your stack; otherwise you're compiling from source with ninja and a working C++ toolchain. Confirm your GPU's compute capability meets the minimum before you spend an afternoon debugging a failed build.

How do I install FlashAttention 2?

Match your CUDA toolkit, PyTorch build, and Python version to an available wheel first, since mismatches are the most common install failure. If no wheel fits, build flash-attn from source, which requires ninja, a compatible compiler, and enough RAM to survive the compile. Pin all three versions together in production rather than trusting "latest."

How do I install FlashAttention 3?

FA3 targets Hopper-class GPUs, so verify your hardware qualifies by checking the repo's FA3-specific directory and release notes before installing. If it doesn't, fall back to FlashAttention 2 or let PyTorch's SDPA select an appropriate fused backend automatically.

References

1. A Persistent-State Dataflow Accelerator for Memory-Bound Linear Attention Decode on FPGA - Neelesh Gupta, Peter Wang, Rajgopal Kannan et al. (2026)