Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
review_content: summary: This paper studies an important topic for current Transformer architectures: accelerating the speed and memory consumption for self attention operations. The proposed method leverages tiling and recomputation to reduce the memory access of HBM, avoiding the memory bound issues in standard self attention kernels. The authors also extend the proposed Fast-Attention with block-wise sparsity, and  is the fastest implementation among exsiting approximate attention methods. The experiments show the superior performance w.r.t. memory and speed against existing solutions. The paper is quite solid with technical details, yet it is still well-written and friendly to the general audience. The idea to exploit IO awareness should be inspiring for the community to develop more efficient algorithms and kernels. 
strengths_and_weaknesses: - The paper studies an important research question: designing fast and memory-cheap attention kernels for the general transformer architectures. The authors conduct extensive experiments and achieve impressive results w.r.t. training time and memory consumption across multiple tasks and network architectures.

- The proposed Fast-Attention can be readily combined with block-sparsity to allow longer sequences, and the results are still faster than existing approximate attention designs. 

- The writing is clear and self-contained, rigorous and yet friendly to the general audience. Though I am not an expert in low-level designs on CUDA but I find it rather pleasing to read the aritcle.
questions: 
- What are the separative gains of reduced HBM access and tiling? In Algorithm 0, the Q,K,V matrices are loaded to HBM by blocks. I wonder whether similar techniques (tiling) are applied in standard attention implementation?

- The baselines in Table 2 and 3 do not fuse the attention kernel. I wonder if there are any public baselines with fused kernels for fair comparisons?

- As a suggestion, it is a bit miss-leading to show that $\Theta(N^2 d^2 M^{-1})$ requires less memory access than $\Theta(Nd + N^2)$, which may be attributed to tiling (due to the presence of $M^{-1}$). The main reduction of HMB access should come from not storing the attention score $S$, which takes $\Theta(N^2)$ accesses. 

limitations: The authors have listed and discussed the potention directions for the limitations.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: The paper leverages fusion and tiling techniques to reduce the overall memory consumption of attention mechanism. The results corroborate with some of the existing work on how fusion/tiling would help to reduce the memory footprint and runtime.
strengths_and_weaknesses: **Strengths**

- The paper consider both forward pass and backward pass in the results that makes it more interesting compared to recent fusion/tiling approaches.

- The paper compares with the hand-tuned MLPerf results (although it is old by now) and show up to 14% improvement for one transformer model.

**Weaknesses**

- Fusion/tiling is not new and has been extensively explored in the literature. It is not clear what the contributions of this work are over previous methods.

- The paper seems to merely provide an engineering solution, merging multiple known approaches (fusion, tiling, algebraic aggregation) --- some already applied to transformer models.

- The paper needs to perform head-to-head comparison with existing fusion techniques for transformers.
questions: (1) Fusion and tiling are well-established techniques in the literature, which is also highlighted in the paper. The block-level Softmax calculation is also not new. What are the main contributions of FlashAttention over existing methods?

