HackerPapers

arXiv:2205.14135cs.LG

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

Abstract

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3$\times$ speedup on GPT-2 (seq. length 1K), and 2.4$\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

arXiv comments
Journal reference
DOI
Subjects
cs.LGMachine Learning

On Hacker News

Linked by 11 identifiable commenters in 11 comments across 11 threads, from Dec 2022 to Apr 2025. Found in 4 HN submissions from Jun 2022 to Jan 2024. Open submission discussions.

Comments linking this paper, per year.

Discussed in comments11 threads

  1. 400x AI Acceleration 1 comment
    • Next Monday, 28 Apr 2025, we are launching the public permanently free unlimited, no payment, no catch of Purem – a revolutionary CPU-accelerated AI kernel that matches the performance of leading in-house infrastructures like OpenAI, Tesla, and Meta. A major leap for macOS users…

      elegantlyView on HNParent context

    • For those who have no idea what's being discussed, quick background. Discussing: Transformer [1] memory issues and approximate attention [2] in machine learning training. Specifically: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. [3] As a side…

      araesView on HNParent context

    • FlashAttention(2)[0] reduces context-length space complexity to linear. Compute is still O(n^2) in length though, AFAIK, so we'd expect these long sequence lengths to take some time to compute. I'm a bit out of my depth, but I think ultra-long exact-attention work like this also…

      kristjanssonView on HNParent context

    • I understand that it's a different approach, but I would still have expected this paper to at least mention FlashAttention [1] since they both leverage flash memory. [1] https://arxiv.org/abs/2205.14135

      scribuView on HNParent context

    • Good question. Yes, the 10GB available for batching is in the HBM. In a single forward pass, you move the entire model from HBM -> SRAM exactly once. In a batched forward pass, this is still the case, so you end up doing more compute for the same amount of memory movement. You…

      varunshenoyView on HNParent context

  2. GPT-4 details leaked? 1 comment
  3. [dead] 1 comment
    • Introducing three new open-source PaLM models trained at a context length of 8k on C4. Open-sourcing LLMs is a necessity for the fair and equitable democratization of AI. The models of sizes 150m, 410m, and 1b are available to download and use here…

      EnricoShippoleView on HNParent context

    • https://arxiv.org/abs/2205.14135 - Section 5 suggests that the biggest limitation is that custom CUDA kernels need to be coded on a per-GPU architecture basis.

      lhlView on HNParent context

  4. The Coming of Local LLMs 1 comment
    • Sharing a comment from a similar line of discussion: """ Complexity is quadratic in sequence length. For 512 tokens it is 262K, but for 4000 tokens it becomes 16M and goes OOM on a single GPU. We need about 100K-1M tokens to load whole books at once. Since 2017 there have been…

      dpflanView on HNParent context

    • Complexity is quadratic in sequence length. For 512 tokens it is 262K, but for 4000 tokens it becomes 16M and goes OOM on a single GPU. We need about 100K-1M tokens to load whole books at once. Since 2017 there have been hundreds of attempts to bring O(N^2) to O(N), but none of…

      visargaView on HNParent context

Submitted to Hacker News

Connected through shared HN threads, not inferred similarity or a recommendation.