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

May combine two claims.

For large sequence lengths N, s is often set to N^{-1/2} or N^{-1} log N, resulting in Θ(N sqrt N) or Θ(N log N) IO complexity.

“For large sequence lengths 𝑁 , 𝑠 is often set to 𝑁 1 2 [11] or 𝑁 1 log 𝑁 [3, 17, 92], resulting in Θ 𝑁 𝑁 or Θ 𝑁 log 𝑁 IO complexity.”

• cites [11] [11] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019. - sourceDisputed

  The citation to Child et al. is attached to the logarithmic-complexity option. Instead, Section 4.3 (p. 4) chooses a stride close to √N, and Section 4.2 (p. 4) gives O(N√N) computation. The paper does not give the attributed Θ(N log N) IO complexity. (LLM, conf=0.85)

  Checked by a reasoning model — may be wrong.

  Experimental — accuracy not yet measured
• cites [3] [3] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 , 2020. - sourceCould not verify

  Section 3.1 (p. 3) gives sliding-window attention a computation complexity of O(n × w) for a fixed window size. Our copy does not say that sparsity s is set to N⁻¹ log N or give a Θ(N log N) IO complexity; a search for those terms found no supporting passage. (LLM, conf=0.73)

  Experimental — accuracy not yet measured
• cites [17] [17] Tri Dao, Beidi Chen, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Ré. Pixelated butterfly: Simple and efficient sparse training for neural network models. In International Conference on Learning Representations (ICLR) , 2022. - sourcePartly supported

  Section 3.2 (p. 5) says flat butterfly matrices have O(n log k) nonzeros; with maximum stride k = n, this supports an O(N log N) sparsity pattern. The Problem Formulation (p. 17) relates memory accesses to nonzero blocks. The paper does not state that s is often set to either claimed value, or establish the stated Θ IO complexities, including the Θ(N√N) case. (LLM, conf=0.78)

  Experimental — accuracy not yet measured
• cites [92] [92] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems , 33, 2020. - sourcePartly supported

  Section 2 (pp. 3–4) describes sparse attention and its reduced complexity, including attention to a limited number of random keys. It does not specify a sparsity fraction of (log N)/N or establish Θ(N log N) IO complexity; the Θ(N√N) alternative is attached to a different citation in the citing sentence. (LLM, conf=0.76)

  Experimental — accuracy not yet measured

May combine two claims.

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).

“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”

• Also stated in: § 1 Introduction · plain

  “We observe a 0.7 improvement in perplexity on GPT-2 and 6.4 points of lift from modeling longer sequences on long-document classification [13].”

• Also stated in: § 4 Experiments · plain

  “Modeling longer sequences yields 6.4 points of lift on two longdocument classification tasks.”

• cites [13] [13] Xiang Dai, Ilias Chalkidis, Sune Darkner, and Desmond Elliott. Revisiting transformer-based models for long document classification. arXiv preprint arXiv:2204.06683 , 2022. - sourcePartly supported

  The Introduction and Conclusions report that processing more text can improve long-document classification, and §3.1 describes Longformer as a way to process longer sequences. Our copy does not attribute those gains to FlashAttention or block-sparse FlashAttention, and it does not report the claimed 0.7 GPT-2 perplexity improvement or 6.4-point classification lift (searched those terms and figures). Page numbers were not available in our copy. (LLM, conf=0.86)

  Experimental — accuracy not yet measured

May combine two claims.

Approximate attention methods reduce compute complexity by trading off model quality but often do not achieve wall-clock speedup.

“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.”

• Also stated in: § 1 Introduction · plain

  “Although these methods reduce the compute requirements to linear or near-linear in sequence length, many of them do not display wall-clock speedup against standard attention and have not gained wide adoption.”

