FlashAttention 如何運作：分塊、SRAM 重用，以及快 10 倍的原因

How FlashAttention Actually Works: Tiling, SRAM Reuse, and Why It's 10x Faster

When Transformers process long sequences, attention is the undisputed computational bottleneck. Standard attention scales quadratically ($O(N^2)$) in both FLOPs and memory. But as sequence lengths grow past 8k, 32k, or 128k tokens, GPUs do not usually run out of compute power first. They run out of memory bandwidth.

Even though modern GPUs like the NVIDIA A100 can perform hundreds of teraflops, they spend most of their time waiting for data to move between slow High Bandwidth Memory (HBM) and fast on-chip Static RAM (SRAM).

Standard attention materializes the massive $N \times N$ attention score matrix in HBM. For a sequence of 16,000 tokens, that single intermediate matrix takes over 1 gigabyte of memory just to write and read back repeatedly.

FlashAttention changes this by making the attention algorithm IO-aware. It does not approximate attention, and it does not drop mathematical precision. It simply reorganizes the execution using tiling and online softmax so that intermediate matrices never touch slow HBM.

The GPU Memory Wall: HBM vs. SRAM

To understand why FlashAttention is so fast, you need to look at the physical memory hierarchy of modern GPUs:

High Bandwidth Memory (HBM): 40GB to 80GB in capacity, but bounded at roughly 1.5 TB/s to 2.0 TB/s bandwidth. This is where large tensors live.
On-Chip SRAM: Extremely fast (roughly 19 TB/s bandwidth), but tiny (only about 192KB per Streaming Multiprocessor on an A100).

In standard attention implemented in PyTorch, the GPU performs matrix multiplications, writes the $Q K^T$ matrix to HBM, reads it back to compute softmax, writes the softmax matrix to HBM, reads it back to apply dropout/masking, writes it back again, and finally multiplies by $V$.

Over 90% of execution time is spent reading and writing tensors to HBM. FlashAttention eliminates these round-trips entirely.

How Tiling Works in FlashAttention

Tiling breaks large matrices into smaller blocks that fit completely inside the GPU's fast on-chip SRAM.

Instead of computing the entire $N \times N$ attention matrix at once, FlashAttention loads blocks of keys ($K$) and values ($V$) from HBM into SRAM. Then, it loops over blocks of queries ($Q$) in SRAM, computing partial attention outputs directly on-chip.

However, standard softmax poses a major mathematical hurdle for tiling: softmax is globally normalized across the entire row.

If you process a row in chunks, you do not know the complete denominator (the sum of exponents across all tokens in the sequence) until you have seen every single token.

The Secret Weapon: Online Softmax

FlashAttention solves this by using online softmax (or incremental softmax rescaling).

Instead of waiting for the full row sum, online softmax maintains running statistics as it processes each block:

Running Maximum ($m$): Tracks the maximum value seen so far in the row to prevent numerical overflow during exponentiation.
Running Scaling Factor ($l$): Tracks the cumulative sum of exponents.

When a new block is processed in SRAM, FlashAttention computes the local exponent sum, rescales the previous output using the updated running maximum and scaling factor, and accumulates the result without ever writing intermediate $N \times N$ matrices to HBM.

The Backward Pass: Recomputation

In standard training, storing the massive attention matrix from the forward pass for backward gradient computation requires quadratic memory ($O(N^2)$), which limits context lengths.

FlashAttention avoids storing this matrix entirely during the forward pass. During the backward pass, it reloads the inputs from HBM, recomputes the attention scores on-chip in SRAM on the fly, and calculates gradients.

Even though recomputation adds extra floating-point operations (FLOPs), it dramatically reduces memory reads and writes. Wall-clock time actually drops by 2x to 4x because memory bandwidth is the true bottleneck.

Practical Takeaways

Exact Computation: FlashAttention is mathematically identical to standard attention. There is zero approximation loss.
Linear Memory Scaling: Memory complexity drops from quadratic $O(N^2)$ to linear $O(N)$ with respect to sequence length.
IO Awareness: Performance is dictated by memory bandwidth (HBM vs SRAM), not just raw FLOP counts.
Ecosystem Standard: FlashAttention is now foundational to modern LLM training and inference frameworks (vLLM, Hugging Face, PyTorch, Megatron-LM)