(2) There are recent proposals to leverage fusing opportunities in transformers, such as [FLAT](https://arxiv.org/abs/2107.06419) and [Self-Attention Does Not Need O(n$^2$) Memory](https://arxiv.org/abs/2112.05682), how is FlashAttention different from these work?

(3) Which version of PyTorch did you use in your implementations? I am wondering if newer versions of PyTorch have already added some supports for fusion.

(4) I understand the limitation to study the scalability of the FlashAttention method over larger number of GPUs. However, I suspect that as we increase the number of GPUs the benefit of FlashAttention over existing methods may plateau. Do you have any intuition on the performance projection of FlashAttention as number of GPUs increase?

(5) What is the overhead of recompute step compared to the end-to-end training runtime? 

(6) In Figure 2-Middle, I suspect if you include the data for larger block sizes the number of HBM accesses would increase due to the limited on-chip resources. Is my understanding correct? 

(7) In addition, in the most recent release of [MLPerf](https://mlcommons.org/en/training-normal-20/), the training runtime has been improved across all the models. I understand that the NeurIPS submission deadline is before the most recent MLPerf release, but I am wondering if you could provide some insights on whether FlashAttention would outperform the SOTA implementations?

(8) This is more a general comment about transformer models and the contribution of each layer to the overall runtime. Even for reasonably large sequence lengths, FF layers are still the main contributor to the end-to-end runtime. I am wondering if we are targeting the right problem to reduce the runtime of transformer models (I understand that memory consumption is still a large part of self-attention layers, but I doubt that this statement is correct for their runtime).

**Minor**

- There is a missing reference in p. 23. 
- In some of the tables, the accuracy of FlashAttention is lower than SOTA. Nothing major, but would be good to make sure that in all cases, the final accuracy matches the SOTA implementation. This would provide a better head-to-head comparison.
limitations: The authors summarizes some of the main limitations of this work in a dedicates section. As authors clearly mentioned, this is not a one-fit-all solutions for all the platforms and models. In this current state, the engineering efforts to adopt such approach may not be transferrable to other platforms and ML libraries.
review_title: null

final_score: rating: 8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
review_content: summary: This paper puts forward the hardware performance perspective in the design of attention algorithms. Specifically, it proposes to use established techniques such as tiling and re-computation to reduce the amount of high-bandwidth memory accesses for exact and approximate attention in transformers which is shown to be the main bottleneck in existing methods. Evaluation demonstrates that the proposed optimization enables training BERT and GPT-2 style architectures faster and achieves new state-of-the-art in the long range benchmark. Apart from training efficiency, the proposed attention improves performance in language modeling, long-document classification and it is the first result that is better than random in Path-X challenge.
strengths_and_weaknesses: **Strengths**
- In terms of quality, the exposition of ideas was clear and provided sufficient background for modern hardware which was necessary for the reader to follow. Writing overall was of high quality and the experimentation/proofs thorough. 
- The paper brings awareness about reducing the number of memory accesses for computing exact attention which has been considered not scalable because of its quadratic complexity and memory access handling of existing implementations.
- The key idea of the algorithm is to compute attention sequentially by loading into on-chip memory blocks of Q, K, V and store the partial result along with easy-to-store statistics required to compute the final result into standard accelerator memory. This is based on the observation that modern accelerators have improved a lot in speed and access to standard memory can become a bottleneck for several standard operations used in deep learning (softmax, batch norm, etc). This kind of optimization is applicable to approximate attention as well and could help improve the efficiency of other operators in the future. 
-  Experiments are well designed and provide evidence of the superiority of the proposed algorithm for both exact and approximate attention methods. It is quite important that this is the first attention that is exact and can extract information from such long contexts in high resolution and improves over all existing alternatives in tasks that require longer context. Should next generation of deep learning frameworks adopt such optimizations the impact can be big for scaling transformers to very long contexts without trading off performance. 

**Weaknesses** 
- Framing and motivation presents previous efforts as unaware of hardware performance and that they fail to speed up over existing transformers in terms of wall-clock time which is not universally true. 
-  The main concern is how much hardware/deep learning framework/experiment-specific the low-level optimization results are. Also, the memory and speed numbers are thoroughly reported for attention layer but they are mostly missing for full models (except training speed).
- The code is provided which is helpful but implementations for big pretrained models such as BERT and GPT-2 are not available which might be an issue for reproducing results.
questions: What is the runtime memory and decoding speed during training/inference for the models and baselines in Tables 1, 2, 4 and 5? An investigation beyond the runtime and memory of forward/backward pass of an individual attention layer is missing.  

The introduction does not cover the complete view of prior work. Some emphasized the fact that intermediate results are often written into GPU memory and proposed unbiased approximations for iterative computation of attention e.g. kernel-based ones to avoid that and make attention compute-bound instead of memory-bound. Also, reducing the objective to reduce quadratic compute complexity is still important for scaling to extremely long sequences as this paper shows (e.g. Table 13) which is not reflected well in the text.  

Do the authors plan to provide implementations for GPT-2 and BERT or the provided kernel would work?  Since it was highlighted that implementing a custom kernel is not straightforward this might be a concern reproducibility. 

The algorithm appears to be applicable to pretrained models with non-optimized exact attention and does not necessarily require training the model from scratch. Is that the case and is there any evidence to support it?  If not please clarify this in the text. 

For GPT-2 style models, could the authors explain how the causal masking is handled? It would be useful to include a description in section 3 for completeness and extra clarity. 
limitations: Yes. Limitations require a bit more discussion (see comments above).
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: ## Summary:
Problem: The paper addresses the problem of large memory requirements with high compute complexities of transformer models.
Solution: the paper attributes the above problem to the missing principle of IO-awareness, where the auditing is missing between different levels of GPU memory hierarchy. Which is a fair finding since the transformers are not designed with that in mind and moreover the GPUs at that time may not be having such complex memory hierarchies either or they may be evolving by then. The proposed solution is FLASHATTENTION
Results: The paper reports some exciting improvements to the speedup while also helping to improve the quality of models with the ability to process longer sequences. 

## Distinction from state-of-the-art:
The paper clearly distinguishes from the state-of-the-art methods in the sub-domain and positions within the body of the knowledge. 

strengths_and_weaknesses: ## Pros:
- The paper is written well, easy to understand and follow the concepts.

- The novelty of the paper is in identifying the underlying hardware causes of long time consumption by transformers and proposing the optimizations in the form of efficient data movement.

- The paper is a good fit for machine learning hardware software codesign in many respects.

## Cons:

- Line 54, typo, “out algorithm “ should be “our algorithm”

- One of the recent improvements are warp level matrix multiplications named as wmma in tensor core api. Since transformers involve a number of cascaded matrix multiplications starting from the input layer, wondering if you can compare the proposed approach with some of those mma/wmma based optimizations.

- I understand the focus of the paper is to perform efficient computations on GPU architectures. On the contrary, what if we  are operating on a resource constraint device (only in terms of memory but not that much on compute) where the SRAM size is not as big as that on the GPUs? Then, this approach looks to perform worse than the standard attention due to the N^2d^2. However, this can be out of the scope but can be listed as a limitation.

questions: Please refer to the strengths and weaknesses above.
limitations: The paper describes the limitations to a satisfactory level.
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: The paper proposes an algorithm to compute exact attention that reduces memory data transfers. 
They tile computation of softmax, which is a major bottleneck in running transformers. 
And they recompute attention matrix to avoid storing it in memory.
They also propose block-sparse attention in which they eliminate zero blocks in model with high sparsity.
This improves training time and produces more accurate models because it can model longer sequence.
They provide IO-complexity for the algorithm.
Previous methods that change the attention module and use approximate function of softmax. 
The proposed method was evaluated on transformer models (e.g: BERT, GPT2). They show improved training throughput and ability to use longer sequence lengths and showcase this in Path-X and Path-256 challenges.

strengths_and_weaknesses: Pros:
The paper tackles an important problem and demonstrates algorithm optimizations to improve training runtime and long-range accuracy of relevant transformer models.

The claims are supported by evaluations

Overall the paper is clear

Cons:
The techniques: tiling and recomputation are known algorithms that is applied into the new transformers workload. 
questions: Some suggestions and questions that could improve the paper are:

Remains to be seem if the method is relevant and if can be applied to mobile variation of transformers (mobile-ViT) and Vision transformers.

The effectiveness of the method is bounded by hardware constraints (SRAM size). It would interesting to see a roofline plot to demonstrate the compute-bound and memory access trade-off with and without flash-attention.  

Evaluation or comment on other memory technologies other than HBM and GPUs to expand coverage of this work

Compilers (e.g: TVM) have mechanisms for tiling and adding passes. It may be interesting to see if other optimizations in TVM could be applied here and further improve Flashattention
limitations: Yes. the paper appropriately discuss limitations and social impact
review_title: null

Meta-review: metareview: The authors study the problem of improving wall-clock time of performing exact and approximate attention in Transformer networks. Towards this, they identified the number of HBM access as a limiting factor. The authors propose FlashAttention which employs tiling and fusion operations to efficiently implement the attention operation on (GPU) SRAM while minimizing the number of HBM accesses. Even though the proposed approach has higher FLOP coun wrt. the standard attention implementation, it leads to reduced wall clock time due to the smaller number of HBM accesses. The proposed FlashAttention can be easily modified to obtain a sparse attention method, namely block-sparse FlashAttention. The paper is very well written and the authors demonstrate the utility of the proposed FlashAttention/block-sparse FlashAttenion on a wide range of settings. On many benchmarks, the proposed method allows for a longer context length, which leads to higher performance as compared to the existing approaches in the literature. 

Overall, the findings of the paper are very interesting and impactful. All the reviewers are quite positive about the paper. During the rebuttal phase the authors have included additional experiments that further strengthen the contributions and value of the paper. Some of the reviewers pointed out that the paper builds on well-known techniques like tiling and fusion. However, given that the paper makes significant improvements on a timely and actively studied problem of making the attention operation efficient and the results in the paper have the potential to inspire many follow-up explorations, I recommend that the paper be accepted in NeurIPS 2022.
confidence: Certain
Decision: Accept