• cites [51] [51] Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML) , 2020. - sourcePartly supported

  The Introduction (p. 1) says LSH attention reduces complexity from O(L²) to O(L log L), and the synthetic-task analysis (p. 5) reports accuracy losses in some settings. But the paper does not support the claim that approximate attention methods often fail to achieve wall-clock speedups; its Experiments section (p. 9) reports that LSH attention remains fast as sequence length increases while regular attention slows. (LLM, conf=0.87)

  Experimental — accuracy not yet measured
• cites [74] [74] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics , 9: 53-68, 2021. - sourcePartly supported

  The Abstract says sparse attention reduces compute complexity, and §6.1 reports that local attention was slightly worse than full attention on CIFAR-10. Section 6.3 supports a qualified wall-clock concern: Routing Transformer was slower than local attention on the reported TPU tests. But local attention was faster than full attention, and some Routing Transformer configurations were faster too, so the paper does not establish that approximate attention methods generally fail to achieve wall-clock speedups. (LLM, conf=0.84)

  Experimental — accuracy not yet measured
• cites [12] [12] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations (ICLR) , 2020. - sourcePartly supported

  The Introduction and Related Work (p. 1) describes more tractable attention methods and limitations that can affect their representation power or accuracy. Section 2.2 (p. 3) shows reduced time complexity for the paper’s approximate attention mechanism. But the paper does not say that approximate methods often fail to achieve wall-clock speedups; Section 4.1 (p. 6) instead reports lower wall-clock time for its own method. (LLM, conf=0.84)

  Experimental — accuracy not yet measured
• cites [50] [50] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning , pages 5156-5165. PMLR, 2020. - sourcePartly supported

  The Introduction (p. 1) says sparse attention and locality-sensitive hashing reduce computational complexity but do not speed up autoregressive inference. That supports the speed claim only for that setting: the paper also says these methods can be trained efficiently on long sequences, and it does not establish a general trade-off in model quality. (LLM, conf=0.86)

  Experimental — accuracy not yet measured
• cites [84] [84] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020. - sourcePartly supported

  The Introduction (p. 1) says sparse attention reduces complexity but incurs a performance drop, while Reformer’s efficiency gains appear only for sequences longer than 2048 tokens. The claim overstates the wall-clock point for sparse attention: the paper reports a limited but real 20% speedup, not no speedup. (LLM, conf=0.85)

  Experimental — accuracy not yet measured

May combine two claims.

Rabe and Staats [66] focuses on reducing total memory footprint while FlashAttention focuses on reducing memory accesses.

“The first major difference is that Rabe and Staats [66] focuses on the reducing the total memory footprint (maximum amount of GPU memory required) while FlashAttention focuses on reducing memory accesses (the number of memory reads/writes).”

• cites Rabe and Staats [66] sourceSupported

May combine two claims.

In prior work, all transformer models have either run out of memory or only achieved random performance on Path-X and Path-256.

“In prior work, all transformer models have either run out of memory, or only achieved random performance [80].”

• cites [80] [80] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations , 2020. - sourcePartly supported

  The Quantitative Results (p. 7) say that all models evaluated on Path-X failed, achieving at best 50% accuracy, consistent with chance performance. The paper does not establish this for all transformer models or for Path-256, and the cited passages do not say that any Path-X model ran out of memory. (LLM, conf=0.88)

  Experimental — accuracy not yet measured

Probably not a substantial claim: may not pass S3 (load-bearing).

LRA accuracy results are highly dependent on the tuning procedure.

“3 LRA accuracy results are known to be highly dependent on the tuning procedure [90].”

• cites [90] [90] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nystöm-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence. AAAI Conference on Artificial Intelligence , volume 35, page 14138, 2021. - sourceCould not verify

  The Long Range Arena Benchmark section (p. 8) describes the evaluation setup and reports accuracy, but does not say that accuracy is highly dependent on the tuning procedure. Searches for LRA tuning and hyperparameter sensitivity found a sensitivity discussion only for MRPC, a different task (Fine-tuning on Downstream NLP tasks, p. 7). (LLM, conf=0.76)

  Experimental — accuracy not yet measured

May combine two claims.

