TALK.md — FlashAttention: Press, Q&A, and Misconceptions

Press & Blog Coverage

  1. IEEE Spectrum — “How FlashAttention Makes Transformers Faster” — article covering Tri Dao’s submission to the MLPerf 2.0 benchmark using FlashAttention. Referenced on the official GitHub repository README.
    • URL: https://spectrum.ieee.org/flashattention (referenced in GitHub README)
  2. Tri Dao’s Blog — “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning” (July 2023) — detailed blog post explaining the improvements in FlashAttention-2 over the original.
    • URL: https://tridao.me/blog/2023/flash2/
  3. Tri Dao’s Blog — “FlashAttention-3: Faster Attention on Hopper GPUs” (May 2024) — blog post on FlashAttention-3 optimizations for H100 GPUs.
    • URL: https://tridao.me/blog/2024/flash3/
  4. PyTorch Blog — “PyTorch 2.0 includes scaled_dot_product_attention with FlashAttention” — PyTorch 2.0 release notes (March 2023) describe integration of FlashAttention as a native operator.
    • URL: https://pytorch.org/blog/pytorch-2.0-release/
  5. HuggingFace Blog — Multiple posts on using FlashAttention with Transformers library for faster training and inference.
    • URL: https://huggingface.co/docs/transformers/en/perf_torch_compile (references FlashAttention integration)
  6. NeurIPS 2022 — Paper accepted and presented at NeurIPS 2022 (Advances in Neural Information Processing Systems 35).

Interview-Style Q&A

Q1: Why did approximate attention methods fail to deliver wall-clock speedups?

A: Most approximate attention methods (Linformer, Performer, etc.) focused on reducing FLOP count, but on modern GPUs, attention is memory-bound — the bottleneck is reading/writing the N×N matrix to HBM, not the arithmetic. Reducing FLOPs without reducing memory accesses doesn’t translate to real speedup. FlashAttention’s key insight was being IO-aware: reducing HBM accesses is what actually matters.

Q2: How does FlashAttention compute exact softmax without seeing the full row?

A: It uses an online/running softmax trick. For each block of the attention matrix, you compute the local maximum and local sum of exponentials. When combining blocks, you rescale: if the new block’s maximum is larger, you rescale the old running sum down by exp(old_max - new_max), and vice versa. This maintains numerical stability and gives the exact same result as computing softmax on the full row at once.

Q3: Isn’t recomputation in the backward pass wasteful?

A: Counterintuitively, no. The backward pass needs the attention matrix P. Standard attention reads P back from HBM (an O(N²) read). FlashAttention recomputes P from Q, K in SRAM — this takes extra FLOPs but avoids the HBM read entirely. Since the operation is memory-bound (not compute-bound), the extra computation is cheaper than the memory access it replaces. The net result is faster, not slower.

Q4: Why can’t you just fuse the operations with a compiler?

A: Standard kernel fusion (like PyTorch’s torch.compile) fuses elementwise operations but cannot fuse across the QK^T → softmax → PV chain because softmax requires a reduction over the full row, and the intermediate N×N matrix must be materialized for the backward pass. FlashAttention’s tiling + online softmax is an algorithmic change that no compiler can derive automatically — it requires restructuring the computation itself.

Q5: What’s the relationship between FlashAttention and PyTorch’s scaled_dot_product_attention?

A: PyTorch 2.0’s F.scaled_dot_product_attention is a dispatch operator that internally calls FlashAttention (or memory-efficient attention) when available. It automatically selects the best backend based on tensor properties. This makes FlashAttention transparent to most users — they get the speedup without changing their code.

Common Misconceptions

  1. “FlashAttention is an approximation.” — No. FlashAttention computes the exact same result as standard attention. It reorganizes the computation to reduce memory accesses, but the output is mathematically identical (within floating-point precision).

  2. “FlashAttention reduces the O(N²) compute complexity.” — No. FlashAttention still does O(N²) FLOPs. What it reduces is memory accesses (IO complexity), which is the actual bottleneck on modern GPUs. The compute is the same; the memory traffic is much less.

  3. “FlashAttention only works on GPUs.” — The algorithm is conceptually hardware-agnostic (it’s about memory hierarchy). The implementation is CUDA-specific, but the tiling principle applies to any system with a memory hierarchy. CPU implementations exist (e.g., in OpenBLAS-style approaches).

  4. “FlashAttention and FlashAttention-2 are the same thing.” — FlashAttention-2 (2023) improves parallelism (better GPU thread block assignment) and work partitioning (reducing non-matmul FLOPs), achieving ~2× speedup over FlashAttention-1. The core tiling idea is the same, but the GPU scheduling is significantly better.

Real Citations

  1. Dao, T., Fu, D.Y., Ermon, S., Rudra, A., Ré, C. (2022). “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” Advances in Neural Information Processing Systems (NeurIPS 2022). arXiv:2205.14135.

  2. Dao, T. (2023). “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.” arXiv:2307.08691.

  3. Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., Dao, T. (2024). “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision.” arXiv:2407.08608.

  4. Vaswani, A., et al. (2017). “Attention Is All You Need.” NeurIPS 2017. (The original Transformer paper that FlashAttention optimizes.)

  5. Rabe, M.N., Staats, C. (2021). “Self-attention Does Not Need O(N²) Memory.” arXiv:2112.05682. (Independent work on memory-efficient attention that FlashAttention builds upon.)

  6. Hong, X., et al. (2022). “MLPerf Training 2.0 Results.” (FlashAttention was used in MLPerf 2.0 submissions achieving training speed records.)