FlashAttention is faster than standard attention (2-4×) while Rabe and Staats [66] is around the same speed or slightly slower than standard attention.

“As a result, FlashAttention is faster than standard attention (2-4 ) while Rabe and Staats [66] is around the same speed or slightly slower than standard attention.”

• cites Rabe and Staats [66] sourceSupported

Probably not a substantial claim: may not pass S3 (load-bearing). May combine two claims.

Even though structured matrices are efficient in theory, they have not seen wide adoption since it is hard to translate their efficiency to wall-clock speedup since dense unconstrained matrix multiply has very optimized implementation, a phenomenon known as the hardware lottery.

“However, even though structured matrices are efficient in theory, they have not seen wide adoption since it is hard to translate their efficiency to wall-clock speedup since dense unconstrained matrix multiply has very optimize implementation, a phenomenon known as the hardware lottery [41].”

• cites [41] [41] Sara Hooker. The hardware lottery. arXiv preprint arXiv:2009.06489 , 2020. - sourcePartly supported

  Split into 2 atoms — 0 confirmed, 2 partially supported, 0 disputed, 0 misattributed, 0 unverifiable. Expand the claim to see per-atom verdicts.

  Experimental — accuracy not yet measured

May combine two claims.

FlashAttention simplifies the backward pass analytically, only recomputes the attention matrix and does not recompute the temporary output of each block, reducing memory requirement for the backward pass and yielding speedup compared to Rabe and Staats [66].

“FlashAttention instead simplifies the backward pass analytically (Appendices B.2 and B.4). It only recomputes the attention matrix and does not recompute the temporary output of each block. This reduces the memory requirement for the backward pass and yields speedup.”

• cites Rabe and Staats [66] sourcePartly supported

  Split into 2 atoms — 0 confirmed, 2 partially supported, 0 disputed, 0 misattributed, 0 unverifiable. Expand the claim to see per-atom verdicts.

  Experimental — accuracy not yet measured

May combine two claims.

GPT-2 small and medium using FlashAttention achieve up to 3x speedup compared to Huggingface and up to 1.7x compared to Megatron-LM.

“Table 2 shows up to 3 endto-end speedup compared to Huggingface and 1.7 speedup compared to Megatron-LM.”

• Also stated in: § 4.1 Faster Models with FlashAttention · plain

  “GPT-2 small and medium using FlashAttention achieve up to 3 speed up compared to Huggingface implementation and up to 1.7 compared to Megatron-LM.”

• cites [87] [87] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, Online, October 2020. Association for Computational Linguistics. URL https://www sourceCould not verify

  The Deployment section reports a nearly 4× speedup using ONNX, not FlashAttention. Searches for the claimed 3× and 1.7× end-to-end comparisons found no such results in our copy; Related Work mentions Megatron-LM but does not benchmark against it. (LLM, conf=0.77)

  Experimental — accuracy not yet measured
• cites [77] [77] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053 , 2019. - sourceSupported

May combine two claims.

FlashAttention has smaller total memory requirement compared to Rabe and Staats [66] because it incrementally updates the output after processing each block, needing only one copy instead of K copies for K blocks.

“FlashAttention instead incrementally updates the output (Algorithm 1 line 12) after processing each block, so only one copy of the output is needed (instead of 𝐾 copies for 𝐾 blocks). This means that FlashAttention has smaller total memory requirement compared to Rabe and Staats [66].”

• cites Rabe and Staats [66] sourcePartly supported

  Section 4 (p. 4) says the practical implementation keeps one summary per key/value chunk before rescaling them, producing √n summaries and O(√n) memory overhead. That supports the comparison with storing K intermediate results, but the paper does not establish the claimed smaller total memory requirement for FlashAttention. It also describes an O(log n)-memory self-attention algorithm in the Abstract (p. 1), so the comparison cannot be applied to all of its methods. (LLM, conf=0.84)

  Experimental — accuracy not yet measured

May combine two claims.

FlashAttention yields the fastest single-node BERT training speed known to the authors, 15% faster than Nvidia MLPerf 1.1.

“FlashAttention yields the fastest single-node BERT training speed that we know of. We train a BERT-large [22] model with FlashAttention on Wikipedia. Table 1 compares our training time to the implementation from Nvidia that set the training speed record for MLPerf 1.1 [58]. Our implementation is 15% faster.”

• cites [22] [22] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. 2019. - sourceNot checked — no full text of the source
• cites [58] [58] Peter Mattson, Christine Cheng, Gregory Diamos, Cody Coleman, Paulius Micikevicius, David Patterson, Hanlin Tang, Gu-Yeon Wei, Peter Bailis, Victor Bittorf, et al. Mlperf training benchmark. Proceedings of Machine Learning and Systems , 2:336-349, 2020. - sourceCould not verify

  The citation is attached to the claim that an NVIDIA implementation set the MLPerf 1.1 training-speed record, not to the citing authors’ 15% comparison. The Results (p. 9) discuss only MLPerf v0.5 and v0.6; searches for MLPerf 1.1, NVIDIA, and BERT found no passage establishing the cited record. (LLM, conf=0.76)

  Experimental — accuracy not yet measured

Classified as probably not a substantial claim.

The backward pass of attention can be done without quadratic extra memory by applying gradient checkpointing to the memory-efficient forward pass.

“Rabe and Staats [66] suggests that the backward pass can be done without quadratic extra memory by applying gradient checkpointing to the memory-efficient forward pass.”

• cites Rabe and Staats [66] sourceSupported

May combine two claims.

Block-sparse FlashAttention is 2-4x faster than even FlashAttention, scaling up to sequence length of 64k.

“As a proof of concept, we implement block-sparse FlashAttention , a sparse attention algorithm that is 2-4 faster than even FlashAttention , scaling up to sequence length of 64k.”

May combine two claims.

Algorithm 1 returns O = softmax(QK^T)V with O(N^2 d) FLOPs and requires O(N) additional memory beyond inputs and output.

“Theorem 1. Algorithm 1 returns O softmax QK V with 𝑂 𝑁 2 𝑑 FLOPs and requires 𝑂 𝑁 additional memory beyond inputs and output.”

May combine two claims.

FlashAttention is optimal for a range of SRAM sizes; no exact attention algorithm can asymptotically improve on the number of HBM accesses over all SRAM sizes.

“and is optimal for a range of SRAM sizes.”

• Also stated in: § 1 Introduction · plain

  “Moreover, we provide a lower bound, showing that no exact attention algorithm can asymptotically improve on the number of HBM accesses over all SRAM sizes.”

• Also stated in: § 3.2 Analysis: IO Complexity of FlashAttention · plain

  “Proposition 3. Let 𝑁 be the sequence length, 𝑑 be the head dimension, and 𝑀 be size of SRAM with 𝑑 𝑀 𝑁𝑑 . There does not exist an algorithm to compute exact attention with 𝑜 𝑁 2 𝑑 2 𝑀 1 HBM accesses for all 𝑀 in the range 𝑑 𝑁𝑑 .”

May combine two claims.

FlashAttention runs faster (up to 7.6x on GPT-2) and uses less memory-linear in sequence length-than standard attention, even with increased FLOPs due to recomputation.

“Even with the increased FLOPs due to recomputation, our algorithm both runs faster (up to 7.6x on GPT-2 [67], Figure 1 right) and uses less memory -linear in sequence length-than standard attention, thanks to the massively reduced amount of HBM access.”

• cites [67] [67] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019. -Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

• AssumptionRests on an assumption: All gradient checkpointing implementations known to the authors trade speed for memory.

May combine two claims.

FlashAttention is up to 3x faster than standard attention across common sequence lengths from 128 to 2K and scales up to 64K.

“FlashAttention is up to 3 faster than the standard attention implementation across common sequence lengths from 128 to 2K and scales up to 64K.”

May combine two claims.

Up to sequence length of 512, FlashAttention is both faster and more memory-efficient than any existing attention method.

“Up to sequence length of 512, FlashAttention is both faster and more memory-efficient than any existing attention method, whereas for sequence length beyond 1K, some approximate attention methods (e.g., Linformer) start to become faster.”

May combine two claims.

FlashAttention is up to 20× more memory efficient than exact attention baselines, and is more memory-efficient than the approximate attention baselines.

“FlashAttention is up to 20 more memory efficient than exact attention baselines, and is more memory-efficient than the approximate attention baselines.”

May combine two claims.

Standard attention backward pass requires Θ(Nd + N^2) HBM accesses, while FlashAttention backward pass requires Θ(N^2 d^2 M^{-1}) HBM accesses.

“Theorem 5. Let 𝑁 be the sequence length, 𝑑 be the head dimension, and 𝑀 be size of SRAM with 𝑑 𝑀 𝑁𝑑 . Standard attention (Algorithm 0) backward pass requires Θ 𝑁𝑑 𝑁 2 HBM accesses, while FlashAttention backward pass (Algorithm 4) requires Θ 𝑁 2 𝑑 2 𝑀 1 HBM accesses.”

• Also stated in: § C Proofs · plain

  “At each step of the standard attention backward pass, one needs to load inputs of size 𝑁𝑑 or 𝑁 2 from HBM, and needs to write the outputs of size 𝑁 2 or 𝑁𝑑 to HBM. This incurs Θ 𝑁𝑑 𝑁 2 HBM accesses.”

• Also stated in: § C Proofs · plain

  “Similar to Theorem 2, we see that each element of K and V is loaded from HBM once. Each element of dK and dV is only written to HBM once. We make 𝑇 𝑐 passes over Q O dO , each pass loading all of Q O dO to HBM. We also make 𝑇 𝑐 passes over dQ , each pass reading/writing all of dQ from/to HBM. Therefore the number of HBM accesses is Θ 𝑁𝑑 𝑁𝑑𝑇 𝑐 Θ 𝑁𝑑𝑇 𝑐 .”

Probably not a substantial claim: may not pass S3 (load-bearing), S4 (based).

Our IO-aware implementation of attention is optimal within constants for computing attention on a single GPU.

“Our IO-aware implementation of attention is optimal within constants for computing attention on a single GPU.”

May combine two claims.

FlashAttention backward pass performs O(N^2) FLOPs and only requires O(N) extra memory beyond inputs, output, output gradient, and input gradients.

“We see that similar to the forward pass, the backward pass performs 𝑂 𝑁 2 FLOPs and only requires 𝑂 𝑁 extra memory beyond inputs, output, output gradient, and input gradients.”

Classified as probably not a substantial claim.

The runtimes of many approximate/sparse attention mechanisms grow linearly with sequence length.

“The runtimes of many approximate/sparse attention mechanisms grow linearly with sequence length, but FlashAttention still runs faster than approximate and sparse attention for short sequences due to fewer memory accesses.”

May combine two claims.

FlashAttention backward pass does not need to store the dropout mask of size O(N^2) from the forward pass; saving PRNG states and regenerating the dropout mask in the backward pass allows only O(N) extra memory.

“We do not need to store the dropout mask of size 𝑂 𝑁 2 from the forward pass. Instead, we can save the pseudo-random number generator states from the forward pass and re-generate the dropout mask in the backward pass. This allows us to only use 𝑂 𝑁 extra memory.”

May combine two claims.

GPT-2 with FlashAttention and context length 4K is still 30% faster than GPT-2 from Megatron with context length 1K, while achieving 0.7 better perplexity.

“Table 4 shows that that GPT-2 with FlashAttention and context length 4K is still 30% faster than GPT-2 from Megatron with context length 1K, while achieving 0.7 better perplexity.”

May combine two claims.

Block-sparse FlashAttention enables Transformers to scale to sequence length 64K, achieving 63.1 accuracy on Path-256.

“Additionally, block-sparse FlashAttention enables the Transformers to scale to sequence length 64K, achieving 63.1 accuracy 4 on Path-256.”

May combine two claims.

All other algorithms except for Linformer run out of memory on an A100 GPU before 64K, and FlashAttention is still 2× more efficient than Linformer.

“All other algorithms except for Linformer run out of memory on an A100 GPU before 64K, and FlashAttention is still 2 more efficient than Linformer.”

May combine two claims.

FlashAttention and block-sparse FlashAttention have the same memory footprint, which grows linearly with sequence length.

“FlashAttention and block-sparse FlashAttention have the same memory footprint, which grows linearly with sequence length.”

May combine two claims.

Block-sparse FlashAttention achieves 2.8x speedup on the LRA benchmark while performing on par with standard attention.

“On the LRA benchmark, block-sparse FlashAttention achieves 2.8 speedup, while performing on par with standard attention (Section 4).”

May combine two claims.

FlashAttention is as numerically stable as the baselines and produces the same training/validation curves.

“Appendix E includes plots of the validation perplexity throughout training, confirming that FlashAttention is as numerically stable as the baselines and produces the same training / validation curves.”

May combine two claims.

Both MIMIC-III and ECtHR contain very long text documents; the average number of tokens in MIMIC is 2,395 tokens, and the longest document contains 14,562 tokens, while the average and longest numbers in ECtHR are 2,197 and 49,392, respectively.

“Both of these datasets contain very long text documents; the average number of tokens in MIMIC is 2,395 tokens, and the longest document contains 14,562 tokens, while the average and longest numbers in ECtHR are 2,197 and 49,392, respectively.”

May combine two claims.

The runtime of block-sparse FlashAttention scales linearly in sequence length and is faster than all existing approximate attention baselines.

“We confirm that runtime of block-sparse FlashAttention scales linearly in seq. length and is faster than all existing approximate attention baselines.”

May combine two claims.

Attention forward and backward passes can be computed with O(n) extra memory, linear instead of quadratic in the sequence length.

“We first derive the forward and backward passes of attention and show that they can be computed in a memory-efficient manner (requiring extra memory linear instead of quadratic in the sequence length).”

• Also stated in: § B.1 Memory-efficient forward pass · plain

  “Therefore the forward pass can be computed with 𝑂 𝑛 extra memory:”

• Also stated in: § B.2 Memory-efficient backward pass · plain

  “Therefore the backward pass can also be computed with 𝑂 𝑛 extra memory:”

May combine two claims.

The memory footprint of FlashAttention scales linearly with sequence length and is up to 3x faster than standard attention for common sequence lengths up to 2K.

“We confirm that the memory footprint of FlashAttention scales linearly with seq. length and is up to 3 faster than standard attention for common seq. lengths (up to 2K).”

May combine two claims.

On an RTX 3090 GPU with batch size 12 and 12 attention heads, FlashAttention speedup over standard PyTorch attention is between 2.5-4.5x, slightly higher than on an A100.

“We observe slightly higher speedups on the RTX 3090 (between 2.5-4.5 ), since the memory bandwidth on an RTX 3090 is lower than on an A100 (roughly 900 GB/s vs. 1.5 TB/s).”

May combine two claims.

FlashAttention can support much longer sequences (e.g., up to length 64K), more head dimensions (16, 32, 64, 128), and broader GPU types (all Turing and Ampere GPUs at the time of writing).

“As a result, we can support much longer sequences (e.g., up to length 64K). We also support more head dimensions (16, 32, 64, 128) and broader GPU types (all Turing and Ampere GPUs at the time of writing).”

May combine two claims.

Training GPT-2 small takes between 2.7-9.5 days, and training GPT-2 medium takes between 6.9-21.0 days on 8 A100-40GB GPUs.

“We train the model on 8 A100-40GB GPUs, and we measure the wall-clock training time. Training GPT-2 small takes between 2.7-9.5 days, and training GPT-2 medium takes between 6.9-21.0 days (Table 2).”

May combine two claims.

For block-sparse FlashAttention, the number of HBM accesses is scaled by s, the fraction of nonzero blocks in the block-sparsity mask, but for small s the result O must still be written.

“For the block-sparse case, notice that we only need to load blocks corresponding to nonzero blocks. As a result, the number of HBM accesses are scaled by 𝑠 , the fraction of nonzero blocks in the block-sparsity mask. However, for small values of 𝑠 , we would still need to write the result O 𝑁 𝑑 . Therefore the number of HBM accesses is”

Probably not a substantial claim: may not pass S3 (load-bearing).

Longformer and BigBird do not support sequence lengths longer than 8092.

“Longformer and BigBird do not support sequence lengths longer than 8092.”

Probably not a substantial claim: may not pass S3 (load-bearing).

Local Attention's implementation only supports FP32.

“We use FP16 for all measurements, except for Local Attention, whose implementation only supports FP32.”

Probably not a substantial claim: may not pass S3 (load-bearing).

After hyperparameter tuning, almost all attention methods achieve similar accuracy on all five LRA tasks.

“After hyperparameter tuning, almost all of the attention methods achieve similar accuracy on all of the five LRA tasks.”

Probably not a substantial claim: may not pass S3 (load-bearing).

Block-Sparse (OpenAI) does not support sequence lengths longer than 4096.

“Block-Sparse (OpenAI) does not support sequence lengths longer than 4096.”

May combine two claims.

Compared to FMHA, overall runtime of FlashAttention is about 4% slower for sequence length 128, 8% faster for sequence length 256, and 5% faster for sequence length 512.

“Compared to FMHA, the overall runtime of FlashAttention is about 4% slower for sequence length 128, 8% faster for sequence length 256, and 5% faster for sequence length 512.”

Probably not a substantial claim: may not pass S3 (load-bearing).

The Megatron implementation does not support sequence lengths longer than 2048.

“The Megatron implementation does not support sequence lengths longer than 2048.”

May combine two claims.

FlashAttention is slightly faster than FMHA in the forward pass and slightly slower than FMHA in the backward pass.

“Generally FlashAttention is slightly faster than FMHA in the forward pass and slightly slower than FMHA in the backward pass.”

May combine two claims.

FMHA stores the attention matrix softmax mask QK^T to HBM in the forward pass for gradient computation, so it does not offer substantial memory saving.

“In the forward pass, it stores the attention matrix softmax mask QK 𝑇 to HBM to be used in gradient computation. As a result, it does not offer substantial memory saving (though for shorter sequences memory footprint is often not a primary concern).”

May combine two claims.

FlashAttention behaves the same as the baseline HuggingFace implementation, with validation perplexity curves of the two implementations almost lying on top of each other.

“We see that FlashAttention behaves the same as the baseline implementation and the validation perplexity curves of the two implementations almost lie on top of each other.”

Probably not a substantial claim: may not pass S3 (load-bearing). May combine two claims.

For Path-X, additional fine-tuning adds roughly 4 points of accuracy to FlashAttention, but the model starts overfitting afterwards.

“For Path-X, we take the best performing checkpoint (according to val accuracy), and additionally fine-tune it for 200 epochs with the same warmup and learning rate (this adds roughly 4 points of accuracy to FlashAttention for Path-X, but the model starts overfitting afterwards).”

May combine two claims.

Transformers are slow and memory-hungry on long sequences because self-attention time and memory complexity are quadratic in sequence length.

“Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length.”

• Also stated in: § 1 Introduction · plain

  “since the self-attention module at their heart has time and memory complexity quadratic in sequence length.”

Probably not a substantial claim: may not pass S3 (load-bearing).

IO-aware algorithms have been critical for memory-bound operations such as database joins, image processing, and numerical linear algebra.

“IO-aware algorithms have been critical for similar memory-bound operations, when reading and writing data can account for a large portion of the runtime-such as database joins [71], image processing [70], numerical linear algebra [4], and more [40, 85].”

• cites [71] [71] Raghu Ramakrishnan, Johannes Gehrke, and Johannes Gehrke. Database management systems , volume 3. McGraw-Hill New York, 2003. -Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

• cites [70] [70] Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines. Acm Sigplan Notices , 48(6):519-530, 2013. - sourceNot checked — no full text of the source
• cites [4] [4] L Susan Blackford, Antoine Petitet, Roldan Pozo, Karin Remington, R Clint Whaley, James Demmel, Jack Dongarra, Iain Duff, Sven Hammarling, Greg Henry, et al. An updated set of basic linear algebra subprograms (blas). ACM Transactions on Mathematical Software , 28(2):135-151, 2002. -Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

May combine two claims.

FlashAttention is implemented in CUDA to achieve fine-grained control over memory access and fuse all attention operations into one GPU kernel.

“We implement FlashAttention in CUDA to achieve fine-grained control over memory access and fuse all the attention operations into one GPU kernel.”

• AssumptionRests on an assumption: Common Python interfaces to deep learning such as PyTorch and Tensorflow do not allow fine-grained control of memory access.

May combine two claims.

Building IO-aware implementations of attention requires writing the attention algorithm in a considerably lower-level language than PyTorch, and requires significant engineering effort.

“This requires writing the attention algorithm in a considerably lower-level language than PyTorch, and requires significant engineering effort.”

Probably not a substantial claim: may not pass S3 (load-bearing).

Implementations may also not be transferrable across GPU architectures.

“Implementations may also not be transferrable across GPU architectures.”

Probably not a substantial claim: may not pass S3 (load-bearing).

Our current approach to building IO-aware implementations of attention requires writing a new CUDA kernel for each new attention implementation.

“Our current approach to building IO-aware implementations of attention requires writing a new CUDA kernel for each new attention implementation.”

Probably not a substantial claim: may not pass S3 (load-bearing), S4 (based). May combine two claims.

The discrepancies in performance between MIMIC-III and ECtHR may be due to subtle distribution shifts: MIMIC-III contains specialized medical text and thus may be more susceptible to a distribution shift in document length, whereas ECtHR contains general language.

“The discrepancies may be due to subtle distribution shifts: MIMIC-III contains specialized medical text and thus may be more susceptible to a distribution shift in the document length, whereas ECtHR contains general language.”

Probably not a substantial claim: may not pass S3 (load-bearing), S4 (based). May combine two claims.

Path-256 requires longer sequences but has relatively shorter paths than Path-X, so it is easier to obtain a higher accuracy.

“Path-256 requires longer sequences but has relatively shorter paths than Path-X, so it is easier to obtain a higher accuracy.”

May combine two claims.

FlashAttention implements both forward and backward passes on GPUs, reduces HBM accesses, and leads to both faster runtime and smaller memory footprint.

“We describe the FlashAttention algorithm to implement both the forward and the backward passes on GPUs that reduces HBM accesses, leading to both faster runtime and smaller memory footprint.”

• AssumptionRests on an assumption: Attention time and memory scale quadratically in the sequence length.

May combine two claims.

When the head dimension increases, each block requires more memory, so smaller block sizes are needed to fit into SRAM.

“Each block requires more memory, so we need to use smaller block sizes to fit into SRAM.”

Probably not a substantial claim: may not pass S3 (load-bearing). May combine two claims.

Because the attention matrix is a function of low-rank QK^T, one can repeatedly load inputs Q,K and recompute the needed block of the attention matrix, significantly reducing HBM access.

“As a result, we can repeatedly load the inputs Q K and recompute the block of the attention matrix that we need, significantly reducing HBM access.”

Probably not a substantial claim: may not pass S3 (load-bearing).

Block Sparse, Longformer, and BigBird did not successfully run the backward pass with masking due to a bug in external libraries.

“These methods did not successfully run the backward pass with masking due to a bug in external libraries, so we measured them without masking to be generous.”

Probably not a substantial claim: may not pass S3 (load-bearing), S4 (based).

Memory footprint does not vary between runs.

“We only measure memory footprint once, since it does not vary between runs.”